From 03f8fca0a3cc4a1d2059b27804ee67dd8be739f0 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Sat, 23 Aug 2025 21:48:18 -0700
Subject: [PATCH 01/20] extend defVar to be able to write field directly to
 NetCDF

---
 ext/OceananigansNCDatasetsExt.jl | 12 ++++++++++++
 test/test_netcdf_writer.jl       | 17 +++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/ext/OceananigansNCDatasetsExt.jl b/ext/OceananigansNCDatasetsExt.jl
index a80c1467fd1..5045581cf45 100644
--- a/ext/OceananigansNCDatasetsExt.jl
+++ b/ext/OceananigansNCDatasetsExt.jl
@@ -1,6 +1,7 @@
 module OceananigansNCDatasetsExt
 
 using NCDatasets
+import NCDatasets: defVar
 
 using Dates: AbstractTime, UTC, now
 using Printf: @sprintf
@@ -44,6 +45,17 @@ const f = Face()
 const BoussinesqSeawaterBuoyancy = SeawaterBuoyancy{FT, <:BoussinesqEquationOfState, T, S} where {FT, T, S}
 const BuoyancyBoussinesqEOSModel = BuoyancyForce{<:BoussinesqSeawaterBuoyancy, g} where {g}
 
+function defVar(ds, name, field::AbstractField;
+                time_dependent=false,
+                dimension_name_generator = trilocation_dim_name,
+                kwargs...)
+    FT = Array{eltype(field)}(field)
+    dims = field_dimensions(field, dimension_name_generator)
+    all_dims = time_dependent ? (dims..., "time") : dims
+
+    defVar(ds, name, FT, all_dims; kwargs...)
+end
+
 #####
 ##### Utils
 #####
diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl
index bdbe04be7a3..ed6a17ef589 100644
--- a/test/test_netcdf_writer.jl
+++ b/test/test_netcdf_writer.jl
@@ -2726,6 +2726,21 @@ function test_netcdf_buoyancy_force(arch)
     return nothing
 end
 
+function test_netcdf_single_field_defvar()
+    grid = RectilinearGrid(size=(4, 4, 4), extent=(1, 1, 1))
+    c = CenterField(grid)
+    filepath = "test_single_field_defvar.nc"
+    isfile(filepath) && rm(filepath)
+
+    ds = NCDataset(filepath, "c")
+    defVar(ds, "c", c, time_dependent=false)
+
+    close(ds)
+    rm(filepath)
+
+    return nothing
+end
+
 for arch in archs
     @testset "NetCDF output writer [$(typeof(arch))]" begin
         @info "  Testing NetCDF output writer [$(typeof(arch))]..."
@@ -2775,5 +2790,7 @@ for arch in archs
         test_netcdf_free_surface_mixed_output(arch)
 
         test_netcdf_buoyancy_force(arch)
+
+        test_netcdf_single_field_defvar()
     end
 end

From 4c8a461fb2011c6ac73dec692940b556a5bc6116 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Mon, 25 Aug 2025 20:14:23 -0700
Subject: [PATCH 02/20] added function to create dimensions

---
 ext/OceananigansNCDatasetsExt.jl | 76 ++++++++++++++++++++++++++++----
 test/test_netcdf_writer.jl       | 22 ++++++++-
 2 files changed, 88 insertions(+), 10 deletions(-)

diff --git a/ext/OceananigansNCDatasetsExt.jl b/ext/OceananigansNCDatasetsExt.jl
index 5045581cf45..b2dde05a087 100644
--- a/ext/OceananigansNCDatasetsExt.jl
+++ b/ext/OceananigansNCDatasetsExt.jl
@@ -7,11 +7,11 @@ using Dates: AbstractTime, UTC, now
 using Printf: @sprintf
 
 using Oceananigans: initialize!, prettytime, pretty_filesize, AbstractModel
-using Oceananigans.AbstractOperations: KernelFunctionOperation
+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.Grids: Center, Face, Flat, AbstractGrid, RectilinearGrid, LatitudeLongitudeGrid, StaticVerticalDiscretization,
+using Oceananigans.Fields: Reduction, reduced_dimensions, reduced_location, location
+using Oceananigans.Grids: Center, Face, Flat, nodes, AbstractGrid, RectilinearGrid, LatitudeLongitudeGrid, StaticVerticalDiscretization,
                           topology, halo_size, xspacings, yspacings, zspacings, λspacings, φspacings,
                           parent_index_range, ξnodes, ηnodes, rnodes, validate_index, peripheral_node
 using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBottom, GFBIBG, GridFittedBoundary, PartialCellBottom, PCBIBG
@@ -45,6 +45,9 @@ const f = Face()
 const BoussinesqSeawaterBuoyancy = SeawaterBuoyancy{FT, <:BoussinesqEquationOfState, T, S} where {FT, T, S}
 const BuoyancyBoussinesqEOSModel = BuoyancyForce{<:BoussinesqSeawaterBuoyancy, g} where {g}
 
+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,
@@ -53,9 +56,56 @@ function defVar(ds, name, field::AbstractField;
     dims = field_dimensions(field, dimension_name_generator)
     all_dims = time_dependent ? (dims..., "time") : dims
 
+    # Validate that all dimensions exist and match the field
+    validate_field_dimensions(ds, field, all_dims, name, dimension_name_generator)
+
     defVar(ds, name, FT, all_dims; kwargs...)
 end
 
+#####
+##### Dimension validation
+#####
+
+"""
+    validate_field_dimensions(ds, field::AbstractField, all_dims, name)
+
+Validates that all dimensions in `all_dims` exist in the NetCDF dataset `ds` and match
+the expected dimensions for the given `field`. Creates missing dimensions if needed.
+
+Arguments:
+- `ds`: NetCDF dataset
+- `field`: AbstractField being written  
+- `all_dims`: Tuple of dimension names to validate
+- `name`: Variable name (for error messages)
+"""
+function validate_field_dimensions(ds, field::AbstractField, all_dims, name, 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)]
+
+    for (i, dim_name) in enumerate(spatial_dims)
+        expected_size = size(field)[i]
+        if dim_name ∉ keys(ds.dim)
+            # Create missing dimension
+            dim_data = nodes(field)[i]
+            defVar(ds, dim_name, dim_data, (dim_name,), attrib=dimension_attributes[dim_name])
+        else
+            # Validate existing dimension size
+            dataset_dim_size = length(ds[dim_name])
+            if dataset_dim_size != expected_size
+                throw(ArgumentError("Dimension '$dim_name' for field '$name' has size $dataset_dim_size " *
+                                    "in the dataset, but field requires size $expected_size"))
+            end
+        end
+    end
+
+    # Create time dimension if needed
+    if "time" in all_dims && "time" ∉ keys(ds.dim)
+        create_time_dimension!(ds)
+    end
+    
+    return nothing
+end
+
 #####
 ##### Utils
 #####
@@ -75,6 +125,20 @@ function collect_dim(ξ, ℓ, T, N, H, inds, with_halos)
     end
 end
 
+function create_time_dimension!(dataset)
+    # 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",))
+end
+
+# function create_spatial_dimensions!(dataset, dims, attributes_dict; kwargs...)
+#     for (dim_name, dim_data) in dims
+#         defVar(dataset, dim_name, dim_data, (dim_name,), attrib=attributes_dict[dim_name]; kwargs...)
+#     end
+# end
+
 #####
 ##### Dimension name generators
 #####
@@ -1146,11 +1210,7 @@ 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_time_dimension!(dataset)
 
         # Create spatial dimensions as variables whose dimensions are themselves.
         # Each should already have a default attribute.
diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl
index ed6a17ef589..838c4c9d8bc 100644
--- a/test/test_netcdf_writer.jl
+++ b/test/test_netcdf_writer.jl
@@ -2727,15 +2727,33 @@ function test_netcdf_buoyancy_force(arch)
 end
 
 function test_netcdf_single_field_defvar()
-    grid = RectilinearGrid(size=(4, 4, 4), extent=(1, 1, 1))
+    N = 4
+    grid = RectilinearGrid(size=(N, N, N), extent=(1, 1, 1))
     c = CenterField(grid)
+    set!(c, (x, y, z) -> x + y + z)
+
     filepath = "test_single_field_defvar.nc"
     isfile(filepath) && rm(filepath)
 
     ds = NCDataset(filepath, "c")
-    defVar(ds, "c", c, time_dependent=false)
+    defVar(ds, "c", c, attrib=Dict("long_name" => "Center field", "units" => "kg/m³"))
+
+    # Write an abstract operation
+    c² = c^2
+    defVar(ds, "c²", c², attrib=Dict("long_name" => "Center field squared"))
+
+    # Write a reduction
+    c̄ = Average(c, dims=3)
+    defVar(ds, "c̄", c̄, attrib=Dict("long_name" => "Average center field", "units" => "kg/m²"))
 
     close(ds)
+    ds = NCDataset(filepath, "r")
+    @test "c" ∈ keys(ds)
+    @test all(ds["c"] .== interior(c))
+    @test ds["c"].attrib["long_name"] == "Center field"
+    @test ds["c"].attrib["units"] == "kg/m³"
+    close(ds)
+
     rm(filepath)
 
     return nothing

From ab770d13db744fa2cd23baf066da57c2b56be70a Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Wed, 27 Aug 2025 06:50:05 -0700
Subject: [PATCH 03/20] add some tests

---
 ext/OceananigansNCDatasetsExt.jl | 39 +++++++-------
 test/test_netcdf_writer.jl       | 88 ++++++++++++++++++++++++++++++++
 2 files changed, 107 insertions(+), 20 deletions(-)

