Passes grid argument to NetCDFOutputWriter

docs/src/model_setup/output_writers.md

@@ -158,6 +158,37 @@ NetCDFOutputWriter scheduled on IterationInterval(1):
 └── file size: 17.8 KiB
 ```
 
+`NetCDFOutputWriter` can also be configured for `outputs` that are interpolated or regridded to a different grid than `model.grid`.
+To use this functionality, the `output_grid` must be passed explicitly when constructing `NetCDFOutputWriter` along with the regridded / interpolated `outputs`.
+
+```jldoctest
+using Oceananigans
+using Oceananigans.Fields: interpolate!
+
+grid = RectilinearGrid(size=(1, 1, 8), extent=(1, 1, 1));
+model = NonhydrostaticModel(; grid)
+simulation = Simulation(model, Δt=1, stop_iteration=1)
+
+coarse_grid = RectilinearGrid(size=(grid.Nx, grid.Ny, grid.Nz÷2), extent=(grid.Lx, grid.Ly, grid.Lz))
+coarse_u = Field{Face, Center, Center}(coarse_grid)
+
+interpolate_u(model) = interpolate!(coarse_u, model.velocities.u)
+outputs = (; u = interpolate_u)
+
+simulation.output_writers[:coarse_u] = NetCDFOutputWriter(model, outputs, coarse_grid;
+                                                          filename = "coarse_u.nc",
+                                                          schedule = IterationInterval(1))
+
+# output
+NetCDFOutputWriter scheduled on IterationInterval(1):
+├── filepath: ./coarse_u.nc
+├── dimensions: zC(4), zF(5), xC(1), yF(1), xF(1), yC(1), time(0)
+├── 1 outputs: u
+└── array type: Array{Float64}
+├── file_splitting: NoFileSplitting
+└── file size: 14.6 KiB
+```
+
 See [`NetCDFOutputWriter`](@ref) for more information.
 
 ## JLD2 output writer

src/Fields/interpolate.jl

@@ -354,5 +354,5 @@ function interpolate!(to_field::Field, from_field::AbstractField)
 
     fill_halo_regions!(to_field)
 
-    return nothing
+    return to_field
 end

src/OutputWriters/netcdf_output_writer.jl

@@ -10,7 +10,8 @@ using Oceananigans.Utils: versioninfo_with_gpu, oceananigans_versioninfo, pretty
 using Oceananigans.TimeSteppers: float_or_date_time
 using Oceananigans.Fields: reduced_dimensions, reduced_location, location, validate_indices
 
-mutable struct NetCDFOutputWriter{D, O, T, A, FS} <: AbstractOutputWriter
+mutable struct NetCDFOutputWriter{G, D, O, T, A, FS} <: AbstractOutputWriter
+    grid :: G
     filepath :: String
     dataset :: D
     outputs :: O
@@ -162,7 +163,7 @@ function add_schedule_metadata!(global_attributes, schedule::AveragedTimeInterva
 end
 
 """
-    NetCDFOutputWriter(model, outputs; filename, schedule
+    NetCDFOutputWriter(model, outputs, grid; filename, schedule
                                           dir = ".",
                                    array_type = Array{Float64},
                                       indices = nothing,
@@ -180,7 +181,7 @@ Construct a `NetCDFOutputWriter` that writes `(label, output)` pairs in `outputs
 be a `Dict`) to a NetCDF file, where `label` is a string that labels the output and `output` is
 either a `Field` (e.g. `model.velocities.u`) or a function `f(model)` that
 returns something to be written to disk. Custom output requires the spatial `dimensions` (a
-`Dict`) to be manually specified (see examples).
+`Dict`) or `grid` to be manually specified (see examples).
 
 Keyword arguments
 =================
@@ -354,20 +355,53 @@ NetCDFOutputWriter scheduled on IterationInterval(1):
 ├── file_splitting: NoFileSplitting
 └── file size: 17.8 KiB
 ```
+
+`NetCDFOutputWriter` can also be configured for `outputs` that are interpolated or regridded to a different grid than `model.grid`.
+To use this functionality, the `output_grid` must be passed explicitly when constructing `NetCDFOutputWriter` along with the regridded / interpolated `outputs`.
+
+```jldoctest
+using Oceananigans
+using Oceananigans.Fields: interpolate!
+
+grid = RectilinearGrid(size=(1, 1, 8), extent=(1, 1, 1));
+model = NonhydrostaticModel(; grid)
+simulation = Simulation(model, Δt=1, stop_iteration=1)
+
+coarse_grid = RectilinearGrid(size=(grid.Nx, grid.Ny, grid.Nz÷2), extent=(grid.Lx, grid.Ly, grid.Lz))
+coarse_u = Field{Face, Center, Center}(coarse_grid)
+
+interpolate_u(model) = interpolate!(coarse_u, model.velocities.u)
+outputs = (; u = interpolate_u)
+
+simulation.output_writers[:coarse_u] = NetCDFOutputWriter(model, outputs, coarse_grid;
+                                                          filename = "coarse_u.nc",
+                                                          schedule = IterationInterval(1))
+
+# output
+NetCDFOutputWriter scheduled on IterationInterval(1):
+├── filepath: ./coarse_u.nc
+├── dimensions: zC(4), zF(5), xC(1), yF(1), xF(1), yC(1), time(0)
+├── 1 outputs: u
+└── array type: Array{Float64}
+├── file_splitting: NoFileSplitting
+└── file size: 14.6 KiB
+```
 """
-function NetCDFOutputWriter(model, outputs; filename, schedule,
-                                          dir = ".",
-                                   array_type = Array{Float64},
-                                      indices = (:, :, :),
-                                   with_halos = false,
+function NetCDFOutputWriter(model, outputs, grid=model.grid;
+                            filename,
+                            schedule,
+                            dir = ".",
+                            array_type = Array{Float64},
+                            indices = (:, :, :),
+                            with_halos = false,
                             global_attributes = Dict(),
                             output_attributes = Dict(),
-                                   dimensions = Dict(),
-                           overwrite_existing = nothing,
-                                 deflatelevel = 0,
-                                         part = 1,
-                               file_splitting = NoFileSplitting(),
-                                      verbose = false)
+                            dimensions = Dict(),
+                            overwrite_existing = nothing,
+                            deflatelevel = 0,
+                            part = 1,
+                            file_splitting = NoFileSplitting(),
+                            verbose = false)
     mkpath(dir)
     filename = auto_extension(filename, ".nc")
     filepath = joinpath(dir, filename)
@@ -394,7 +428,7 @@ function NetCDFOutputWriter(model, outputs; filename, schedule,
     # with LagrangianParticles output (see the end of the file).
     # We shouldn't support this in the future; we should require users to 'name' LagrangianParticles output.
     outputs = dictify(outputs)
-    outputs = Dict(string(name) => construct_output(outputs[name], model.grid, indices, with_halos) for name in keys(outputs))
+    outputs = Dict(string(name) => construct_output(outputs[name], grid, indices, with_halos) for name in keys(outputs))
 
     output_attributes = dictify(output_attributes)
     global_attributes = dictify(global_attributes)
@@ -414,9 +448,11 @@ function NetCDFOutputWriter(model, outputs; filename, schedule,
                                                      dimensions,
                                                      overwrite_existing,
                                                      deflatelevel,
+                                                     grid,
                                                      model)
 
-    return NetCDFOutputWriter(filepath,
+    return NetCDFOutputWriter(grid,
+                              filepath,
                               dataset,
                               outputs,
                               schedule,
@@ -446,22 +482,34 @@ get_default_dimension_attributes(grid::ImmersedBoundaryGrid) =
 ##### Variable definition
 #####
 
+materialize_output(func, model) = func(model)
+materialize_output(field::AbstractField, model) = field
+materialize_output(particles::LagrangianParticles, model) = particles
+
 """ Defines empty variables for 'custom' user-supplied `output`. """
-function define_output_variable!(dataset, output, name, array_type, deflatelevel, output_attributes, dimensions)
+function define_output_variable!(dataset, output, name, array_type,
+                                 deflatelevel, attrib, dimensions)
+
     name ∉ keys(dimensions) && error("Custom output $name needs dimensions!")
 
-    defVar(dataset, name, eltype(array_type), (dimensions[name]..., "time"),
-           deflatelevel=deflatelevel, attrib=output_attributes)
+    dims = dimensions[name]
+    FT = eltype(array_type)
+    defVar(dataset, name, FT, (dims..., "time"); deflatelevel, attrib)
 
     return nothing
 end
 
 
 """ Defines empty field variable. """
-define_output_variable!(dataset, output::AbstractField, name, array_type, deflatelevel, output_attributes, dimensions) =
-    defVar(dataset, name, eltype(array_type),
-           (netcdf_spatial_dimensions(output)..., "time"),
-           deflatelevel=deflatelevel, attrib=output_attributes)
+function define_output_variable!(dataset, output::AbstractField, name, array_type,
+                                 deflatelevel, attrib, dimensions)
+
+    dims = netcdf_spatial_dimensions(output)
+    FT = eltype(array_type)
+    defVar(dataset, name, FT, (dims..., "time"); deflatelevel, attrib)
+
+    return nothing
+end
 
 """ Defines empty field variable for `WindowedTimeAverage`s over fields. """
 define_output_variable!(dataset, output::WindowedTimeAverage{<:AbstractField}, args...) =
@@ -580,12 +628,17 @@ end
 #####
 
 """ Defines empty variable for particle trackting. """
-function define_output_variable!(dataset, output::LagrangianParticles, name, array_type, deflatelevel, output_attributes, dimensions)
+function define_output_variable!(dataset, output::LagrangianParticles, name, array_type,
+                                 deflatelevel, output_attributes, dimensions)
+
     particle_fields = eltype(output.properties) |> fieldnames .|> string
+    T = eltype(array_type)
+
     for particle_field in particle_fields
-        defVar(dataset, particle_field, eltype(array_type),
-               ("particle_id", "time"), deflatelevel=deflatelevel)
+        defVar(dataset, particle_field, T, ("particle_id", "time"); deflatelevel)
     end
+
+    return nothing
 end
 
 dictify(outputs::LagrangianParticles) = Dict("particles" => outputs)
@@ -633,6 +686,7 @@ function initialize_nc_file!(filepath,
                              dimensions,
                              overwrite_existing,
                              deflatelevel,
+                             grid,
                              model)
 
     mode = overwrite_existing ? "c" : "a"
@@ -648,12 +702,12 @@ function initialize_nc_file!(filepath,
     # schedule::AveragedTimeInterval
     schedule, outputs = time_average_outputs(schedule, outputs, model)
 
-    dims = default_dimensions(outputs, model.grid, indices, with_halos)
+    dims = default_dimensions(outputs, grid, indices, with_halos)
 
     # Open the NetCDF dataset file
     dataset = NCDataset(filepath, mode, attrib=global_attributes)
 
-    default_dimension_attributes = get_default_dimension_attributes(model.grid)
+    default_dimension_attributes = get_default_dimension_attributes(grid)
 
     # Define variables for each dimension and attributes if this is a new file.
     if mode == "c"
@@ -669,7 +723,7 @@ function initialize_nc_file!(filepath,
 
         # Creates an unlimited dimension "time"
         defDim(dataset, "time", Inf)
-        defVar(dataset, "time", eltype(model.grid), ("time",), attrib=time_attrib)
+        defVar(dataset, "time", eltype(grid), ("time",), attrib=time_attrib)
 
         # Use default output attributes for known outputs if the user has not specified any.
         # Unknown outputs get an empty tuple (no output attributes).
@@ -681,7 +735,14 @@ function initialize_nc_file!(filepath,
 
         for (name, output) in outputs
             attributes = try output_attributes[name]; catch; Dict(); end
-            define_output_variable!(dataset, output, name, array_type, deflatelevel, attributes, dimensions)
+            materialized = materialize_output(output, model) 
+            define_output_variable!(dataset,
+                                    materialized,
+                                    name,
+                                    array_type,
+                                    deflatelevel,
+                                    attributes,
+                                    dimensions)
         end
 
         sync(dataset)
@@ -704,4 +765,5 @@ initialize_nc_file!(ow::NetCDFOutputWriter, model) =
                         ow.dimensions,
                         ow.overwrite_existing,
                         ow.deflatelevel,
+                        ow.grid,
                         model)