Move velocity halo filling into `rk3_substep!` and `ab2_step!` for `HydrostaticFreeSurfaceModel` by simone-silvestri · Pull Request #5352 · CliMA/Oceananigans.jl

simone-silvestri · 2026-03-02T08:36:33Z

this PR is preparatory for #5351. It moves the fill_halo_regions! of the u and v velocities out of the update_state! and before the tracer update, interleaving them with the computation of the tracer tendency.
This allows us to correct the barotropic mode after the fill halo regions and moves the correction also inside the halo regions.

With this strategy, even in distributed computations, we get a correct w field also inside the halos at the end of the timestep. In addition, tracers which need the new u and v velocity component (like CATKE) for timestepping, have complete access to the velocity fields correct also in the halos, such that we can use ansynchronous computation - communication also for CATKE.

This PR additionally excludes the periphery from the velocity timestepping like we do in the non hydrostatic model

Indepth description:

Restructures the HydrostaticFreeSurfaceModel time-stepping algorithm by moving fill_halo_regions! for u, v from update_state! into the time-stepping kernels (rk_substep! / ab2_step!), interleaving it with tendency computation. Preparatory for #5351 (open boundaries).

Motivation

In the old algorithm, velocity halos were filled at the start of the next substep (inside update_state!). This had three consequences for distributed architectures:

Tracer tendencies used stale velocity halos — closures like CATKE that need u, v during tracer time-stepping had no access to correctly filled halos, preventing asynchronous communication-computation overlap.
Barotropic mode correction happened before the halo fill — the correction was only valid in the interior, so w (diagnosed from continuity) was incorrect in the halos at the end of the time step. This matters for distributed computations.
The old code had an explicit # TODO: fill halo regions for horizontal velocities should be here before the tracer update.

Algorithm change

Shown for the Split Runge-Kutta timestepper; AB2 follows the same pattern.

Before (velocity halos filled in update_state!, at the start of the next stage):

for each RK stage:
  rk_substep!:
    1. compute_momentum_flux_bcs!
    2. step_free_surface!
    3. compute_transport_velocities!
    ┌─ @apply_regionally ────────────────────────────────┐
    │ 4. compute_tracer_tendencies!    ← stale u,v halos │
    │ 5. rk_substep_velocities!        (full domain)     │
    │ 6. correct_barotropic_mode!      ← interior only   │
    │ 7. rk_substep_tracers!                             │
    └────────────────────────────────────────────────────┘
  update_state!:
    fill_halo_regions!(u, v)          ← halos filled here
    fill_halo_regions!(tracers)
    compute diagnostics + momentum tendencies

After (velocity halos filled within the substep, before tracer tendencies):

for each RK stage:
  rk_substep!:
    1. compute_momentum_flux_bcs!
    2. step_free_surface!
    3. compute_transport_velocities!
    4. rk_substep_velocities!          (exclude periphery)
    5. fill_halo_regions!(u, v)       ← halos filled here
    ┌─ @apply_regionally ───────────────────────────────────────────────────────────────────────────────┐
    │ 6. compute_tracer_tendencies!    ← correct u,v halos (also synchronizes halo pass for velocities) │
    │ 7. correct_barotropic_mode!      ← also in halos                                                  │
    │ 8. rk_substep_tracers!                                                                            │
    └───────────────────────────────────────────────────────────────────────────────────────────────────┘
  update_state!:
    fill_halo_regions!(tracers)       ← only tracers
    compute diagnostics + momentum tendencies

Other changes

Velocity stepping excludes the periphery (exclude_periphery=true), consistent with the NonhydrostaticModel.
Initialization fix: since update_state! no longer fills velocity halos, the first call (iteration 0) would compute momentum tendencies from unfilled halos. Fixed by adding fill_halo_regions!(prognostic_fields(model)) in maybe_initialize_state!.

Future work

After this PR, momentum tendency computation can be detangled from update_state! and moved to the beginning of the time step, further simplifying the algorithm.

…ananigans.jl into ss/prepare-for-open-boundaries

glwagner · 2026-03-17T17:47:02Z

this PR is preparatory for #5351. It moves the fill_halo_regions! of the u and v velocities inside the time-stepping scheme, interleaving them with the computation of the tracer tendency. This allows us to correct the barotropic mode after the fill halo regions and moves the correction also inside the halo regions.

