Skip to content

Add SimulationState object to make checkpointing easier #655

@cncastillo

Description

@cncastillo

Feature Request

For checkpointing, we want a: "minimal dynamic state needed to continue integrating the same Bloch system".

This could be achieved by

struct SimulationState{T, ArrT<:AbtractArray{T}, SpinSt<:SpinStateRepresentation{T}}
    last_block_idx::Int # Needed to skip simulated blocks
    sig::ArrT           # Acquired signal so far (filled with 0s for remaining points)
    Xt::SpinSt          # Current spin state
    sim_params::Dict{String,Any}
end

Then, if we create a simulate method than can be called like this:

seq = load_seq(...)
obj = load_phatom(...) 
sys = Scanner(...)
spin_state_file_path = "spin_state.jld2"
if file_exists(spin_state_file_path)
    sim_state = @jld2_load(spin_state_file_path)
    # verify that the hash of seq, obj, sys, are the same as the ones stored
    ...
else
    sim_state = sim_state_init(sim_params) # initializes sig and Xt, last_block_idx=0
end
sig = simulate(seq, obj, sys, sim_state)

In run_sim_time_iter the call back would look like this:

function run_sim_time_iter(...; callbacks=SpinSavingCallback[])
    for block_idx in 1:Nblocks
        # 1. run GPU kernel over this block (updates Xt and sig)
        ...

        # 2. host-side "SavingCallback" analogue
        for cb in callbacks
                cb.fun(...)
        end
    end
end

This proposal is a work in progress.

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions