Unable to output η using NetCDFWriter

[HXC303]

I’m trying to output the free surface η using NetCDFWriter, but I encountered an error. Below is an example:

using Oceananigans
using NCDatasets

grid = RectilinearGrid(size = (20, 40, 100), 
                       extent=(4e5, 8e5, 2e3), 
                       topology = (Periodic, Bounded, Bounded))

model = HydrostaticFreeSurfaceModel(;           grid,
                            buoyancy = BuoyancyTracer(),
                             tracers = :b)

b0=rand(Float64, 100)
using LinearAlgebra
b0=repeat(reshape(b0,(1,1,100)),20,40,1);

set!(model, b=b0)

simulation = Simulation(model, Δt=120, stop_time=24*60*60)

fields = Dict("η" => model.free_surface.η);
simulation.output_writers[:field_writerETA] =
    NetCDFWriter(model, fields, filename="eta.nc", 
    schedule=TimeInterval(60*60), array_type = Array{Float64})

run!(simulation)

This results in the error:

ArgumentError: view indices (1:20, 1:40, 1:101) do not intersect field indices (Colon(), Colon(), 101:101)

I also tried specifying: indices=(:,:,grid.Nz+1)，but then got:

ArgumentError: view indices (1:1, 1:1, 101:100) do not intersect field indices (Colon(), Colon(), 101:101)

I attempted with_halos=true, but that led to:

