Skip to content

Add split runge kutta to reactant extension#5237

Merged
jlk9 merged 25 commits intomainfrom
jlk9/add-split-runge-kutta-reactant-extension
Feb 13, 2026
Merged

Add split runge kutta to reactant extension#5237
jlk9 merged 25 commits intomainfrom
jlk9/add-split-runge-kutta-reactant-extension

Conversation

@jlk9
Copy link
Copy Markdown
Collaborator

@jlk9 jlk9 commented Jan 30, 2026

This adds an additional method for time_step! to OceananigansReactantExt for ReactantModel{<:SplitRungeKuttaTimeStepper}. Its only difference from the method for non-Reactant models with SplitRungeKuttaTimeStepper is commenting out the conditional statement:

#=
if model.clock.iteration == 0
    update_state!(model, callbacks)
end
=#

since this breaks Reactant, and the Reactant methods for first_time_step! automatically call update_state! anyway. A method for Reactant models with Adams-Bashforth timesteppers has already been implemented for the same reason.

@jlk9
Copy link
Copy Markdown
Collaborator Author

jlk9 commented Jan 30, 2026

NOTE: this PR already includes the small change from #5236, so maybe merge that one first.

@codecov
Copy link
Copy Markdown

codecov bot commented Jan 30, 2026

Codecov Report

❌ Patch coverage is 0% with 12 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.48%. Comparing base (a63cb5c) to head (7154ccb).

Files with missing lines Patch % Lines
ext/OceananigansReactantExt/TimeSteppers.jl 0.00% 12 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #5237      +/-   ##
==========================================
- Coverage   73.51%   73.48%   -0.04%     
==========================================
  Files         395      395              
  Lines       22194    22206      +12     
==========================================
+ Hits        16316    16318       +2     
- Misses       5878     5888      +10     
Flag Coverage Δ
buildkite 68.76% <0.00%> (-0.04%) ⬇️
julia 68.76% <0.00%> (-0.04%) ⬇️

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

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@glwagner
Copy link
Copy Markdown
Member

we need a test here, in addition to the PR that this depends on.

@glwagner
Copy link
Copy Markdown
Member

@dkytezab this is probably of interest to you.

@giordano
Copy link
Copy Markdown
Collaborator

giordano commented Feb 2, 2026

Its only difference from the method for non-Reactant models with SplitRungeKuttaTimeStepper is commenting out the conditional statement:

Can't this be fixed by @trace? Like

@trace if model.clock.iteration == 0
update_state!(model, callbacks; compute_tendencies=true)
end

@glwagner glwagner changed the title Add split runge kutta reactant extension Add split runge kutta to reactant extension Feb 2, 2026
@dkytezab
Copy link
Copy Markdown
Collaborator

dkytezab commented Feb 2, 2026

We should also add support for RK3 inside of OceananigansReactantExt. We should be able to get by adopting code from BreezeReactantExt's support for SSPRK3.

