Pass Discrete Sequence to GPU

KomaMRIBase/src/KomaMRIBase.jl

@@ -87,7 +87,7 @@ export addblock!, @addblock, @addblocks
 export Grad, RF, ADC, Delay, Duration, QuaternionRot
 export dur, get_block_start_times, get_samples
 export RFuse, Excitation, Refocusing, Inversion, Saturation, Preparation, Other, Undefined
-export DiscreteSequence
+export DiscreteSequence, AbstractDiscreteSequence
 export discretize, get_adc_phase_compensation, get_adc_sampling_times
 export is_Gx_on, is_Gy_on, is_Gz_on, is_RF_on, is_ADC_on
 export times, ampls, freqs

KomaMRIBase/src/datatypes/simulation/DiscreteSequence.jl

@@ -1,3 +1,5 @@
+abstract type AbstractDiscreteSequence end
+
 """
     seqd = DiscreteSequence(Gx, Gy, Gz, B1, Δf, ψ, ADC, t, Δt)
 
@@ -18,16 +20,19 @@ times. DiscreteSequence is the struct used for simulation.
 # Returns
 - `seqd`: (`::DiscreteSequence`) DiscreteSequence struct
 """
-struct DiscreteSequence{T<:Real}
-    Gx::AbstractVector{T}
-    Gy::AbstractVector{T}
-    Gz::AbstractVector{T}
-    B1::AbstractVector{Complex{T}}
-    Δf::AbstractVector{T}
-    ψ::AbstractVector{T}
-    ADC::AbstractVector{Bool}
-    t::AbstractVector{T}
-    Δt::AbstractVector{T}
+struct DiscreteSequence{T<:Real,
+                        AT<:AbstractVector{T},
+                        ACT<:AbstractVector{Complex{T}},
+                        AB<:AbstractVector{Bool}} <: AbstractDiscreteSequence
+    Gx::AT
+    Gy::AT
+    Gz::AT
+    B1::ACT
+    Δf::AT
+    ψ::AT
+    ADC::AB
+    t::AT
+    Δt::AT
 end
 
 Base.length(seq::DiscreteSequence) = length(seq.Δt)

KomaMRICore/src/KomaMRICore.jl

@@ -19,6 +19,7 @@ include("datatypes/Spinor.jl")
 include("other/DiffusionModel.jl")
 include("callbacks/Callback.jl")
 # Simulator
+include("simulation/SpatialEncoding.jl")
 include("simulation/GPUFunctions.jl")
 include("simulation/Functors.jl")
 include("simulation/SimulatorCore.jl")

KomaMRICore/src/simulation/Functors.jl

@@ -1,4 +1,4 @@
-import Adapt: adapt, adapt_storage
+import Adapt: adapt, adapt_storage, @adapt_structure
 import Functors: @functor, functor, fmap, isleaf
 
 #Aux. funcitons to check if the variable we want to move to the GPU is numeric