diff --git a/ext/OceananigansNCDatasetsExt.jl b/ext/OceananigansNCDatasetsExt.jl
index b2dde05a087..1c21c6d4c47 100644
--- a/ext/OceananigansNCDatasetsExt.jl
+++ b/ext/OceananigansNCDatasetsExt.jl
@@ -82,21 +82,7 @@ function validate_field_dimensions(ds, field::AbstractField, all_dims, name, dim
     dimension_attributes = default_dimension_attributes(field.grid, dimension_name_generator)
     spatial_dims = all_dims[1:end-(("time" in all_dims) ? 1 : 0)]
 
-    for (i, dim_name) in enumerate(spatial_dims)
-        expected_size = size(field)[i]
-        if dim_name ∉ keys(ds.dim)
-            # Create missing dimension
-            dim_data = nodes(field)[i]
-            defVar(ds, dim_name, dim_data, (dim_name,), attrib=dimension_attributes[dim_name])
-        else
-            # Validate existing dimension size
-            dataset_dim_size = length(ds[dim_name])
-            if dataset_dim_size != expected_size
-                throw(ArgumentError("Dimension '$dim_name' for field '$name' has size $dataset_dim_size " *
-                                    "in the dataset, but field requires size $expected_size"))
-            end
-        end
-    end
+    create_spatial_dimensions!(ds, spatial_dims, dimension_attributes)
 
     # Create time dimension if needed
     if "time" in all_dims && "time" ∉ keys(ds.dim)
@@ -133,11 +119,23 @@ function create_time_dimension!(dataset)
     defVar(dataset, "time", Float64, ("time",))
 end
 
-# function create_spatial_dimensions!(dataset, dims, attributes_dict; kwargs...)
-#     for (dim_name, dim_data) in dims
-#         defVar(dataset, dim_name, dim_data, (dim_name,), attrib=attributes_dict[dim_name]; kwargs...)
-#     end
-# end
+function create_spatial_dimensions!(dataset, dims, attributes_dict; kwargs...)
+    for (i, dim_name) in enumerate(dims)
+        expected_size = size(field)[i]
+        if dim_name ∉ keys(ds.dim)
+            # Create missing dimension
+            dim_data = nodes(field)[i]
+            defVar(dataset, dim_name, dim_data, (dim_name,), attrib=attributes_dict[dim_name]; kwargs...)
+        else
+            # Validate existing dimension size
+            dataset_dim_size = length(ds[dim_name])
+            if dataset_dim_size != expected_size
+                throw(ArgumentError("Dimension '$dim_name' for field '$name' has size $dataset_dim_size " *
+                                    "in the dataset, but field requires size $expected_size"))
+            end
+        end
+    end
+end
 
 #####
 ##### Dimension name generators
@@ -1218,6 +1216,7 @@ function initialize_nc_file(model,
             defVar(dataset, dim_name, array_type(dim_array), (dim_name,),
                    deflatelevel=deflatelevel, attrib=output_attributes[dim_name])
         end
+        # create_spatial_dimensions!(dataset, dims; deflatelevel=1, attrib=output_attributes)
 
         time_independent_vars = Dict()
 
diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl
index 838c4c9d8bc..54d7c58928c 100644
--- a/test/test_netcdf_writer.jl
+++ b/test/test_netcdf_writer.jl
@@ -2759,6 +2759,93 @@ function test_netcdf_single_field_defvar()
     return nothing
 end
 
+function test_netcdf_field_dimension_validation()
+    grid = RectilinearGrid(size=(4, 4, 4), extent=(1, 1, 1))
+    c = CenterField(grid)
+
+    # Test 1: Successful validation with proper dimensions
+    filepath_success = "test_dimension_validation_success.nc"
+    isfile(filepath_success) && rm(filepath_success)
+
+    # Create NetCDF file with proper dimensions
+    ds = NCDataset(filepath_success, "c")
+
+    # Define dimensions first
+    defDim(ds, "x_caa", 4)
+    defDim(ds, "y_aca", 4)
+    defDim(ds, "z_aac", 4)
+    defVar(ds, "x_caa", Float64, ("x_caa",))
+    defVar(ds, "y_aca", Float64, ("y_aca",))
+    defVar(ds, "z_aac", Float64, ("z_aac",))
+
+    # This should succeed
+    defVar(ds, "c", c, time_dependent=false)
+
+    close(ds)
+    rm(filepath_success)
+
+    # Test 2: Missing dimension should throw error
+    filepath_missing = "test_dimension_validation_missing.nc"
+    isfile(filepath_missing) && rm(filepath_missing)
+
+    ds_missing = NCDataset(filepath_missing, "c")
+
+    # Define only some dimensions
+    defDim(ds_missing, "x_caa", 4)
+    defDim(ds_missing, "y_aca", 4)
+    defVar(ds_missing, "x_caa", Float64, ("x_caa",))
+    defVar(ds_missing, "y_aca", Float64, ("y_aca",))
+    # Note: z_aac dimension is missing
+
+    @test_throws ArgumentError defVar(ds_missing, "c", c, time_dependent=false)
+
+    close(ds_missing)
+    rm(filepath_missing)
+
+    # Test 3: Wrong dimension size should throw error
+    filepath_wrong_size = "test_dimension_validation_wrong_size.nc"
+    isfile(filepath_wrong_size) && rm(filepath_wrong_size)
+
+    ds_wrong = NCDataset(filepath_wrong_size, "c")
+
+    # Define dimensions with wrong sizes
+    defDim(ds_wrong, "x_caa", 4)
+    defDim(ds_wrong, "y_aca", 4)
+    defDim(ds_wrong, "z_aac", 8)  # Wrong size! Should be 4
+    defVar(ds_wrong, "x_caa", Float64, ("x_caa",))
+    defVar(ds_wrong, "y_aca", Float64, ("y_aca",))
+    defVar(ds_wrong, "z_aac", Float64, ("z_aac",))
+
+    @test_throws ArgumentError defVar(ds_wrong, "c", c, time_dependent=false)
+
+    close(ds_wrong)
+    rm(filepath_wrong_size)
+
+    # Test 4: Test with time-dependent field
+    filepath_time = "test_dimension_validation_time.nc"
+    isfile(filepath_time) && rm(filepath_time)
+
+    ds_time = NCDataset(filepath_time, "c")
+
+    # Define dimensions including time
+    defDim(ds_time, "x_caa", 4)
+    defDim(ds_time, "y_aca", 4)
+    defDim(ds_time, "z_aac", 4)
+    defDim(ds_time, "time", Inf)
+    defVar(ds_time, "x_caa", Float64, ("x_caa",))
+    defVar(ds_time, "y_aca", Float64, ("y_aca",))
+    defVar(ds_time, "z_aac", Float64, ("z_aac",))
+    defVar(ds_time, "time", Float64, ("time",))
+
+    # This should succeed
+    defVar(ds_time, "c", c, time_dependent=true)
+
+    close(ds_time)
+    rm(filepath_time)
+
+    return nothing
+end
+
 for arch in archs
     @testset "NetCDF output writer [$(typeof(arch))]" begin
         @info "  Testing NetCDF output writer [$(typeof(arch))]..."
@@ -2810,5 +2897,6 @@ for arch in archs
         test_netcdf_buoyancy_force(arch)
 
         test_netcdf_single_field_defvar()
+        # test_netcdf_field_dimension_validation()
     end
 end

From 776cd7e4ce62c8605164585f7cb4fc30e0bc2396 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Wed, 27 Aug 2025 08:45:36 -0700
Subject: [PATCH 04/20] fix test

---
 ext/OceananigansNCDatasetsExt.jl | 28 ++++++++++++++--------------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/ext/OceananigansNCDatasetsExt.jl b/ext/OceananigansNCDatasetsExt.jl
index 1c21c6d4c47..476328bd4d2 100644
--- a/ext/OceananigansNCDatasetsExt.jl
+++ b/ext/OceananigansNCDatasetsExt.jl
@@ -57,7 +57,7 @@ function defVar(ds, name, field::AbstractField;
     all_dims = time_dependent ? (dims..., "time") : dims
 
     # Validate that all dimensions exist and match the field
-    validate_field_dimensions(ds, field, all_dims, name, dimension_name_generator)
+    validate_or_create_field_dimensions!(ds, field, all_dims, name, dimension_name_generator)
 
     defVar(ds, name, FT, all_dims; kwargs...)
 end
@@ -67,7 +67,7 @@ end
 #####
 
 """
-    validate_field_dimensions(ds, field::AbstractField, all_dims, name)
+    validate_or_create_field_dimensions!(ds, field::AbstractField, all_dims, name)
 
 Validates that all dimensions in `all_dims` exist in the NetCDF dataset `ds` and match
 the expected dimensions for the given `field`. Creates missing dimensions if needed.
@@ -78,11 +78,12 @@ Arguments:
 - `all_dims`: Tuple of dimension names to validate
 - `name`: Variable name (for error messages)
 """
-function validate_field_dimensions(ds, field::AbstractField, all_dims, name, dimension_name_generator)
+function validate_or_create_field_dimensions!(ds, field::AbstractField, all_dims, name, 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)]
 
-    create_spatial_dimensions!(ds, spatial_dims, dimension_attributes)
+    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)
@@ -119,19 +120,18 @@ function create_time_dimension!(dataset)
     defVar(dataset, "time", Float64, ("time",))
 end
 
-function create_spatial_dimensions!(dataset, dims, attributes_dict; kwargs...)
-    for (i, dim_name) in enumerate(dims)
-        expected_size = size(field)[i]
-        if dim_name ∉ keys(ds.dim)
+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
-            dim_data = nodes(field)[i]
-            defVar(dataset, dim_name, dim_data, (dim_name,), attrib=attributes_dict[dim_name]; kwargs...)
+            defVar(dataset, dim_name, array_type(dim_array), (dim_name,), attrib=attributes_dict[dim_name], kwargs...)
         else
             # Validate existing dimension size
-            dataset_dim_size = length(ds[dim_name])
-            if dataset_dim_size != expected_size
-                throw(ArgumentError("Dimension '$dim_name' for field '$name' has size $dataset_dim_size " *
-                                    "in the dataset, but field requires size $expected_size"))
+            dataset_dim_length = length(dataset[dim_name])
+            expected_length = length(dim_array)
+            if dataset_dim_length != expected_length
+                throw(ArgumentError("Dimension '$dim_name' has size $dataset_dim_length " *
+                                    "in the dataset, but field requires size $expected_length"))
             end
         end
     end

From 21302b39dfe1704a4a9f61adc24d3306035a6ceb Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Wed, 27 Aug 2025 09:02:28 -0700
Subject: [PATCH 05/20] better function signature

---
 ext/OceananigansNCDatasetsExt.jl | 35 ++++++++++++++++----------------
 1 file changed, 17 insertions(+), 18 deletions(-)

diff --git a/ext/OceananigansNCDatasetsExt.jl b/ext/OceananigansNCDatasetsExt.jl
index 476328bd4d2..9cd057bef00 100644
--- a/ext/OceananigansNCDatasetsExt.jl
+++ b/ext/OceananigansNCDatasetsExt.jl
@@ -57,7 +57,7 @@ function defVar(ds, name, field::AbstractField;
     all_dims = time_dependent ? (dims..., "time") : dims
 
     # Validate that all dimensions exist and match the field
-    validate_or_create_field_dimensions!(ds, field, all_dims, name, dimension_name_generator)
+    create_field_dimensions!(ds, field, all_dims, dimension_name_generator)
 
     defVar(ds, name, FT, all_dims; kwargs...)
 end
@@ -67,18 +67,18 @@ end
 #####
 
 """
-    validate_or_create_field_dimensions!(ds, field::AbstractField, all_dims, name)
+    create_field_dimensions!(ds, field::AbstractField, all_dims, dimension_name_generator)
 
-Validates that all dimensions in `all_dims` exist in the NetCDF dataset `ds` and match
-the expected dimensions for the given `field`. Creates missing dimensions if needed.
+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 validate
-- `name`: Variable name (for error messages)
+- `all_dims`: Tuple of dimension names to create/validate
+- `dimension_name_generator`: Function to generate dimension names
 """
-function validate_or_create_field_dimensions!(ds, field::AbstractField, all_dims, name, dimension_name_generator)
+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)]
 
