Enable Reactant AD through sequence discretization

KomaMRIBase/Project.toml

@@ -19,9 +19,11 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 
 [weakdeps]
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [extensions]
+KomaMRIBaseReactantExt = "Reactant"
 KomaMRIBaseUnitfulExt = "Unitful"
 
 [compat]
@@ -32,6 +34,7 @@ Interpolations = "0.13, 0.14, 0.15, 0.16"
 MAT = "0.10, 0.11, 0.12"
 MRIBase = "0.4"
 Parameters = "0.12, 0.13"
+Reactant = "0.2"
 Reexport = "1"
 Unitful = "1"
 julia = "1.9"

KomaMRIBase/ext/KomaMRIBaseReactantExt.jl

@@ -0,0 +1,12 @@
+module KomaMRIBaseReactantExt
+
+using KomaMRIBase
+using Reactant
+
+import KomaMRIBase: RF, is_on
+
+# A traced RF waveform has a statically known shape, but its values cannot be
+# inspected on the host to determine whether the event is active.
+is_on(rf::RF{<:Reactant.AnyTracedRVector}) = !isempty(rf.A)
+
+end

KomaMRIBase/src/datatypes/sequence/RF.jl

@@ -259,6 +259,8 @@ function rf_center(rf::RF)
     weights = abs.(ampls(rf))
     isempty(weights) && return 0.0
     total = sum(weights)
-    iszero(total) && return 0.0
-    return sum(weights .* (times(rf) .- rf.delay)) / total
+    active = !iszero(total)
+    denominator = ifelse(active, total, one(total))
+    center = sum(weights .* (times(rf) .- rf.delay)) / denominator
+    return ifelse(active, center, zero(center))
 end

KomaMRIBase/src/datatypes/sequence/SequenceEventWaveforms.jl

@@ -47,10 +47,13 @@ function ampls(rf::BlockPulseRF; freq_in_phase=false)
 end
 
 function ampls(rf::RF; freq_in_phase=false)
-    A = collect(cis(rf.ϕ) .* rf.A)
+    A = cis(rf.ϕ) .* rf.A
     is_on(rf) || return similar(A, 0)
-    length(A) == 1 && (A = [only(A), only(A)])
-    A = [zero(eltype(A)); A; zero(eltype(A))]
+    length(A) == 1 && (A = A[[1, 1]])
+    out = similar(A, length(A) + 2)
+    fill!(out, zero(eltype(A)))
+    out[2:(end - 1)] .= A
+    A = out
     if freq_in_phase
         t  = times(rf)
         Δf = (t=times(rf, :Δf)[2:(end - 1)], A=freqs(rf)[2:(end - 1)])

KomaMRIBase/src/discretization/SequenceSampling.jl

@@ -29,15 +29,15 @@ function append_adc_start_padding!(out, values, first_t)
     return out
 end
 
-function same_boundary_sample(out, values)
-    return all(last(dst) == first(src) for (dst, src) in zip(table_columns(out), table_columns(values)))
-end
-
 function append_sampled_block!(out, values, t0)
     isempty(values.t) && return out
     first_t = t0 + first(values.t)
     first_row = firstindex(values.t)
-    isempty(out.t) ? append_adc_start_padding!(out, values, first_t) : if first_t == last(out.t) && same_boundary_sample(out, values)
+    if isempty(out.t)
+        append_adc_start_padding!(out, values, first_t)
+    elseif all(dst -> dst isa Vector, table_columns(out)) &&
+           first_t == last(out.t) &&
+           all(last(dst) == first(src) for (dst, src) in zip(table_columns(out), table_columns(values)))
         first_row += 1
     else
         push!(out.Δt, first_t - last(out.t))
@@ -47,20 +47,24 @@ function append_sampled_block!(out, values, t0)
     for t in view(values.t, rows)
         push!(out.t, t0 + t)
     end
-    foreach((dst, src) -> append!(dst, view(src, rows)), table_columns(out), table_columns(values))
+    foreach((dst, src) -> append!(dst, dst isa Vector ? view(src, rows) : src[rows]), table_columns(out), table_columns(values))
     append!(out.Δt, values.Δt)
     append!(out.excitation_bool, values.excitation_bool)
     return out
 end
 
 # -- 6.3. Sample the full sequence ------------------------------------------
 function sample_sequence(seq; motion=NoMotion(), sampling_rule=MaxStepSizeRule(1e-3, 5e-5), freq_in_phase=false)