[ Info: Initializing simulation...
DimensionMismatch: new dimensions (26, 46, 107, 1) must be consistent with array size 1196

How can I correctly output the η variable with NetCDFWriter?

[ali-ramadhan]

Hi @HXC303! Apologies for the oversight here. We recently refactored NetCDFWriter and I need to improve how it handles windowed fields like model.free_surface.η.

For now you can output Average(model.free_surface.η, dims=3). The tests show examples of how to do this. In particular see test_netcdf_free_surface_only_output and test_netcdf_free_surface_mixed_output:

Oceananigans.jl/test/test_netcdf_writer.jl

Lines 2530 to 2645 in a84b3f8

	function test_netcdf_free_surface_only_output(arch)
	Nλ, Nφ, Nz = 8, 8, 4
	Hλ, Hφ, Hz = 3, 4, 2
	grid = LatitudeLongitudeGrid(arch;
	topology = (Bounded, Bounded, Bounded),
	size = (Nλ, Nφ, Nz),
	halo = (Hλ, Hφ, Hz),
	longitude = (-1, 1),
	latitude = (-1, 1),
	z = (-100, 0)
	)
	model = HydrostaticFreeSurfaceModel(; grid,
	closure = ScalarDiffusivity(ν=4e-2, κ=4e-2),
	buoyancy = SeawaterBuoyancy(),
	tracers = (:T, :S)
	)
	Nt = 5
	simulation = Simulation(model, Δt=0.1, stop_iteration=Nt)
	# Kind of a hack because we want η to be a ReducedField.
	outputs = (;
	η = Average(model.free_surface.η, dims=3)
	)
	Arch = typeof(arch)
	filepath_with_halos = "test_free_surface_with_halos_$Arch.nc"
	isfile(filepath_with_halos) && rm(filepath_with_halos)
	simulation.output_writers[:with_halos] =
	NetCDFWriter(model, outputs;
	filename = filepath_with_halos,
	schedule = IterationInterval(1),
	with_halos = true)
	filepath_no_halos = "test_free_surface_no_halos_$Arch.nc"
	isfile(filepath_no_halos) && rm(filepath_no_halos)
	simulation.output_writers[:no_halos] =
	NetCDFWriter(model, outputs;
	filename = filepath_no_halos,
	schedule = IterationInterval(1),
	with_halos = false)
	run!(simulation)
	ds_h = NCDataset(filepath_with_halos)
	@test haskey(ds_h, "η")
	@test dimsize(ds_h["η"]) == (λ_caa=Nλ + 2Hλ, φ_aca=Nφ + 2Hφ, time=Nt + 1)
	close(ds_h)
	rm(filepath_with_halos)
	ds_n = NCDataset(filepath_no_halos)
	@test haskey(ds_n, "η")
	@test dimsize(ds_n["η"]) == (λ_caa=Nλ, φ_aca=Nφ, time=Nt + 1)
	close(ds_n)
	rm(filepath_no_halos)
	return nothing
	end
	function test_netcdf_free_surface_mixed_output(arch)
	Nλ, Nφ, Nz = 8, 8, 4
	Hλ, Hφ, Hz = 3, 4, 2
	grid = LatitudeLongitudeGrid(arch;
	topology = (Bounded, Bounded, Bounded),
	size = (Nλ, Nφ, Nz),
	halo = (Hλ, Hφ, Hz),
	longitude = (-1, 1),
	latitude = (-1, 1),
	z = (-100, 0)
	)
	model = HydrostaticFreeSurfaceModel(; grid,
	closure = ScalarDiffusivity(ν=4e-2, κ=4e-2),
	buoyancy = SeawaterBuoyancy(),
	tracers = (:T, :S)
	)
	Nt = 5
	simulation = Simulation(model, Δt=0.1, stop_iteration=Nt)
	# Kind of a hack because we want η to be a ReducedField.
	free_surface_outputs = (;
	η = Average(model.free_surface.η, dims=3)
	)
	outputs = merge(model.velocities, model.tracers, free_surface_outputs)
	Arch = typeof(arch)
	filepath_with_halos = "test_mixed_free_surface_with_halos_$Arch.nc"
	isfile(filepath_with_halos) && rm(filepath_with_halos)
	simulation.output_writers[:with_halos] =
	NetCDFWriter(model, outputs;
	filename = filepath_with_halos,
	schedule = IterationInterval(1),
	with_halos = true)
	filepath_no_halos = "test_mixed_free_surface_no_halos_$Arch.nc"
	isfile(filepath_no_halos) && rm(filepath_no_halos)
	simulation.output_writers[:no_halos] =
	NetCDFWriter(model, outputs;
	filename = filepath_no_halos,
	schedule = IterationInterval(1),
	with_halos = false)
	run!(simulation)

[glwagner]

Now that's what I call a workaround

What is the challenge here? Is it that we can't properly compute the intersection of indices?

[ali-ramadhan]

I don't think there's really a challenge, just an oversight on my part in that I forgot to fix this behavior by dispatching on WindowedField in PR #4046 which should take care of fields like model.free_surface.η. Instead I left the hack in the test and it kept passing lol.

I suppose this was antithetical to test-driven development. I should have let the test keep failing so that I actually fix the code!

[glwagner]

Well, that is good news then that ultimately it's a simple thing to fix!

[navidcy]

(#3260 might be related; at least the initial motivation was the same: trying to save free surface wasn't working)

[glwagner]

Does #4401 fix it?

[glwagner]

Ok it doesn't fix it for me -- basically NetCDFWriter does not seem to support windowed fields, if I am not mistaken. cc @tomchor @ali-ramadhan

[glwagner]

this looks bad to me:

Oceananigans.jl/ext/OceananigansNCDatasetsExt.jl

Lines 57 to 65 in bd708ef

	function collect_dim(ξ, ℓ, T, N, H, inds, with_halos)
	if with_halos
	return collect(ξ)
	else
	inds = validate_index(inds, ℓ, T, N, H)
	inds = restrict_to_interior(inds, ℓ, T, N)
	return collect(view(ξ, inds))
	end
	end

shouldn't we be using output_indices from the output writers? Otherwise things are not guaranteed to be consistent (note repeated calls / code dupcliation):

Oceananigans.jl/src/OutputWriters/output_construction.jl

Lines 33 to 43 in bd708ef

	function output_indices(output::Union{AbstractField, Reduction}, grid, indices, with_halos)
	indices = validate_indices(indices, location(output), grid)
	if !with_halos # Maybe chop those indices
	loc = map(instantiate, location(output))
	topo = map(instantiate, topology(grid))
	indices = map(restrict_to_interior, indices, loc, topo, size(grid))
	end
	return indices
	end

[ali-ramadhan]

I'm starting to look into this issue. Indeed not supporting WindowedField is an oversight! But it seems like the free surface field is not a WindowedField?

using Oceananigans
using Oceananigans.Fields: WindowedField

grid = LatitudeLongitudeGrid(CPU();
    topology = (Bounded, Bounded, Bounded),
    size = (8, 8, 4),
    halo = (3, 4, 2),
    longitude = (-1, 1),
    latitude = (-1, 1),
    z = (-100, 0)
)

model = HydrostaticFreeSurfaceModel(; grid,
    closure = ScalarDiffusivity(ν=4e-2, κ=4e-2),
    buoyancy = SeawaterBuoyancy(),
    tracers = (:T, :S)
)

julia> model.free_surface.η isa WindowedField
false

PS: Agree that we should be using output_indices instead of collect_dim for NetCDFWriter. I don't remember if I missed output_indices or was attempting a workaround of something, but it should be easy to switch over and check to see if all the tests still pass.

[ali-ramadhan]

Based on these definitions:

Oceananigans.jl/src/Fields/field.jl

Lines 21 to 35 in dfe0b6e

	struct Field{LX, LY, LZ, O, G, I, D, T, B, S, F} <: AbstractField{LX, LY, LZ, G, T, 3}
	grid :: G
	data :: D
	boundary_conditions :: B
	indices :: I
	operand :: O
	status :: S
	boundary_buffers :: F
	# Inner constructor that does not validate _anything_!
	function Field{LX, LY, LZ}(grid::G, data::D, bcs::B, indices::I, op::O, status::S, buffers::F) where {LX, LY, LZ, G, D, B, O, S, I, F}
	T = eltype(data)
	return new{LX, LY, LZ, O, G, I, D, T, B, S, F}(grid, data, bcs, indices, op, status, buffers)
	end
	end

Oceananigans.jl/src/Fields/field.jl

Lines 357 to 358 in dfe0b6e

	const WindowedData = OffsetArray{<:Any, <:Any, <:SubArray}
	const WindowedField = Field{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:WindowedData}

shouldn't a WindowedField be defined as

const WindowedField = Field{<:Any, <:WindowedData, <:Any, <:Any, <:Any, <:Any, <:Any} 

Looking through the tests, no test explicitly tests for this and it seems like reductions on WindowedFields still works correctly.

---

EDIT: This post is wrong. It's Field{LX, LY, LZ, O, G, I, D, ...} so

const WindowedField = Field{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:WindowedData} 

is correct.

[ali-ramadhan]

Hmmm, maybe the distinction is that the free surface has at least one UnitRange in its indices instead of all Colon. The data is still a regular OffsetArray for both.

julia> typeof(model.velocities.u.data)
OffsetArrays.OffsetArray{Float64, 3, Array{Float64, 3}}

julia> typeof(model.free_surface.η.data)
OffsetArrays.OffsetArray{Float64, 3, Array{Float64, 3}}

but

julia> typeof(model.velocities.u.indices)
Tuple{Colon, Colon, Colon}

julia> typeof.(model.free_surface.η.indices)
(Colon, Colon, UnitRange{Int64})

So if model.free_surface.η is not currently a WindowedField then should we be updating the definition of a WindowedField to consider the the type of indices? Or is model.free_surface.η a different kind of field that warrants a new definition?

If we can decide what type of field model.free_surface.η is, then I think fixing this issue is just a matter of dispatching on it in fetch_output to account for the field.indices.

[glwagner]

I think you're right, it looks like WindowedField is defined incorrectly! It is making an assumption that the windowed field describes a view into another, non-windowed field. But we can also directly create windowed fields (the case for the free surface). So we really need to look at indices as a definition of windowedness.

[ali-ramadhan]

I don't really see WindowedField being used anywhere in the code except for the definition, a broadcasting method, and these tests:

Oceananigans.jl/test/test_field.jl

Lines 477 to 492 in dfe0b6e

	@info " Test reductions on WindowedFields [$(typeof(arch)), $FT]..."
	grid = RectilinearGrid(arch, FT, size=(2, 3, 4), x=(0, 1), y=(0, 1), z=(0, 1))
	c = CenterField(grid)
	Random.seed!(42)
	set!(c, rand(size(c)...))
	windowed_c = view(c, :, 2:3, 1:2)
	for fun in (sum, maximum, minimum)
	@test fun(c) ≈ fun(interior(c))
	@test fun(windowed_c) ≈ fun(interior(windowed_c))
	end
	@test mean(c) ≈ CUDA.@allowscalar mean(interior(c))
	@test mean(windowed_c) ≈ CUDA.@allowscalar mean(interior(windowed_c))

Should we define a new type of Field to classify model.free_surface.η or should we change the definition of a WindowedField to be something like

const XWindowedIndices = Tuple{<:UnitRange, Colon, Colon}
const YWindowedIndices = Tuple{Colon, <:UnitRange, Colon}
const ZWindowedIndices = Tuple{Colon, Colon, <:UnitRange}
const XYWindowedIndices = Tuple{<:UnitRange, <:UnitRange, Colon}
const XZWindowedIndices = Tuple{<:UnitRange, Colon, <:UnitRange}
const YZWindowedIndices = Tuple{Colon, <:UnitRange, <:UnitRange}
const XYZWindowedIndices = Tuple{<:UnitRange, <:UnitRange, <:UnitRange}
const WindowedIndices = Union{XWindowedIndices, YWindowedIndices, ZWindowedIndices, XYWindowedIndices, XZWindowedIndices, YZWindowedIndices, XYZWindowedIndices}
const WindowedField = Field{<:Any, <:Any, <:Any, <:Any, <:Any, <:WindowedIndices}

?

I guess either way something like this definition using WindowedIndices is what we need to classify the free surface field?

[glwagner]

If we don't need the definition then its best not to introduce needless code, perhaps?

[ali-ramadhan]

Hmmm, the way I see it it's either dispatch on a definition or do if statements on field.indices to treat fields like model.free_surface.η in fetch_output differently. I know we tend to prefer multiple dispatch.

[glwagner]

That makes sense, just trying to understand. Where in fetch_output do we need to dispatch on the type of field.indices?

[glwagner]

It may be easier / cleaner to special case the non-windowed case with Tuple{Colon, Colon, Colon} and then treat all other cases as windowed