@@ -113,11 +113,13 @@ function collect_dim(ξ, ℓ, T, N, H, inds, with_halos)
 end
 
 function create_time_dimension!(dataset)
-    # 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",))
+    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",))
+    end
 end
 
 function create_spatial_dimensions!(dataset, dims, attributes_dict; array_type=Array{Float32}, kwargs...)
@@ -126,12 +128,9 @@ function create_spatial_dimensions!(dataset, dims, attributes_dict; array_type=A
             # Create missing dimension
             defVar(dataset, dim_name, array_type(dim_array), (dim_name,), attrib=attributes_dict[dim_name], kwargs...)
         else
-            # Validate existing dimension size
-            dataset_dim_length = length(dataset[dim_name])
-            expected_length = length(dim_array)
-            if dataset_dim_length != expected_length
-                throw(ArgumentError("Dimension '$dim_name' has size $dataset_dim_length " *
-                                    "in the dataset, but field requires size $expected_length"))
+            # Validate existing dimension
+            if dataset[dim_name] != dim_array
+                throw(ArgumentError("Dimension '$dim_name' already exists in dataset but is different from expected."))
             end
         end
     end

From 6efd39cb65d847d21ee008bff4f62e8faff294d6 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Fri, 19 Sep 2025 13:31:20 -0300
Subject: [PATCH 06/20] import node

---
 ext/OceananigansNCDatasetsExt.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ext/OceananigansNCDatasetsExt.jl b/ext/OceananigansNCDatasetsExt.jl
index 6516e08e9a6..afb642cbed1 100644
--- a/ext/OceananigansNCDatasetsExt.jl
+++ b/ext/OceananigansNCDatasetsExt.jl
@@ -16,7 +16,7 @@ using Oceananigans.Fields: Reduction, reduced_dimensions, reduced_location, loca
 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

From fe7e746ccded1163efe2901c8febb57668dff948 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Fri, 19 Sep 2025 13:36:41 -0300
Subject: [PATCH 07/20] add some more tests

---
 test/test_netcdf_writer.jl | 27 +++++++++++++++++++++------
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl
index 54d7c58928c..02dc162417a 100644
--- a/test/test_netcdf_writer.jl
+++ b/test/test_netcdf_writer.jl
@@ -2736,24 +2736,39 @@ function test_netcdf_single_field_defvar()
     isfile(filepath) && rm(filepath)
 
     ds = NCDataset(filepath, "c")
-    defVar(ds, "c", c, attrib=Dict("long_name" => "Center field", "units" => "kg/m³"))
+    c_long_name = "Center field"
+    c_units = "kg/m³"
+    defVar(ds, "c", c, attrib=Dict("long_name" => c_long_name, "units" => c_units))
 
     # Write an abstract operation
     c² = c^2
-    defVar(ds, "c²", c², attrib=Dict("long_name" => "Center field squared"))
+    c²_long_name = "Center field squared"
+    defVar(ds, "c²", c², attrib=Dict("long_name" => c²_long_name))
 
     # Write a reduction
     c̄ = Average(c, dims=3)
-    defVar(ds, "c̄", c̄, attrib=Dict("long_name" => "Average center field", "units" => "kg/m²"))
+    c̄_long_name = "Average center field"
+    c̄_units = "kg/m³"
+    defVar(ds, "c̄", c̄, attrib=Dict("long_name" => c̄_long_name, "units" => c̄_units))
 
     close(ds)
     ds = NCDataset(filepath, "r")
+
     @test "c" ∈ keys(ds)
     @test all(ds["c"] .== interior(c))
-    @test ds["c"].attrib["long_name"] == "Center field"
-    @test ds["c"].attrib["units"] == "kg/m³"
-    close(ds)
+    @test ds["c"].attrib["long_name"] == c_long_name
+    @test ds["c"].attrib["units"] == c_units
+
+    @test "c²" ∈ keys(ds)
+    @test all(ds["c²"] .== interior(c²))
+    @test ds["c²"].attrib["long_name"] == c²_long_name
 
+    @test "c̄" ∈ keys(ds)
+    @test all(ds["c̄"] .== interior(Field(c̄)))
+    @test ds["c̄"].attrib["long_name"] == c̄_long_name
+    @test ds["c̄"].attrib["units"] == c̄_units
+
+    close(ds)
     rm(filepath)
 
     return nothing

From 59431281b89db677f1ad122a71457aff478c1150 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Fri, 19 Sep 2025 18:18:05 -0300
Subject: [PATCH 08/20] fix test of dimension validation

---
 test/test_netcdf_writer.jl | 79 +++++++++-----------------------------
 1 file changed, 18 insertions(+), 61 deletions(-)

diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl
index 02dc162417a..f24df2f49d7 100644
--- a/test/test_netcdf_writer.jl
+++ b/test/test_netcdf_writer.jl
@@ -2779,85 +2779,42 @@ function test_netcdf_field_dimension_validation()
     c = CenterField(grid)
 
     # Test 1: Successful validation with proper dimensions
-    filepath_success = "test_dimension_validation_success.nc"
-    isfile(filepath_success) && rm(filepath_success)
+    filepath = "test_dimension_validation_success.nc"
+    isfile(filepath) && rm(filepath)
 
     # Create NetCDF file with proper dimensions
-    ds = NCDataset(filepath_success, "c")
+    ds = NCDataset(filepath, "c")
 
     # Define dimensions first
-    defDim(ds, "x_caa", 4)
-    defDim(ds, "y_aca", 4)
-    defDim(ds, "z_aac", 4)
-    defVar(ds, "x_caa", Float64, ("x_caa",))
-    defVar(ds, "y_aca", Float64, ("y_aca",))
-    defVar(ds, "z_aac", Float64, ("z_aac",))
-
-    # This should succeed
+    defVar(ds, "x_caa", xnodes(grid, Center()), ("x_caa",))
+    # We "forget" to define the y dimension. Should work regardless
+    defVar(ds, "z_aac", znodes(grid, Center()), ("z_aac",))
+
+    # Write variable to disk after dimensions are created and match the
     defVar(ds, "c", c, time_dependent=false)
+    @test ds["x_caa"][:] == xnodes(grid, Center())
+    @test ds["y_aca"][:] == ynodes(grid, Center())
+    @test ds["z_aac"][:] == znodes(grid, Center())
 
     close(ds)
-    rm(filepath_success)
-
-    # Test 2: Missing dimension should throw error
-    filepath_missing = "test_dimension_validation_missing.nc"
-    isfile(filepath_missing) && rm(filepath_missing)
-
-    ds_missing = NCDataset(filepath_missing, "c")
-
-    # Define only some dimensions
-    defDim(ds_missing, "x_caa", 4)
-    defDim(ds_missing, "y_aca", 4)
-    defVar(ds_missing, "x_caa", Float64, ("x_caa",))
-    defVar(ds_missing, "y_aca", Float64, ("y_aca",))
-    # Note: z_aac dimension is missing
-
-    @test_throws ArgumentError defVar(ds_missing, "c", c, time_dependent=false)
-
-    close(ds_missing)
-    rm(filepath_missing)
+    rm(filepath)
 
-    # Test 3: Wrong dimension size should throw error
+    # Test that wrong dimension should throw error
     filepath_wrong_size = "test_dimension_validation_wrong_size.nc"
     isfile(filepath_wrong_size) && rm(filepath_wrong_size)
 
     ds_wrong = NCDataset(filepath_wrong_size, "c")
 
     # Define dimensions with wrong sizes
-    defDim(ds_wrong, "x_caa", 4)
-    defDim(ds_wrong, "y_aca", 4)
-    defDim(ds_wrong, "z_aac", 8)  # Wrong size! Should be 4
-    defVar(ds_wrong, "x_caa", Float64, ("x_caa",))
-    defVar(ds_wrong, "y_aca", Float64, ("y_aca",))
-    defVar(ds_wrong, "z_aac", Float64, ("z_aac",))
+    defVar(ds_wrong, "x_caa", xnodes(grid, Center()), ("x_caa",))
+    defVar(ds_wrong, "y_aca", ynodes(grid, Center()), ("y_aca",))
+    defVar(ds_wrong, "z_aac", znodes(grid, Face()), ("z_aac",)) # ⚠ wrongly write faces instead of center
 
-    @test_throws ArgumentError defVar(ds_wrong, "c", c, time_dependent=false)
+    @test_throws ArgumentError defVar(ds_wrong, "c", c)
 
     close(ds_wrong)
     rm(filepath_wrong_size)
 
-    # Test 4: Test with time-dependent field
-    filepath_time = "test_dimension_validation_time.nc"
-    isfile(filepath_time) && rm(filepath_time)
-
-    ds_time = NCDataset(filepath_time, "c")
-
-    # Define dimensions including time
-    defDim(ds_time, "x_caa", 4)
-    defDim(ds_time, "y_aca", 4)
-    defDim(ds_time, "z_aac", 4)
-    defDim(ds_time, "time", Inf)
-    defVar(ds_time, "x_caa", Float64, ("x_caa",))
-    defVar(ds_time, "y_aca", Float64, ("y_aca",))
-    defVar(ds_time, "z_aac", Float64, ("z_aac",))
-    defVar(ds_time, "time", Float64, ("time",))
-
-    # This should succeed
-    defVar(ds_time, "c", c, time_dependent=true)
-
-    close(ds_time)
-    rm(filepath_time)
-
     return nothing
 end
 
@@ -2912,6 +2869,6 @@ for arch in archs
         test_netcdf_buoyancy_force(arch)
 
         test_netcdf_single_field_defvar()
-        # test_netcdf_field_dimension_validation()
+        test_netcdf_field_dimension_validation()
     end
 end

From 7ec7b2986345adb691e75c10423f89c368201cd0 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Fri, 19 Sep 2025 22:32:20 -0300
Subject: [PATCH 09/20] add attributes to time dimension

---
 ext/OceananigansNCDatasetsExt.jl | 6 +++---
 test/test_netcdf_writer.jl       | 1 -
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/ext/OceananigansNCDatasetsExt.jl b/ext/OceananigansNCDatasetsExt.jl
index bcd6735448e..932937d5733 100644
--- a/ext/OceananigansNCDatasetsExt.jl
+++ b/ext/OceananigansNCDatasetsExt.jl
@@ -116,13 +116,13 @@ function collect_dim(ξ, ℓ, T, N, H, inds, with_halos)
     end
 end
 
-function create_time_dimension!(dataset)
+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",))
+        defVar(dataset, "time", Float64, ("time",), attrib=attrib)
     end
 end
 
@@ -1219,7 +1219,7 @@ 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_time_dimension!(dataset)
+        create_time_dimension!(dataset, attrib=time_attrib)
 
         # Create spatial dimensions as variables whose dimensions are themselves.
         # Each should already have a default attribute.
diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl
index f24df2f49d7..e03cc385158 100644
--- a/test/test_netcdf_writer.jl
+++ b/test/test_netcdf_writer.jl
@@ -33,7 +33,6 @@ function test_datetime_netcdf_output(arch)
                                                       filename = filepath,
                                                       schedule = IterationInterval(1),
                                                       include_grid_metrics = false)
