Add Magnus-based methods & fix simulation block splitting issues

KomaMRIBase/src/datatypes/simulation/DiscreteSequence.jl

@@ -40,31 +40,31 @@ Base.getindex(seq::DiscreteSequence, i::Integer) = begin
                      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.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.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)

KomaMRIBase/src/timing/TimeStepCalculation.jl

@@ -102,8 +102,7 @@ function get_variable_times(seq; Δt=1e-3, Δt_rf=1e-5, motion=NoMotion())
 			delay, T = y.delay, y.T
 			t1 = t0 + delay
 			t2 = t1 + sum(T)
-			rf0 = t0 + get_RF_center(y) #get_RF_center includes delays
-			taux = points_from_key_times([t1, t1 + ϵ, rf0, t2 - ϵ, t2]; dt=Δt_rf)
+			taux = points_from_key_times([t1 - ϵ, t1, t1 + ϵ, t2 - ϵ, t2, t2 + ϵ]; dt=Δt_rf)
             append!(t_block, taux)
 		end
 		if is_GR_on(s)

KomaMRIBase/test/runtests.jl

@@ -291,7 +291,7 @@ 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]
 
         T, N = 1.0, 4
         seq = RF(1.0e-6, 1.0)

KomaMRICore/src/KomaMRICore.jl

@@ -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

KomaMRICore/src/simulation/SimMethods/Bloch/cpu/BlochCPU.jl

@@ -26,9 +26,9 @@ function prealloc(sim_method::Bloch, backend::KA.CPU, obj::Phantom{T}, M::Mag{T}
             similar(M.xy),
             similar(M.z)
         ),