Once we get around FFT planning issues (see #5244) this will be the furthest upstream barrier to AD of NonhydrostaticModel.

@jlk9
Copy link
Copy Markdown
Collaborator Author

jlk9 commented Feb 3, 2026

Its only difference from the method for non-Reactant models with SplitRungeKuttaTimeStepper is commenting out the conditional statement:

Can't this be fixed by @trace? Like

@trace if model.clock.iteration == 0
update_state!(model, callbacks; compute_tendencies=true)
end

Using @trace here (or in SplitRungeKutta3) doesn't work. It produces a NoFieldMatchError.

@glwagner
Copy link
Copy Markdown
Member

glwagner commented Feb 4, 2026

Its only difference from the method for non-Reactant models with SplitRungeKuttaTimeStepper is commenting out the conditional statement:

Can't this be fixed by @trace? Like

@trace if model.clock.iteration == 0
update_state!(model, callbacks; compute_tendencies=true)
end

Using @trace here (or in SplitRungeKutta3) doesn't work. It produces a NoFieldMatchError.

Does this deserve an issue in Reactant?

@jlk9
Copy link
Copy Markdown
Collaborator Author

jlk9 commented Feb 5, 2026

@giordano I created an issue for using @trace on the update_state! conditional here.

@glwagner I also added test_reactant_model_correctness calls for Reactant models with SplitRungeKutta3 to test_reactant.jl and test_reactant_latitude_longitude_grid.jl.

@jlk9
Copy link
Copy Markdown
Collaborator Author

jlk9 commented Feb 13, 2026

@glwagner Request to review

@jlk9 jlk9 requested a review from glwagner February 13, 2026 15:37
Copy link
Copy Markdown
Member

@glwagner glwagner left a comment

Choose a reason for hiding this comment

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

This is good. I think we should note, this is a temporary fix until we have more comprehensive support for nested tracing (eg we need more traced_type_inner, possibly). @dkytezab if you agree with what I just said, then @jlk9 let's merge this.

@dkytezab
Copy link
Copy Markdown
Collaborator

I'm with @glwagner, note #5270.

@glwagner
Copy link
Copy Markdown
Member

@jlk9 summary is you can merge this and snatch that cred, just know we will hopefully implement the alternative interface on #5270 very soon so we don't have copy-pasted time_step! in extensions.

@jlk9 jlk9 merged commit da22d71 into main Feb 13, 2026
75 of 77 checks passed
@jlk9 jlk9 deleted the jlk9/add-split-runge-kutta-reactant-extension branch February 13, 2026 19:52
glwagner referenced this pull request Feb 20, 2026
* Support boolean comparison operations on fields and infer eltype from operations

Register comparison operators (>, <, >=, <=) as binary operations on fields,
and use Base.promote_op to infer the result type of all AbstractOperation
constructors instead of hardcoding eltype(grid). Field(op) now allocates data
matching the operation's eltype, so boolean operations produce Bool arrays.

Closes #5307

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Fix doctest and operator precedence in comparison tests

- Suppress Set output in @binary doctest (non-deterministic ordering)
- Wrap comparison `isa` tests in parentheses to fix precedence
  (`u > v isa T` parses as `u > (v isa T)`, not `(u > v) isa T`)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Fall back to eltype(grid) when promote_op fails for KernelFunctionOperation

Base.promote_op can return Union{} or non-concrete types (like Any) for
complex kernel functions (e.g. ForcingKernelFunction, buoyancy_perturbation).
This caused MethodErrors when trying to allocate data arrays. Fall back to
eltype(grid) when promote_op cannot infer a concrete return type.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Bump ambiguity threshold from 345 to 355 for comparison operators

Adding >, <, >=, <= as binary operations on fields introduces 10 new
method ambiguities with Base's definitions.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Refactor operation constructors: accept T as parameter, infer at higher level

- All operation inner constructors now accept `::Type{T}` as an optional
  final argument with a sensible default, making them clean passthroughs
- BinaryOperation, UnaryOperation, MultiaryOperation: default T via
  `Base.promote_op` (reliable for simple eltypes)
- KernelFunctionOperation inner constructor takes `(kernel_function, grid,
  arguments::Tuple, T=eltype(grid))` — no argument splatting, usable as
  internal interface. Outer convenience constructor splats + infers T.
- adapt_structure, on_architecture, and multi-region getregion pass
  `eltype(op)` explicitly, avoiding redundant re-inference

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Remove promote_op from KernelFunctionOperation constructor

Kernel functions are too complex for Base.promote_op to reliably infer
return types. For example, SeawaterPolynomials.ρ with a Float32 grid
but Float64 EOS coefficients causes promote_op to return Float64,
changing the KFO eltype from Float32 to Float64 and breaking
Float32 seawater density tests.

Revert to always using eltype(grid) for KFO, as on main. The
promote_op eltype inference is still used for Binary, Unary, and
Multiary operations where it works reliably.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Remove promote_op from KFO outer constructor, use eltype(grid) default

promote_op inferred Float64 for TEOS10 density on Float32 grids
(because EOS polynomial coefficients are Float64), causing precision
mismatches in seawater density tests. The inner constructor already
defaults to eltype(grid); users can override T explicitly via the
inner constructor when needed.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Retrigger CI

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
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.

4 participants