-
     run!(simulation)
 
     ds = NCDataset(filepath)

From 0dfd9e7bc2124c904b1e74c8169d6e514497857a Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Sat, 20 Sep 2025 07:28:27 -0300
Subject: [PATCH 10/20] use function to create_time_dimensions

---
 ext/OceananigansNCDatasetsExt.jl | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/ext/OceananigansNCDatasetsExt.jl b/ext/OceananigansNCDatasetsExt.jl
index 932937d5733..e5f0516687b 100644
--- a/ext/OceananigansNCDatasetsExt.jl
+++ b/ext/OceananigansNCDatasetsExt.jl
@@ -126,11 +126,21 @@ function create_time_dimension!(dataset; attrib=nothing)
     end
 end
 
+
+"""
+    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...)
+            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
@@ -1220,14 +1230,7 @@ function initialize_nc_file(model,
             Dict("long_name" => "Time", "units" => "seconds")
 
         create_time_dimension!(dataset, 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_spatial_dimensions!(dataset, dims; deflatelevel=1, attrib=output_attributes)
+        create_spatial_dimensions!(dataset, dims, output_attributes; deflatelevel=1, array_type)
 
         time_independent_vars = Dict()
 

From 6230e5137d070deeba51ffea26afa3e1d37c9bb0 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Sat, 20 Sep 2025 11:03:55 -0300
Subject: [PATCH 11/20] test multiple grids in the same netCDF file

---
 ext/OceananigansNCDatasetsExt.jl   | 64 ++++++++------------------
 src/OutputWriters/netcdf_writer.jl | 74 ++++++++++++++++++++++++++++++
 test/test_netcdf_writer.jl         | 63 +++++++++++++++++++++++++
 3 files changed, 157 insertions(+), 44 deletions(-)

diff --git a/ext/OceananigansNCDatasetsExt.jl b/ext/OceananigansNCDatasetsExt.jl
index e5f0516687b..8f616b38ce9 100644
--- a/ext/OceananigansNCDatasetsExt.jl
+++ b/ext/OceananigansNCDatasetsExt.jl
@@ -42,7 +42,7 @@ 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, loc2letter, trilocation_location_string, suffixed_dim_name_generator
 
 const c = Center()
 const f = Face()
@@ -151,39 +151,32 @@ function create_spatial_dimensions!(dataset, dims, attributes_dict; array_type=A
 end
 
 #####
-##### Dimension name generators
+##### Specialized dimension name generators for specific grid types
 #####
 
-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
+# Override loc2letter for specific types
+loc2letter(::Face, full=true) = "f"
+loc2letter(::Center, full=true) = "c"
+loc2letter(::Nothing, full=true) = full ? "a" : ""
 
-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
+# Specialized suffixed_dim_name_generator for StaticVerticalDiscretization
+suffixed_dim_name_generator(var_name, ::StaticVerticalDiscretization, LX, LY, LZ, dim::Val{:z}; connector="_", location_letters) = var_name * connector * location_letters
 
-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
+# Specialized trilocation_location_string for specific grid types
+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"
 
-suffixed_dim_name_generator(var_name, ::StaticVerticalDiscretization, LX, LY, LZ, dim::Val{:z}; connector="_", location_letters) = var_name * connector * location_letters
+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"
 
-loc2letter(::Face, full=true) = "f"
-loc2letter(::Center, full=true) = "c"
-loc2letter(::Nothing, full=true) = full ? "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)
+
+# Specialized trilocation_dim_name for ImmersedBoundaryGrid
+trilocation_dim_name(var_name, grid::ImmersedBoundaryGrid, args...) = trilocation_dim_name(var_name, grid.underlying_grid, args...)
 
+# Minimal location string functions (for other dimension name generators)
 minimal_location_string(::RectilinearGrid, LX, LY, LZ, ::Val{:x}) = loc2letter(LX, false)
 minimal_location_string(::RectilinearGrid, LX, LY, LZ, ::Val{:y}) = loc2letter(LY, false)
 
@@ -200,23 +193,6 @@ 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...)
-
-
 #####
 ##### Gathering of grid dimensions
 #####
diff --git a/src/OutputWriters/netcdf_writer.jl b/src/OutputWriters/netcdf_writer.jl
index fb323e8335a..4c55a156d3e 100644
--- a/src/OutputWriters/netcdf_writer.jl
+++ b/src/OutputWriters/netcdf_writer.jl
@@ -4,6 +4,80 @@
 ##### NetCDFWriter functionality is implemented in ext/OceananigansNCDatasetsExt
 #####
 
+using Oceananigans.Grids: topology, Flat
+
+#####
+##### Dimension name generators
+#####
+
+"""
+    suffixed_dim_name_generator(var_name, grid, LX, LY, LZ, dim; connector="_", location_letters)
+
+Generate dimension names by suffixing the variable name with location letters.
+"""
+function suffixed_dim_name_generator(var_name, grid, LX, LY, LZ, dim; connector="_", location_letters)
+    # Check if topology is Flat for the given dimension
+    topo = topology(grid)
+
+    if dim == Val(:x)
+        TX = topo[1]
+        if TX == Flat || isnothing(LX)
+            return ""
+        else
+            return "$(var_name)" * connector * location_letters
+        end
+    elseif dim == Val(:y)
+        TY = topo[2]
+        if TY == Flat || isnothing(LY)
+            return ""
+        else
+            return "$(var_name)" * connector * location_letters
+        end
+    elseif dim == Val(:z)
+        TZ = topo[3]
+        if TZ == Flat || isnothing(LZ)
+            return ""
+        else
+            return "$(var_name)" * connector * location_letters
+        end
+    end
+
+    return "$(var_name)" * connector * location_letters
+end
+
+# Special case for StaticVerticalDiscretization (will be overridden in extension)
+function suffixed_dim_name_generator(var_name, coordinate, LX, LY, LZ, dim::Val{:z}; connector="_", location_letters)
+    return var_name * connector * location_letters
+end
+
+"""
+    loc2letter(location, full=true)
+
+Convert location types to letter representations.
+"""
+loc2letter(::Any, full=true) = full ? "a" : ""
+
+"""
+    trilocation_location_string(grid, LX, LY, LZ, dim)
+
+Generate location strings for trilocation naming convention.
+"""
+function trilocation_location_string(grid, LX, LY, LZ, dim)
+    # Default implementation - will be overridden in extension for specific grid types
+    return loc2letter(LX, true) * loc2letter(LY, true) * loc2letter(LZ, true)
+end
+
+"""
+    trilocation_dim_name(var_name, grid, LX, LY, LZ, dim)
+
+Generate dimension names using trilocation naming convention.
+This is the default dimension name generator for NetCDF output.
+"""
+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, args...) = trilocation_dim_name(var_name, grid, args...)
+
 mutable struct NetCDFWriter{G, D, O, T, A, FS, DN} <: AbstractOutputWriter
     grid :: G
     filepath :: String
diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl
index e03cc385158..5ddf66d9d96 100644
--- a/test/test_netcdf_writer.jl
+++ b/test/test_netcdf_writer.jl
@@ -2817,6 +2817,68 @@ function test_netcdf_field_dimension_validation()
     return nothing
 end
 
+function test_netcdf_multiple_grids_defvar()
+    # Create two different grids with different sizes
+    grid1 = RectilinearGrid(size=(4, 4, 4), extent=(1, 1, 1))
+    grid2 = RectilinearGrid(size=(6, 8, 5), extent=(2, 3, 1.5))
+
+    # Create fields on each grid
+    c1 = CenterField(grid1)
+    c2 = CenterField(grid2)
+
+    # Set different values for each field
+    set!(c1, (x, y, z) -> x + y + z)
+    set!(c2, (x, y, z) -> 2 * (x^2 + y^2 + z^2))
+
+    filepath = "test_multiple_grids_defvar.nc"
+    isfile(filepath) && rm(filepath)
+
+    # Open NetCDF file for writing
+    ds = NCDataset(filepath, "c")
+
+    # Define variables from grid1 (4x4x4)
+    suffixed_dim_name_grid1(args...) = trilocation_dim_name(args...) * "_grid1"
+    defVar(ds, "c1", c1, attrib=Dict("long_name" => "Field from grid 1"), dimension_name_generator=suffixed_dim_name_grid1)
+
+    # Define variables from grid2 (6x8x5) - this should create new dimension names
+    suffixed_dim_name_grid2(args...) = trilocation_dim_name(args...) * "_grid2"
+    defVar(ds, "c2", c2, attrib=Dict("long_name" => "Field from grid 2"), dimension_name_generator=suffixed_dim_name_grid2)
+
+    close(ds)
+
+    # Verify the file was created correctly
+    ds = NCDataset(filepath, "r")
+
+    # Check that both fields are present
+    @test "c1" ∈ keys(ds)
+    @test "c2" ∈ keys(ds)
+
+    # Check dimensions for grid1 fields (4x4x4)
+    @test size(ds["c1"]) == (4, 4, 4)
+
+    # Check dimensions for grid2 field (6x8x5)
+    @test size(ds["c2"]) == (6, 8, 5)
+
+    # Check that the coordinate variables were created with appropriate names
+    # Grid1 coordinates (default names)
+    @test "x_caa_grid1" ∈ keys(ds)
+    @test "y_aca_grid1" ∈ keys(ds)
+    @test "z_aac_grid1" ∈ keys(ds)
+
+    @test "x_caa_grid2" ∈ keys(ds)
+    @test "y_aca_grid2" ∈ keys(ds)
+    @test "z_aac_grid2" ∈ keys(ds)
+
+    # Check that data values match
+    @test all(ds["c1"] .== interior(c1))
+    @test all(ds["c2"] .== interior(c2))
+
+    close(ds)
+    rm(filepath)
+
+    return nothing
+end
+
 for arch in archs
     @testset "NetCDF output writer [$(typeof(arch))]" begin
         @info "  Testing NetCDF output writer [$(typeof(arch))]..."
@@ -2869,5 +2931,6 @@ for arch in archs
 
         test_netcdf_single_field_defvar()
         test_netcdf_field_dimension_validation()
+        test_netcdf_multiple_grids_defvar()
     end
 end

From 8c9c7863cb5180f2e31363d429da318ee23382f9 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Sat, 20 Sep 2025 16:01:03 -0300
Subject: [PATCH 12/20] redo move to netcdf_writer.jl

---
 ext/OceananigansNCDatasetsExt.jl   | 60 +++++--------------
 src/OutputWriters/netcdf_writer.jl | 94 +++++++++++++-----------------
 2 files changed, 53 insertions(+), 101 deletions(-)

diff --git a/ext/OceananigansNCDatasetsExt.jl b/ext/OceananigansNCDatasetsExt.jl
index 8f616b38ce9..800aef4348f 100644
--- a/ext/OceananigansNCDatasetsExt.jl
+++ b/ext/OceananigansNCDatasetsExt.jl
@@ -42,13 +42,23 @@ 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, trilocation_dim_name, loc2letter, trilocation_location_string, suffixed_dim_name_generator
+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...)
 
@@ -62,7 +72,6 @@ function defVar(ds, name, field::AbstractField;
 
     # Validate that all dimensions exist and match the field
     create_field_dimensions!(ds, field, all_dims, dimension_name_generator)
-
     defVar(ds, name, FT, all_dims; kwargs...)
 end
 
@@ -150,49 +159,6 @@ function create_spatial_dimensions!(dataset, dims, attributes_dict; array_type=A
     end
 end
 
-#####
-##### Specialized dimension name generators for specific grid types
-#####
-
-# Override loc2letter for specific types
-loc2letter(::Face, full=true) = "f"
-loc2letter(::Center, full=true) = "c"
-loc2letter(::Nothing, full=true) = full ? "a" : ""
-
-# Specialized suffixed_dim_name_generator for StaticVerticalDiscretization
-suffixed_dim_name_generator(var_name, ::StaticVerticalDiscretization, LX, LY, LZ, dim::Val{:z}; connector="_", location_letters) = var_name * connector * location_letters
-
-# Specialized trilocation_location_string for specific grid types
-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)
-
-# Specialized trilocation_dim_name for ImmersedBoundaryGrid
-trilocation_dim_name(var_name, grid::ImmersedBoundaryGrid, args...) = trilocation_dim_name(var_name, grid.underlying_grid, args...)
-
-# Minimal location string functions (for other dimension name generators)
-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...)
-
-
 #####
 ##### Gathering of grid dimensions
 #####
@@ -476,8 +442,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)
diff --git a/src/OutputWriters/netcdf_writer.jl b/src/OutputWriters/netcdf_writer.jl
index 4c55a156d3e..e1a8a90bcba 100644
--- a/src/OutputWriters/netcdf_writer.jl
+++ b/src/OutputWriters/netcdf_writer.jl
@@ -4,79 +4,65 @@
 ##### NetCDFWriter functionality is implemented in ext/OceananigansNCDatasetsExt
 #####
 
-using Oceananigans.Grids: topology, Flat
+using Oceananigans.Grids: topology, Flat, StaticVerticalDiscretization
+using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
 
 #####
 ##### Dimension name generators
 #####
 
-"""
-    suffixed_dim_name_generator(var_name, grid, LX, LY, LZ, dim; connector="_", location_letters)
-
-Generate dimension names by suffixing the variable name with location letters.
-"""
-function suffixed_dim_name_generator(var_name, grid, LX, LY, LZ, dim; connector="_", location_letters)
-    # Check if topology is Flat for the given dimension
-    topo = topology(grid)
-
-    if dim == Val(:x)
-        TX = topo[1]
-        if TX == Flat || isnothing(LX)
-            return ""
-        else
-            return "$(var_name)" * connector * location_letters
-        end
-    elseif dim == Val(:y)
-        TY = topo[2]
-        if TY == Flat || isnothing(LY)
-            return ""
-        else
-            return "$(var_name)" * connector * location_letters
-        end
-    elseif dim == Val(:z)
-        TZ = topo[3]
-        if TZ == Flat || isnothing(LZ)
-            return ""
-        else
-            return "$(var_name)" * connector * location_letters
-        end
+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
 
-    return "$(var_name)" * connector * location_letters
+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
 
-# Special case for StaticVerticalDiscretization (will be overridden in extension)
-function suffixed_dim_name_generator(var_name, coordinate, LX, LY, LZ, dim::Val{:z}; connector="_", location_letters)
-    return var_name * connector * location_letters
+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
 
-"""
-    loc2letter(location, full=true)
+suffixed_dim_name_generator(var_name, ::StaticVerticalDiscretization, LX, LY, LZ, dim::Val{:z}; connector="_", location_letters) = var_name * connector * location_letters
 
