Skip to content

Enable Reactant AD through sequence discretization - #924

Open
KookiesNKareem wants to merge 5 commits into
branch/reactant-parallel-adfrom
branch/reactant-discretize-ad
Open

Enable Reactant AD through sequence discretization#924
KookiesNKareem wants to merge 5 commits into
branch/reactant-parallel-adfrom
branch/reactant-discretize-ad

Conversation

@KookiesNKareem

Copy link
Copy Markdown
Collaborator

Summary

  • make single-block sequence discretization traceable through Reactant, including RF waveform interpolation and ADC-at-start padding
  • add a narrow KomaMRIBase Reactant extension for traced RF arrays while preserving the existing native discretization path
  • differentiate an acquired-signal loss through discretize, get_sim_ranges, and the BlochSimple time iterator

Validation

  • focused CPU Reactant + Enzyme ADC test: 16/16
  • KomaMRIBase: 767/767
  • full unfiltered KomaMRICore CPU suite: 448/448
  • KomaMRIFiles: 24448/24448
  • Titan RTX CUDA: existing parallel BlochSimple Reactant + Enzyme test 4/4; new discretize + ADC + iterator test 2/2
  • CUDA gradient matched a fifth-order central finite-difference baseline at rtol=1e-8, atol=1e-10
  • git diff --check

Stacked on #922.

@cncastillo
cncastillo marked this pull request as ready for review July 31, 2026 20:15
@cncastillo
cncastillo self-requested a review as a code owner July 31, 2026 20:15
@KookiesNKareem
KookiesNKareem marked this pull request as draft August 1, 2026 19:01
@KookiesNKareem
KookiesNKareem force-pushed the branch/reactant-discretize-ad branch from db10f1e to 09f0468 Compare August 2, 2026 01:59
@codecov

codecov Bot commented Aug 2, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 88.24%. Comparing base (9dfc437) to head (122b9e7).

Additional details and impacted files

Impacted file tree graph

@@                       Coverage Diff                       @@
##           branch/reactant-parallel-ad     #924      +/-   ##
===============================================================
- Coverage                        91.57%   88.24%   -3.33%     
===============================================================
  Files                              146      146              
  Lines                             7534     7495      -39     
===============================================================
- Hits                              6899     6614     -285     
- Misses                             635      881     +246     
Flag Coverage Δ
base 90.99% <100.00%> (+0.03%) ⬆️
core 72.19% <100.00%> (-17.60%) ⬇️
files 95.15% <ø> (ø)
komamri 90.51% <ø> (ø)
plots 91.76% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
KomaMRIBase/src/datatypes/sequence/RF.jl 90.90% <100.00%> (+0.18%) ⬆️
...e/src/datatypes/sequence/SequenceEventWaveforms.jl 88.46% <100.00%> (+0.27%) ⬆️
KomaMRIBase/src/discretization/SequenceSampling.jl 92.42% <100.00%> (+0.48%) ⬆️
...RIBase/src/discretization/WaveformInterpolation.jl 100.00% <100.00%> (ø)
KomaMRICore/src/simulation/SimulatorCore.jl 88.82% <100.00%> (-4.11%) ⬇️

... and 23 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@KookiesNKareem
KookiesNKareem force-pushed the branch/reactant-discretize-ad branch from c4bc57f to eca8bbf Compare August 2, 2026 15:59
@KookiesNKareem
KookiesNKareem marked this pull request as ready for review August 2, 2026 18:53

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As mentioned in other PRs probably not needed.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

problem is that when reactant is compiling with the traced array it doesn't know the numerical values so it can't evaluate is_on the same way

Comment on lines +20 to +30
@inline _scalar_getindex(x, i) = x[i]
@inline _scalar_setindex!(x, value, i) = setindex!(x, value, i)

function _zero_pad_samples(A)
out = similar(A, length(A) + 2)
fill!(out, zero(eltype(A)))
for i in eachindex(A)
_scalar_setindex!(out, _scalar_getindex(A, i), i + 1)
end
return out
end

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remove.

Comment on lines +64 to +65
length(A) == 1 && (A = A[[1, 1]])
A = _zero_pad_samples(A)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This one is kind of sketchy, probably not needed. What about, could the problem be the [el1 ; ... ; eln] in another PR you changes something to [el1, ..., eln]

Comment on lines 85 to 89
if length(seq) == 1
values = sample_sequence_block(seq, 1; sampling_rule, motion_times=block_global_event_times(T0, 1, global_event_times), freq_in_phase)
return _with_adc_start_padding(values)
end
out = DiscreteSequence()

@cncastillo cncastillo Aug 10, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why define a special case for length(seq) == 1? When is this used? Also _with_adc_start_padding and _prepend_sample are probably not needed..

Comment thread KomaMRIBase/src/discretization/WaveformInterpolation.jl
Comment thread KomaMRICore/test/test_files/ad_utils.jl Outdated
Comment thread KomaMRICore/test/reactant_cuda.jl Outdated
Comment thread KomaMRICore/test/runtests.jl
@KookiesNKareem
KookiesNKareem force-pushed the branch/reactant-discretize-ad branch from 7f6640c to 7ea4dab Compare August 10, 2026 16:53
@KookiesNKareem
KookiesNKareem force-pushed the branch/reactant-discretize-ad branch from 1c8ddb5 to e8bbd9c Compare August 19, 2026 23:45
@KookiesNKareem
KookiesNKareem force-pushed the branch/reactant-discretize-ad branch from 98cacf3 to 96387ed Compare August 21, 2026 04:57
@KookiesNKareem
KookiesNKareem force-pushed the branch/reactant-discretize-ad branch from 96387ed to 122b9e7 Compare August 24, 2026 21:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants