Pass Discrete Sequence to GPU

KomaMRIBase/src/KomaMRIBase.jl

@@ -41,7 +41,7 @@ export γ    # gyro-magnetic ratio [Hz/T]
 export Scanner, Sequence, Phantom
 export Grad, RF, ADC, Delay
 export dur, get_block_start_times, get_samples
-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)
 
@@ -17,15 +19,18 @@ 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}
-    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
+    ADC::AB
+    t::AT
+    Δt::AT
 end
 
 Base.length(seq::DiscreteSequence) = length(seq.Δt)

KomaMRICore/src/KomaMRICore.jl

@@ -18,6 +18,7 @@ include("rawdata/ISMRMRD.jl")
 include("datatypes/Spinor.jl")
 include("other/DiffusionModel.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
@@ -115,5 +115,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

@@ -66,8 +66,8 @@ function run_spin_precession!(
     Mxy = prealloc.M.xy
     ΔBz = prealloc.ΔBz
     fill!(ϕ, zero(T))
-    @. Bz_old = x[:,1] * seq.Gx[1] + y[:,1] * seq.Gy[1] + z[:,1] * seq.Gz[1] + ΔBz
-
+    Bz_old .= get_Bz(seq, x, y, z, 1) .+ ΔBz
+    
     # Fill sig[1] if needed
     ADC_idx = 1
     if (seq.ADC[1])
@@ -81,7 +81,7 @@ function run_spin_precession!(
         t_seq += seq.Δt[seq_idx-1]
 
         #Effective Field
-        @. Bz_new = x * seq.Gx[seq_idx] + y * seq.Gy[seq_idx] + z * seq.Gz[seq_idx] + ΔBz
+        Bz_new .= get_Bz(seq, x, y, z, seq_idx) .+ ΔBz
 
         #Rotation
         @. ϕ += (Bz_old + Bz_new) * T(-π * γ) * seq.Δt[seq_idx-1]
@@ -136,24 +136,24 @@ function run_spin_excitation!(
     Maux_z = prealloc.M.z
 
     #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 .= 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
-        @. φ = T(-π * γ) * (B * s.Δt) # TODO: Use trapezoidal integration here (?),  this is just Forward Euler 
+        @. φ = T(-π * γ) * (B * seq.Δt[i]) # TODO: Use trapezoidal integration here (?),  this is just Forward Euler 
         @. α = cos(φ) - Complex{T}(im) * (Bz / B) * sin(φ)
-        @. β = -Complex{T}(im) * (s.B1 / B) * sin(φ)
+        @. β = -Complex{T}(im) * (seq.B1[i] / B) * sin(φ)
         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.ρ)
     end
     #Acquired signal
     #sig .= -1.4im #<-- This was to test if an ADC point was inside an RF block

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

@@ -20,7 +20,7 @@ end
 
 function run_spin_precession!(
     p::Phantom{T},
-    seq::DiscreteSequence{T},
+    seq::AbstractDiscreteSequence,
     sig::AbstractArray{Complex{T}},
     M::Mag{T},
     sim_method::Bloch,
@@ -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, UInt32(length(seq.Δt)+1),
         Val(!(p.motion isa NoMotion)), Val(supports_warp_reduction(backend)),
         ndrange=(cld(length(M.xy), groupsize) * groupsize)
     )
@@ -53,7 +53,7 @@ end
 
 function run_spin_excitation!(
     p::Phantom{T},
-    seq::DiscreteSequence{T},
+    seq::AbstractDiscreteSequence,
     sig::AbstractArray{Complex{T}},
     M::Mag{T},
     sim_method::Bloch,
@@ -68,7 +68,7 @@ function run_spin_excitation!(
     excitation_kernel!(backend, groupsize)(
         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)),
+        seq, UInt32(length(seq.Δt)),
         Val(!(p.motion isa NoMotion)),
         ndrange=(cld(length(M.xy), groupsize) * groupsize)
     )

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

@@ -3,7 +3,7 @@
 @kernel unsafe_indices=true inbounds=true function excitation_kernel!(
     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,
+    seq::AbstractDiscreteSequence, N_Δt,
     ::Val{MOTION}
 ) where {T, MOTION}
 
@@ -27,10 +27,10 @@
                 x, y, z = get_spin_coordinates(p_x, p_y, p_z, i, s_idx)
             end
 
-            Bz = (x * s_Gx[s_idx] + y * s_Gy[s_idx] + z * s_Gz[s_idx]) + ΔBz - s_Δf[s_idx] / T(γ)
-            B1_r, B1_i = reim(s_B1[s_idx])
+            Bz = get_Bz(seq, x, y, z, s_idx) + ΔBz - seq.Δf[s_idx] / T(γ)
+            B1_r, B1_i = reim(seq.B1[s_idx])
             B = sqrt(B1_r^2 + B1_i^2 + Bz^2)
-            Δt = s_Δt[s_idx]
+            Δt = seq.Δt[s_idx]
             φ = T(-π * γ) * B * Δt
             sin_φ, cos_φ = sincos(φ)
             α_r = 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},
 ) where {T, MOTION, USE_WARP_REDUCTION}
 
@@ -37,11 +37,11 @@
         Δ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
-    if s_ADC[1]
+    if seq.ADC[1]
         sig_r, sig_i = reduce_signal!(sig_r, sig_i, sig_group_r, sig_group_i, i_l, N, T, Val(USE_WARP_REDUCTION))
         if i_l == 1u32
             sig_output[i_g, 1] = complex(sig_r, sig_i)
@@ -56,13 +56,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
 
-        if s_idx < s_length && s_ADC[s_idx]
+        if s_idx < s_length && seq.ADC[s_idx]
             if i <= N_spins
                 ΔT2 = exp(-t / T2)
                 cis_ϕ_i, cis_ϕ_r = sincos(ϕ)

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

@@ -31,7 +31,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/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,
@@ -73,7 +73,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

@@ -80,7 +80,7 @@ separating the spins of the phantom `obj` in `Nthreads`.
 """
 function run_spin_precession_parallel!(
     obj::Phantom{T},
-    seq::DiscreteSequence{T},
+    seq::AbstractDiscreteSequence,
     sig::AbstractArray{Complex{T}},
     Xt::SpinStateRepresentation{T},
     sim_method::SimulationMethod,
@@ -122,7 +122,7 @@ different number threads to excecute the process.
 """
 function run_spin_excitation_parallel!(
     obj::Phantom{T},
-    seq::DiscreteSequence{T},
+    seq::AbstractDiscreteSequence,
     sig::AbstractArray{Complex{T}},
     Xt::SpinStateRepresentation{T},
     sim_method::SimulationMethod,
@@ -170,7 +170,7 @@ take advantage of CPU parallel processing.
 """
 function run_sim_time_iter!(
     obj::Phantom,
-    seq::DiscreteSequence,
+    seq::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::DiscreteSequence, x, y, z, s_idx)
+    return @. x * seq.Gx[s_idx] + y * seq.Gy[s_idx] + z * seq.Gz[s_idx]
+end