With this strategy, even in distributed computations, we get a correct w field also inside the halos at the end of the timestep. In addition, tracers which need the new u and v velocity component (like CATKE) for timestepping, have complete access to the velocity fields correct also in the halos, such that we can use ansynchronous computation - communication also for CATKE.

This PR additionally excludes the periphery from the velocity timestepping like we do in the non hydrostatic model

Does this mean that u, v halos are not correct at the end of a time-step? Can you make the description more verbose (sketch out the whole algorithm --- the description assumes that one is intimately familiar with the algorithm already). Also please change the PR title to be more specific.

simone-silvestri · 2026-03-17T18:01:59Z

Does this mean that u, v halos are not correct at the end of a time-step? Can you make the description more verbose (sketch out the whole algorithm --- the description assumes that one is intimately familiar with the algorithm already). Also please change the PR title to be more specific.

I have given to claude the task to be nice and descriptive :)

glwagner · 2026-03-17T18:06:02Z

I'm still confused about the title. Prior to this PR, velocity halo filling was obviously part of the time-stepping loop. What does this PR actually do?

simone-silvestri · 2026-03-17T18:11:48Z

It does the halo pass before computing tracer tendencies. This has 2 main advantages for distributed computations:

tracer update has access to a complete velocity field
we can compute the vertical velocity also inside the halos in update_state!

It is also necessary for open boundary conditions because the barotropic correction needs to act also in the halos as the barotropic component of the BC needs to be consistent with what is used to update the barotropic velocity.
This is actually the initial motivation that lead to this PR, however I realized that this restructuring actually cleans up a lot of the distributed logic which is great.

glwagner · 2026-03-17T18:57:19Z

How does this algorithm ensure that the u, v halos are filled before the first time-step? This is the idea behind putting auxiliary state updates into update_state! --- this way, update_state! is called both at the end of a time-step as well as before the first step.

simone-silvestri · 2026-03-18T08:07:42Z

Yeah, that is a bit tricky, but we can put all the necessary things to do before the first time step in initialization_update_state! which can be as redundant as we want since we call it only one time.

…ananigans.jl into ss/prepare-for-open-boundaries

Copilot

Pull request overview

This PR refactors HydrostaticFreeSurfaceModel time-stepping so u, v halo filling occurs inside the RK/AB2 substeps (before tracer tendency evaluation), enabling barotropic-mode correction to apply in halo regions and aligning hydrostatic behavior with nonhydrostatic periphery-exclusion.

Changes:

Move fill_halo_regions! for horizontal velocities into RK3/AB2 stepping routines, interleaved with tracer tendency computation.
Adjust initialization to ensure prognostic halos are filled when update_state! no longer fills u, v halos.
Update split-explicit correction / transport-velocity computation and distributed buffer communication sequencing.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
src/TimeSteppers/TimeSteppers.jl	Adds `initialization_update_state!` hook and uses it at iteration 0.
src/Models/Models.jl	Imports/defines fallback `initialization_update_state!` signature with callbacks.
src/MultiRegion/multi_region_models.jl	Updates CubedSphere initialization hook signature to accept callbacks.
src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl	Removes `u, v` halo fill from `update_state!`; adjusts immersed masking helper.
src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk_step.jl	Reorders explicit RK substep operations; fills velocity halos before tracer tendencies; excludes periphery.
src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl	Mirrors RK changes for AB2; excludes periphery.
src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl	Uses kernel-parameter helpers; modifies correction masking; adjusts transport-velocity computation flow.
src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_buffers.jl	Reorders what gets synchronized for distributed buffer computations.
src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl	Changes transport-velocity field construction and initialization update flow.
src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl	Tweaks transport-velocity halo fill behavior.
src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl	Updates transport-velocity plumbing for prescribed-velocity models.
src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl	Makes `compute_transport_velocities!` a declared interface (no-op removed).
src/Models/interleave_communication_and_computation.jl	Minor whitespace-only change.
ext/OceananigansReactantExt/Models.jl	Updates Reactant HFSM initialization hook signature to accept callbacks.

Comments suppressed due to low confidence (2)

src/MultiRegion/multi_region_models.jl:90

This initialization_update_state! method now accepts callbacks but does not pass them through to update_state! (or otherwise execute UpdateStateCallsite callbacks). If callers rely on callbacks running at iteration 0 (via maybe_initialize_state!), they will be skipped for CubedSphereModel. Consider forwarding callbacks to update_state!(model, callbacks) or explicitly applying the callback filtering used elsewhere.