@@ -116,5 +116,6 @@ adapt_storage(T::Type{<:Real}, xs::MotionList) = MotionList(paramtype.(T, xs.mot
 @functor Spinor
 # DiscreteSequence
 @functor DiscreteSequence
+@adapt_structure DiscreteSequence
 
 export gpu, cpu, f32, f64

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

@@ -47,7 +47,7 @@ that they can be re-used from block to block.
 """
 function run_spin_precession!(
     p::Phantom{T},
-    seq::DiscreteSequence{T},
+    seq::AbstractDiscreteSequence,
     sig::AbstractArray{Complex{T}},
     M::Mag{T},
     sim_method::Bloch,
@@ -70,7 +70,7 @@ function run_spin_precession!(
         #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[i + 1] + y * seq.Gy[i + 1] + z * seq.Gz[i + 1] + ΔBz
+        Bz_new .= get_Bz(seq, x, y, z, i + 1) .+ ΔBz
         #Rotation
         @. ϕ += (Bz_old + Bz_new) * T(-π * γ) * seq.Δt[i]
         block_time += seq.Δt[i]
@@ -103,7 +103,7 @@ optimized for the CPU. Uses preallocation for all arrays to reduce memory usage.
 """
 function run_spin_excitation!(
     p::Phantom{T},
-    seq::DiscreteSequence{T},
+    seq::AbstractDiscreteSequence,
     sig::AbstractArray{Complex{T}},
     M::Mag{T},
     sim_method::Bloch,
@@ -132,7 +132,7 @@ function run_spin_excitation!(
         #Motion
         x, y, z = get_spin_coords(p.motion, p.x, p.y, p.z, seq.t[i])
         #Effective field
-        @. 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)
+        Bz .= get_Bz(seq, x, y, z, i) .+ Δ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

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

@@ -27,7 +27,7 @@ end
 
 function run_spin_precession!(
     p::Phantom{T},
-    seq::DiscreteSequence{T},
+    seq::AbstractDiscreteSequence,
     sig::AbstractArray{Complex{T}},
     M::Mag{T},
     sim_method::SM,
@@ -44,7 +44,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)),
+        seq, UInt32(length(seq.t)),
         Val(!(p.motion isa NoMotion)), Val(supports_warp_reduction(backend)), Val(has_adc),
         sim_method,
         ndrange=(cld(length(M.xy), groupsize) * groupsize)
@@ -64,7 +64,7 @@ end
 
 function run_spin_excitation!(
     p::Phantom{T},
-    seq::DiscreteSequence{T},
+    seq::AbstractDiscreteSequence,
     sig::AbstractArray{Complex{T}},
     M::Mag{T},
     sim_method::SM,
@@ -81,7 +81,7 @@ function run_spin_excitation!(
         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, seq.ψ, seq.ADC, UInt32(length(seq.t)),
+        seq, UInt32(length(seq.t)),
         Val(!(p.motion isa NoMotion)), Val(supports_warp_reduction(backend)), Val(has_adc),
         sim_method,
         ndrange=(cld(length(M.xy), groupsize) * groupsize)

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

@@ -4,7 +4,7 @@
     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), @Const(s_ψ), @Const(s_ADC), s_length,
+    seq::AbstractDiscreteSequence, s_length,
     ::Val{MOTION}, ::Val{USE_WARP_REDUCTION}, ::Val{HAS_ADC},
     sim_method::SM
 ) where {T, MOTION, USE_WARP_REDUCTION, HAS_ADC, SM <: BlochLikeSimMethods}
@@ -41,16 +41,16 @@
         T2 = p_T2[i]
         # Rotating frame -> RF frame
         # M * exp(-i * ψ)
-        ψ_start = s_ψ[1]
+        ψ_start = seq.ψ[1]
         if !iszero(ψ_start)
             sin_ψ, cos_ψ = sincos(ψ_start)
             Mxy_r, Mxy_i = Mxy_r * cos_ψ + Mxy_i * sin_ψ, Mxy_i * cos_ψ - Mxy_r * sin_ψ
         end
 
         # Calculate initial B
         x, y, z = get_spin_coordinates(p_x, p_y, p_z, i, 1)
-        Bx_prev, By_prev = reim(s_B1[1])
-        Bz_prev = x * s_Gx[1] + y * s_Gy[1] + z * s_Gz[1] + ΔBz - s_Δf[1] / T(γ)
+        Bx_prev, By_prev = reim(seq.B1[1])
+        Bz_prev = get_Bz(seq, x, y, z, 1) + ΔBz - seq.Δf[1] / T(γ)
     end
 
     ADC_idx = 1u32
@@ -60,10 +60,10 @@
             if MOTION
                 x, y, z = get_spin_coordinates(p_x, p_y, p_z, i, s_idx)
             end
-            Bx_next, By_next = reim(s_B1[s_idx])
-            Bz_next = (x * s_Gx[s_idx] + y * s_Gy[s_idx] + z * s_Gz[s_idx]) + ΔBz - s_Δf[s_idx] / T(γ)
+            Bx_next, By_next = reim(seq.B1[s_idx])
+            Bz_next = get_Bz(seq, x, y, z, s_idx) + ΔBz - seq.Δf[s_idx] / T(γ)
 
-            Δt = s_Δt[s_idx - 1]
+            Δt = seq.Δt[s_idx - 1]
 
             # Spinor rotation
             θx, θy, θz = effective_rotation_vector(Bx_prev, By_prev, Bz_prev, Bx_next, By_next, Bz_next, Δt, sim_method)
@@ -103,7 +103,7 @@
         end
 
         # Acquire Signal
-        if HAS_ADC && s_ADC[s_idx]
+        if HAS_ADC && seq.ADC[s_idx]
             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)
@@ -117,7 +117,7 @@
     if active
         # RF frame -> Rotating frame
         # M * exp(i * ψ)
-        ψ_end = s_ψ[s_length]
+        ψ_end = seq.ψ[s_length]
         if !iszero(ψ_end)
             sin_ψ, cos_ψ = sincos(ψ_end)
             Mxy_r, Mxy_i = Mxy_r * cos_ψ - Mxy_i * sin_ψ, Mxy_r * sin_ψ + Mxy_i * cos_ψ

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

@@ -4,7 +4,7 @@
     sig_output::AbstractMatrix{Complex{T}}, 
     M_xy, 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_ADC), s_length,
+    seq::AbstractDiscreteSequence, s_length,
     ::Val{MOTION}, ::Val{USE_WARP_REDUCTION}, ::Val{HAS_ADC},
     sim_method::BlochLikeSimMethods
 ) where {T, MOTION, USE_WARP_REDUCTION, HAS_ADC}
