Magnus-based methods, RF reference frame, and block splitting (#612)

* Port simulation block timing fixes

* Handle RF frequency modulation frame

* Bump package compat for RF frame changes

* Update simulation docs and benchmarks

* Refine RF frame phase display

* Restore legacy Pulseq RF use fallback

* Fix get_samples doc attachment

Author: cncastillo

KomaMRIBase/Project.toml

@@ -1,7 +1,7 @@
 name = "KomaMRIBase"
 uuid = "d0bc0b20-b151-4d03-b2a4-6ca51751cb9c"
 
-version = "0.10.1"
+version = "0.11.0"
 authors = ["Carlos Castillo Passi <cncastillo@uc.cl>"]
 
 [workspace]

KomaMRIBase/src/datatypes/Sequence.jl

@@ -706,9 +706,11 @@ function get_samples(seq::Sequence, range; events=[:rf, :gr, :adc], freq_in_phas
         t_Δf = reduce(vcat, [T0[i] .+ times(seq.RF[1,i], :Δf)  for i in range])
         A_rf = reduce(vcat, [ampls(seq.RF[1,i]; freq_in_phase) for i in range])
         A_Δf = reduce(vcat, [freqs(seq.RF[1,i])                for i in range])
+        A_ψ  = reduce(vcat, [rf_frame_phase(seq.RF[1,i])       for i in range])
         rf_samples = (
             rf  = fill_if_empty((t = t_rf, A = A_rf)),
-            Δf  = fill_if_empty((t = t_Δf, A = A_Δf))
+            Δf  = fill_if_empty((t = t_Δf, A = A_Δf)),
+            ψ   = fill_if_empty((t = t_Δf, A = A_ψ))
 		)
     end
     # Gradients
@@ -742,6 +744,8 @@ function get_samples(seq::Sequence, range; events=[:rf, :gr, :adc], freq_in_phas
 end
 get_samples(seq::Sequence; kwargs...) = get_samples(seq, 1:length(seq); kwargs...)
 
+wrap_to_pi(x) = mod(x + π, 2π) - π
+
 """
     Gx, Gy, Gz = get_grads(seq, t::Vector)
     Gx, Gy, Gz = get_grads(seq, t::Matrix)
@@ -769,9 +773,9 @@ function get_grads(seq, t::Union{Vector, Matrix})
 end
 
 """
-    B1, Δf_rf  = get_rfs(seq::Sequence, t)
+    B1, Δf_rf, ψ  = get_rfs(seq::Sequence, t)
 
-Returns the RF pulses and the delta frequency.
+Returns the RF pulses, delta frequency, and RF rotating-frame phase.
 
 # Arguments
 - `seq`: (`::Sequence`) Sequence struct
@@ -780,13 +784,15 @@ Returns the RF pulses and the delta frequency.
 # Returns
 - `B1`: (`1-row ::Matrix{ComplexF64}`, `[T]`) vector of RF pulses
 - `Δf_rf`: (`1-row ::Matrix{Float64}`, `[Hz]`) delta frequency vector
+- `ψ`: (`1-row ::Matrix{Float64}`, `[rad]`) RF rotating-frame phase vector
 """
 function get_rfs(seq, t::Union{Vector, Matrix})
     rf_samples = get_samples(seq; events=[:rf])
     for event in rf_samples
         Interpolations.deduplicate_knots!(event.t; move_knots=true)
     end
-    return Tuple(linear_interpolation(event..., extrapolation_bc=0.0).(t) for event in rf_samples)
+    B1, Δf, ψ = Tuple(linear_interpolation(event..., extrapolation_bc=0.0).(t) for event in rf_samples)
+    return B1, Δf, wrap_to_pi.(ψ)
 end
 
 """

KomaMRIBase/src/datatypes/simulation/DiscreteSequence.jl

@@ -1,5 +1,5 @@
 """
-    seqd = DiscreteSequence(Gx, Gy, Gz, B1, Δf, ADC, t, Δt)
+    seqd = DiscreteSequence(Gx, Gy, Gz, B1, Δf, ψ, ADC, t, Δt)
 
 A sampled version of a Sequence struct, containing vectors for event amplitudes at specified
 times. DiscreteSequence is the struct used for simulation.
@@ -10,6 +10,7 @@ times. DiscreteSequence is the struct used for simulation.
 - `Gz`: (`::AbstractVector{T<:Real}`, `[T/m]`) z-gradient vector
 - `B1`: (`::AbstractVector{Complex{T<:Real}}`, `[T]`) RF amplitude vector
 - `Δf`: (`::AbstractVector{T<:Real}`, `[Hz]`) RF carrier frequency displacement vector
+- `ψ`: (`::AbstractVector{T<:Real}`, `[rad]`) RF rotating-frame phase vector
 - `ADC`: (`::AbstractVector{Bool}`) ADC sample vector
 - `t`: (`::AbstractVector{T<:Real}`, `[s]`) time vector
 - `Δt`: (`::AbstractVector{T<:Real}`, `[s]`) delta time vector
@@ -23,6 +24,7 @@ struct DiscreteSequence{T<:Real}
     Gz::AbstractVector{T}
     B1::AbstractVector{Complex{T}}
     Δf::AbstractVector{T}
+    ψ::AbstractVector{T}
     ADC::AbstractVector{Bool}
     t::AbstractVector{T}
     Δt::AbstractVector{T}
@@ -35,36 +37,39 @@ Base.getindex(seq::DiscreteSequence, i::Integer) = begin
                      seq.Gz[i, :],
                      seq.B1[i, :],
                      seq.Δf[i, :],
+                     seq.ψ[i, :],
                      seq.ADC[i, :],
                      seq.t[i, :],
                      seq.Δt[i, :])
 end
 Base.getindex(seq::DiscreteSequence, i::UnitRange) = begin
-    DiscreteSequence(seq.Gx[i.start:i.stop+1],
-                     seq.Gy[i.start:i.stop+1],
-                     seq.Gz[i.start:i.stop+1],
-                     seq.B1[i.start:i.stop+1],
-                     seq.Δf[i.start:i.stop+1],
+    DiscreteSequence(seq.Gx[i],
+                     seq.Gy[i],
+                     seq.Gz[i],
+                     seq.B1[i],
+                     seq.Δf[i],
+                     seq.ψ[i],
                      seq.ADC[i],
-                     seq.t[i.start:i.stop+1],
-                     seq.Δt[i])
+                     seq.t[i],
+                     seq.Δt[i.start:i.stop-1])
 end
 Base.view(seq::DiscreteSequence, i::UnitRange) = begin
-    @views DiscreteSequence(seq.Gx[i.start:i.stop+1],
-                     seq.Gy[i.start:i.stop+1],
-                     seq.Gz[i.start:i.stop+1],
-                     seq.B1[i.start:i.stop+1],
-                     seq.Δf[i.start:i.stop+1],
+    @views DiscreteSequence(seq.Gx[i],
+                     seq.Gy[i],
+                     seq.Gz[i],
+                     seq.B1[i],
+                     seq.Δf[i],
+                     seq.ψ[i],
                      seq.ADC[i],
-                     seq.t[i.start:i.stop+1],
-                     seq.Δt[i])
+                     seq.t[i],
+                     seq.Δt[i.start:i.stop-1])
 end
 Base.iterate(seq::DiscreteSequence) = (seq[1], 2)
 Base.iterate(seq::DiscreteSequence, i) = (i <= length(seq)) ? (seq[i], i+1) : nothing
 
-is_GR_on(seq::DiscreteSequence) =  sum(abs.([seq.Gx[1:end-1]; seq.Gy[1:end-1]; seq.Gz[1:end-1]])) != 0
-is_RF_on(seq::DiscreteSequence) =  sum(abs.(seq.B1[1:end-1])) != 0
-is_ADC_on(seq::DiscreteSequence) = sum(abs.(seq.ADC[1:end-1])) != 0
+is_GR_on(seq::DiscreteSequence) =  sum(abs.([seq.Gx; seq.Gy; seq.Gz])) != 0
+is_RF_on(seq::DiscreteSequence) =  sum(abs.(seq.B1)) != 0
+is_ADC_on(seq::DiscreteSequence) = sum(abs.(seq.ADC)) != 0
 is_GR_off(seq::DiscreteSequence) =  !is_GR_on(seq)
 is_RF_off(seq::DiscreteSequence) =  !is_RF_on(seq)
 is_ADC_off(seq::DiscreteSequence) = !is_ADC_on(seq)
@@ -87,12 +92,12 @@ based on simulation parameters.
 """
 function discretize(seq::Sequence; sampling_params=default_sampling_params(), motion=NoMotion())
     t, Δt      = get_variable_times(seq; Δt=sampling_params["Δt"], Δt_rf=sampling_params["Δt_rf"], motion=motion)
-    B1, Δf     = get_rfs(seq, t)
+    B1, Δf, ψ  = get_rfs(seq, t)
     Gx, Gy, Gz = get_grads(seq, t)
     tadc       = get_adc_sampling_times(seq)
     tadc_set = Set(tadc)
     ADCflag = [tt in tadc_set for tt in t] # Displaced 1 dt, sig[i]=S(ti+dt)
-    seqd       = DiscreteSequence(Gx, Gy, Gz, complex.(B1), Δf, ADCflag, t, Δt)
+    seqd       = DiscreteSequence(Gx, Gy, Gz, complex.(B1), Δf, ψ, ADCflag, t, Δt)
     return seqd
 end
 

KomaMRIBase/src/timing/KeyValuesCalculation.jl

@@ -53,6 +53,28 @@ function freqs(rf::RF)
     end
     return _shape_samples(rf.Δf)
 end
+
+"""
+    ψ = rf_frame_phase(rf::RF)
+
+RF rotating-frame phase samples from frequency modulation, with `ψ` zeroed at the RF center.
+"""
+function rf_frame_phase(rf::RF)
+    f = freqs(rf)
+    isempty(f) && return Float64[]
+    t = times(rf, :Δf)
+    Interpolations.deduplicate_knots!(t; move_knots=true)
+    center_time = rf.delay + get_RF_center(rf)
+    t_aug = sort!([t; center_time])
+    f_aug = linear_interpolation(t, f, extrapolation_bc=Interpolations.Flat()).(t_aug)
+    ψ_aug = -2π .* [0.0; cumtrapz(diff(t_aug), f_aug)]
+    ψ = ψ_aug[searchsortedfirst.(Ref(t_aug), t)]
+    ψ_center = ψ_aug[searchsortedfirst(t_aug, center_time)]
+    ψ .-= ψ_center
+    ψ[1] = zero(eltype(ψ))
+    ψ[end] = zero(eltype(ψ))
+    return ψ
+end
 function ampls(adc::ADC)
     if !is_on(adc)
         return Bool[]

KomaMRIBase/src/timing/TimeStepCalculation.jl

@@ -99,11 +99,11 @@ function get_variable_times(seq; Δt=1e-3, Δt_rf=1e-5, motion=NoMotion())
 		t0 = T0[i]
 		if is_RF_on(s)
 			y = s.RF[1]
-			delay, T, center = y.delay, y.T, y.center
+			delay, T = y.delay, y.T
 			t1 = t0 + delay
 			t2 = t1 + sum(T)
-			rf0 = t0 + delay + center
-			taux = points_from_key_times([t1, t1 + ϵ, rf0, t2 - ϵ, t2]; dt=Δt_rf)
+			tc = t0 + delay + get_RF_center(y)
+			taux = points_from_key_times(sort([t1, t1 + ϵ, tc, t2 - ϵ, t2]); dt=Δt_rf)
             append!(t_block, taux)
 		end
 		if is_GR_on(s)
@@ -112,7 +112,7 @@ function get_variable_times(seq; Δt=1e-3, Δt_rf=1e-5, motion=NoMotion())
 			if is_Gy_on(s) append!(active_gradients, s.GR.y) end
 			if is_Gz_on(s) append!(active_gradients, s.GR.z) end
 			for y = active_gradients
-				ts = _gradient_interpolation_samples(y).t .+ t0
+				ts = times(y) .+ t0
 				taux = points_from_key_times([ts[1] + ϵ; ts; ts[end] - ϵ]; dt=Δt)
                 append!(t_block, taux)
 			end

KomaMRIBase/test/runtests.jl

@@ -539,6 +539,26 @@ using TestItems, TestItemRunner
             @test KomaMRIBase.get_char_from_RF_use(use) == char
             @test KomaMRIBase.get_RF_use_from_char(Val(char)) == use
         end
+
+        # Constant frequency offset: ψ is centered so the RF center has zero frame phase.
+        rf = RF([1.0, 2.0, 1.0] .* 1e-6, 1e-3, 1000.0, 0.0; ϕ=0.3)
+        @test KomaMRIBase.rf_frame_phase(rf) ≈ [0, π, -π, 0]
+
+        # Frequency-modulated RF: ψ is generated from the integrated Δf waveform and centered.
+        rf_fm = RF([1.0, 2.0, 1.0] .* 1e-6, 1e-3, [0.0, 1000.0, 0.0], 0.0)
+        ψ = KomaMRIBase.rf_frame_phase(rf_fm)
+        @test ψ[argmin(abs.(times(rf_fm, :Δf) .- (rf_fm.delay + rf_fm.center)))] ≈ 0
+
+        # Discretization carries ψ so simulators can handle split RF blocks independently.
+        seq = Sequence()
+        seq += rf
+        seqd = discretize(seq; sampling_params=Dict{String,Any}("Δt"=>1e-4, "Δt_rf"=>1e-4))
+        center_idx = findfirst(==(rf.delay + rf.center), seqd.t)
+        @test hasproperty(seqd, :ψ)
+        @test !isnothing(center_idx)
+        @test iszero(seqd.ψ[center_idx])
+        @test length(seqd.ψ) == length(seqd.t)
+        @test !all(iszero, seqd.ψ)
     end
 
     @testset "Delay" begin
@@ -726,7 +746,8 @@ using TestItems, TestItemRunner
         seqd = KomaMRIBase.discretize(seq)
         i1, i2 = rand(1:Int(floor(0.5*length(seqd)))), rand(Int(ceil(0.5*length(seqd))):length(seqd))
         @test seqd[i1].t ≈ [t[i1]]
-        @test seqd[i1:i2-1].t ≈ t[i1:i2]
+        @test seqd[i1:i2].t ≈ t[i1:i2]
+        @test seqd[i1:i2].Δt ≈ Δt[i1:i2-1]
 
         T, N = 1.0, 4
         seq = Sequence()

KomaMRICore/Project.toml

@@ -1,7 +1,7 @@
 name = "KomaMRICore"
 uuid = "4baa4f4d-2ae9-40db-8331-a7d1080e3f4e"
 authors = ["Carlos Castillo Passi <cncastillo@uc.cl>"]
-version = "0.10.0"
+version = "0.11.0"
 
 [deps]
 AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1"
@@ -32,7 +32,7 @@ Adapt = "3, 4"
 CUDA = "3, 4, 5, 6.0"
 Functors = "0.4, 0.5"
 KernelAbstractions = "0.9.34"
-KomaMRIBase = "0.10"
+KomaMRIBase = "0.11"
 Metal = "1.2"
 ProgressMeter = "1"
 Reexport = "1"

KomaMRICore/src/KomaMRICore.jl

@@ -1,4 +1,4 @@
-module KomaMRICore
+module KomaMRICore 
 
 # General
 import Base.*, Base.abs
@@ -17,6 +17,7 @@ include("rawdata/ISMRMRD.jl")
 # Datatypes
 include("datatypes/Spinor.jl")
 include("other/DiffusionModel.jl")
+include("callbacks/Callback.jl")
 # Simulator
 include("simulation/GPUFunctions.jl")
 include("simulation/Functors.jl")
@@ -30,5 +31,7 @@ export Mag
 export simulate, simulate_slice_profile
 # Spinors
 export Spinor, Rx, Ry, Rz, Q, Un
+# Callback
+export Callback
 
 end

KomaMRICore/src/callbacks/Callback.jl

@@ -0,0 +1,16 @@
+struct Callback{F}
+    enabled::Bool
+    every::Int
+    fun::F
+end
+
+Callback(every::Int, fun::F) where {F} = Callback(every > 0, every, fun)
+
+function (cb::Callback)(progress_info, simulation_blocks_info, device_data, sim_params)
+    if cb.enabled && (progress_info.block - 1) % cb.every == 0
+        cb.fun(progress_info, simulation_blocks_info, device_data, sim_params)
+    end
+end
+
+# Included Callbacks
+include("progress_callback.jl")

KomaMRICore/src/callbacks/progress_callback.jl

@@ -0,0 +1,14 @@
+function progressbar_callback(Nblocks)
+    progress_bar = Progress(Nblocks; desc="Running simulation...")
+    function progressbar_callback_fun(progress_info, simulation_blocks_info, device_data, sim_params)
+        next!(
+            progress_bar;
+            showvalues=[
+                (:simulated_blocks, progress_info.block), 
+                (:rf_blocks, progress_info.rfs), 
+                (:acq_samples, progress_info.samples - 1)
+            ],
+        )
+    end
+    return progressbar_callback_fun
+end