-Convert location types to letter representations.
-"""
-loc2letter(::Any, full=true) = full ? "a" : ""
+loc2letter(::Face, full=true) = "f"
+loc2letter(::Center, full=true) = "c"
+loc2letter(::Nothing, full=true) = full ? "a" : ""
 
-"""
-    trilocation_location_string(grid, LX, LY, LZ, dim)
+minimal_location_string(::RectilinearGrid, LX, LY, LZ, ::Val{:x}) = loc2letter(LX, false)
+minimal_location_string(::RectilinearGrid, LX, LY, LZ, ::Val{:y}) = loc2letter(LY, false)
 
-Generate location strings for trilocation naming convention.
-"""
-function trilocation_location_string(grid, LX, LY, LZ, dim)
-    # Default implementation - will be overridden in extension for specific grid types
-    return loc2letter(LX, true) * loc2letter(LY, true) * loc2letter(LZ, true)
-end
+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_dim_name(var_name, grid, LX, LY, LZ, dim)
+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)
 
-Generate dimension names using trilocation naming convention.
-This is the default dimension name generator for NetCDF output.
-"""
 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, args...) = trilocation_dim_name(var_name, grid, args...)
+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

From 20d579da9a7e3fe6ebe0630c5068e864f23fbb23 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Sat, 20 Sep 2025 16:01:15 -0300
Subject: [PATCH 13/20] fix test

