"Dimension-aware" behavior of set! and interior with Flat directions?

[glwagner]

EDIT: We can't set the locations of all Field to Nothing. So I'm changing this issue topic to a discussion of whether we should support 1- and 2-argument functions for set! (when 2 or 1 directions are Flat), and also dropdims the Flat dimensions in interior.

Previously:
This would allow us to use two- or one-argument function specification in set!, as well as two- and two-dimensional indexing.

It might also make sense to dropdims Nothing locations when outputting raw field data... (it makes broadcasting with reduced fields, etc more difficult, but probably makes most activities easier, like plotting)... ?

This would automagically solve @francispoulin's pain provided that Flat vertical topologies are enforced in the constructor for ShallowWaterModel.

[francispoulin]

Great idea and I needless to say I'm in favour.

[glwagner]

@navidcy @ali-ramadhan any comments? This would change examples and API, because if Flat directions have Nothing locations then users will specify function initial conditions with 2D or 1D functions, rather than always using 3D functions even in 2D domains as it works now.

[navidcy]

I think this makes absolute sense. If one sets up a 2D grid then they can't be imposing 3D initial conditions/forcing functions. :)

This change will allow this if I understood correctly, right?

[glwagner]

Correct. Flat dimensions will have to be omitted from initial conditions and forcing functions if we make this change.

[glwagner]

@simone-silvestri thoughts about this? Major API change so maybe we want to do it sooner rather than later.

I'm not sure if this will affect ensemble models. I don't think it will but I'm not sure.

We probably want to wait until after #2246 .

[simone-silvestri]

Hmmm, will the underlying data of a field be 2D and/or 1D in that case?

I think that will be extremely intrusive in the code.
Might be nice in terms of API but will give rise to a lot of problems in the future and endless dispatch on the Flat dimension....

In conclusion... if you want to make it an API change (i.e. a change only in the location which will require an extension in set! and forcings) I am ok, but under the hood I would like to keep 3D arrays as data for the fields

[glwagner]

All data is always 3D, even when a field has Nothing location. See

Oceananigans.jl/src/Grids/new_data.jl

Lines 58 to 65 in b3ddbc8

	function new_data(FT, grid, loc)
	arch = architecture(grid)
	Tx = total_length(loc[1], topology(grid, 1), grid.Nx, grid.Hx)
	Ty = total_length(loc[2], topology(grid, 2), grid.Ny, grid.Hy)
	Tz = total_length(loc[3], topology(grid, 3), grid.Nz, grid.Hz)
	underlying_data = zeros(FT, arch, Tx, Ty, Tz)
	return offset_data(underlying_data, grid, loc)
	end

and

Oceananigans.jl/src/Grids/grid_utils.jl

Line 110 in b3ddbc8

@inline total_length(::Type{Nothing}, topo, N, H=0) = 1

In other words, fields with Nothing location have size 1.

This change therefore likely affects set! primarily. The inner kernels don't know about field location, and we already do endless dispatch on Flat:

There is an exception for Flat topology, because in this case N>1 is used for ensemble models; thus

Oceananigans.jl/src/Grids/grid_utils.jl

Lines 111 to 113 in b3ddbc8

	@inline total_length(::Type{Nothing}, ::Type{Flat}, N, H=0) = N
	@inline total_length(::Type{Face}, ::Type{Flat}, N, H=0) = N
	@inline total_length(::Type{Center}, ::Type{Flat}, N, H=0) = N

My main concern in fact for this change is how it might affect ensemble models. It shouldn't matter but we may have used a few cheats here and there regarding field location in Flat directions...

An alternative strategy, which might actually be better (now that I think about it), is to predicate alternative behavior for functions like set! and interior! on the grid topology, rather than changing the field location. We can just check for the case that topology(grid, i) === Flat and size(grid, i) == 1. We can also support both 3D and ND function input (eg try 3D, and then resort to ND?)

[glwagner]

I guess the way this goes down is via these functions:

Oceananigans.jl/src/Fields/function_field.jl

Line 64 in b3ddbc8

@inline (f::FunctionField)(x...) = call_func(f.clock, f.parameters, f.func, x...)

Oceananigans.jl/src/Fields/function_field.jl

Lines 57 to 59 in b3ddbc8

	# For setting ReducedField
	@inline call_func(::Nothing, ::Nothing, func, x, y) = func(x, y)
	@inline call_func(::Nothing, ::Nothing, func, x) = func(x)

that's why

Oceananigans.jl/examples/geostrophic_adjustment.jl

Lines 73 to 77 in b3ddbc8

	ηⁱ(x, y) = 2 * ηᵍ(x)
	# We set the initial condition to ``vᵍ`` and ``ηⁱ``,
	set!(model, v=vᵍ, η=ηⁱ)

is 2D for the free surface.

This'll affect #2246 .

[glwagner]

Ok, now I realize that we actually can't do this because it will probably prevent GPU compilation. ReducedField keep their location on the GPU, but we can't keep the locations of all fields on the GPU.

I'll change this issue to a discussion of behavior changes for set! and interior with Flat directions.

[glwagner]

We changed set!.