Fix broken LagrangianAveraged DynamicSmagorinsky regression tests (#5506)

* Fix broken LagrangianAveraged DynamicSmagorinsky regression tests

Two bugs prevented the regression test from exercising the Lagrangian-
averaged dynamic Smagorinsky closure at all:

1. Duplicate nested `for` loop in `test_nonhydrostatic_regression.jl`
   caused the inner loop to shadow the outer one, running each closure
   16 times instead of 4 (with no effect on correctness, since the
   inner variable always won).

2. `ocean_large_eddy_simulation_regression_test.jl` restores model
   state via direct `.=` assignment rather than `set!`, so
   `initialize_closure_fields!` was never called with the loaded
   velocities. For LagrangianAveraging this left 𝒥ᴹᴹ = 0, which
   permanently disables the closure (ϵ = 0 → 𝒥ᴹᴹ never accumulates).
   Now calls `initialize_closure_fields!` after state is loaded so the
   closure is properly bootstrapped from the checkpoint velocities.

Note: the LagrangianDynamicSmagorinsky regression test will correctly
fail until reference data is regenerated with the fixed code.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* comment out tests we're not testing

* better fix

* Restore previous_compute_time Ref for bit-exact LagAvg DynamicSmag regression

The regression test loader previously only restored `Field`-typed leaves
in `model.closure_fields`, skipping the `previous_compute_time::RefValue`
scalar. After `set!(model, ...)`, that Ref was left at `clock.time = 0`
(because `model.clock.time` is updated only later). At the first time
step from the loaded checkpoint, `Δt_lagrangian = clock.time - 0` was
catastrophically large, making `ϵ ≈ 1` in `_lagrangian_average_LM_MM!`
and effectively discarding the carefully-restored 𝒥 history.

Extends the closure-restoration loop to also copy `Base.RefValue` leaves,
so `previous_compute_time` is set to the value from the checkpoint. With
this, the LagrangianAveragedDynamicSmagorinsky regression test reproduces
the reference iteration to machine epsilon (~1e-15) over 10 time steps.

Also:
- Auto-detect new (`simulation/model/...`) vs legacy (root) checkpoint
  paths in `get_fields_from_checkpoint`.
- Add `load_interior(data, target_size)` helper that auto-detects how
  many halo layers are saved and unwraps OffsetArrays — removes the
  brittle hardcoded `[2:end-1, 2:end-1, 2:end-1]` slicing.
- Return `closure_fields` (and `pNHS`, harmlessly) from
  `get_fields_from_checkpoint`.
- Fix the regression-data-generation block (which had a stale
  `simulation.stop_iteration = ...` referring to a `simulation` that
  was never constructed) so it can be uncommented to regenerate
  reference data.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* force redownload of the data if the files have changed sizes

* make it work on GPU

* uncomment temporarily commented out regions

* remove unnused function

* Make rayleigh_benard regression's get_fields_from_checkpoint destructure explicit

`get_fields_from_checkpoint` now returns 5 values (`solution, Gⁿ, G⁻,
closure_fields, pNHS`). Rayleigh-Bénard only needs the first three;
Julia's destructuring iterates and silently discards extras, so the
old `solution, Gⁿ, G⁻ = get_fields_from_checkpoint(...)` works, but
it's not self-documenting. Bind the unused values to `_` so the call
site explicitly shows that two extras are being ignored.

Addresses Copilot review comment on #5506.

---------

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: Gregory L. Wagner <wagner.greg@gmail.com>

Author: 3 people

test/data_dependencies.jl

@@ -46,11 +46,28 @@ dh = DataDep("regression_truth_data",
 
 DataDeps.register(dh)
 
-# Invalidate stale DataDeps cache if new files are missing
+# Invalidate stale DataDeps cache when any expected reference file is missing
+# or has an unexpected size. Persistent CI caches (e.g. buildkite agents)
+# may hold reference data from previous artifact uploads; expected sizes are
+# the authoritative way to detect that.
+#
+# Sizes refer to the current files in glwagner/OceananigansArtifacts.jl
+# (data_for_regression_tests/). When regenerating reference data, update
+# these alongside the new upload.
+const _expected_reference_sizes = Dict(
+    "ocean_large_eddy_simulation_DynamicSmagorinsky_directional_iteration10000.jld2" => 713110,
+    "ocean_large_eddy_simulation_DynamicSmagorinsky_directional_iteration10010.jld2" => 713110,
+    "ocean_large_eddy_simulation_DynamicSmagorinsky_lagrangian_iteration10000.jld2"  => 1244676,
+    "ocean_large_eddy_simulation_DynamicSmagorinsky_lagrangian_iteration10010.jld2"  => 1244676,
+)
+
 dd_path = try; datadep"regression_truth_data"; catch; nothing; end
 if dd_path !== nothing
-    expected = joinpath(dd_path, "ocean_large_eddy_simulation_DynamicSmagorinsky_directional_iteration10000.jld2")
-    if !isfile(expected)
+    cache_stale = any(_expected_reference_sizes) do (filename, expected_size)
+        path = joinpath(dd_path, filename)
+        !isfile(path) || filesize(path) != expected_size
+    end
+    if cache_stale
         @info "Regression truth data cache is stale, re-downloading..."
         rm(dd_path; recursive=true)
         datadep"regression_truth_data"

test/regression_tests/ocean_large_eddy_simulation_regression_test.jl

@@ -2,6 +2,22 @@ using Oceananigans.TurbulenceClosures: AnisotropicMinimumDissipation, Lagrangian
 using Oceananigans.TimeSteppers: update_state!
 using Oceananigans.DistributedComputations: cpu_architecture, partition
 
+# Extract a window matching `target_size` from the center of `data`,
+# auto-detecting and stripping halo layers. Works for any saved halo
+# count (e.g. data saved as `parent(field)` with halo=1 is 18³,
+# halo=2 is 20³; both reduce to the same 16³ interior). Also unwraps
+# OffsetArrays via `parent` (which is what new-format Checkpointer
+# files contain). `parent` keeps the underlying storage backend
+# (Array or CuArray) intact — important for GPU correctness.
+function load_interior(data, target_size)
+    arr = data isa AbstractArray ? parent(data) : data
+    sz = size(arr)
+    Hx = (sz[1] - target_size[1]) ÷ 2
+    Hy = (sz[2] - target_size[2]) ÷ 2
+    Hz = (sz[3] - target_size[3]) ÷ 2
+    return arr[Hx+1:Hx+target_size[1], Hy+1:Hy+target_size[2], Hz+1:Hz+target_size[3]]
+end
+
 function run_ocean_large_eddy_simulation_regression_test(arch, grid_type, closure)
     if first(closure) isa SmagorinskyLilly
         name = "ocean_large_eddy_simulation_SmagorinskyLilly"
@@ -73,7 +89,7 @@ function run_ocean_large_eddy_simulation_regression_test(arch, grid_type, closur
     u₀(x, y, z) = sqrt(abs(Qᵘ)) * 1e-3 * Ξ(z)
     set!(model, u=u₀, w=u₀, T=T₀, S=35)
 
-    simulation.stop_iteration = spinup_steps-test_steps
+    simulation = Simulation(model, Δt=Δt, stop_iteration=spinup_steps-test_steps)
     run!(simulation)
 
     checkpointer = Checkpointer(model, schedule = IterationInterval(test_steps), prefix = name,
@@ -93,35 +109,31 @@ function run_ocean_large_eddy_simulation_regression_test(arch, grid_type, closur
     datadep_path = "regression_truth_data/" * name * "_iteration$spinup_steps.jld2"
     initial_filename = @datadep_str datadep_path
 
-    solution₀, Gⁿ₀, G⁻₀ = get_fields_from_checkpoint(initial_filename)
+    solution₀, Gⁿ₀, G⁻₀, closure₀, pNHS₀ = get_fields_from_checkpoint(initial_filename)
 
     Nz = grid.Nz
 
     cpu_arch = cpu_architecture(architecture(grid))
 
-    u₀ = partition(ArrayType(solution₀.u[2:end-1, 2:end-1, 2:end-1]), cpu_arch, size(u))
-    v₀ = partition(ArrayType(solution₀.v[2:end-1, 2:end-1, 2:end-1]), cpu_arch, size(v))
-    w₀ = partition(ArrayType(solution₀.w[2:end-1, 2:end-1, 2:end-1]), cpu_arch, size(w))
-    T₀ = partition(ArrayType(solution₀.T[2:end-1, 2:end-1, 2:end-1]), cpu_arch, size(T))
-    S₀ = partition(ArrayType(solution₀.S[2:end-1, 2:end-1, 2:end-1]), cpu_arch, size(S))
-
-    Gⁿu₀ = partition(ArrayType(Gⁿ₀.u)[2:end-1, 2:end-1, 2:end-1], cpu_arch, size(u))
-    Gⁿv₀ = partition(ArrayType(Gⁿ₀.v)[2:end-1, 2:end-1, 2:end-1], cpu_arch, size(v))
-    Gⁿw₀ = partition(ArrayType(Gⁿ₀.w)[2:end-1, 2:end-1, 2:end-1], cpu_arch, size(w))
-    GⁿT₀ = partition(ArrayType(Gⁿ₀.T)[2:end-1, 2:end-1, 2:end-1], cpu_arch, size(T))
-    GⁿS₀ = partition(ArrayType(Gⁿ₀.S)[2:end-1, 2:end-1, 2:end-1], cpu_arch, size(S))
-
-    G⁻u₀ = partition(ArrayType(G⁻₀.u)[2:end-1, 2:end-1, 2:end-1], cpu_arch, size(u))
-    G⁻v₀ = partition(ArrayType(G⁻₀.v)[2:end-1, 2:end-1, 2:end-1], cpu_arch, size(v))
-    G⁻w₀ = partition(ArrayType(G⁻₀.w)[2:end-1, 2:end-1, 2:end-1], cpu_arch, size(w))
-    G⁻T₀ = partition(ArrayType(G⁻₀.T)[2:end-1, 2:end-1, 2:end-1], cpu_arch, size(T))
-    G⁻S₀ = partition(ArrayType(G⁻₀.S)[2:end-1, 2:end-1, 2:end-1], cpu_arch, size(S))
-
-    interior(model.velocities.u) .= u₀
-    interior(model.velocities.v) .= v₀
-    interior(model.velocities.w) .= w₀
-    interior(model.tracers.T)    .= T₀
-    interior(model.tracers.S)    .= S₀
+    u₀ = partition(ArrayType(load_interior(solution₀.u, size(u))), cpu_arch, size(u))
+    v₀ = partition(ArrayType(load_interior(solution₀.v, size(v))), cpu_arch, size(v))
+    w₀ = partition(ArrayType(load_interior(solution₀.w, size(w))), cpu_arch, size(w))
+    T₀ = partition(ArrayType(load_interior(solution₀.T, size(T))), cpu_arch, size(T))
+    S₀ = partition(ArrayType(load_interior(solution₀.S, size(S))), cpu_arch, size(S))
+
+    Gⁿu₀ = partition(load_interior(ArrayType(Gⁿ₀.u), size(u)), cpu_arch, size(u))
+    Gⁿv₀ = partition(load_interior(ArrayType(Gⁿ₀.v), size(v)), cpu_arch, size(v))
+    Gⁿw₀ = partition(load_interior(ArrayType(Gⁿ₀.w), size(w)), cpu_arch, size(w))
+    GⁿT₀ = partition(load_interior(ArrayType(Gⁿ₀.T), size(T)), cpu_arch, size(T))
+    GⁿS₀ = partition(load_interior(ArrayType(Gⁿ₀.S), size(S)), cpu_arch, size(S))
+
+    G⁻u₀ = partition(load_interior(ArrayType(G⁻₀.u), size(u)), cpu_arch, size(u))
+    G⁻v₀ = partition(load_interior(ArrayType(G⁻₀.v), size(v)), cpu_arch, size(v))
+    G⁻w₀ = partition(load_interior(ArrayType(G⁻₀.w), size(w)), cpu_arch, size(w))
+    G⁻T₀ = partition(load_interior(ArrayType(G⁻₀.T), size(T)), cpu_arch, size(T))
+    G⁻S₀ = partition(load_interior(ArrayType(G⁻₀.S), size(S)), cpu_arch, size(S))
+
+    set!(model, u=u₀, v=v₀, w=w₀, T=T₀, S=S₀)
 
     interior(model.timestepper.Gⁿ.u) .= Gⁿu₀
     interior(model.timestepper.Gⁿ.v) .= Gⁿv₀
@@ -135,6 +147,40 @@ function run_ocean_large_eddy_simulation_regression_test(arch, grid_type, closur
     interior(model.timestepper.G⁻.T) .= G⁻T₀
     interior(model.timestepper.G⁻.S) .= G⁻S₀
 
+    # Restore closure prognostic state (𝒥ᴸᴹ, 𝒥ᴹᴹ, ..., previous_compute_time)
+    # so the closure picks up where the spinup left off rather than bootstrapping
+    # fresh. Saved structure: Tuple (one entry per closure) of NamedTuples whose
+    # leaves are either (data = OffsetArray,) Field structs or scalar Refs.
+    if closure₀ !== nothing
+        for (saved, current) in zip(closure₀, model.closure_fields)
+            saved isa NamedTuple || continue
+            for name in keys(saved)
+                hasproperty(current, name) || continue
+                target = getproperty(current, name)
+                leaf = saved[name]
+                if target isa Oceananigans.Fields.Field
+                    arr = leaf isa NamedTuple && haskey(leaf, :data) ? leaf.data :
+                          leaf isa AbstractArray ? leaf : nothing
+                    arr === nothing && continue
+                    # ArrayType(...) transfers to the field's backend (CPU/GPU)
+                    # before broadcast-assigning to the GPU/CPU interior.
+                    interior(target) .= ArrayType(load_interior(arr, size(target)))
+                elseif target isa Base.RefValue
+                    # e.g. previous_compute_time — restore the scalar so
+                    # Δt_lagrangian = clock.time - previous_compute_time
+                    # matches the running simulation on the very next step.
+                    target[] = leaf
+                end
+            end
+        end
+    end
+
+    # Restore non-hydrostatic pressure so the next step's velocity
+    # correction starts from the same pressure field as the reference.
+    if pNHS₀ !== nothing
+        interior(model.pressures.pNHS) .= ArrayType(load_interior(pNHS₀, size(model.pressures.pNHS)))
+    end
+
     model.clock.time = spinup_steps * Δt
     model.clock.iteration = spinup_steps
 
@@ -148,19 +194,19 @@ function run_ocean_large_eddy_simulation_regression_test(arch, grid_type, closur
     datadep_path = "regression_truth_data/" * name * "_iteration$(spinup_steps+test_steps).jld2"
     final_filename = @datadep_str datadep_path
 
-    solution₁, Gⁿ₁, G⁻₁ = get_fields_from_checkpoint(final_filename)
+    solution₁, Gⁿ₁, G⁻₁, _, _ = get_fields_from_checkpoint(final_filename)
 
     test_fields = @allowscalar (u = Array(interior(model.velocities.u)),
                                 v = Array(interior(model.velocities.v)),
                                 w = Array(interior(model.velocities.w)[:, :, 1:nz]),
                                 T = Array(interior(model.tracers.T)),
                                 S = Array(interior(model.tracers.S)))
 
-    u₁ = partition(Array(solution₁.u)[2:end-1, 2:end-1, 2:end-1], cpu_arch, size(u))
-    v₁ = partition(Array(solution₁.v)[2:end-1, 2:end-1, 2:end-1], cpu_arch, size(v))
-    w₁ = partition(Array(solution₁.w)[2:end-1, 2:end-1, 2:end-2], cpu_arch, size(test_fields.w))
-    T₁ = partition(Array(solution₁.T)[2:end-1, 2:end-1, 2:end-1], cpu_arch, size(T))
-    S₁ = partition(Array(solution₁.S)[2:end-1, 2:end-1, 2:end-1], cpu_arch, size(S))
+    u₁ = partition(load_interior(solution₁.u, size(u)), cpu_arch, size(u))
+    v₁ = partition(load_interior(solution₁.v, size(v)), cpu_arch, size(v))
+    w₁ = partition(load_interior(solution₁.w, size(test_fields.w)), cpu_arch, size(test_fields.w))
+    T₁ = partition(load_interior(solution₁.T, size(T)), cpu_arch, size(T))
+    S₁ = partition(load_interior(solution₁.S, size(S)), cpu_arch, size(S))
 
     @show size(test_fields.w), size(w₁)
 

test/regression_tests/rayleigh_benard_regression_test.jl

@@ -104,7 +104,7 @@ function run_rayleigh_benard_regression_test(arch, grid_type)
     datadep_path = "regression_truth_data/" * prefix * "_iteration$spinup_steps.jld2"
     initial_filename = @datadep_str datadep_path
 
-    solution₀, Gⁿ₀, G⁻₀ = get_fields_from_checkpoint(initial_filename)
+    solution₀, Gⁿ₀, G⁻₀, _, _ = get_fields_from_checkpoint(initial_filename)
 
     cpu_arch = cpu_architecture(architecture(grid))
 
@@ -156,7 +156,7 @@ function run_rayleigh_benard_regression_test(arch, grid_type)
     datadep_path = "regression_truth_data/" * prefix * "_iteration$(spinup_steps+test_steps).jld2"
     final_filename = @datadep_str datadep_path
 
-    solution₁, Gⁿ₁, G⁻₁ = get_fields_from_checkpoint(final_filename)
+    solution₁, Gⁿ₁, G⁻₁, _, _ = get_fields_from_checkpoint(final_filename)
 
     test_fields =  @allowscalar (u = Array(interior(model.velocities.u)),
                                  v = Array(interior(model.velocities.v)),

test/test_nonhydrostatic_regression.jl

@@ -8,37 +8,51 @@ using Oceananigans.TurbulenceClosures: LagrangianAveraging
 function get_fields_from_checkpoint(filename)
     file = jldopen(filename)
 
-    tracers = keys(file["tracers"])
+    # Auto-detect format: new Checkpointer saves under "simulation/model/...",
+    # legacy reference data has fields at the file root.
+    base = haskey(file, "simulation") ? "simulation/model/" : ""
+
+    tracers = keys(file["$(base)tracers"])
     tracers = Tuple(Symbol(c) for c in tracers)
 
-    velocity_fields = (u = file["velocities/u/data"],
-                       v = file["velocities/v/data"],
-                       w = file["velocities/w/data"])
+    velocity_fields = (u = file["$(base)velocities/u/data"],
+                       v = file["$(base)velocities/v/data"],
+                       w = file["$(base)velocities/w/data"])
 
     tracer_fields =
-        NamedTuple{tracers}(Tuple(file["tracers/$c/data"] for c in tracers))
+        NamedTuple{tracers}(Tuple(file["$(base)tracers/$c/data"] for c in tracers))
 
-    current_tendency_velocity_fields = (u = file["timestepper/Gⁿ/u/data"],
-                                        v = file["timestepper/Gⁿ/v/data"],
-                                        w = file["timestepper/Gⁿ/w/data"])
+    current_tendency_velocity_fields = (u = file["$(base)timestepper/Gⁿ/u/data"],
+                                        v = file["$(base)timestepper/Gⁿ/v/data"],
+                                        w = file["$(base)timestepper/Gⁿ/w/data"])
 
     current_tendency_tracer_fields =
-        NamedTuple{tracers}(Tuple(file["timestepper/Gⁿ/$c/data"] for c in tracers))
+        NamedTuple{tracers}(Tuple(file["$(base)timestepper/Gⁿ/$c/data"] for c in tracers))
 
-    previous_tendency_velocity_fields = (u = file["timestepper/G⁻/u/data"],
-                                         v = file["timestepper/G⁻/v/data"],
-                                         w = file["timestepper/G⁻/w/data"])
+    previous_tendency_velocity_fields = (u = file["$(base)timestepper/G⁻/u/data"],
+                                         v = file["$(base)timestepper/G⁻/v/data"],
+                                         w = file["$(base)timestepper/G⁻/w/data"])
 
     previous_tendency_tracer_fields =
-        NamedTuple{tracers}(Tuple(file["timestepper/G⁻/$c/data"] for c in tracers))
+        NamedTuple{tracers}(Tuple(file["$(base)timestepper/G⁻/$c/data"] for c in tracers))
+
+    # Closure prognostic fields. Only present in new-format checkpoints.
+    # JLD2 fully deserializes these into the original Tuple/NamedTuple
+    # structure, so we just return the loaded object as-is.
+    closure_fields = haskey(file, "$(base)closure_fields") ? file["$(base)closure_fields"] : nothing
+
+    # Non-hydrostatic pressure — without restoring this, the next step's
+    # velocity correction (u -= Δt * ∇p) starts from zero pressure and
+    # diverges from the running-state reference at O(Δt) per step.
+    pNHS = haskey(file, "$(base)pressures/pNHS/data") ? file["$(base)pressures/pNHS/data"] : nothing
 
     close(file)
 
     solution = merge(velocity_fields, tracer_fields)
     Gⁿ = merge(current_tendency_velocity_fields, current_tendency_tracer_fields)
     G⁻ = merge(previous_tendency_velocity_fields, previous_tendency_tracer_fields)
 
-    return solution, Gⁿ, G⁻
+    return solution, Gⁿ, G⁻, closure_fields, pNHS
 end
 
 include("regression_tests/thermal_bubble_regression_test.jl")
@@ -70,17 +84,12 @@ include("regression_tests/ocean_large_eddy_simulation_regression_test.jl")
                 dyn_smag_directional = (DynamicSmagorinsky(averaging=(1, 2)),)
                 dyn_smag_lagrangian = (DynamicSmagorinsky(averaging=LagrangianAveraging()),)
 
-                for (closurename, closure) in [("AnisotropicMinimumDissipation", amd_closure),
-                                                ("SmagorinskyLilly", smag_closure),
-                                                ("DirectionalDynamicSmagorinsky", dyn_smag_directional),
-                                                ("LagrangianDynamicSmagorinsky", dyn_smag_lagrangian)]
                 for (closurename, closure) in [("AnisotropicMinimumDissipation", amd_closure),
                                                ("SmagorinskyLilly", smag_closure),
                                                ("DirectionalDynamicSmagorinsky", dyn_smag_directional),
                                                ("LagrangianDynamicSmagorinsky", dyn_smag_lagrangian)]
-                        @info "  Testing oceanic large eddy simulation regression [$A, $closurename, $grid_type grid]"
-                        run_ocean_large_eddy_simulation_regression_test(arch, grid_type, closure)
-                    end
+                    @info "  Testing oceanic large eddy simulation regression [$A, $closurename, $grid_type grid]"
+                    run_ocean_large_eddy_simulation_regression_test(arch, grid_type, closure)
                 end
             end
         end