Add methods to write a field directly to NetCDF (#4730)

Co-authored-by: Copilot <[email protected]>

Author: tomchor

ext/OceananigansNCDatasetsExt.jl

@@ -1,21 +1,22 @@
 module OceananigansNCDatasetsExt
 
 using NCDatasets
+import NCDatasets: defVar
 
 using Dates: AbstractTime, UTC, now
 using Printf: @sprintf
 using OrderedCollections: OrderedDict
 
 using Oceananigans: initialize!, prettytime, pretty_filesize, AbstractModel
-using Oceananigans.Architectures: CPU, GPU
-using Oceananigans.AbstractOperations: KernelFunctionOperation
+using Oceananigans.Architectures: CPU, GPU, on_architecture
+using Oceananigans.AbstractOperations: KernelFunctionOperation, AbstractOperation
 using Oceananigans.BuoyancyFormulations: BuoyancyForce, BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState
 using Oceananigans.Fields
-using Oceananigans.Fields: reduced_dimensions, reduced_location, location
+using Oceananigans.Fields: Reduction, reduced_dimensions, reduced_location, location
 using Oceananigans.Grids: Center, Face, Flat, Periodic, Bounded,
                           AbstractGrid, RectilinearGrid, LatitudeLongitudeGrid, StaticVerticalDiscretization,
                           topology, halo_size, xspacings, yspacings, zspacings, λspacings, φspacings,
-                          parent_index_range, ξnodes, ηnodes, rnodes, validate_index, peripheral_node,
+                          parent_index_range, nodes, ξnodes, ηnodes, rnodes, validate_index, peripheral_node,
                           constructor_arguments
 using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBottom, GFBIBG, GridFittedBoundary, PartialCellBottom, PCBIBG
 using Oceananigans.Models: ShallowWaterModel, LagrangianParticles
@@ -41,13 +42,71 @@ using Oceananigans.OutputWriters:
     show_array_type
 
 import Oceananigans: write_output!
-import Oceananigans.OutputWriters: NetCDFWriter, write_grid_reconstruction_data!, convert_for_netcdf, materialize_from_netcdf, reconstruct_grid
+import Oceananigans.OutputWriters:
+    NetCDFWriter,
+    write_grid_reconstruction_data!,
+    convert_for_netcdf,
+    materialize_from_netcdf,
+    reconstruct_grid,
+    trilocation_dim_name
 
 const c = Center()
 const f = Face()
 const BoussinesqSeawaterBuoyancy = SeawaterBuoyancy{FT, <:BoussinesqEquationOfState, T, S} where {FT, T, S}
 const BuoyancyBoussinesqEOSModel = BuoyancyForce{<:BoussinesqSeawaterBuoyancy, g} where {g}
 
+#####
+##### Extend defVar to be able to write fields to NetCDF directly
+#####
+
+defVar(ds, name, op::AbstractOperation; kwargs...) = defVar(ds, name, Field(op); kwargs...)
+defVar(ds, name, op::Reduction; kwargs...) = defVar(ds, name, Field(op); kwargs...)
+
+function defVar(ds, name, field::AbstractField;
+                time_dependent=false,
+                dimension_name_generator = trilocation_dim_name,
+                kwargs...)
+    field_cpu = on_architecture(CPU(), field) # Need to bring field to CPU in order to write it to NetCDF
+    field_data = Array{eltype(field)}(field_cpu)
+    dims = field_dimensions(field, dimension_name_generator)
+    all_dims = time_dependent ? (dims..., "time") : dims
+
+    # Validate that all dimensions exist and match the field
+    create_field_dimensions!(ds, field, all_dims, dimension_name_generator)
+    defVar(ds, name, field_data, all_dims; kwargs...)
+end
+
+#####
+##### Dimension validation
+#####
+
+"""
+    create_field_dimensions!(ds, field::AbstractField, all_dims, dimension_name_generator)
+
+Creates all dimensions for the given `field` in the NetCDF dataset `ds`. If the dimensions
+already exist, they are validated to match the expected dimensions for the given `field`.
+
+Arguments:
+- `ds`: NetCDF dataset
+- `field`: AbstractField being written  
+- `all_dims`: Tuple of dimension names to create/validate
+- `dimension_name_generator`: Function to generate dimension names
+"""
+function create_field_dimensions!(ds, field::AbstractField, all_dims, dimension_name_generator)
+    dimension_attributes = default_dimension_attributes(field.grid, dimension_name_generator)
+    spatial_dims = all_dims[1:end-(("time" in all_dims) ? 1 : 0)]
+
+    spatial_dims_dict = Dict(dim_name => dim_data for (dim_name, dim_data) in zip(spatial_dims, nodes(field)))
+    create_spatial_dimensions!(ds, spatial_dims_dict, dimension_attributes; array_type=Array{eltype(field)})
+
+    # Create time dimension if needed
+    if "time" in all_dims && "time" ∉ keys(ds.dim)
+        create_time_dimension!(ds)
+    end
+    
+    return nothing
+end
+
 #####
 ##### Utils
 #####
@@ -67,73 +126,42 @@ function collect_dim(ξ, ℓ, T, N, H, inds, with_halos)
     end
 end
 
-#####
-##### Dimension name generators
-#####
-
-function suffixed_dim_name_generator(var_name, grid::AbstractGrid{FT, TX}, LX, LY, LZ, dim::Val{:x}; connector="_", location_letters) where {FT, TX}
-    if TX == Flat || isnothing(LX)
-        return ""
-    else
-        return "$(var_name)" * connector * location_letters
+function create_time_dimension!(dataset; attrib=nothing)
+    if "time" ∉ keys(dataset.dim)
+        # Create an unlimited dimension "time"
+        # Time should always be Float64 to be extra safe from rounding errors.
+        # See: https://github.com/CliMA/Oceananigans.jl/issues/3056
+        defDim(dataset, "time", Inf)
+        defVar(dataset, "time", Float64, ("time",), attrib=attrib)
     end
 end
 
-function suffixed_dim_name_generator(var_name, grid::AbstractGrid{FT, TX, TY}, LX, LY, LZ, dim::Val{:y}; connector="_", location_letters) where {FT, TX, TY}
-    if TY == Flat || isnothing(LY)
-        return ""
-    else
-        return "$(var_name)" * connector * location_letters
-    end
-end
 
-function suffixed_dim_name_generator(var_name, grid::AbstractGrid{FT, TX, TY, TZ}, LX, LY, LZ, dim::Val{:z}; connector="_", location_letters) where {FT, TX, TY, TZ}
-    if TZ == Flat || isnothing(LZ)
-        return ""
-    else
-        return "$(var_name)" * connector * location_letters
+"""
+    create_spatial_dimensions!(dataset, dims, attributes_dict; array_type=Array{Float32}, kwargs...)
+
+Create spatial dimensions in the NetCDF dataset and define corresponding variables to store
+their coordinate values. Each dimension variable has itself as its sole dimension (e.g., the
+`x` variable has dimension `x`). The dimensions are created if they don't exist, and validated
+against provided arrays if they do exist. An error is thrown if the dimension already exists
+but is different from the provided array.
+"""
+function create_spatial_dimensions!(dataset, dims, attributes_dict; array_type=Array{Float32}, kwargs...)
+    for (i, (dim_name, dim_array)) in enumerate(dims)
+        if dim_name ∉ keys(dataset.dim)
+            # Create missing dimension
+            defVar(dataset, dim_name, array_type(dim_array), (dim_name,), attrib=attributes_dict[dim_name]; kwargs...)
+        else
+            # Validate existing dimension
+            if dataset[dim_name] != dim_array
+                throw(ArgumentError("Dimension '$dim_name' already exists in dataset but is different from expected.\n" *
+                                    "  Actual:   $(dataset[dim_name]) (length=$(length(dataset[dim_name])))\n" *
+                                    "  Expected: $(dim_array) (length=$(length(dim_array)))"))
+            end
+        end
     end
 end
 
-suffixed_dim_name_generator(var_name, ::StaticVerticalDiscretization, LX, LY, LZ, dim::Val{:z}; connector="_", location_letters) = var_name * connector * location_letters
-
-loc2letter(::Face, full=true) = "f"
-loc2letter(::Center, full=true) = "c"
-loc2letter(::Nothing, full=true) = full ? "a" : ""
-
-minimal_location_string(::RectilinearGrid, LX, LY, LZ, ::Val{:x}) = loc2letter(LX, false)
-minimal_location_string(::RectilinearGrid, LX, LY, LZ, ::Val{:y}) = loc2letter(LY, false)
-
-minimal_location_string(::LatitudeLongitudeGrid, LX, LY, LZ, ::Val{:x}) = loc2letter(LX, false) * loc2letter(LY, false)
-minimal_location_string(::LatitudeLongitudeGrid, LX, LY, LZ, ::Val{:y}) = loc2letter(LX, false) * loc2letter(LY, false)
-
-minimal_location_string(grid::AbstractGrid,             LX, LY, LZ, dim::Val{:z}) = minimal_location_string(grid.z, LX, LY, LZ, dim)
-minimal_location_string(::StaticVerticalDiscretization, LX, LY, LZ, dim::Val{:z}) = loc2letter(LZ, false)
-minimal_location_string(grid,                           LX, LY, LZ, dim)          = loc2letter(LX, false) * loc2letter(LY, false) * loc2letter(LZ, false)
-
-minimal_dim_name(var_name, grid, LX, LY, LZ, dim) =
-    suffixed_dim_name_generator(var_name, grid, LX, LY, LZ, dim, connector="_", location_letters=minimal_location_string(grid, LX, LY, LZ, dim))
-
-minimal_dim_name(var_name, grid::ImmersedBoundaryGrid, args...) = minimal_dim_name(var_name, grid.underlying_grid, args...)
-
-
-
-trilocation_location_string(::RectilinearGrid, LX, LY, LZ, ::Val{:x}) = loc2letter(LX) * "aa"
-trilocation_location_string(::RectilinearGrid, LX, LY, LZ, ::Val{:y}) = "a" * loc2letter(LY) * "a"
-
-trilocation_location_string(::LatitudeLongitudeGrid, LX, LY, LZ, ::Val{:x}) = loc2letter(LX) * loc2letter(LY) * "a"
-trilocation_location_string(::LatitudeLongitudeGrid, LX, LY, LZ, ::Val{:y}) = loc2letter(LX) * loc2letter(LY) * "a"
-
-trilocation_location_string(grid::AbstractGrid,             LX, LY, LZ, dim::Val{:z}) = trilocation_location_string(grid.z, LX, LY, LZ, dim)
-trilocation_location_string(::StaticVerticalDiscretization, LX, LY, LZ, dim::Val{:z}) = "aa" * loc2letter(LZ)
-trilocation_location_string(grid,                           LX, LY, LZ, dim)          = loc2letter(LX) * loc2letter(LY) * loc2letter(LZ)
-
-trilocation_dim_name(var_name, grid, LX, LY, LZ, dim) =
-    suffixed_dim_name_generator(var_name, grid, LX, LY, LZ, dim, connector="_", location_letters=trilocation_location_string(grid, LX, LY, LZ, dim))
-
-trilocation_dim_name(var_name, grid::ImmersedBoundaryGrid, args...) = trilocation_dim_name(var_name, grid.underlying_grid, args...)
-
-
 #####
 ##### Gathering of grid dimensions
 #####
@@ -417,8 +445,8 @@ function field_dimensions(field::AbstractField, grid::LatitudeLongitudeGrid, dim
     LΛ, LΦ, LZ = location(field)
 
     λ_dim_name = dim_name_generator("λ", grid, LΛ(), nothing, nothing, Val(:x))
-    φ_dim_name = dim_name_generator("φ",  grid, nothing, LΦ(), nothing, Val(:y))
-    z_dim_name = dim_name_generator("z",         grid, nothing, nothing, LZ(), Val(:z))
+    φ_dim_name = dim_name_generator("φ", grid, nothing, LΦ(), nothing, Val(:y))
+    z_dim_name = dim_name_generator("z", grid, nothing, nothing, LZ(), Val(:z))
 
     λ_dim_name = isempty(λ_dim_name) ? tuple() : tuple(λ_dim_name)
     φ_dim_name = isempty(φ_dim_name) ? tuple() : tuple(φ_dim_name)
@@ -1146,18 +1174,8 @@ function initialize_nc_file(model,
             Dict("long_name" => "Time", "units" => "seconds since 2000-01-01 00:00:00") :
             Dict("long_name" => "Time", "units" => "seconds")
 
-        # Create an unlimited dimension "time"
-        # Time should always be Float64 to be extra safe from rounding errors.
-        # See: https://github.com/CliMA/Oceananigans.jl/issues/3056
-        defDim(dataset, "time", Inf)
-        defVar(dataset, "time", Float64, ("time",), attrib=time_attrib)
-
-        # Create spatial dimensions as variables whose dimensions are themselves.
-        # Each should already have a default attribute.
-        for (dim_name, dim_array) in dims
-            defVar(dataset, dim_name, array_type(dim_array), (dim_name,),
-                   deflatelevel=deflatelevel, attrib=output_attributes[dim_name])
-        end
+        create_time_dimension!(dataset, attrib=time_attrib)
+        create_spatial_dimensions!(dataset, dims, output_attributes; deflatelevel=1, array_type)
 
         time_independent_vars = Dict()
 

src/OutputWriters/netcdf_writer.jl

@@ -4,6 +4,66 @@
 ##### NetCDFWriter functionality is implemented in ext/OceananigansNCDatasetsExt
 #####
 
+using Oceananigans.Grids: topology, Flat, StaticVerticalDiscretization
+using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
+
+#####
+##### Dimension name generators
+#####
+
+function suffixed_dim_name_generator(var_name, grid::AbstractGrid{FT, TX}, LX, LY, LZ, dim::Val{:x}; connector="_", location_letters) where {FT, TX}
+    if TX == Flat || isnothing(LX)
+        return ""
+    else
+        return "$(var_name)" * connector * location_letters
+    end
+end
+
+function suffixed_dim_name_generator(var_name, grid::AbstractGrid{FT, TX, TY}, LX, LY, LZ, dim::Val{:y}; connector="_", location_letters) where {FT, TX, TY}
+    if TY == Flat || isnothing(LY)
+        return ""
+    else
+        return "$(var_name)" * connector * location_letters
+    end
+end
+
+function suffixed_dim_name_generator(var_name, grid::AbstractGrid{FT, TX, TY, TZ}, LX, LY, LZ, dim::Val{:z}; connector="_", location_letters) where {FT, TX, TY, TZ}
+    if TZ == Flat || isnothing(LZ)
+        return ""
+    else
+        return "$(var_name)" * connector * location_letters
+    end
+end
+
+suffixed_dim_name_generator(var_name, ::StaticVerticalDiscretization, LX, LY, LZ, dim::Val{:z}; connector="_", location_letters) = var_name * connector * location_letters
+
+loc2letter(::Face, full=true) = "f"
+loc2letter(::Center, full=true) = "c"
+loc2letter(::Nothing, full=true) = full ? "a" : ""
+
+minimal_location_string(::RectilinearGrid, LX, LY, LZ, ::Val{:x}) = loc2letter(LX, false)
+minimal_location_string(::RectilinearGrid, LX, LY, LZ, ::Val{:y}) = loc2letter(LY, false)
+
+minimal_dim_name(var_name, grid, LX, LY, LZ, dim) =
+minimal_dim_name(var_name, grid::ImmersedBoundaryGrid, args...) = minimal_dim_name(var_name, grid.underlying_grid, args...)
+
+trilocation_location_string(::RectilinearGrid, LX, LY, LZ, ::Val{:x}) = loc2letter(LX) * "aa"
+trilocation_location_string(::RectilinearGrid, LX, LY, LZ, ::Val{:y}) = "a" * loc2letter(LY) * "a"
+
+trilocation_location_string(::LatitudeLongitudeGrid, LX, LY, LZ, ::Val{:x}) = loc2letter(LX) * loc2letter(LY) * "a"
+trilocation_location_string(::LatitudeLongitudeGrid, LX, LY, LZ, ::Val{:y}) = loc2letter(LX) * loc2letter(LY) * "a"
+
+trilocation_location_string(grid::AbstractGrid,             LX, LY, LZ, dim::Val{:z}) = trilocation_location_string(grid.z, LX, LY, LZ, dim)
+trilocation_location_string(::StaticVerticalDiscretization, LX, LY, LZ, dim::Val{:z}) = "aa" * loc2letter(LZ)
+trilocation_location_string(grid,                           LX, LY, LZ, dim)          = loc2letter(LX) * loc2letter(LY) * loc2letter(LZ)
+
+trilocation_dim_name(var_name, grid, LX, LY, LZ, dim) =
+    suffixed_dim_name_generator(var_name, grid, LX, LY, LZ, dim, connector="_", location_letters=trilocation_location_string(grid, LX, LY, LZ, dim))
+
+trilocation_dim_name(var_name, grid::ImmersedBoundaryGrid, args...) = trilocation_dim_name(var_name, grid.underlying_grid, args...)
+
+
+
 mutable struct NetCDFWriter{G, D, O, T, A, FS, DN} <: AbstractOutputWriter
     grid :: G
     filepath :: String