function Models.initialization_update_state!(model::CubedSphereModel, callbacks=[])

    # Update the state of the model
    update_state!(model)

src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_buffers.jl:109

This branch assumes transport velocities do not need synchronization for ExplicitFreeSurface because transport_velocities === velocities, but HydrostaticFreeSurfaceModel now typically constructs distinct transport_velocities fields. Either restore aliasing for ExplicitFreeSurface/Nothing, or adjust this logic so buffer-region tracer tendencies always see transport-velocity halos/w̃ consistent with the synchronized u,v and free surface.

    # We do not need to synchronize the transport velocities
    # for an ExplicitFreeSurface (`transport_velocities === velocities`)
    if !(model.free_surface isa ExplicitFreeSurface)
        ũ, ṽ, _ = model.transport_velocities
        synchronize_communication!(ũ)

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk_step.jl

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl

...ydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl

…liMA/Oceananigans.jl into glw/remove-initialization-update-state

…-state' into ss/prepare-for-open-boundaries

…oundaries

Copilot

Pull request overview

This PR restructures HydrostaticFreeSurfaceModel time-stepping so u,v halo filling happens inside the RK/AB2 stepping kernels (before tracer tendency evaluation), enabling barotropic corrections to be applied consistently inside halos and excluding the periphery from velocity stepping (to match the nonhydrostatic model).

Changes:

Move fill_halo_regions! for horizontal velocities from update_state! into RK/AB2 stepping flows and apply exclude_periphery=true for velocity stepping kernels.
Refactor transport velocity handling (construction + updating) and update split-explicit/implicit transport velocity kernels to use surface_kernel_parameters / volume_kernel_parameters.
Adjust distributed-buffer synchronization logic and extend immersed-boundary utilities for “immersed periphery” handling.

Reviewed changes

Copilot reviewed 15 out of 15 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
`test/test_metal.jl`	Extends a MetalGPU immersed-boundary run length to validate evolution over more steps.
`src/TimeSteppers/TimeSteppers.jl`	Adds a new (currently unused) `initialization_update_state!` interface hook.
`src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl`	Stops filling `u,v` halos in `update_state!`; introduces `mask_immersed_horizontal_velocities!`.
`src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl`	Updates prescribed-velocity specializations for new masking/transport-velocity signatures.
`src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl`	Updates implicit transport-velocity kernel launch parameters and streamlines transport-velocity update.
`src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk_step.jl`	Moves velocity stepping + halo fill into RK substep and applies `exclude_periphery=true`.
`src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl`	Moves velocity stepping + halo fill into AB2 step and applies `exclude_periphery=true`.
`src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl`	Changes how `transport_velocities` are constructed/updated; adds fallback transport-velocity updater.
`src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_buffers.jl`	Adjusts async distributed synchronization for tracer/momentum buffer computation.
`src/Models/HydrostaticFreeSurfaceModels/barotropic_pressure_correction.jl`	Switches implicit barotropic correction launch to `volume_kernel_parameters(...)`.
`src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl`	Uses kernel-parameter helpers; applies immersed-periphery logic to barotropic corrections and transport velocities.
`src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl`	Turns `compute_transport_velocities!` into a declared interface function.
`src/ImmersedBoundaries/immersed_boundary_interface.jl`	Adds non-IBG fallback for `immersed_peripheral_node`.
`equally_partitioned_grids.jl`	Adds an MPI/distributed grid inspection script.
`.claude/rules/julia-repl-rules.md`	Adds internal guidance for using a Julia MCP REPL when available.

Comments suppressed due to low confidence (1)

src/Models/HydrostaticFreeSurfaceModels/barotropic_pressure_correction.jl:41

