Support NetCDF output on any AbstractVerticalCoordinate

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*
+# 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,19 +433,22 @@ 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 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))
 
     Δ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 (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")

src/OutputWriters/netcdf_writer.jl

@@ -22,12 +22,15 @@ 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 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)
 

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!
 
@@ -3688,6 +3688,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))
@@ -3885,9 +3915,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