@@ -36,7 +36,7 @@
         ΔBz = p_ΔBz[i]
         T2 = p_T2[i]
         x, y, z = get_spin_coordinates(p_x, p_y, p_z, i, 1)
-        Bz_prev = x * s_Gx[1] + y * s_Gy[1] + z * s_Gz[1] + ΔBz
+        Bz_prev = get_Bz(seq, x, y, z, 1) + ΔBz
     end
 
     ADC_idx = 1u32
@@ -47,13 +47,13 @@
                 x, y, z = get_spin_coordinates(p_x, p_y, p_z, i, s_idx)
             end
 
-            Δt = s_Δt[s_idx-1]
+            Δt = seq.Δt[s_idx-1]
             t += Δt
-            Bz_next = x * s_Gx[s_idx] + y * s_Gy[s_idx] + z * s_Gz[s_idx] + ΔBz
+            Bz_next = get_Bz(seq, x, y, z, s_idx) + ΔBz
             ϕ += (Bz_prev + Bz_next) * T(-π * γ) * Δt
         end
         # Acquire Signal
-        if HAS_ADC && s_ADC[s_idx]
+        if HAS_ADC && seq.ADC[s_idx]
             sig_r = zero(T)
             sig_i = zero(T)
             if active

KomaMRICore/src/simulation/SimMethods/BlochDict/BlochDict.jl

@@ -36,7 +36,7 @@ precession.
 """
 function run_spin_precession!(
     p::Phantom{T},
-    seq::DiscreteSequence{T},
+    seq::AbstractDiscreteSequence,
     sig::AbstractArray{Complex{T}},
     M::Mag{T},
     sim_method::BlochDict,

KomaMRICore/src/simulation/SimMethods/BlochMagnus/cpu/BlochMagnusCPU.jl

@@ -52,7 +52,7 @@ end
 # Use Bloch implementation for precession
 function run_spin_precession!(
     p::Phantom{T},
-    seq::DiscreteSequence{T},
+    seq::AbstractDiscreteSequence,
     sig::AbstractArray{Complex{T}},
     M::Mag{T},
     sim_method::BlochMagnus,
@@ -66,7 +66,7 @@ end
 # This part changes a bit more
 function run_spin_excitation!(
     p::Phantom{T},
-    seq::DiscreteSequence{T},
+    seq::AbstractDiscreteSequence,
     sig::AbstractArray{Complex{T}},
     M::Mag{T},
     sim_method::BlochMagnus,

KomaMRICore/src/simulation/SimMethods/BlochSimple/BlochSimple.jl

@@ -23,7 +23,7 @@ precession.
 """
 function run_spin_precession!(
     p::Phantom{T},
-    seq::DiscreteSequence{T},
+    seq::AbstractDiscreteSequence,
     sig::AbstractArray{Complex{T}},
     M::Mag{T},
     sim_method::SimulationMethod,
@@ -72,7 +72,7 @@ It gives rise to a rotation of `M0` with an angle given by the efective magnetic
 """
 function run_spin_excitation!(
     p::Phantom{T},
-    seq::DiscreteSequence{T},
+    seq::AbstractDiscreteSequence,
     sig::AbstractArray{Complex{T}},
     M::Mag{T},
     sim_method::SimulationMethod,

KomaMRICore/src/simulation/SimulatorCore.jl

@@ -62,7 +62,7 @@ end
 
 function run_spin_precession_parallel!(
     obj::Phantom{T},
-    seq::DiscreteSequence{T},
+    seq::AbstractDiscreteSequence,
     sig::AbstractArray{Complex{T}},
     Xt::SpinStateRepresentation{T},
     sim_method::SimulationMethod,
@@ -84,7 +84,7 @@ end
 
 function run_spin_excitation_parallel!(
     obj::Phantom{T},
-    seq::DiscreteSequence{T},
+    seq::AbstractDiscreteSequence,
     sig::AbstractArray{Complex{T}},
     Xt::SpinStateRepresentation{T},
     sim_method::SimulationMethod,
@@ -131,7 +131,7 @@ take advantage of CPU parallel processing.
 """
 function run_sim_time_iter!(
     obj::Phantom,
-    seqd::DiscreteSequence,
+    seqd::AbstractDiscreteSequence,
     sig::AbstractArray{Complex{T}},
     Xt::SpinStateRepresentation{T},
     sim_method::SimulationMethod,

KomaMRICore/src/simulation/SpatialEncoding.jl

@@ -0,0 +1,3 @@
+@inline function get_Bz(seq::AbstractDiscreteSequence, x, y, z, s_idx)
+    return @. x * seq.Gx[s_idx] + y * seq.Gy[s_idx] + z * seq.Gz[s_idx]
+end