-        similar(obj.x),
-        similar(obj.x),
-        similar(obj.x),
+        zeros(T, size(obj.x)),
+        zeros(T, size(obj.x)),
+        zeros(T, size(obj.x)),
         Spinor(
             similar(M.xy),
             similar(M.xy)
@@ -53,60 +53,45 @@ function run_spin_precession!(
     sim_method::Bloch,
     groupsize,
     backend::KA.CPU,
-    prealloc::BlochCPUPrealloc
+    prealloc::PreallocResult{T}
 ) where {T<:Real}
-    #Simulation
-    #Motion
-    x, y, z = get_spin_coords(p.motion, p.x, p.y, p.z, seq.t[1])
-
-    #Initialize arrays
+    #Rename arrays
     Bz_old = prealloc.Bz_old
     Bz_new = prealloc.Bz_new
     ϕ = prealloc.ϕ
     Mxy = prealloc.M.xy
     ΔBz = prealloc.ΔBz
+    #Initialize
     fill!(ϕ, zero(T))
-    @. Bz_old = x[:,1] * seq.Gx[1] + y[:,1] * seq.Gy[1] + z[:,1] * seq.Gz[1] + ΔBz
-
-    # Fill sig[1] if needed
-    ADC_idx = 1
-    if (seq.ADC[1])
-        sig[1] = sum(M.xy)
-        ADC_idx += 1
-    end
-
-    t_seq = zero(T) # Time
-    for seq_idx=2:length(seq.t)
-        x, y, z = get_spin_coords(p.motion, p.x, p.y, p.z, seq.t[seq_idx])
-        t_seq += seq.Δt[seq_idx-1]
-
+    block_time = zero(T)
+    sample = 1
+    #Simulation
+    for i in eachindex(seq.Δt)
+        #Motion
+        x, y, z = get_spin_coords(p.motion, p.x, p.y, p.z, seq.t[i + 1])
         #Effective Field
-        @. Bz_new = x * seq.Gx[seq_idx] + y * seq.Gy[seq_idx] + z * seq.Gz[seq_idx] + ΔBz
-
+        @. Bz_new = x * seq.Gx[i + 1] + y * seq.Gy[i + 1] + z * seq.Gz[i + 1] + ΔBz
         #Rotation
-        @. ϕ += (Bz_old + Bz_new) * T(-π * γ) * seq.Δt[seq_idx-1]
-
+        @. ϕ += (Bz_old + Bz_new) * T(-π * γ) * seq.Δt[i]
+        block_time += seq.Δt[i]
         #Acquired Signal
-        if seq_idx <= length(seq.ADC) && seq.ADC[seq_idx]
-            @. Mxy = exp(-t_seq / p.T2) * M.xy * cis(ϕ)
-
+        if seq.ADC[i + 1]
+            #Update signal
+            @. Mxy = exp(-block_time / p.T2) * M.xy * cis(ϕ)
             #Reset Spin-State (Magnetization). Only for FlowPath
-            outflow_spin_reset!(Mxy, seq.t[seq_idx], p.motion)
-
-            sig[ADC_idx] = sum(Mxy) 
-            ADC_idx += 1
+            outflow_spin_reset!(Mxy, seq.t[i + 1], p.motion)
+            #Acquired signal
+            sig[sample] = sum(Mxy) 
+            sample += 1
         end
-
-        Bz_old, Bz_new = Bz_new, Bz_old
+        #Update simulation state
+        Bz_old .= Bz_new
     end
-
     #Final Spin-State
-    @. M.xy = M.xy * exp(-t_seq / p.T2) * cis(ϕ)
-    @. M.z = M.z * exp(-t_seq / p.T1) + p.ρ * (T(1) - exp(-t_seq / p.T1))
-
+    @. M.xy = M.xy * exp(-block_time / p.T2) * cis(ϕ)
+    @. M.z = M.z * exp(-block_time / p.T1) + p.ρ * (T(1) - exp(-block_time / p.T1))
     #Reset Spin-State (Magnetization). Only for FlowPath
     outflow_spin_reset!(M,  seq.t', p.motion; replace_by=p.ρ)
-
     return nothing
 end
 
@@ -126,36 +111,40 @@ function run_spin_excitation!(
     backend::KA.CPU,
     prealloc::BlochCPUPrealloc
 ) where {T<:Real}
+    #Rename arrays 
     Bz = prealloc.Bz_old
     B = prealloc.Bz_new
-    φ = prealloc.ϕ
+    φ_half = prealloc.ϕ
     α = prealloc.Rot.α
     β = prealloc.Rot.β
     ΔBz = prealloc.ΔBz
     Maux_xy = prealloc.M.xy
     Maux_z = prealloc.M.z
-
+    #Initialize
+    sample = 1
     #Simulation
-    for s in seq #This iterates over seq, "s = seq[i,:]"
+    for i in eachindex(seq.Δt)
         #Motion
-        x, y, z = get_spin_coords(p.motion, p.x, p.y, p.z, s.t)
+        x, y, z = get_spin_coords(p.motion, p.x, p.y, p.z, seq.t[i])
         #Effective field
-        @. Bz = (s.Gx * x + s.Gy * y + s.Gz * z) + ΔBz - s.Δf / T(γ) # ΔB_0 = (B_0 - ω_rf/γ), Need to add a component here to model scanner's dB0(x,y,z)
-        @. B = sqrt(abs(s.B1)^2 + abs(Bz)^2)
+        @. Bz = (seq.Gx[i] * x + seq.Gy[i] * y + seq.Gz[i] * z) + ΔBz - seq.Δf[i] / T(γ) # ΔB_0 = (B_0 - ω_rf/γ), Need to add a component here to model scanner's dB0(x,y,z)
+        @. B = sqrt(abs(seq.B1[i])^2 + abs(Bz)^2)
         @. B[B == 0] = eps(T)
         #Spinor Rotation
-        @. φ = T(-π * γ) * (B * s.Δt) # TODO: Use trapezoidal integration here (?),  this is just Forward Euler 
-        @. α = cos(φ) - Complex{T}(im) * (Bz / B) * sin(φ)
-        @. β = -Complex{T}(im) * (s.B1 / B) * sin(φ)
+        @. φ_half = T(-π * γ) * (B * seq.Δt[i]) # TODO: Use trapezoidal integration here (?),  this is just Forward Euler
+        @. α = cos(φ_half) - Complex{T}(im) * (Bz / B) * sin(φ_half)
+        @. β = -Complex{T}(im) * (seq.B1[i] / B) * sin(φ_half)
         mul!(Spinor(α, β), M, Maux_xy, Maux_z)
         #Relaxation
-        @. M.xy = M.xy * exp(-s.Δt / p.T2)
-        @. M.z = M.z * exp(-s.Δt / p.T1) + p.ρ * (T(1) - exp(-s.Δt / p.T1))
-
+        @. M.xy = M.xy * exp(-seq.Δt[i] / p.T2)
+        @. M.z = M.z * exp(-seq.Δt[i] / p.T1) + p.ρ * (T(1) - exp(-seq.Δt[i] / p.T1))
         #Reset Spin-State (Magnetization). Only for FlowPath
-        outflow_spin_reset!(M,  s.t, p.motion; replace_by=p.ρ)
+        outflow_spin_reset!(M,  seq.t[i + 1, :], p.motion; replace_by=p.ρ)
+        #Acquire signal
+        if seq.ADC[i + 1] # ADC at the end of the time step
+            sig[sample] = sum(M.xy) 
+            sample += 1
+        end
     end
-    #Acquired signal
-    #sig .= -1.4im #<-- This was to test if an ADC point was inside an RF block
     return nothing
 end

KomaMRICore/src/simulation/SimMethods/Bloch/gpu/BlochGPU.jl

@@ -36,7 +36,7 @@ function run_spin_precession!(
         pre.sig_output,
         M.xy, M.z,
         x, y, z, pre.ΔBz, p.T1, p.T2, p.ρ, UInt32(length(M.xy)),
-        seq.Gx, seq.Gy, seq.Gz, seq.Δt, seq.ADC, UInt32(length(seq.Δt)+1),
+        seq.Gx, seq.Gy, seq.Gz, seq.Δt, seq.ADC, UInt32(length(seq.t)),
         Val(!(p.motion isa NoMotion)), Val(supports_warp_reduction(backend)),
         ndrange=(cld(length(M.xy), groupsize) * groupsize)
     )
@@ -66,13 +66,18 @@ function run_spin_excitation!(
 
     #Excitation
     excitation_kernel!(backend, groupsize)(
+        pre.sig_output,
         M.xy, M.z,
         x, y, z, pre.ΔBz, p.T1, p.T2, p.ρ, UInt32(length(M.xy)),
-        seq.Gx, seq.Gy, seq.Gz, seq.Δt, seq.Δf, seq.B1, UInt32(length(seq.Δt)),
-        Val(!(p.motion isa NoMotion)),
+        seq.Gx, seq.Gy, seq.Gz, seq.Δt, seq.Δf, seq.B1, seq.ADC, UInt32(length(seq.Δt)),
+        Val(!(p.motion isa NoMotion)), Val(supports_warp_reduction(backend)),
         ndrange=(cld(length(M.xy), groupsize) * groupsize)
     )
 
+    #Signal
+    AK.reduce(+, view(pre.sig_output,:,1:length(sig)); init=zero(Complex{T}), dims=1, temp=view(pre.sig_output_final,:,1:length(sig)))
+    sig .= transpose(view(pre.sig_output_final,:,1:length(sig)))
+
     #Reset Spin-State (Magnetization). Only for FlowPath
     outflow_spin_reset!(M,  seq.t', p.motion; replace_by=p.ρ) # TODO: reset state inside kernel
 

KomaMRICore/src/simulation/SimMethods/Bloch/gpu/ExcitationKernel.jl

@@ -1,18 +1,24 @@
 ## COV_EXCL_START
 
 @kernel unsafe_indices=true inbounds=true function excitation_kernel!(
+    sig_output::AbstractMatrix{Complex{T}}, 
     M_xy::AbstractVector{Complex{T}}, M_z, 
-    @Const(p_x), @Const(p_y), @Const(p_z), @Const(p_ΔBz), @Const(p_T1), @Const(p_T2), @Const(p_ρ), N_Spins,
-    @Const(s_Gx), @Const(s_Gy), @Const(s_Gz), @Const(s_Δt), @Const(s_Δf), @Const(s_B1), N_Δt,
-    ::Val{MOTION}
-) where {T, MOTION}
+    @Const(p_x), @Const(p_y), @Const(p_z), @Const(p_ΔBz), @Const(p_T1), @Const(p_T2), @Const(p_ρ), N_spins,
+    @Const(s_Gx), @Const(s_Gy), @Const(s_Gz), @Const(s_Δt), @Const(s_Δf), @Const(s_B1), @Const(s_ADC), N_Δt,
+    ::Val{MOTION}, ::Val{USE_WARP_REDUCTION},
+) where {T, MOTION, USE_WARP_REDUCTION}
 
     @uniform N = @groupsize()[1]
     i_l = @index(Local, Linear)
     i_g = @index(Group, Linear)
     i = (i_g - 1u32) * UInt32(N) + i_l
 
-    if i <= N_Spins
+    sig_group_r = @localmem T USE_WARP_REDUCTION ? 32 : N
+    sig_group_i = @localmem T USE_WARP_REDUCTION ? 32 : N
+    sig_r = zero(T)
+    sig_i = zero(T)
+
+    if i <= N_spins
         x, y, z = get_spin_coordinates(p_x, p_y, p_z, i, 1)
         ΔBz = p_ΔBz[i]
         Mxy_r, Mxy_i = reim(M_xy[i])
@@ -21,6 +27,7 @@
         T1 = p_T1[i]
         T2 = p_T2[i]
 
+        ADC_idx = 1u32
         s_idx = 1u32
         while (s_idx <= N_Δt)
             if MOTION
@@ -61,6 +68,15 @@
             Mxy_r = Mxy_new_r * ΔT2
             Mxy_i = Mxy_new_i * ΔT2
             Mz = Mz_new * ΔT1 + ρ * (T(1) - ΔT1)
+
+            # Acquire Signal
+            if s_idx <= N_Δt && s_ADC[s_idx + 1]
+                sig_r, sig_i = reduce_signal!(Mxy_r, Mxy_i, sig_group_r, sig_group_i, i_l, N, T, Val(USE_WARP_REDUCTION))
+                if i_l == 1u32
+                    sig_output[i_g, ADC_idx] = complex(sig_r, sig_i)
+                end
+                ADC_idx += 1u32
+            end
             s_idx += 1u32
         end