---
 test/test_netcdf_writer.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl
index 5ddf66d9d96..510f68e8a4a 100644
--- a/test/test_netcdf_writer.jl
+++ b/test/test_netcdf_writer.jl
@@ -11,6 +11,7 @@ using SeawaterPolynomials.SecondOrderSeawaterPolynomials: RoquetEquationOfState
 
 using Oceananigans: Clock
 using Oceananigans.Models.HydrostaticFreeSurfaceModels: VectorInvariant
+using Oceananigans.OutputWriters: trilocation_dim_name
 
 function test_datetime_netcdf_output(arch)
     grid = RectilinearGrid(arch, size=(1, 1, 1), extent=(1, 1, 1))

From 5b951167abcc77ac0632d190b02900854279ddf5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Chor?= <tomaschor@gmail.com>
Date: Sat, 20 Sep 2025 17:29:18 -0700
Subject: [PATCH 14/20] Update ext/OceananigansNCDatasetsExt.jl

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 ext/OceananigansNCDatasetsExt.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/ext/OceananigansNCDatasetsExt.jl b/ext/OceananigansNCDatasetsExt.jl
index 800aef4348f..f44f016e9cb 100644
--- a/ext/OceananigansNCDatasetsExt.jl
+++ b/ext/OceananigansNCDatasetsExt.jl
@@ -153,7 +153,9 @@ function create_spatial_dimensions!(dataset, dims, attributes_dict; array_type=A
         else
             # Validate existing dimension
             if dataset[dim_name] != dim_array
-                throw(ArgumentError("Dimension '$dim_name' already exists in dataset but is different from expected."))
+                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

From 441e5d842aac23ea734c99141a8ea7d4515075c5 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Tue, 30 Sep 2025 08:22:45 -0300
Subject: [PATCH 15/20] test also on GPUs

---
 test/test_netcdf_writer.jl | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl
index 510f68e8a4a..9e7ae44cd1f 100644
--- a/test/test_netcdf_writer.jl
+++ b/test/test_netcdf_writer.jl
@@ -2726,9 +2726,9 @@ function test_netcdf_buoyancy_force(arch)
     return nothing
 end
 
-function test_netcdf_single_field_defvar()
+function test_netcdf_single_field_defvar(arch)
     N = 4
-    grid = RectilinearGrid(size=(N, N, N), extent=(1, 1, 1))
+    grid = RectilinearGrid(arch, size=(N, N, N), extent=(1, 1, 1))
     c = CenterField(grid)
     set!(c, (x, y, z) -> x + y + z)
 
@@ -2774,8 +2774,8 @@ function test_netcdf_single_field_defvar()
     return nothing
 end
 
-function test_netcdf_field_dimension_validation()
-    grid = RectilinearGrid(size=(4, 4, 4), extent=(1, 1, 1))
+function test_netcdf_field_dimension_validation(arch)
+    grid = RectilinearGrid(arch, size=(4, 4, 4), extent=(1, 1, 1))
     c = CenterField(grid)
 
     # Test 1: Successful validation with proper dimensions
@@ -2818,10 +2818,10 @@ function test_netcdf_field_dimension_validation()
     return nothing
 end
 
-function test_netcdf_multiple_grids_defvar()
+function test_netcdf_multiple_grids_defvar(arch)
     # Create two different grids with different sizes
-    grid1 = RectilinearGrid(size=(4, 4, 4), extent=(1, 1, 1))
-    grid2 = RectilinearGrid(size=(6, 8, 5), extent=(2, 3, 1.5))
+    grid1 = RectilinearGrid(arch, size=(4, 4, 4), extent=(1, 1, 1))
+    grid2 = RectilinearGrid(arch, size=(6, 8, 5), extent=(2, 3, 1.5))
 
     # Create fields on each grid
     c1 = CenterField(grid1)
@@ -2930,8 +2930,8 @@ for arch in archs
 
         test_netcdf_buoyancy_force(arch)
 
-        test_netcdf_single_field_defvar()
-        test_netcdf_field_dimension_validation()
-        test_netcdf_multiple_grids_defvar()
+        test_netcdf_single_field_defvar(arch)
+        test_netcdf_field_dimension_validation(arch)
+        test_netcdf_multiple_grids_defvar(arch)
     end
 end

From 72154933a30103b363330a4fa59c44fa9de39d14 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Tue, 30 Sep 2025 08:49:44 -0300
Subject: [PATCH 16/20] add tests for immersed boundaries

---
 test/test_netcdf_writer.jl | 30 +++++++++++++++++++++---------
 1 file changed, 21 insertions(+), 9 deletions(-)

diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl
index 9e7ae44cd1f..afd89d8ce3b 100644
--- a/test/test_netcdf_writer.jl
+++ b/test/test_netcdf_writer.jl
@@ -2676,7 +2676,6 @@ function test_netcdf_free_surface_mixed_output(arch)
 end
 
 function test_netcdf_buoyancy_force(arch)
-
     Nx, Nz = 8, 8
     Hx, Hz = 2, 3
     Lx, H  = 2, 1
@@ -2697,7 +2696,6 @@ function test_netcdf_buoyancy_force(arch)
                       RoquetEquationOfState(:SimplestRealistic))
 
     for eos in Boussinesq_eos
-
         model = NonhydrostaticModel(; grid,
                                     closure = ScalarDiffusivity(ν=4e-2, κ=4e-2),
                                     buoyancy = SeawaterBuoyancy(equation_of_state=eos),
@@ -2726,9 +2724,11 @@ function test_netcdf_buoyancy_force(arch)
     return nothing
 end
 
-function test_netcdf_single_field_defvar(arch)
+function test_netcdf_single_field_defvar(arch; immersed=false)
     N = 4
     grid = RectilinearGrid(arch, size=(N, N, N), extent=(1, 1, 1))
+    immersed && (grid = ImmersedBoundaryGrid(grid, GridFittedBoundary(-1/2)))
+
     c = CenterField(grid)
     set!(c, (x, y, z) -> x + y + z)
 
@@ -2774,8 +2774,11 @@ function test_netcdf_single_field_defvar(arch)
     return nothing
 end
 
-function test_netcdf_field_dimension_validation(arch)
-    grid = RectilinearGrid(arch, size=(4, 4, 4), extent=(1, 1, 1))
+function test_netcdf_field_dimension_validation(arch; immersed=false)
+    N = 4
+    grid = RectilinearGrid(arch, size=(N, N, N), extent=(1, 1, 1))
+    immersed && (grid = ImmersedBoundaryGrid(grid, GridFittedBoundary(-1/2)))
+
     c = CenterField(grid)
 
     # Test 1: Successful validation with proper dimensions
@@ -2818,11 +2821,17 @@ function test_netcdf_field_dimension_validation(arch)
     return nothing
 end
 
-function test_netcdf_multiple_grids_defvar(arch)
+function test_netcdf_multiple_grids_defvar(arch; immersed=false)
     # Create two different grids with different sizes
     grid1 = RectilinearGrid(arch, size=(4, 4, 4), extent=(1, 1, 1))
     grid2 = RectilinearGrid(arch, size=(6, 8, 5), extent=(2, 3, 1.5))
 
+    if immersed
+        # Create different immersed boundaries for each grid
+        grid1 = ImmersedBoundaryGrid(grid1, GridFittedBoundary(-1/2))
+        grid2 = ImmersedBoundaryGrid(grid2, GridFittedBoundary(-0.8))
+    end
+
     # Create fields on each grid
     c1 = CenterField(grid1)
     c2 = CenterField(grid2)
@@ -2930,8 +2939,11 @@ for arch in archs
 
         test_netcdf_buoyancy_force(arch)
 
-        test_netcdf_single_field_defvar(arch)
-        test_netcdf_field_dimension_validation(arch)
-        test_netcdf_multiple_grids_defvar(arch)
+        test_netcdf_single_field_defvar(arch, immersed=false)
+        test_netcdf_single_field_defvar(arch, immersed=true)
+        test_netcdf_field_dimension_validation(arch, immersed=false)
+        test_netcdf_field_dimension_validation(arch, immersed=true)
+        test_netcdf_multiple_grids_defvar(arch, immersed=false)
+        test_netcdf_multiple_grids_defvar(arch, immersed=true)
     end
 end

From 7fcb2e68af31ca7a0d47546e610f2d0aafd386e3 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Tue, 30 Sep 2025 09:51:40 -0300
Subject: [PATCH 17/20] use GridFittedBottom

---
 test/test_netcdf_writer.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl
index afd89d8ce3b..39d22fd9d40 100644
--- a/test/test_netcdf_writer.jl
+++ b/test/test_netcdf_writer.jl
@@ -2727,7 +2727,7 @@ end
 function test_netcdf_single_field_defvar(arch; immersed=false)
     N = 4
     grid = RectilinearGrid(arch, size=(N, N, N), extent=(1, 1, 1))
-    immersed && (grid = ImmersedBoundaryGrid(grid, GridFittedBoundary(-1/2)))
+    immersed && (grid = ImmersedBoundaryGrid(grid, GridFittedBottom(-1/2)))
 
     c = CenterField(grid)
     set!(c, (x, y, z) -> x + y + z)
@@ -2777,7 +2777,7 @@ end
 function test_netcdf_field_dimension_validation(arch; immersed=false)
     N = 4
     grid = RectilinearGrid(arch, size=(N, N, N), extent=(1, 1, 1))
-    immersed && (grid = ImmersedBoundaryGrid(grid, GridFittedBoundary(-1/2)))
+    immersed && (grid = ImmersedBoundaryGrid(grid, GridFittedBottom(-1/2)))
 
     c = CenterField(grid)
 
@@ -2828,8 +2828,8 @@ function test_netcdf_multiple_grids_defvar(arch; immersed=false)
 
     if immersed
         # Create different immersed boundaries for each grid
-        grid1 = ImmersedBoundaryGrid(grid1, GridFittedBoundary(-1/2))
-        grid2 = ImmersedBoundaryGrid(grid2, GridFittedBoundary(-0.8))
+        grid1 = ImmersedBoundaryGrid(grid1, GridFittedBottom(-1/2))
+        grid2 = ImmersedBoundaryGrid(grid2, GridFittedBottom(-0.8))
     end
 
     # Create fields on each grid

From 2a4566000a9b1b07f9ab6a4f18b0ec1b62f219dc Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Tue, 30 Sep 2025 14:05:58 +0000
Subject: [PATCH 18/20] bring field to CPU before writing to NetCDF

---
 ext/OceananigansNCDatasetsExt.jl |  7 ++++---
 test/test_netcdf_writer.jl       | 10 +++++-----
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/ext/OceananigansNCDatasetsExt.jl b/ext/OceananigansNCDatasetsExt.jl
index f44f016e9cb..4d915d5219b 100644
--- a/ext/OceananigansNCDatasetsExt.jl
+++ b/ext/OceananigansNCDatasetsExt.jl
@@ -8,7 +8,7 @@ using Printf: @sprintf
 using OrderedCollections: OrderedDict
 
 using Oceananigans: initialize!, prettytime, pretty_filesize, AbstractModel
-using Oceananigans.Architectures: CPU, GPU
+using Oceananigans.Architectures: CPU, GPU, on_architecture
 using Oceananigans.AbstractOperations: KernelFunctionOperation, AbstractOperation
 using Oceananigans.BuoyancyFormulations: BuoyancyForce, BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState
 using Oceananigans.Fields
@@ -66,13 +66,14 @@ function defVar(ds, name, field::AbstractField;
                 time_dependent=false,
                 dimension_name_generator = trilocation_dim_name,
                 kwargs...)
-    FT = Array{eltype(field)}(field)
+    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, FT, all_dims; kwargs...)
+    defVar(ds, name, field_data, all_dims; kwargs...)
 end
 
 #####
diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl
index 39d22fd9d40..3059d2a2b3d 100644
--- a/test/test_netcdf_writer.jl
+++ b/test/test_netcdf_writer.jl
@@ -2755,16 +2755,16 @@ function test_netcdf_single_field_defvar(arch; immersed=false)
     ds = NCDataset(filepath, "r")
 
     @test "c" ∈ keys(ds)
-    @test all(ds["c"] .== interior(c))
+    @test ds["c"] == Array(interior(c))
     @test ds["c"].attrib["long_name"] == c_long_name
     @test ds["c"].attrib["units"] == c_units
 
     @test "c²" ∈ keys(ds)
-    @test all(ds["c²"] .== interior(c²))
+    @test ds["c²"] == Array(interior(Field(c²)))
     @test ds["c²"].attrib["long_name"] == c²_long_name
 
     @test "c̄" ∈ keys(ds)
-    @test all(ds["c̄"] .== interior(Field(c̄)))
+    @test all(ds["c̄"] .== Array(interior(Field(c̄))))
     @test ds["c̄"].attrib["long_name"] == c̄_long_name
     @test ds["c̄"].attrib["units"] == c̄_units
 
@@ -2880,8 +2880,8 @@ function test_netcdf_multiple_grids_defvar(arch; immersed=false)
     @test "z_aac_grid2" ∈ keys(ds)
 
     # Check that data values match
-    @test all(ds["c1"] .== interior(c1))
-    @test all(ds["c2"] .== interior(c2))
+    @test ds["c1"] == Array(interior(c1))
+    @test ds["c2"] == Array(interior(c2))
 
     close(ds)
     rm(filepath)

From e3d44bbb6f7bf2ee8a39d4cb1ba750fd11223bf6 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Tue, 30 Sep 2025 17:27:50 +0000
Subject: [PATCH 19/20] add tests for LatLonGrid

---
 test/test_netcdf_writer.jl | 105 ++++++++++++++++++++++---------------
 1 file changed, 63 insertions(+), 42 deletions(-)

diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl
index 3059d2a2b3d..1acacba9d25 100644
--- a/test/test_netcdf_writer.jl
+++ b/test/test_netcdf_writer.jl
@@ -12,6 +12,7 @@ using SeawaterPolynomials.SecondOrderSeawaterPolynomials: RoquetEquationOfState
 using Oceananigans: Clock
 using Oceananigans.Models.HydrostaticFreeSurfaceModels: VectorInvariant
 using Oceananigans.OutputWriters: trilocation_dim_name
+using Oceananigans.Grids: ξname, ηname, rname, ξnodes, ηnodes
 
 function test_datetime_netcdf_output(arch)
     grid = RectilinearGrid(arch, size=(1, 1, 1), extent=(1, 1, 1))
@@ -2724,10 +2725,9 @@ function test_netcdf_buoyancy_force(arch)
     return nothing
 end
 
-function test_netcdf_single_field_defvar(arch; immersed=false)
-    N = 4
-    grid = RectilinearGrid(arch, size=(N, N, N), extent=(1, 1, 1))
-    immersed && (grid = ImmersedBoundaryGrid(grid, GridFittedBottom(-1/2)))
+function test_netcdf_single_field_defvar(grid; immersed=false)
+    # Create immersed boundary grid if requested
+    grid = immersed ? ImmersedBoundaryGrid(grid, GridFittedBottom(-1/2)) : grid
 
     c = CenterField(grid)
     set!(c, (x, y, z) -> x + y + z)
@@ -2774,11 +2774,9 @@ function test_netcdf_single_field_defvar(arch; immersed=false)
     return nothing
 end
 
-function test_netcdf_field_dimension_validation(arch; immersed=false)
-    N = 4
-    grid = RectilinearGrid(arch, size=(N, N, N), extent=(1, 1, 1))
-    immersed && (grid = ImmersedBoundaryGrid(grid, GridFittedBottom(-1/2)))
-
+function test_netcdf_field_dimension_validation(grid; immersed=false)
+    # Create immersed boundary grid if requested
+    grid = immersed ? ImmersedBoundaryGrid(grid, GridFittedBottom(-1/2)) : grid
     c = CenterField(grid)
 
     # Test 1: Successful validation with proper dimensions
@@ -2788,16 +2786,21 @@ function test_netcdf_field_dimension_validation(arch; immersed=false)
     # Create NetCDF file with proper dimensions
     ds = NCDataset(filepath, "c")
 
+    # Get dimension names in a way that works for multiple grid types
+    ξ_center_name = string(ξname(grid)) * "_caa"
+    η_center_name = string(ηname(grid)) * "_aca"
+    r_center_name = string(rname(grid)) * "_aac"
+
     # Define dimensions first
-    defVar(ds, "x_caa", xnodes(grid, Center()), ("x_caa",))
-    # We "forget" to define the y dimension. Should work regardless
-    defVar(ds, "z_aac", znodes(grid, Center()), ("z_aac",))
+    defVar(ds, ξ_center_name, ξnodes(grid, Center()), (ξ_center_name,))
+    # We "forget" to define the η dimension. Should work regardless
+    defVar(ds, r_center_name, rnodes(grid, Center()), (r_center_name,))
 
     # Write variable to disk after dimensions are created and match the
     defVar(ds, "c", c, time_dependent=false)
-    @test ds["x_caa"][:] == xnodes(grid, Center())
-    @test ds["y_aca"][:] == ynodes(grid, Center())
-    @test ds["z_aac"][:] == znodes(grid, Center())
+    @test ds[ξ_center_name][:] == ξnodes(grid, Center())
+    @test ds[η_center_name][:] == ηnodes(grid, Center())
+    @test ds[r_center_name][:] == rnodes(grid, Center())
 
     close(ds)
     rm(filepath)
@@ -2809,9 +2812,9 @@ function test_netcdf_field_dimension_validation(arch; immersed=false)
     ds_wrong = NCDataset(filepath_wrong_size, "c")
 
     # Define dimensions with wrong sizes
-    defVar(ds_wrong, "x_caa", xnodes(grid, Center()), ("x_caa",))
-    defVar(ds_wrong, "y_aca", ynodes(grid, Center()), ("y_aca",))
-    defVar(ds_wrong, "z_aac", znodes(grid, Face()), ("z_aac",)) # ⚠ wrongly write faces instead of center
+    defVar(ds_wrong, ξ_center_name, ξnodes(grid, Center()), (ξ_center_name,))
+    defVar(ds_wrong, η_center_name, ηnodes(grid, Center()), (η_center_name,))
+    defVar(ds_wrong, r_center_name, rnodes(grid, Face()), (r_center_name,)) # ⚠ wrongly write faces instead of center
 
     @test_throws ArgumentError defVar(ds_wrong, "c", c)
 
@@ -2821,11 +2824,8 @@ function test_netcdf_field_dimension_validation(arch; immersed=false)
     return nothing
 end
 