-    out = DiscreteSequence()
-    sizehint = max(8length(seq), 5sum(seq.ADC.N))
-    foreach(x -> Base.sizehint!(x, sizehint), (out.t, table_columns(out)..., out.excitation_bool, out.Δt))
+    length(seq) == 0 && return DiscreteSequence()
     T0 = get_block_start_times(seq)
     global_event_times = merge_sampling_times(sequence_boundary_sampling_times(seq), motion_sampling_times(seq, motion))
-    for block in 1:length(seq)
+    values = sample_sequence_block(seq, 1; sampling_rule, motion_times=block_global_event_times(T0, 1, global_event_times), freq_in_phase)
+    columns = (table_columns(values)..., values.excitation_bool, values.t, values.Δt)
+    out = DiscreteSequence(map(x -> similar(x, 0), columns)...)
+    sizehint = max(8length(seq), 5sum(seq.ADC.N))
+    foreach(x -> x isa Vector && Base.sizehint!(x, sizehint), (out.t, table_columns(out)..., out.excitation_bool, out.Δt))
+    append_sampled_block!(out, values, T0[1])
+    for block in 2:length(seq)
         values = sample_sequence_block(seq, block; sampling_rule, motion_times=block_global_event_times(T0, block, global_event_times), freq_in_phase)
         append_sampled_block!(out, values, T0[block])
     end

KomaMRIBase/src/discretization/WaveformInterpolation.jl

@@ -6,8 +6,9 @@
 # sampling grid.
 
 function linear_interpolate_samples(samples, t; default=zero(eltype(samples.A)), interpolate=true)
-    out = Vector{typeof(default)}(undef, length(t))
-    isempty(samples.t) && return fill!(out, default)
+    out = similar(samples.A, typeof(default), length(t))
+    fill!(out, default)
+    (isempty(samples.t) || isempty(samples.A)) && return out
     last_sample = min(lastindex(samples.t), lastindex(samples.A))
     sample = firstindex(samples.t)
     i = firstindex(t)
@@ -27,16 +28,17 @@ function linear_interpolate_samples(samples, t; default=zero(eltype(samples.A)),
             end
             for k in i:j
                 l = interpolate ? min(sample + k - i, sample_end) : sample_end - (j - k)
-                out[k] = l >= sample ? samples.A[l] : default
+                l >= sample && (out[k] = samples.A[l])
             end
             sample = sample_end + 1
-        elseif !interpolate || ti < first(samples.t) || sample > last_sample
-            out[i:j] .= default
-        else
+        elseif interpolate && ti >= first(samples.t) && sample <= last_sample
             lo_time, hi_time = samples.t[sample - 1], samples.t[sample]
             w = (ti - lo_time) / (hi_time - lo_time)
-            value = samples.A[sample - 1] + (samples.A[sample] - samples.A[sample - 1]) * w
-            out[i:j] .= value
+            lo = samples.A[sample - 1]
+            value = lo + (samples.A[sample] - lo) * w
+            for k in i:j
+                out[k] = value
+            end
         end
         i = j + 1
     end

KomaMRICore/src/simulation/SimulatorCore.jl

@@ -278,8 +278,12 @@ end
 function get_sim_ranges(seqd::DiscreteSequence; max_block_length=Inf, max_rf_block_length=Inf, eval_intervals_per_step=1)
     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]
-    stops = [starts[2:end] .- 1; lastindex(seqd.Δt)]
+    starts = Int[firstindex(seqd.Δt)]
+    for i in (firstindex(seqd.excitation_bool) + 1):lastindex(seqd.excitation_bool)
+        seqd.excitation_bool[i] == seqd.excitation_bool[i - 1] || push!(starts, i)
+    end
+    stops = [starts[i] - 1 for i in 2:length(starts)]
+    push!(stops, lastindex(seqd.Δt))
     for (start, stop) in zip(starts, stops)
         is_excitation = seqd.excitation_bool[start]
         max_length = is_excitation ? max_rf_block_length : max_block_length