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
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions