Merge pull request #2865 from CliMA/tc/reduced-field-netcdf

Fixes writing of 3d-reduced `Field`s to NetCDF

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.79.2"
+version = "0.79.3"
 
 [deps]
 AMGX = "c963dde9-0319-47f5-bf0c-b07d3c80ffa6"
@@ -39,9 +39,9 @@ StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 
 [compat]
+AMGX = "0.1.3"
 Adapt = "3"
 AlgebraicMultigrid = "0.5"
-AMGX = "0.1.3"
 CUDA = "3.8, 3.9"
 CUDAKernels = "0.4.7"
 Crayons = "4"

src/OutputWriters/netcdf_output_writer.jl

@@ -437,7 +437,7 @@ function save_output!(ds, output, model, ow, time_index, name)
     data = fetch_and_convert_output(output, model, ow)
     data = drop_output_dims(output, data)
     colons = Tuple(Colon() for _ in 1:ndims(data))
-    ds[name][colons..., time_index] = data
+    ds[name][colons..., time_index:time_index] = data
     return nothing
 end
 

test/test_computed_field.jl

@@ -390,8 +390,8 @@ for arch in archs
                 end
 
                 ϵ(x, y, z) = 2rand() - 1
-                u, v, w = model.velocities
                 set!(model, u=ϵ, v=ϵ)
+                u, v, w = model.velocities
                 ζ_op = KernelFunctionOperation{Face, Face, Center}(ζ₃ᶠᶠᶜ, grid, computed_dependencies=(u, v))
 
                 ζ = Field(ζ_op) # identical to `VerticalVorticityField`
@@ -400,15 +400,15 @@ for arch in archs
 
                 ζxy = Field(ζ_op, indices=(:, :, 1))
                 compute!(ζxy)
-                @test ζxy == view(ζ, :, :, 1)
+                @test all(interior(ζxy, :, :, 1) .== interior(ζ, :, :, 1))
 
                 ζxz = Field(ζ_op, indices=(:, 1, :))
                 compute!(ζxz)
-                @test ζxz == view(ζ, :, 1, :)
+                @test all(interior(ζxz, :, 1, :) .== interior(ζ, :, 1, :))
 
                 ζyz = Field(ζ_op, indices=(1, :, :))
                 compute!(ζyz)
-                @test ζyz == view(ζ, 1, :, :)
+                @test all(interior(ζyz, 1, :, :) .== interior(ζ, 1, :, :))
             end
 
             @testset "Operations with computed Fields [$A, $G]" begin
@@ -417,7 +417,7 @@ for arch in archs
             end
 
             @testset "Horizontal averages of operations [$A, $G]" begin
-                @info "      Testing horizontal averges..."
+                @info "      Testing horizontal averages..."
                 @test horizontal_average_of_plus(model)
                 @test horizontal_average_of_minus(model)
                 @test horizontal_average_of_times(model)
@@ -427,12 +427,12 @@ for arch in archs
             end
 
             @testset "Zonal averages of operations [$A, $G]" begin
-                @info "      Testing zonal averges..."
+                @info "      Testing zonal averages..."
                 @test zonal_average_of_plus(model)
             end
 
             @testset "Volume averages of operations [$A, $G]" begin
-                @info "      Testing volume averges..."
+                @info "      Testing volume averages..."
                 @test volume_average_of_times(model)
             end
 

test/test_netcdf_output_writer.jl

@@ -504,10 +504,53 @@ function test_netcdf_function_output(arch)
     return nothing
 end
 
+function test_netcdf_spatial_average(arch)
+    topo = (Periodic, Periodic, Periodic)
+    domain = (x=(0, 1), y=(0, 1), z=(0, 1))
+    grid = RectilinearGrid(arch, topology=topo, size=(4, 4, 4); domain...)
+
+    model = NonhydrostaticModel(grid = grid,
+                                timestepper = :RungeKutta3,
+                                tracers = (:c,), 
+                                coriolis = nothing, 
+                                buoyancy = nothing, 
+                                closure = nothing)
+    set!(model, c=1)
+
+    Δt = 1/64 # Nice floating-point number
+    simulation = Simulation(model, Δt=Δt, stop_time=50Δt)
+
+    ∫c_dx = Field(Average(model.tracers.c, dims=(1)))
+    ∫∫c_dxdy = Field(Average(model.tracers.c, dims=(1, 2)))
+    ∫∫∫c_dxdydz = Field(Average(model.tracers.c, dims=(1, 2, 3)))
+
+    volume_avg_nc_filepath = "volume_averaged_field_test.nc"
+
+    simulation.output_writers[:averages] = NetCDFOutputWriter(model, (; ∫c_dx, ∫∫c_dxdy, ∫∫∫c_dxdydz),
+                                                              array_type = Array{Float64},
+                                                              verbose = true,
+                                                              filename = volume_avg_nc_filepath,
+                                                              schedule = TimeInterval(10Δt))
+    run!(simulation)
+
+    ds = NCDataset(volume_avg_nc_filepath)
+
+    for (n, t) in enumerate(ds["time"])
+        @test all(ds["∫c_dx"][:,:, n] .≈ 1)
+        @test all(ds["∫∫c_dxdy"][:, n] .≈ 1)
+        @test all(ds["∫∫∫c_dxdydz"][n] .≈ 1)
+    end
+
+    close(ds)
+
+    return nothing
+end
+
+
 function test_netcdf_time_averaging(arch)
     topo = (Periodic, Periodic, Periodic)
     domain = (x=(0, 1), y=(0, 1), z=(0, 1))
-    grid = RectilinearGrid(topology=topo, size=(4, 4, 4); domain...)
+    grid = RectilinearGrid(arch, topology=topo, size=(4, 4, 4); domain...)
 
     λ1(x, y, z) = x + (1 - y)^2 + tanh(z)
     λ2(x, y, z) = x + (1 - y)^2 + tanh(4z)
@@ -795,6 +838,7 @@ for arch in archs
         test_thermal_bubble_netcdf_output_with_halos(arch)
         test_netcdf_function_output(arch)
         test_netcdf_output_alignment(arch)
+        test_netcdf_spatial_average(arch)
         test_netcdf_time_averaging(arch)
         test_netcdf_vertically_stretched_grid_output(arch)
         test_netcdf_regular_lat_lon_grid_output(arch)