Add periodic FlowPath magnetization remapping

KomaMRIBase/src/motion/Motion.jl

@@ -145,7 +145,7 @@ function path(dx, dy, dz, time=TimeRange(t_start=zero(eltype(dx)), t_end=eps(elt
 end
 
 """
-    fp = flowpath(dx, dy, dz, spin_reset, time, spins)
+    fp = flowpath(dx, dy, dz, spin_reset, time, spins; cycle_map=nothing)
 
 # Arguments
 - `dx`: (`::AbstractArray{T<:Real}`, `[m]`) displacements in x
@@ -155,6 +155,9 @@ end
 - `time`: (`::TimeCurve{T<:Real}`) time information about the motion
 - `spins`: (`::AbstractSpinSpan`) spin indexes affected by the motion
 
+# Keywords
+- `cycle_map`: (`::Union{Nothing,AbstractVector{Int}}`) optional periodic magnetization map
+
 # Returns
 - `fp`: (`::Motion`) Motion struct with [`FlowPath`](@ref) action
 
@@ -170,8 +173,10 @@ julia> fp = flowpath(
        )
 ```
 """
-function flowpath(dx, dy, dz, spin_reset, time=TimeRange(t_start=zero(eltype(dx)), t_end=eps(eltype(dx))), spins=AllSpins())
-    return Motion(FlowPath(dx, dy, dz, spin_reset), time, spins)
+function flowpath(dx, dy, dz, spin_reset, time=TimeRange(t_start=zero(eltype(dx)), t_end=eps(eltype(dx))), spins=AllSpins(); cycle_map=nothing)
+    !isnothing(cycle_map) && !time.periodic &&
+        throw(ArgumentError("cycle_map requires a periodic TimeCurve."))
+    return Motion(FlowPath(dx, dy, dz, spin_reset; cycle_map), time, spins)
 end
 
 """ Compare two Motions """
@@ -219,6 +224,20 @@ end
 # Auxiliary functions
 times(m::Motion) = times(m.time)
 is_composable(m::Motion) = is_composable(m.action)
+cycle_map(::AbstractAction) = nothing
+cycle_map(action::FlowPath) = action.cycle_map
+cycle_remap(m::Motion) = isnothing(cycle_map(m.action)) ? nothing : m
+
+function cycle_remap_times(motion, t_max)
+    m = cycle_remap(motion)
+    t = typeof(float(t_max))[]
+    isnothing(m) && return t
+    add_cycle_end_times!(t, m.time.t_start, m.time.t_end, m.time.periods)
+    period = sum((m.time.t_end - m.time.t_start) .* m.time.periods)
+    extend_periodic!(t, t_max, period, Val(m.time.periodic))
+    filter!(x -> m.time.t_start < x <= t_max, t)
+    return sort!(unique!(t))
+end
 
 """
     add_key_time_points!(t, motion)
@@ -234,6 +253,7 @@ function add_key_time_points!(t, a, t_start::T, t_end::T, periods, periodic) whe
     t_max = maximum(t)
     add_period_times!(aux, t_start, t_end, periods)
     add_reset_times!(aux, a, t_start, t_end, periods)
+    add_cycle_remap_times!(aux, a, t_start, t_end, periods)
     extend_periodic!(aux, t_max, period, Val(periodic))
     append!(t, aux[aux .<= t_max])
     return nothing
@@ -266,6 +286,18 @@ function add_period_times!(t, t_start, t_end, periods)
     return nothing
 end
 
+function add_cycle_end_times!(t, t_start, t_end, periods)
+    period_times = times([t_start, t_end], t_start, t_end, periods)
+    append!(t, @view(period_times[2:2:end]))
+    return nothing
+end
+
+add_cycle_remap_times!(t, ::AbstractAction, t_start, t_end, periods) = nothing
+function add_cycle_remap_times!(t, action::FlowPath, t_start, t_end, periods)
+    isnothing(action.cycle_map) || add_cycle_end_times!(t, t_start, t_end, periods)
+    return nothing
+end
+
 """
     add_reset_times!(t, action, t_start, t_end, periods)
 """ 

KomaMRIBase/src/motion/MotionList.jl

@@ -130,6 +130,12 @@ end
 """ MotionList length """
 Base.length(m::MotionList) = length(m.motions)
 
+function cycle_remap(ml::MotionList)
+    remaps = filter(!isnothing, cycle_remap.(ml.motions))
+    length(remaps) <= 1 || throw(ArgumentError("Only one cycle-remapped FlowPath is supported per phantom."))
+    return isempty(remaps) ? nothing : only(remaps)
+end
+
 function get_spin_coords(
     ml::MotionList{T}, x::AbstractVector{T}, y::AbstractVector{T}, z::AbstractVector{T}, t
 ) where {T<:Real}

KomaMRIBase/src/motion/NoMotion.jl

@@ -50,3 +50,4 @@ function get_spin_coords(
     return x, y, z
 end
 add_key_time_points!(t, ::NoMotion) = nothing
+cycle_remap(::NoMotion) = nothing

KomaMRIBase/src/motion/actions/arbitraryactions/FlowPath.jl

@@ -1,5 +1,5 @@
 @doc raw"""
-    f = FlowPath(dx, dy, dz, spin_reset)
+    f = FlowPath(dx, dy, dz, spin_reset; cycle_map=nothing)
 
 FlowPath struct. This action is the same as `Path`, 
 except that it includes an additional field, called `spin_reset`, 
@@ -10,11 +10,18 @@ state of these spins must be reset during the simulation.
 As with the `dx`, `dy` and `dz` matrices, `spin_reset`
 has a size of (``N_{spins} \times \; N_{discrete\,times}``).
 
+For periodic, non-closed trajectories, `cycle_map` can provide a precomputed
+destination-to-source particle mapping. At each cycle boundary, particle `i`
+receives the magnetization of particle `cycle_map[i]`. The map uses 1-based
+indices local to this `FlowPath`. Koma does not calculate this map. A map is only
+meaningful for the whole particle set, so slicing a `FlowPath` drops it.
+
 # Arguments
 - `dx`: (`::AbstractArray{T<:Real}`, `[m]`) displacements in x
 - `dy`: (`::AbstractArray{T<:Real}`, `[m]`) displacements in y 
 - `dz`: (`::AbstractArray{T<:Real}`, `[m]`) displacements in z
 - `spin_reset`: (`::AbstractArray{Bool}`) reset spin state flags
+- `cycle_map`: (`::Union{Nothing,AbstractVector{Int}}`) optional periodic magnetization map
 
 # Returns
 - `f`: (`::FlowPath`) FlowPath struct
@@ -34,11 +41,24 @@ julia> f = FlowPath(
     dy::AbstractArray{T}
     dz::AbstractArray{T}
     spin_reset::AbstractArray{Bool}
+    cycle_map::Union{Nothing,Vector{Int}} = nothing
+end
+
+FlowPath(dx, dy, dz, spin_reset::BitMatrix; cycle_map=nothing) = FlowPath(dx, dy, dz, collect(spin_reset); cycle_map)
+function FlowPath(dx, dy, dz, spin_reset::AbstractArray{Bool}; cycle_map=nothing)
+    isnothing(cycle_map) ||
+        (eltype(cycle_map) !== Bool && length(cycle_map) == size(dx, 1) && all(in(axes(dx, 1)), cycle_map)) ||
+        throw(ArgumentError("cycle_map must hold one source particle index per FlowPath particle."))
+    return FlowPath(dx, dy, dz, spin_reset, isnothing(cycle_map) ? nothing : Vector{Int}(cycle_map))
 end
 
-FlowPath(dx::AbstractArray{T}, dy::AbstractArray{T}, dz::AbstractArray{T}, spin_reset::BitMatrix) where T<:Real = FlowPath(dx, dy, dz, collect(spin_reset))
+# A cycle_map indexes the whole particle set, so a sub-group cannot carry it.
+_sliced_cycle_map(a::FlowPath, p) = (p isa Colon || p == 1:size(a.dx, 1)) ? a.cycle_map : nothing
+
+Base.getindex(a::FlowPath, p) = FlowPath(a.dx[p, :], a.dy[p, :], a.dz[p, :], a.spin_reset[p, :], _sliced_cycle_map(a, p))
+Base.view(a::FlowPath, p) = @views FlowPath(a.dx[p, :], a.dy[p, :], a.dz[p, :], a.spin_reset[p, :], _sliced_cycle_map(a, p))
 
 function add_reset_times!(t, a::FlowPath, t_start, t_end, periods)
     aux = t_start .+ (t_end - t_start)/(size(a.spin_reset)[2]-1) * (getindex.(findall(a.spin_reset .== 1), 2) .- 1)
     append!(t, times(aux, t_start, t_end, periods) .- MIN_RISE_TIME)
-end
+end

KomaMRIBase/test/runtests.jl

@@ -2116,6 +2116,24 @@ end
         @test xt == ph.x .+ dx
         @test yt == ph.y .+ dy
         @test zt == ph.z .+ dz
+
+        # Periodic magnetization remapping uses a supplied destination-to-source map.
+        Ns = 3
+        trajectory = zeros(Ns, 2)
+        cycle_map = [2, 2, 1]
+        fp = flowpath(trajectory, trajectory, trajectory, falses(Ns, 2), Periodic(1.0, 1.0); cycle_map)
+        @test fp.action.cycle_map == cycle_map
+        @test KomaMRIBase.cycle_remap_times(fp, 2.5) == [1.0, 2.0]
+        @test_throws ArgumentError flowpath(trajectory, trajectory, trajectory, falses(Ns, 2), TimeRange(0.0, 1.0); cycle_map)
+        @test_throws ArgumentError FlowPath(trajectory, trajectory, trajectory, falses(Ns, 2); cycle_map=[1, 2])
+        @test_throws ArgumentError FlowPath(trajectory, trajectory, trajectory, falses(Ns, 2); cycle_map=[1, 2, 4])
+
+        # A cycle_map indexes the whole particle set, so sub-groups drop it.
+        ph = Phantom(x=zeros(Ns), motion=fp)
+        @test ph[:] == ph
+        @test isnothing(ph[1:2].motion.action.cycle_map)
+        @test isnothing(view(ph, 1:2).motion.action.cycle_map)
+        @test ph[1:Ns].motion.action.cycle_map == cycle_map
     end
     @testset "Translate + Rotate" begin
         ph = Phantom(x=[1.0, 1.0, -1.0, -1.0], y=[1.0, -1.0, 1.0, -1.0])

KomaMRICore/src/simulation/Flow.jl

@@ -1,6 +1,36 @@
 spin_coordinates(motion, x, y, z, t) = get_spin_coords(motion, x, y, z, t)
 spin_coordinates(::NoMotion, x, y, z, t) = x, y, z
 
+# Global particle index each particle takes its magnetization from at a cycle boundary.
+function cycle_remap_sources(obj, motion)
+    affected = KomaMRIBase.get_indexing_range(KomaMRIBase.expand(motion.spins, length(obj)))
+    source = collect(eachindex(obj.ρ))
+    source[affected] .= affected[motion.action.cycle_map]
+    return source
+end
+
+# Cycle boundaries are motion key times, so discretize always samples them. The grid
+# value can differ from a freshly computed boundary by a few ulp, since it round-trips
+# through a per-block time offset.
+function cycle_remap_break_indices(seqd, motion)
+    breaks = Int[]
+    for t in KomaMRIBase.cycle_remap_times(motion, last(seqd.t))
+        t < last(seqd.t) || continue
+        tol = max(KomaMRIBase.MAX_STEP_TIME_SNAP_TOL, 4eps(t))
+        i = searchsortedlast(seqd.t, t + tol)
+        i >= firstindex(seqd.t) && abs(seqd.t[i] - t) <= tol ||
+            error("No sampling time within $tol of cycle boundary $t; motion key times are missing from the simulation grid.")
+        push!(breaks, i)
+    end
+    return breaks
+end
+
+function remap_magnetization!(M::Mag, source)
+    M.xy .= M.xy[source]
+    M.z  .= M.z[source]
+    return nothing
+end
+
 outflow_spin_reset_at!(spin_state, t, i, motion; replace_by=0) =
    outflow_spin_reset!(spin_state, t[i, :], motion; replace_by)
 outflow_spin_reset_at!(spin_state, t, i, ::NoMotion; replace_by=0) = nothing

KomaMRICore/src/simulation/Functors.jl

@@ -122,7 +122,7 @@ adapt_storage(T::Type{<:Real}, xs::MotionList) = MotionList(paramtype.(T, xs.mot
 @functor Rotate
 @functor HeartBeat
 @functor Path
-@functor FlowPath
+@functor FlowPath (dx, dy, dz, spin_reset)
 @functor TimeCurve
 # Spinor
 @functor Spinor

KomaMRICore/src/simulation/SimulatorCore.jl

@@ -178,6 +178,8 @@ function run_sim_time_iter!(
     excitation_groupsize=DEFAULT_EXCITATION_GROUPSIZE,
     parts=[1:length(seqd)],
     excitation_bool=ones(Bool, size(parts)),
+    remap_before=falses(length(parts)),
+    remap_sources=nothing,
     sim_params=Dict{String,Any}(),
     callbacks=(),
 ) where {T<:Real}
@@ -190,6 +192,7 @@ function run_sim_time_iter!(
     (excitation_groupsize % 32 == 0) || throw("Groupsize must be a multiple of 32")
 
     for (block, p) in enumerate(parts)
+        remap_before[block] && remap_magnetization!(Xt, remap_sources)
         seqd_block = @view seqd[p]
         # Params
         Nadc = sum(seqd_block.ADC[2:end]) # if ADC[1] == true, that is handled by the previous block
@@ -250,10 +253,10 @@ function split_range(r, max_block_length, eval_intervals_per_step)
     return [i:min(i + block_length, last(r)) for i in first(r):block_length:(last(r) - 1)]
 end
 
-function get_sim_ranges(seqd::DiscreteSequence; max_block_length=Inf, max_rf_block_length=Inf, eval_intervals_per_step=1)
+function get_sim_ranges(seqd::DiscreteSequence; max_block_length=Inf, max_rf_block_length=Inf, eval_intervals_per_step=1, breaks=Int[])
     ranges, ranges_bool = UnitRange{Int}[], Bool[]; isempty(seqd.Δt) && return ranges, ranges_bool
 
-    starts = [firstindex(seqd.Δt); findall(seqd.excitation_bool[2:end] .!= seqd.excitation_bool[1:(end - 1)]) .+ 1]
+    starts = sort!(unique!([firstindex(seqd.Δt); findall(seqd.excitation_bool[2:end] .!= seqd.excitation_bool[1:(end - 1)]) .+ 1; breaks]))
     stops = [starts[2:end] .- 1; lastindex(seqd.Δt)]
     for (start, stop) in zip(starts, stops)
         is_excitation = seqd.excitation_bool[start]
@@ -349,12 +352,17 @@ function simulate(
     end
     # Simulation init
     seqd = discretize(seq; sampling_rule, motion=obj.motion, freq_in_phase=sim_params["freq_in_phase"]) # Sampling of Sequence waveforms
+    remap_motion = KomaMRIBase.cycle_remap(obj.motion)
+    remap_breaks = isnothing(remap_motion) ? Int[] : cycle_remap_break_indices(seqd, obj.motion)
+    remap_sources = isnothing(remap_motion) ? nothing : cycle_remap_sources(obj, remap_motion)
     parts, excitation_bool = get_sim_ranges(
         seqd;
         max_block_length=sim_params["max_block_length"],
         max_rf_block_length=sim_params["max_rf_block_length"],
         eval_intervals_per_step=eval_intervals_per_step(sim_method),
+        breaks=remap_breaks,
     ) # Generating simulation blocks
+    remap_before = [first(p) in remap_breaks for p in parts]
     Nblocks = length(parts)
     t_sim_parts = [seqd.t[p[1]] for p in parts]
     append!(t_sim_parts, seqd.t[end])
@@ -387,6 +395,7 @@ function simulate(
         seqd = seqd |> gpu #DiscreteSequence
         Xt = Xt |> gpu #SpinStateRepresentation
         sig = sig |> gpu #Signal
+        remap_sources = remap_sources |> gpu #Cycle remap indexes
     end
 
     # Simulation
@@ -408,6 +417,8 @@ function simulate(
         excitation_groupsize=sim_params["gpu_groupsize_excitation"],
         parts,
         excitation_bool,
+        remap_before,
+        remap_sources,
         sim_params,
         callbacks=all_callbacks,
     )

KomaMRICore/test/runtests.jl

@@ -846,6 +846,53 @@ end
 end
 
 # --------- Motion-related tests -------------
+@testitem "Periodic FlowPath magnetization remapping" tags=[:core, :motion] begin
+    include("initialize_backend.jl")
+
+    N = 3
+    trajectory = zeros(N, 2)
+    cycle_map = [2, 2, 1]
+    motion = flowpath(trajectory, trajectory, trajectory, falses(N, 2), Periodic(1.0, 1.0); cycle_map)
+    obj = Phantom(x=zeros(N), ρ=[1.0, 2.0, 3.0], T1=fill(Inf, N), T2=fill(Inf, N), motion=motion)
+    seq = Sequence([Grad(0.0, 2.1)])
+
+    seqd = discretize(seq; motion)
+    breaks = KomaMRICore.cycle_remap_break_indices(seqd, motion)
+    ranges, _ = KomaMRICore.get_sim_ranges(seqd; breaks)
+    @test seqd.t[breaks] == [1.0, 2.0]
+    @test all(any(first(r) == i for r in ranges) for i in breaks)
+    @test KomaMRICore.cycle_remap_sources(obj, motion) == cycle_map
+    @test eltype(f32(obj).motion.action.cycle_map) === Int
+
+    # Two boundaries (t = 1, 2 s) apply the map twice: [1,2,3] -> [2,2,1] -> [2,2,2]
+    sim_params = Dict{String,Any}("gpu" => USE_GPU, "return_type" => "state")
+    state = simulate(obj, seq, Scanner(); sim_params, verbose=false)
+    @test state.z ≈ Float32[2, 2, 2]
+end
+
+# Cycle boundaries are motion key times, so discretize promotes them to integration-step
+# boundaries and the forced block break stays aligned with the Magnus node stencil.
+@testitem "Cycle-remap breaks align with Magnus steps" tags=[:core, :motion] begin
+    include("initialize_backend.jl")
+
+    sim_method = BlochMagnusBGL4()
+    seq = PulseDesigner.RF_hard(10e-6, 1e-3, Scanner())
+    trajectory = zeros(1, 2)
+    # 0.33 ms period: boundaries land inside the RF pulse, off the nominal Δt_rf grid
+    motion = flowpath(trajectory, trajectory, trajectory, falses(1, 2), Periodic(0.33e-3, 1.0); cycle_map=[1])
+    sim_params = KomaMRICore.default_sim_params(Dict{String,Any}("sim_method" => sim_method))
+    seqd = discretize(seq; sampling_rule=KomaMRICore.simulation_sampling_rule(sim_method, sim_params), motion)
+
+    breaks = KomaMRICore.cycle_remap_break_indices(seqd, motion)
+    rf_start = first(findall(seqd.excitation_bool))
+    eval_stride = KomaMRICore.eval_intervals_per_step(sim_method)
+    parts, excitation_bool = KomaMRICore.get_sim_ranges(seqd; breaks)
+    @test !isempty(breaks)
+    @test all(iszero((b - rf_start) % eval_stride) for b in breaks)
+    @test all(any(first(p) == b for p in parts) for b in breaks)
+    @test all(excitation_bool[findfirst(p -> first(p) == b, parts)] for b in breaks)
+end
+
 # We compare with the result given by OrdinaryDiffEqTsit5
 @testitem "Motion" tags=[:core, :motion] begin
     using OrdinaryDiffEqTsit5

KomaMRIFiles/src/Phantom/Phantom.jl

@@ -75,7 +75,8 @@ function import_motion_field!(motion_fields::Array, motion::HDF5.Group, name::St
             for subname in fieldnames(subtype_vector[i]) # dx, dy, dz, pitch, roll...
                 key = string(subname)
                 if !(key in ["t_start", "t_end"])
-                    subfield_value = key in keys(field_group) ? read(field_group, key) : read_attribute(field_group, key)
+                    subfield_value = key in keys(field_group) ? read(field_group, key) :
+                        haskey(HDF5.attributes(field_group), key) ? read_attribute(field_group, key) : nothing
                     import_motion_subfield!(motion_subfields, subfield_value, key, T)
                 end
             end
@@ -88,6 +89,11 @@ function import_motion_subfield!(motion_subfields::Array, subfield_value::Union{
     push!(motion_subfields, subfield_value)
     return nothing
 end
+""" Subfields absent from the file (e.g. an unset cycle_map) default to nothing """
+function import_motion_subfield!(motion_subfields::Array, ::Nothing, key::String, T::Type{<:Real})
+    push!(motion_subfields, nothing)
+    return nothing
+end
 function import_motion_subfield!(motion_subfields::Array, subfield_value::String, key::String, T::Type{<:Real})
     if subfield_value in ["true", "false"]
         return push!(motion_subfields, subfield_value == "true" ? true : false)
@@ -181,3 +187,4 @@ end
 function export_motion_subfield!(field_group::HDF5.Group, subfield::CenterOfMass, subname::String)
     field_group[subname] = "CenterOfMass"
 end
+export_motion_subfield!(field_group::HDF5.Group, ::Nothing, subname::String) = nothing

KomaMRIFiles/test/runtests.jl

@@ -57,6 +57,18 @@ end
         obj2 = read_phantom(filename)
         @test obj1 == obj2
     end
+    @testset "Periodic FlowPath remap" begin
+        pth = @__DIR__
+        filename = pth * "/test_files/phantom/flowpath_cyclemap_w.phantom"
+        Ns = 3
+        trajectory = zeros(Ns, 2)
+        cycle_map = [2, 2, 1]
+        obj1 = Phantom(x=zeros(Ns), motion=flowpath(trajectory, trajectory, trajectory, falses(Ns, 2), Periodic(1.0, 1.0); cycle_map))
+        write_phantom(obj1, filename)
+        obj2 = read_phantom(filename)
+        @test obj1 == obj2
+        @test obj2.motion.action.cycle_map == cycle_map
+    end
 end
 
 @testitem "Pulseq" tags=[:files, :pulseq] begin