Passes `grid` argument to `NetCDFOutputWriter`

[tomchor]

Related to https://github.com/CliMA/Oceananigans.jl/issues/3460

[tomchor]

Unfortunately this example

using Oceananigans

grid = RectilinearGrid(size = (1, 1, 8), extent = (1,1,1));
model = NonhydrostaticModel(; grid, closure = ScalarDiffusivity(ν=1e-2))

set!(model, u=(x, y, z,) -> z)

simulation = Simulation(model,
                        Δt=0.5*maximum(zspacings(grid, Center())) / maximum(abs, model.velocities.u),
                        stop_time=20)
# Create regular output
simulation.output_writers[:fullfields] = NetCDFOutputWriter(model, (; model.velocities.u),
                                                            filename = "fullfields.nc",
                                                            schedule = TimeInterval(5),
                                                            overwrite_existing = true,)

# Create interpolated u on coarse grid
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)

using Oceananigans.Fields: interpolate!
update_coarse_u(simulation) = interpolate!(coarse_u, simulation.model.velocities.u)
simulation.callbacks[:update_interp] = Callback(update_coarse_u)

# Create coarse output
simulation.output_writers[:coarsefields] = NetCDFOutputWriter(model, (; coarse_u,), coarse_grid;
                                                              filename="coarsefields.nc",
                                                              schedule=TimeInterval(5),
                                                              overwrite_existing=true,)

run!(simulation)

Throws the following error:

ERROR: LoadError: DimensionMismatch: new dimensions (1, 1, 8, 1) must be consistent with array size 4
Stacktrace:
  [1] (::Base.var"#throw_dmrsa#328")(dims::NTuple{4, Int64}, len::Int64)
    @ Base ./reshapedarray.jl:41
  [2] reshape(a::Array{Float64, 3}, dims::NTuple{4, Int64})
    @ Base ./reshapedarray.jl:45
  [3] setindex_disk!(::NCDatasets.Variable{Float64, 4, NCDatasets.NCDataset{Nothing, Missing}}, ::Array{Float64, 3}, ::Function, ::Vararg{Any})
    @ DiskArrays ~/.julia/packages/DiskArrays/bZBJE/src/diskarray.jl:56
  [4] setindex!
    @ ~/.julia/packages/DiskArrays/bZBJE/src/diskarray.jl:229 [inlined]
  [5] setindex!(::CommonDataModel.CFVariable{…}, ::Array{…}, ::Colon, ::Colon, ::Colon, ::UnitRange{…})
    @ CommonDataModel ~/.julia/packages/CommonDataModel/GGvMn/src/cfvariable.jl:477
  [6] save_output!(ds::NCDatasets.NCDataset{…}, output::Field{…}, model::NonhydrostaticModel{…}, ow::NetCDFOutputWriter{…}, time_index::Int64, name::String)
    @ Oceananigans.OutputWriters ~/repos/Oceananigans.jl/src/OutputWriters/netcdf_output_writer.jl:482
  [7] write_output!(ow::NetCDFOutputWriter{…}, model::NonhydrostaticModel{…})
    @ Oceananigans.OutputWriters ~/repos/Oceananigans.jl/src/OutputWriters/netcdf_output_writer.jl:525
  [8] initialize!(sim::Simulation{…})
    @ Oceananigans.Simulations ~/repos/Oceananigans.jl/src/Simulations/run.jl:212
  [9] time_step!(sim::Simulation{…})
    @ Oceananigans.Simulations ~/repos/Oceananigans.jl/src/Simulations/run.jl:112
 [10] run!(sim::Simulation{…}; pickup::Bool)
    @ Oceananigans.Simulations ~/repos/Oceananigans.jl/src/Simulations/run.jl:97
 [11] run!(sim::Simulation{NonhydrostaticModel{…}, Float64, Float64, OrderedCollections.OrderedDict{…}, OrderedCollections.OrderedDict{…}, OrderedCollections.OrderedDict{…}})
    @ Oceananigans.Simulations ~/repos/Oceananigans.jl/src/Simulations/run.jl:85
 [12] top-level scope
    @ ~/repos/Oceananigans.jl/sandbox/mwe.jl:31
 [13] include(fname::String)
    @ Base.MainInclude ./client.jl:489

