diff --git a/ext/OceananigansNCDatasetsExt/dimensions.jl b/ext/OceananigansNCDatasetsExt/dimensions.jl index cb501cd88b..ce53a8a63c 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* +# 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 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 (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 df74c0a4f4..80e733d27a 100644 --- a/src/OutputWriters/netcdf_writer.jl +++ b/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) diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl index e6998c528e..1a60f9dfc4 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! @@ -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