From 2647265566955e7bde13feafa0ab04bf3f9420ca Mon Sep 17 00:00:00 2001 From: Eliot Quon Date: Mon, 15 Jun 2026 15:31:25 -0700 Subject: [PATCH 1/3] Support NetCDF output on any AbstractVerticalCoordinate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The NCDatasets extension dispatched `vertical_coordinate_name`, `gather_vertical_dimensions`, and `default_vertical_dimension_attributes` only on the concrete `StaticVerticalDiscretization` and `MutableVerticalDiscretization`, so `NetCDFWriter` threw a `MethodError` on a grid whose vertical coordinate is any other `AbstractVerticalCoordinate` subtype (e.g. one defined in a downstream package). Widen the two non-static methods to dispatch on `AbstractVerticalCoordinate` and add a `vertical_coordinate_name(::AbstractVerticalCoordinate) = "r"` fallback. Such coordinates save their reference (Lagrangian) coordinate `r` (the shared `cᵃᵃᶠ`/`cᵃᵃᶜ` fields); physical height is reconstructed at read time from `r` and the coordinate transform, identical to the existing z-star path. `Static` keeps its own method, and `MutableVerticalDiscretization` behaviour is unchanged (`MVD <: AbstractVerticalCoordinate`). Verified by writing a NetCDF file from a `NetCDFWriter` on a grid whose vertical coordinate is a downstream `AbstractVerticalCoordinate` subtype. Co-Authored-By: Claude Opus 4.8 --- ext/OceananigansNCDatasetsExt/dimensions.jl | 25 ++++++++++++--------- src/OutputWriters/netcdf_writer.jl | 3 +++ 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/ext/OceananigansNCDatasetsExt/dimensions.jl b/ext/OceananigansNCDatasetsExt/dimensions.jl index cb501cd88b8..2a8cf6b094a 100644 --- a/ext/OceananigansNCDatasetsExt/dimensions.jl +++ b/ext/OceananigansNCDatasetsExt/dimensions.jl @@ -186,10 +186,12 @@ function gather_vertical_dimensions(coordinate::StaticVerticalDiscretization, TZ zᵃᵃᶜ_name => zᵃᵃᶜ_data) end -# For MutableVerticalDiscretization, the saved 1D coordinate is the *reference* (Lagrangian) -# coordinate `r`. The physical `z = z(r, η, …)` is reconstructible at read time from `r` and -# the time-varying `η` (output separately). -function gather_vertical_dimensions(coordinate::MutableVerticalDiscretization, TZ, Nz, Hz, z_indices, with_halos, dim_name_generator) +# Generic fallback for non-static vertical coordinates (`MutableVerticalDiscretization` +# and any others defined downstream): the saved 1D coordinate is the *reference* +# (Lagrangian) coordinate `r`, stored in the shared `cᵃᵃᶠ`/`cᵃᵃᶜ` fields. The physical +# `z = z(r, …)` is reconstructible at read time from `r` and the coordinate transform +# (e.g. the time-varying free-surface `η`), output separately. +function gather_vertical_dimensions(coordinate::AbstractVerticalCoordinate, TZ, Nz, Hz, z_indices, with_halos, dim_name_generator) rᵃᵃᶠ_name = dim_name_generator("r", coordinate, nothing, nothing, f, Val(:z)) rᵃᵃᶜ_name = dim_name_generator("r", coordinate, nothing, nothing, c, Val(:z)) @@ -431,10 +433,13 @@ function default_vertical_dimension_attributes(coordinate::StaticVerticalDiscret return suffix_grid_keys(vertical_dimension_attributes, grid_index) end -function default_vertical_dimension_attributes(coordinate::MutableVerticalDiscretization, dim_name_generator; grid_index=nothing) - # `r` is the reference (Lagrangian) coordinate. Physical depth `z = z(r, η, …)` is - # reconstructible at read time from `r` and the time-varying free-surface `η` - # written separately (see grid_reconstruction.jl). +# Generic fallback for non-static vertical coordinates: `MutableVerticalDiscretization` +# (z-star / σ) and any other `AbstractVerticalCoordinate` defined downstream. We save the +# reference (Lagrangian) coordinate `r`; physical height `z = z(r, …)` is reconstructible at +# read time from `r` and the coordinate transform (e.g. the time-varying free-surface `η`) — +# see grid_reconstruction.jl. The `StaticVerticalDiscretization` method above handles the +# plain-`z` case. +function default_vertical_dimension_attributes(coordinate::AbstractVerticalCoordinate, dim_name_generator; grid_index=nothing) r = vertical_coordinate_name(coordinate) rᵃᵃᶠ_name = dim_name_generator(r, coordinate, nothing, nothing, f, Val(:z)) rᵃᵃᶜ_name = dim_name_generator(r, coordinate, nothing, nothing, c, Val(:z)) @@ -442,8 +447,8 @@ function default_vertical_dimension_attributes(coordinate::MutableVerticalDiscre Δrᵃᵃᶠ_name = dim_name_generator("Δr", coordinate, nothing, nothing, f, Val(:z)) Δrᵃᵃᶜ_name = dim_name_generator("Δr", coordinate, nothing, nothing, c, Val(:z)) - long_face = "Reference cell-face locations in the vertical (Lagrangian r). Physical depth is reconstructible from r and η." - long_center = "Reference cell-center locations in the vertical (Lagrangian r). Physical depth is reconstructible from r and η." + long_face = "Reference cell-face locations in the vertical (Lagrangian r). Physical height is reconstructible from r and the vertical coordinate transform." + long_center = "Reference cell-center locations in the vertical (Lagrangian r). Physical height is reconstructible from r and the vertical coordinate transform." rᵃᵃᶠ_attrs = Dict("long_name" => long_face, "units" => "m", "axis" => "Z", "positive" => "up") rᵃᵃᶜ_attrs = Dict("long_name" => long_center, "units" => "m", "axis" => "Z", "positive" => "up") diff --git a/src/OutputWriters/netcdf_writer.jl b/src/OutputWriters/netcdf_writer.jl index df74c0a4f48..bb169322082 100644 --- a/src/OutputWriters/netcdf_writer.jl +++ b/src/OutputWriters/netcdf_writer.jl @@ -28,6 +28,9 @@ const MVD = MutableVerticalDiscretization vertical_coordinate_name(::SVD) = "z" vertical_coordinate_name(::MVD) = "r" +# Generic fallback: any other non-static vertical coordinate (e.g. one defined in a +# downstream package) saves its reference (Lagrangian) coordinate `r`. +vertical_coordinate_name(::AbstractVerticalCoordinate) = "r" vertical_coordinate_name(grid::AbstractGrid) = vertical_coordinate_name(grid.z) vertical_coordinate_name(grid::ImmersedBoundaryGrid) = vertical_coordinate_name(grid.underlying_grid) From 5a2cefba432d5a2185ba5cc5768cbc2288111726 Mon Sep 17 00:00:00 2001 From: Eliot Quon Date: Mon, 15 Jun 2026 16:44:30 -0700 Subject: [PATCH 2/3] Test the AbstractVerticalCoordinate vertical-coordinate-name fallback Add a unit test that a custom non-static `AbstractVerticalCoordinate` subtype (standing in for one defined in a downstream package) routes to the generic `r` fallback in `vertical_coordinate_name`, while `StaticVerticalDiscretization` still maps to `z` and `MutableVerticalDiscretization` still maps to `r`. The `r` NetCDF write path itself is covered end-to-end by `test_netcdf_rectilinear_mvd_output`, which now dispatches through the widened `::AbstractVerticalCoordinate` methods. Co-Authored-By: Claude Opus 4.8 --- test/test_netcdf_writer.jl | 37 ++++++++++++++++++++++++++++++++++--- 1 file changed, 34 insertions(+), 3 deletions(-) diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl index 895ee996583..cfd62902359 100644 --- a/test/test_netcdf_writer.jl +++ b/test/test_netcdf_writer.jl @@ -11,7 +11,7 @@ using SeawaterPolynomials.SecondOrderSeawaterPolynomials: RoquetEquationOfState using Oceananigans: Clock using Oceananigans.Models.HydrostaticFreeSurfaceModels: VectorInvariant, ImplicitFreeSurface -using Oceananigans.OutputWriters: trilocation_dim_name +using Oceananigans.OutputWriters: trilocation_dim_name, vertical_coordinate_name using Oceananigans.Grids: ξname, ηname, rname, ξnodes, ηnodes using Oceananigans.Fields: interpolate! @@ -3630,6 +3630,36 @@ function test_netcdf_rectilinear_mvd_output(arch) return nothing end +# Minimal stand-in for a non-static vertical coordinate defined outside Oceananigans +# (e.g. in a downstream package); carries the reference-coordinate fields the NetCDF +# writer reads, shared with `MutableVerticalDiscretization`. +struct ReferenceTestVerticalCoordinate{C, D, E, F} <: Oceananigans.Grids.AbstractVerticalCoordinate + cᵃᵃᶠ :: C + cᵃᵃᶜ :: D + Δᵃᵃᶠ :: E + Δᵃᵃᶜ :: F +end + +function test_netcdf_abstract_vertical_coordinate_name(arch) + # The NetCDF writer labels the vertical axis from the coordinate's *type*: a + # `StaticVerticalDiscretization` writes the physical `z`, while every other + # `AbstractVerticalCoordinate` — `MutableVerticalDiscretization` and coordinates + # defined in downstream packages — writes the reference (Lagrangian) `r`. A custom + # subtype must route to the generic `r` fallback rather than hit a `MethodError`; the + # `r` write path itself is exercised end-to-end by `test_netcdf_rectilinear_mvd_output`. + custom = ReferenceTestVerticalCoordinate(0:3, 0.5:2.5, 1, 1) + @test vertical_coordinate_name(custom) == "r" # generic fallback (under test) + + mvd = Oceananigans.Grids.MutableVerticalDiscretization(collect(0.0:3.0)) + @test vertical_coordinate_name(mvd) == "r" # specific MutableVerticalDiscretization method still wins + + static_grid = RectilinearGrid(arch; size=(1, 1, 2), x=(0, 1), y=(0, 1), z=(0, 1), + topology=(Periodic, Periodic, Bounded)) + @test vertical_coordinate_name(static_grid.z) == "z" # static coordinate still writes z + + return nothing +end + function test_netcdf_tripolar_grid_reconstruction(arch) Nx, Ny, Nz = 20, 16, 3 grid = TripolarGrid(arch, size=(Nx, Ny, Nz), z=(-100, 0)) @@ -3822,9 +3852,10 @@ end test_netcdf_cubed_sphere_panel_immersed_output(arch) end - @testset "MutableVerticalDiscretization output [$A]" begin - @info " Testing MutableVerticalDiscretization output [$A]..." + @testset "Non-static vertical coordinate output [$A]" begin + @info " Testing non-static vertical coordinate output [$A]..." test_netcdf_rectilinear_mvd_output(arch) + test_netcdf_abstract_vertical_coordinate_name(arch) end end end From bbbfdd73c9d99389941ef9e4929b31739ff22cf0 Mon Sep 17 00:00:00 2001 From: Eliot Quon Date: Mon, 15 Jun 2026 17:18:10 -0700 Subject: [PATCH 3/3] Don't call the reference coordinate "Lagrangian" in the generic path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review: the reference coordinate `r` is only Lagrangian for z-star. A general `AbstractVerticalCoordinate` (e.g. a terrain-following σ coordinate) uses `r` as a fixed reference/computational coordinate. Drop "(Lagrangian)" from the now-generic comments and the NetCDF `long_name` attributes. Co-Authored-By: Claude Opus 4.8 --- ext/OceananigansNCDatasetsExt/dimensions.jl | 8 ++++---- src/OutputWriters/netcdf_writer.jl | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/ext/OceananigansNCDatasetsExt/dimensions.jl b/ext/OceananigansNCDatasetsExt/dimensions.jl index 2a8cf6b094a..ce53a8a63ca 100644 --- a/ext/OceananigansNCDatasetsExt/dimensions.jl +++ b/ext/OceananigansNCDatasetsExt/dimensions.jl @@ -188,7 +188,7 @@ end # Generic fallback for non-static vertical coordinates (`MutableVerticalDiscretization` # and any others defined downstream): the saved 1D coordinate is the *reference* -# (Lagrangian) coordinate `r`, stored in the shared `cᵃᵃᶠ`/`cᵃᵃᶜ` fields. The physical +# coordinate `r`, stored in the shared `cᵃᵃᶠ`/`cᵃᵃᶜ` fields. The physical # `z = z(r, …)` is reconstructible at read time from `r` and the coordinate transform # (e.g. the time-varying free-surface `η`), output separately. function gather_vertical_dimensions(coordinate::AbstractVerticalCoordinate, TZ, Nz, Hz, z_indices, with_halos, dim_name_generator) @@ -435,7 +435,7 @@ end # Generic fallback for non-static vertical coordinates: `MutableVerticalDiscretization` # (z-star / σ) and any other `AbstractVerticalCoordinate` defined downstream. We save the -# reference (Lagrangian) coordinate `r`; physical height `z = z(r, …)` is reconstructible at +# reference coordinate `r`; physical height `z = z(r, …)` is reconstructible at # read time from `r` and the coordinate transform (e.g. the time-varying free-surface `η`) — # see grid_reconstruction.jl. The `StaticVerticalDiscretization` method above handles the # plain-`z` case. @@ -447,8 +447,8 @@ function default_vertical_dimension_attributes(coordinate::AbstractVerticalCoord Δrᵃᵃᶠ_name = dim_name_generator("Δr", coordinate, nothing, nothing, f, Val(:z)) Δrᵃᵃᶜ_name = dim_name_generator("Δr", coordinate, nothing, nothing, c, Val(:z)) - long_face = "Reference cell-face locations in the vertical (Lagrangian r). Physical height is reconstructible from r and the vertical coordinate transform." - long_center = "Reference cell-center locations in the vertical (Lagrangian r). Physical height is reconstructible from r and the vertical coordinate transform." + long_face = "Reference cell-face locations in the vertical (reference coordinate r). Physical height is reconstructible from r and the vertical coordinate transform." + long_center = "Reference cell-center locations in the vertical (reference coordinate r). Physical height is reconstructible from r and the vertical coordinate transform." rᵃᵃᶠ_attrs = Dict("long_name" => long_face, "units" => "m", "axis" => "Z", "positive" => "up") rᵃᵃᶜ_attrs = Dict("long_name" => long_center, "units" => "m", "axis" => "Z", "positive" => "up") diff --git a/src/OutputWriters/netcdf_writer.jl b/src/OutputWriters/netcdf_writer.jl index bb169322082..80e733d27ab 100644 --- a/src/OutputWriters/netcdf_writer.jl +++ b/src/OutputWriters/netcdf_writer.jl @@ -22,14 +22,14 @@ const MVD = MutableVerticalDiscretization # - `StaticVerticalDiscretization`: the reference and physical vertical coordinates coincide, # so we use "z". # - `MutableVerticalDiscretization` (z-star, σ-coordinates): the saved 1D coordinate is the -# reference (Lagrangian) coordinate; the physical `z = z(r, η, …)` is reconstructible from +# reference coordinate; the physical `z = z(r, η, …)` is reconstructible from # `r` and the time-varying free surface `η`. We use "r" for the 1D reference coordinate. # vertical_coordinate_name(::SVD) = "z" vertical_coordinate_name(::MVD) = "r" # Generic fallback: any other non-static vertical coordinate (e.g. one defined in a -# downstream package) saves its reference (Lagrangian) coordinate `r`. +# downstream package) saves its reference coordinate `r`. vertical_coordinate_name(::AbstractVerticalCoordinate) = "r" vertical_coordinate_name(grid::AbstractGrid) = vertical_coordinate_name(grid.z) vertical_coordinate_name(grid::ImmersedBoundaryGrid) = vertical_coordinate_name(grid.underlying_grid)