So it's not as trivial as the single change I just made. From glancing at the code we at least have to modify initialize_nc_file!() to take a grid option as well:

https://github.com/CliMA/Oceananigans.jl/blob/6730e6f6b2c8f1695e20b95ef467b5b14fdc4c5f/src/OutputWriters/netcdf_output_writer.jl#L625-L636

plus a couple of other things. Still pretty easy, but more work/time that I have right now.

@iuryt feel free to jump in here and make these changes if you feel it's necessary, since creating a whole separate model can be a bit onerous and wastes precious GPU memory.

[glwagner]

Still pretty easy, but more work/time that I have right now.

Reading the code this is just one more line. Can you be more specific about how much more time this will take?

To clarify, I'm pretty sure the main difficulty of this PR is the documentation, not the code itself. The documentation was always going to take a bit of effort. I say this to give accurate impression of the effort this stuff requires.

[glwagner]

Ok, the above example works by adding grid as a property to the output writer, and then passing grid explicitly to the initialization.

Making this change required a couple minutes. The main technique was to search the file for model.grid, and make the necessary changes to use the user-provided grid instead. In this process I noticed that file initialization, which required the grid, has to occur outside the output writer constructor. This implies that that the "output writer grid" (which is now different from the model grid) must be stored within the output writer. So I added a grid property to the output writer.

These changes took about 5 minutes. However, the main work is still there, to document this and add tests and an example needed.

If there's any other source code changes needed I'm happy to put those in. The documentation will take more effort.

[tomchor]

Still pretty easy, but more work/time that I have right now.

Reading the code this is just one more line. Can you be more specific about how much more time this will take?

I wasn't specific about time because I wasn't sure. I often find myself spending way longer on PRs than I initially anticipate so I'm generally not very good at assessing these things lol

But I'm glad it's working now, pending docs. If the example above works without needing to create a second model, then I don't think anything else is needed on the coding part. @iuryt ?

[glwagner]

I wasn't specific about time because I wasn't sure. I often find myself spending way longer on PRs than I initially anticipate so I'm generally not very good at assessing these things lol

ok ok! Here's a rule of thumb: a bugfix or refactor is the least committment, because you can get away with no new tests or docs.

A new feature takes more time because of documentation. One should expect to spend most of their time on docs and tests. If you're spending a lot of time implementing a feature, then either the feature is very complicated / hard, or your workflow can use some improvement.

Among new features, different types of features require different amount of effort. This case is one of the easier cases --- it's extending functionality without changing existing functionality, also its fairly niche (at this point) so simple documentation will suffice. The work required to test the feature also has already been partially completed (the script you posted).

Other features, like adding new physics will take much longer, because often you'll have to run a full validation case + analysis to assess whether things work as expected. So even if the source code change is small to implement new physics, the validation will take a while and dominate the development effort.

If a new feature also requires changing existing functionality / refactoring, that's going to take the most amount of time, because you will probably also have to change existing tests. And many tests are poorly written, so updating test code is a potential rabbit hole.

[glwagner]

Tests pass, so documentation in your courts @iuryt @tomchor. I think for the docstring we can just say that the grid needs to be provided if outputs are on a different grid than model.

But it could be nice to also put an example in the docs, since I think the whole idea of output on a different grid from the model takes a bit to explain. Eg here: https://clima.github.io/OceananigansDocumentation/stable/model_setup/output_writers/

I can contribute a test if there's desire to put in the documentation.

[iuryt]

I can work on the docs

[iuryt]

Once the wording is finalized, I will edit the docstring.

[iuryt]

Let me know what else is missing or if we just need a third review to merge.

[tomchor]

Let me know what else is missing or if we just need a third review to merge.

Tests still aren't passing.

[glwagner]

Some commits got added that overwrote changes that I made

[glwagner]

Hmm ok somehow changes I made were lost, but I can't figure out how. Anyways I have tried to put them back in.

[iuryt]

Hmm ok somehow changes I made were lost, but I can't figure out how. Anyways I have tried to put them back in.

I was wondering the same. I noticed the error was the same we had before.

[iuryt]

Not sure why this is failing now.

[glwagner]

Might just need a retry -- gave it one. Thanks for the help!

[glwagner]

heck ya