volume_kernel_parameters is referenced here but is not imported/qualified in this file at include time, which will raise UndefVarError. Add the appropriate using Oceananigans.Models: volume_kernel_parameters (or qualify the call) in this file.


    launch!(model.architecture, model.grid, volume_kernel_parameters(model.grid),
            _barotropic_pressure_correction!,
            model.velocities,
            model.grid,
            Δt,

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

equally_partitioned_grids.jl

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk_step.jl

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl

test/test_metal.jl

simone-silvestri added 2 commits February 27, 2026 10:13

this should work

2454792

this should be correct

2fbcbc9

simone-silvestri requested a review from glwagner March 2, 2026 08:36

simone-silvestri added 6 commits March 2, 2026 10:54

Import surface and volume kernel parameters

6f20c82

correct the implicit solver

78e80f5

Merge branch 'main' into ss/prepare-for-open-boundaries

b8477fc

will this work?

83d7d4b

Merge branch 'ss/prepare-for-open-boundaries' of github.com:CliMA/Oce…

49acd83

…ananigans.jl into ss/prepare-for-open-boundaries

this should work now

2c64d3b

simone-silvestri changed the title ~~Prepare for open boundaries~~ Move velocity halo filling into the time-stepping loop for HydrostaticFreeSurfaceModel Mar 17, 2026

simone-silvestri changed the title ~~Move velocity halo filling into the time-stepping loop for HydrostaticFreeSurfaceModel~~ Move velocity halo filling into rk3_substep! and ab2_step! for HydrostaticFreeSurfaceModel Mar 17, 2026

simone-silvestri added 2 commits March 18, 2026 00:45

just a test still stuff to fix

e25919c

masking was the issue

8a1dd66

simone-silvestri added 8 commits March 18, 2026 09:17

new changes

11a3a3c

Merge branch 'main' into ss/prepare-for-open-boundaries

7c3213b

some more bugfixes

08ce256

Merge branch 'ss/prepare-for-open-boundaries' of github.com:CliMA/Oce…

14bd912

…ananigans.jl into ss/prepare-for-open-boundaries

Merge branch 'main' into ss/prepare-for-open-boundaries

b2035c6

we need transport velocities for all free surfaces

d2088cd

Merge branch 'ss/prepare-for-open-boundaries' of github.com:CliMA/Oce…

7fb9235

…ananigans.jl into ss/prepare-for-open-boundaries

no-op for a nothing free surface

92ed34f

simone-silvestri requested a review from Copilot March 19, 2026 11:29

Merge branch 'main' into ss/prepare-for-open-boundaries

d69a0b5

Copilot AI reviewed Mar 19, 2026

View reviewed changes

simone-silvestri added 25 commits March 25, 2026 16:07

fill free surface halos before timestepping

8e1fad6

Merge branch 'glw/remove-initialization-update-state' of github.com:C…

f9fc85f

…liMA/Oceananigans.jl into glw/remove-initialization-update-state

remove stale imports

1131761

add fill halos in reconcile_state!

56d30fd

now just need to fix multi-region

a8e49ff

cubed sphere implementation

2f57d80

Merge remote-tracking branch 'origin/glw/remove-initialization-update…

5870a84

…-state' into ss/prepare-for-open-boundaries

Merge remote-tracking branch 'origin/glw/remove-initialization-update…

d2e7ee5

…-state' into ss/prepare-for-open-boundaries

fix the dispatch

8d89c68

Merge remote-tracking branch 'origin/glw/remove-initialization-update…

b698f2c

…-state' into ss/prepare-for-open-boundaries

barotropic correction for implicit free surface

ec9ecdf

Merge remote-tracking branch 'origin/main' into ss/prepare-for-open-b…

139bf11

…oundaries

add an MC Repl command

1f09dd4

fix all tests

9ac96f9

remove whitespace

6d9ccf2

mask velocities before filling

36ed3b7

remove stale imports

156338c

fix multi-region behavior

8bee07a

remove whitespace

f2b2959

Merge branch 'main' into ss/prepare-for-open-boundaries

a09dc81

Merge branch 'main' into ss/prepare-for-open-boundaries

7bca61e

correct GPU tests

ff80e2b

fix all tests

b77efbf

fix whitespace

1ae40c5

this is better

30ea5cb

simone-silvestri requested a review from Copilot March 30, 2026 12:06

Copilot started reviewing on behalf of simone-silvestri March 30, 2026 12:06 View session

Copilot AI reviewed Mar 30, 2026

View reviewed changes

simone-silvestri added 2 commits March 30, 2026 14:16

remove vestigial file

978a65b

update now

7efddc8

Conversation

simone-silvestri commented Mar 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Algorithm change

Other changes

Future work

Uh oh!

glwagner commented Mar 17, 2026

Uh oh!

simone-silvestri commented Mar 17, 2026

Uh oh!

glwagner commented Mar 17, 2026

Uh oh!

simone-silvestri commented Mar 17, 2026

Uh oh!

glwagner commented Mar 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

simone-silvestri commented Mar 18, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

simone-silvestri commented Mar 2, 2026 •

edited

Loading

glwagner commented Mar 17, 2026 •

edited

Loading