Skip to content

(0.105.0) Implement interface for stepping closure prognostic variables#5274

Merged
glwagner merged 25 commits intomainfrom
step-closure-prognostics
Feb 13, 2026
Merged

(0.105.0) Implement interface for stepping closure prognostic variables#5274
glwagner merged 25 commits intomainfrom
step-closure-prognostics

Conversation

@glwagner
Copy link
Copy Markdown
Member

@glwagner glwagner commented Feb 10, 2026

This PR creates an interface called

step_closure_prognostics!(model.closure_fields, model.closure, model, dt)

which can be extended by closures in order to step forward prognostic fields outside the models default time-stepping mechanism. This is useful for closures that implement off-pattern time-stepping methods, such as Lagrangian advection by DynamicSmagorinsky, or substepping, like CATKEVerticalDiffusivity and TKEDissipationVerticalDiffusivity. CATKEVerticalDiffusivity also evolves a non-advected prognostic variable that represents an Eulerian time-average of the surface buoyancy flux.

The new interface removes state-dependent conditionals in compute_closure_fields!, which complicate preserving the "idempotency" of update_state! (ie repeated calls to update_state! should not change the prognostic state of the model). This makes checkpointing easier, and also simplifies raising via Reactant.

I have also taken the opportunity to clean up the semantics of the closure interface.

It also adds initialize_closure_fields! which is needed to remove a conditional in DynamicSmagorinsky.

Takes over from #5267

@glwagner glwagner requested a review from tomchor February 10, 2026 20:54
glwagner and others added 3 commits February 10, 2026 13:56
…mments

Rename the `diffusivities` local variable / function argument to
`closure_fields` throughout model tendency kernels, update_state
helpers, and TKEBasedVerticalDiffusivities internals. Also renames
struct types CATKEDiffusivityFields → CATKEClosureFields and
TKEDissipationDiffusivityFields → TKEDissipationClosureFields,
and kernel functions compute_CATKE_diffusivities! →
compute_CATKE_closure_fields! and compute_TKEDissipation_diffusivities!
→ compute_TKEDissipation_closure_fields!.

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

Rename the `diffusivities` variable/argument to `closure_fields` across
all closure implementation files for consistency with the earlier rename
of the container and top-level function.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
@giordano giordano added the breaking change 💔 Concerning a change which breaks the API label Feb 11, 2026
glwagner and others added 5 commits February 11, 2026 13:27
- Add initialize_closure_fields! interface for proper closure initialization
- Fix Lagrangian averaging bug: use accumulated time instead of stage Δt
- Remove NaN sentinel pattern from DynamicSmagorinsky
- Remove obsolete time-based conditionals from CATKE and TKE-dissipation
- Add initialize! functions to NonhydrostaticModel and HydrostaticFreeSurfaceModel

Co-Authored-By: Claude Opus 4.6 <[email protected]>
- Remove unused commented buoyancy_flux wrapper from catke_equation.jl
- Remove commented compute_closure_fields! from scalar_diffusivity.jl
- Fix alignment in convective_adjustment_vertical_diffusivity.jl

Co-Authored-By: Claude Opus 4.6 <[email protected]>
- Only step DynamicSmagorinsky at stage 1 to avoid recomputing 3x per iteration with RK3
- Add initialize_closure_fields! fallback for nothing closure

Co-Authored-By: Claude Opus 4.6 <[email protected]>
When multiple closures are used, iterate through each closure
and initialize its corresponding fields.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
@codecov
Copy link
Copy Markdown

codecov bot commented Feb 11, 2026

Codecov Report

❌ Patch coverage is 75.58685% with 52 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.49%. Comparing base (b1b07fe) to head (c9ebd18).
⚠️ Report is 7 commits behind head on main.

Files with missing lines Patch % Lines
...mplementations/Smagorinskys/dynamic_coefficient.jl 46.42% 15 Missing ⚠️
...dVerticalDiffusivities/time_step_catke_equation.jl 54.83% 14 Missing ⚠️
...e_implementations/ri_based_vertical_diffusivity.jl 62.50% 6 Missing ⚠️
...fusivities/tke_dissipation_vertical_diffusivity.jl 85.00% 3 Missing ⚠️
...rbulenceClosures/turbulence_closure_diagnostics.jl 80.00% 2 Missing ⚠️
...bulence_closure_implementations/nothing_closure.jl 50.00% 2 Missing ⚠️
...odels/NonhydrostaticModels/nonhydrostatic_model.jl 66.66% 1 Missing ⚠️
...owWaterModels/shallow_water_diffusion_operators.jl 0.00% 1 Missing ⚠️
src/TurbulenceClosures/TurbulenceClosures.jl 90.90% 1 Missing ⚠️
src/TurbulenceClosures/closure_tuples.jl 85.71% 1 Missing ⚠️
... and 6 more
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #5274      +/-   ##
==========================================
- Coverage   73.51%   73.49%   -0.02%     
==========================================
  Files         395      394       -1     
  Lines       22194    22234      +40     
==========================================
+ Hits        16316    16341      +25     
- Misses       5878     5893      +15     
Flag Coverage Δ
buildkite 68.79% <79.60%> (-0.01%) ⬇️
julia 68.79% <79.60%> (-0.01%) ⬇️

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 Author

this is ready to merge

implicit_solver = model.timestepper.implicit_solver
β = model.timestepper.β[model.clock.stage] # Get the correct β value for the current stage
Δτ = model.clock.last_Δt / β
Δτ = Δt
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Suggested change
Δτ = Δt

we can use Δt

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

wait no, this is wrong. we need to substep CATKE if it has a smaller time-step than Δt

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

@simone-silvestri I added substepping here. Please confirm this makes sense.

Copy link
Copy Markdown
Member

@tomchor tomchor left a comment

Choose a reason for hiding this comment

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

You introduced 4 new test files but I don't see them being called anywhere. Am I missing something?

(Cursor/Claude code sometimes creates these test files that aren't called, so maybe that's what happened here?)

@glwagner
Copy link
Copy Markdown
Member Author

image

results for SplitRungeKutta substepping

@glwagner
Copy link
Copy Markdown
Member Author

What I changed:

  • I removed the MWE dev files (@tomchor)
  • I implemented SplitRungeKutta substepping (@simone-silvestri)
  • I also removed dead CATKE/TKEDissipation features like previous_compute_time (@simone-silvestri)
  • I removed some unused test files (distributed lat lon, test_cuda)
  • I added an unused test file to runtests (background_flux_divergence)

Ready to merge.

@glwagner glwagner merged commit aa1ad5f into main Feb 13, 2026
80 checks passed
@glwagner glwagner deleted the step-closure-prognostics branch February 13, 2026 14:20
@glwagner glwagner changed the title Implement interface for stepping closure prognostic variables (0.105.0) Implement interface for stepping closure prognostic variables Feb 15, 2026
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

breaking change 💔 Concerning a change which breaks the API

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants