Add split runge kutta to reactant extension#5237
Conversation
|
NOTE: this PR already includes the small change from #5236, so maybe merge that one first. |
Codecov Report❌ Patch coverage is
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
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
we need a test here, in addition to the PR that this depends on. |
|
@dkytezab this is probably of interest to you. |
…ectilineargrid test
Can't this be fixed by Oceananigans.jl/ext/OceananigansReactantExt/TimeSteppers.jl Lines 65 to 67 in 97061af |
|
We should also add support for Once we get around FFT planning issues (see #5244) this will be the furthest upstream barrier to AD of |
Using @trace here (or in SplitRungeKutta3) doesn't work. It produces a |
Does this deserve an issue in Reactant? |
|
@glwagner Request to review |
* 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]>
This adds an additional method for
time_step!toOceananigansReactantExtforReactantModel{<:SplitRungeKuttaTimeStepper}. Its only difference from the method for non-Reactant models withSplitRungeKuttaTimeStepperis commenting out the conditional statement:since this breaks Reactant, and the Reactant methods for
first_time_step!automatically callupdate_state!anyway. A method for Reactant models with Adams-Bashforth timesteppers has already been implemented for the same reason.