-function test_netcdf_multiple_grids_defvar(arch; immersed=false)
-    # Create two different grids with different sizes
-    grid1 = RectilinearGrid(arch, size=(4, 4, 4), extent=(1, 1, 1))
-    grid2 = RectilinearGrid(arch, size=(6, 8, 5), extent=(2, 3, 1.5))
-
+function test_netcdf_multiple_grids_defvar(grid1, grid2; immersed=false)
+    # Create immersed boundary grids if requested
     if immersed
         # Create different immersed boundaries for each grid
         grid1 = ImmersedBoundaryGrid(grid1, GridFittedBottom(-1/2))
@@ -2846,11 +2846,11 @@ function test_netcdf_multiple_grids_defvar(arch; immersed=false)
     # Open NetCDF file for writing
     ds = NCDataset(filepath, "c")
 
-    # Define variables from grid1 (4x4x4)
+    # Define variables from grid1
     suffixed_dim_name_grid1(args...) = trilocation_dim_name(args...) * "_grid1"
     defVar(ds, "c1", c1, attrib=Dict("long_name" => "Field from grid 1"), dimension_name_generator=suffixed_dim_name_grid1)
 
-    # Define variables from grid2 (6x8x5) - this should create new dimension names
+    # Define variables from grid2 - this should create new dimension names
     suffixed_dim_name_grid2(args...) = trilocation_dim_name(args...) * "_grid2"
     defVar(ds, "c2", c2, attrib=Dict("long_name" => "Field from grid 2"), dimension_name_generator=suffixed_dim_name_grid2)
 
@@ -2863,21 +2863,29 @@ function test_netcdf_multiple_grids_defvar(arch; immersed=false)
     @test "c1" ∈ keys(ds)
     @test "c2" ∈ keys(ds)
 
-    # Check dimensions for grid1 fields (4x4x4)
-    @test size(ds["c1"]) == (4, 4, 4)
+    # Check dimensions for grid1 fields
+    @test size(ds["c1"]) == size(grid1)
 
-    # Check dimensions for grid2 field (6x8x5)
-    @test size(ds["c2"]) == (6, 8, 5)
+    # Check dimensions for grid2 field
+    @test size(ds["c2"]) == size(grid2)
 
     # Check that the coordinate variables were created with appropriate names
-    # Grid1 coordinates (default names)
-    @test "x_caa_grid1" ∈ keys(ds)
-    @test "y_aca_grid1" ∈ keys(ds)
-    @test "z_aac_grid1" ∈ keys(ds)
+    # Get dimension names in a way that works for multiple grid types
+    ξ_center_name1 = string(ξname(grid1)) * "_caa_grid1"
+    η_center_name1 = string(ηname(grid1)) * "_aca_grid1"
+    r_center_name1 = string(rname(grid1)) * "_aac_grid1"
+
+    @test ξ_center_name1 ∈ keys(ds)
+    @test η_center_name1 ∈ keys(ds)
+    @test r_center_name1 ∈ keys(ds)
 
-    @test "x_caa_grid2" ∈ keys(ds)
-    @test "y_aca_grid2" ∈ keys(ds)
-    @test "z_aac_grid2" ∈ keys(ds)
+    ξ_center_name2 = string(ξname(grid2)) * "_caa_grid2"
+    η_center_name2 = string(ηname(grid2)) * "_aca_grid2"
+    r_center_name2 = string(rname(grid2)) * "_aac_grid2"
+
+    @test ξ_center_name2 ∈ keys(ds)
+    @test η_center_name2 ∈ keys(ds)
+    @test r_center_name2 ∈ keys(ds)
 
     # Check that data values match
     @test ds["c1"] == Array(interior(c1))
@@ -2893,6 +2901,15 @@ for arch in archs
     @testset "NetCDF output writer [$(typeof(arch))]" begin
         @info "  Testing NetCDF output writer [$(typeof(arch))]..."
 
+        # Pre-build grids
+        N = 4
+        rectilinear_grid1 = RectilinearGrid(arch, size=(N, N, N), extent=(1, 1, 1))
+        rectilinear_grid2 = RectilinearGrid(arch, size=(6, 8, 5), extent=(2, 3, 1.5))
+
+        # Pre-build LatitudeLongitudeGrids with different sizes
+        latlon_grid1 = LatitudeLongitudeGrid(arch, size=(N, N, N), longitude=(-180, 180), latitude=(-90, 90), z=(-1, 0))
+        latlon_grid2 = LatitudeLongitudeGrid(arch, size=(8, 6, 4), longitude=(-120, 60), latitude=(-60, 60), z=(-2, 0))
+
         test_datetime_netcdf_output(arch)
         test_timedate_netcdf_output(arch)
 
@@ -2939,11 +2956,15 @@ for arch in archs
 
         test_netcdf_buoyancy_force(arch)
 
-        test_netcdf_single_field_defvar(arch, immersed=false)
-        test_netcdf_single_field_defvar(arch, immersed=true)
-        test_netcdf_field_dimension_validation(arch, immersed=false)
-        test_netcdf_field_dimension_validation(arch, immersed=true)
-        test_netcdf_multiple_grids_defvar(arch, immersed=false)
-        test_netcdf_multiple_grids_defvar(arch, immersed=true)
+        for grids in ((rectilinear_grid1, rectilinear_grid2),
+                      (latlon_grid1, latlon_grid2))
+            grid1, grid2 = grids
+            test_netcdf_single_field_defvar(grid1, immersed=false)
+            test_netcdf_single_field_defvar(grid1, immersed=true)
+            test_netcdf_field_dimension_validation(grid1, immersed=false)
+            test_netcdf_field_dimension_validation(grid1, immersed=true)
+            test_netcdf_multiple_grids_defvar(grid1, grid2, immersed=false)
+            test_netcdf_multiple_grids_defvar(grid1, grid2, immersed=true)
+        end
     end
 end

From 827cd79da88a659143e5f68a73636670c142b651 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Tue, 30 Sep 2025 16:57:49 -0300
Subject: [PATCH 20/20] fewer comments

---
 test/test_netcdf_writer.jl | 19 ++-----------------
 1 file changed, 2 insertions(+), 17 deletions(-)

diff --git a/test/test_netcdf_writer.jl b/test/test_netcdf_writer.jl
index 1acacba9d25..8dfe0431275 100644
--- a/test/test_netcdf_writer.jl
+++ b/test/test_netcdf_writer.jl
@@ -2726,11 +2726,10 @@ function test_netcdf_buoyancy_force(arch)
 end
 
 function test_netcdf_single_field_defvar(grid; immersed=false)
-    # Create immersed boundary grid if requested
     grid = immersed ? ImmersedBoundaryGrid(grid, GridFittedBottom(-1/2)) : grid
 
     c = CenterField(grid)
-    set!(c, (x, y, z) -> x + y + z)
+    set!(c, (x, y, z) -> x + y + z) # Make sure field is non-trivial
 
     filepath = "test_single_field_defvar.nc"
     isfile(filepath) && rm(filepath)
@@ -2775,15 +2774,11 @@ function test_netcdf_single_field_defvar(grid; immersed=false)
 end
 
 function test_netcdf_field_dimension_validation(grid; immersed=false)
-    # Create immersed boundary grid if requested
     grid = immersed ? ImmersedBoundaryGrid(grid, GridFittedBottom(-1/2)) : grid
     c = CenterField(grid)
 
-    # Test 1: Successful validation with proper dimensions
     filepath = "test_dimension_validation_success.nc"
     isfile(filepath) && rm(filepath)
-
-    # Create NetCDF file with proper dimensions
     ds = NCDataset(filepath, "c")
 
     # Get dimension names in a way that works for multiple grid types
@@ -2825,14 +2820,11 @@ function test_netcdf_field_dimension_validation(grid; immersed=false)
 end
 
 function test_netcdf_multiple_grids_defvar(grid1, grid2; immersed=false)
-    # Create immersed boundary grids if requested
     if immersed
-        # Create different immersed boundaries for each grid
         grid1 = ImmersedBoundaryGrid(grid1, GridFittedBottom(-1/2))
         grid2 = ImmersedBoundaryGrid(grid2, GridFittedBottom(-0.8))
     end
 
-    # Create fields on each grid
     c1 = CenterField(grid1)
     c2 = CenterField(grid2)
 
@@ -2842,8 +2834,6 @@ function test_netcdf_multiple_grids_defvar(grid1, grid2; immersed=false)
 
     filepath = "test_multiple_grids_defvar.nc"
     isfile(filepath) && rm(filepath)
-
-    # Open NetCDF file for writing
     ds = NCDataset(filepath, "c")
 
     # Define variables from grid1
@@ -2858,15 +2848,10 @@ function test_netcdf_multiple_grids_defvar(grid1, grid2; immersed=false)
 
     # Verify the file was created correctly
     ds = NCDataset(filepath, "r")
-
-    # Check that both fields are present
     @test "c1" ∈ keys(ds)
     @test "c2" ∈ keys(ds)
 
-    # Check dimensions for grid1 fields
     @test size(ds["c1"]) == size(grid1)
-
-    # Check dimensions for grid2 field
     @test size(ds["c2"]) == size(grid2)
 
     # Check that the coordinate variables were created with appropriate names
@@ -2901,7 +2886,7 @@ for arch in archs
     @testset "NetCDF output writer [$(typeof(arch))]" begin
         @info "  Testing NetCDF output writer [$(typeof(arch))]..."
 
-        # Pre-build grids
+        # Pre-build RectilinearGrids with different sizes
         N = 4
         rectilinear_grid1 = RectilinearGrid(arch, size=(N, N, N), extent=(1, 1, 1))
         rectilinear_grid2 = RectilinearGrid(arch, size=(6, 8, 5), extent=(2, 3, 1.5))