Oceananigans.jl Can we calculate a `Field` at every `N` grid points?

I'd like to compute a Field (ideally in order to write it to a NetCDF file) but only at every N grid points. Something like the following example, which tries to compute u at every 2 grid points in the vertical direction:

using Oceananigans
grid = RectilinearGrid(size = (1, 1, 8), extent = (1,1,1));
model = NonhydrostaticModel(; grid,)
u_slices = Field(model.velocities.u, indices=(:, :, 1:2:grid.Nz))

However, the above code fails since Field currently doesn't accept StepRanges as indices, just (I think) UnitRanges and Ints:

ERROR: LoadError: MethodError: no method matching isinteger(::StepRange{Int64, Int64})

Closest candidates are:
  isinteger(::Integer)
   @ Base number.jl:20
  isinteger(::Complex)
   @ Base complex.jl:148
  isinteger(::Rational)
   @ Base rational.jl:281
  ...

Stacktrace:
 [1] validate_index(idx::StepRange{Int64, Int64}, loc::Center, topo::Bounded, N::Int64, H::Int64)
   @ Oceananigans.Grids ~/repos/Oceananigans.jl/src/Grids/input_validation.jl:196
 [2] map (repeats 3 times)
   @ ./tuple.jl:318 [inlined]
 [3] validate_indices(indices::Tuple{Colon, Colon, StepRange{Int64, Int64}}, loc::Tuple{DataType, DataType, DataType}, topo::Tuple{DataType, DataType, DataType}, sz::Tuple{Int64, Int64, Int64}, halo_sz::Tuple{Int64, Int64, Int64})

Is there a workaround?

CC @iuryt

Feb 06 '24 18:02 tomchor

I think the best way to implement this is to write a function that returns a subsampled array, rather than trying to use the Field machinery.

If you would like to preserve grid information in the file automatically, you can generate the subsampled grid + subsampled Field, and then write a function that computes the subsampled Field.

Another way to implement this is to use interpolate or the conservative regrid! from the fine to coarse field.

I think the concept of a field with subsampled indices is a little tricky. For example, what should we do if we try to index that field in cells that lie between the existing data?

Feb 06 '24 21:02 glwagner

Hi! Thanks for the explanation and reply. Do you have any example that uses interpolate of regrid! I could use as a base for our case?

Feb 14 '24 17:02 iuryt

both are called with the syntax

interpolate!(target_field, source_field)
regrid!(target_field, source_field)

I think there are docstrings but I could be wrong. Can you check the docstring to see if it is sufficient?

Feb 14 '24 20:02 glwagner

both are called with the syntax
interpolate!(target_field, source_field)
regrid!(target_field, source_field)
I think there are docstrings but I could be wrong. Can you check the docstring to see if it is sufficient?

Yes, there are docstrings. Sorry. Just to make sure I understand, can you explain the difference between regridding and interpolating? I noticed that regrid also interpolates.

Feb 14 '24 23:02 iuryt

"regridding" is conservative, ie the volume integral of the target field should be the same as the source field. However because of that, we can only regrid in one direction at a time right now. If you want to regrid in multiple directions, you need to form intermediate fields. General conservative is actually possible, just a bit more difficult.

interpolate! does not conserve global integrals. It's just simple linear interpolation. Usually interpolate! is good enough.

Feb 15 '24 15:02 glwagner

Sorry for the long delay.

I came back to this problem and could make it work like this.

# subset grid for 20 times less vertical nodes
subset_grid = RectilinearGrid(arch,
    size = (grid.Nx, grid.Ny, div(grid.Nz, 20)),
    extent = (grid.Lx, grid.Ly, grid.Lz)
)

subset_fields = NamedTuple( key=>Field(location(output_fields[key]), subset_grid) for key in keys(output_fields) )

for key in keys(subset_fields)
    interpolate!(subset_fields[key], output_fields[key])
end

But this code seems too verbose and I am not sure if the interpolated data will be updated at each save. Does that seem right?

Apr 01 '24 21:04 iuryt

Basically but to make this work with output writers you need

u_subsampled = XFaceField(subset_grid)

function subsample_u(model)
    u = model.velocities.u
    interpolate!(u_subsampled, u)
    return u_subsampled
end

outputs = (; u=subsample_u)

Apr 01 '24 23:04 glwagner

u_subsampled = XFaceField(subset_grid)

function subsample_u(model)
    u = model.velocities.u
    interpolate!(u_subsampled, u)
    return u_subsampled
end

outputs = (; u=subsample_u)

I see, so there is no way to make this seamlessly work with a NamedTuple I already have with all variables I want to output? I would need to create a function for each of the variables?

Apr 02 '24 15:04 iuryt

I'll let @glwagner comment on the possibility of an easier way, but worst case scenario, there's always the possibility of using metaprogramming to automate the creation of these functions. For example as in https://docs.julialang.org/en/v1/manual/metaprogramming/#Code-Generation

Oceananigans uses that approach in a few places too.

Apr 02 '24 16:04 tomchor

Basically but to make this work with output writers you need

u_subsampled = XFaceField(subset_grid)

function subsample_u(model)
    u = model.velocities.u
    interpolate!(u_subsampled, u)
    return u_subsampled
end

outputs = (; u=subsample_u)

I tried

subset_grid = RectilinearGrid(arch,
    size = (grid.Nx, grid.Ny, div(grid.Nz, 20)),
    extent = (grid.Lx, grid.Ly, grid.Lz)
)

u_subsampled = XFaceField(subset_grid)
v_subsampled = YFaceField(subset_grid)
w_subsampled = ZFaceField(subset_grid)

function subsample_u(model)
    u = model.velocities.u
    interpolate!(u_subsampled, u)
    return u_subsampled
end

function subsample_v(model)
    v = model.velocities.v
    interpolate!(v_subsampled, v)
    return v_subsampled
end

function subsample_w(model)
    w = model.velocities.w
    interpolate!(w_subsampled, w)
    return w_subsampled
end



subset_outputs = (; u = subsample_u, v = subsample_v, w = subsample_w)

and

simulation.output_writers[:xyz] = NetCDFOutputWriter(model, subset_outputs,
                                                    schedule = TimeInterval(300),
                                                    filename = "test.nc",
                                                    with_halos = false,
                                                    array_type = Array{Float32},
                                                    )

is returning the following error.

Custom output v needs dimensions!

Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] define_output_variable!(dataset::NCDatasets.NCDataset{Nothing}, output::Function, name::String, array_type::Type, deflatelevel::Int64, output_attributes::Dict{String, String}, dimensions::Dict{Any, Any})
   @ Oceananigans.OutputWriters ~/.julia/packages/Oceananigans/3LHMs/src/OutputWriters/netcdf_output_writer.jl:448
 [3] NetCDFOutputWriter(model::NonhydrostaticModel{Oceananigans.TimeSteppers.RungeKutta3TimeStepper{Float64, NamedTuple{(:u, :v, :w, :b), Tuple{Field{Face, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}, Field{Center, Face, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}, Field{Center, Center, Face, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, Nothing, Nothing, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}, Field{Center, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}}}, Nothing}, AnisotropicMinimumDissipation{Oceananigans.TurbulenceClosures.ExplicitTimeDiscretization, NamedTuple{(:b,), Tuple{Float64}}, Float64, Nothing}, GPU, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Float64, Buoyancy{BuoyancyTracer, Oceananigans.Grids.NegativeZDirection}, FPlane{Float64}, UniformStokesDrift{Nothing, typeof(∂z_uˢ), typeof(Oceananigans.StokesDrifts.zerofunction), typeof(Oceananigans.StokesDrifts.zerofunction), typeof(Oceananigans.StokesDrifts.zerofunction)}, NamedTuple{(:u, :v, :w), Tuple{Field{Face, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Float64}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}, Field{Center, Face, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}, Field{Center, Center, Face, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Open, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Open, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}}}, NamedTuple{(:b,), Tuple{Field{Center, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Gradient, Float64}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Float64}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}}}, NamedTuple{(:pHY′, :pNHS), Tuple{Field{Center, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}, Field{Center, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}}}, NamedTuple{(:u, :v, :w, :b), Tuple{Oceananigans.Forcings.ContinuousForcing{Face, Center, Center, NamedTuple{(:Nx, :Ny, :Nz, :Lx, :Ly, :Lz, :constant_wind_stress, :surface_buoyancy_flux, :amplitude, :wavelength, :wavenumber, :frequency, :Uˢ, :N², :N2, :Qᵘ, :u_star, :La_t, :Qʰ, :Qᵇ, :initial_mixed_layer_depth, :w_star, :Λ, :sponge_depth, :sponge_bottom, :sponge_top, :σ), Tuple{Int64, Int64, Int64, Int64, Int64, Int64, Int64, Int64, Float64, Int64, Float64, Float64, Float64, Float64, Float64, Float64, Float64, Float64, Int64, Int64, Int64, Vararg{Float64, 6}}}, typeof(sponge_u), Tuple{Symbol}, Tuple{Int64}, Tuple{typeof(Oceananigans.Operators.identity4)}}, Oceananigans.Forcings.ContinuousForcing{Center, Face, Center, NamedTuple{(:Nx, :Ny, :Nz, :Lx, :Ly, :Lz, :constant_wind_stress, :surface_buoyancy_flux, :amplitude, :wavelength, :wavenumber, :frequency, :Uˢ, :N², :N2, :Qᵘ, :u_star, :La_t, :Qʰ, :Qᵇ, :initial_mixed_layer_depth, :w_star, :Λ, :sponge_depth, :sponge_bottom, :sponge_top, :σ), Tuple{Int64, Int64, Int64, Int64, Int64, Int64, Int64, Int64, Float64, Int64, Float64, Float64, Float64, Float64, Float64, Float64, Float64, Float64, Int64, Int64, Int64, Vararg{Float64, 6}}}, typeof(sponge_v), Tuple{Symbol}, Tuple{Int64}, Tuple{typeof(Oceananigans.Operators.identity5)}}, Oceananigans.Forcings.ContinuousForcing{Center, Center, Face, NamedTuple{(:Nx, :Ny, :Nz, :Lx, :Ly, :Lz, :constant_wind_stress, :surface_buoyancy_flux, :amplitude, :wavelength, :wavenumber, :frequency, :Uˢ, :N², :N2, :Qᵘ, :u_star, :La_t, :Qʰ, :Qᵇ, :initial_mixed_layer_depth, :w_star, :Λ, :sponge_depth, :sponge_bottom, :sponge_top, :σ), Tuple{Int64, Int64, Int64, Int64, Int64, Int64, Int64, Int64, Float64, Int64, Float64, Float64, Float64, Float64, Float64, Float64, Float64, Float64, Int64, Int64, Int64, Vararg{Float64, 6}}}, typeof(sponge_w), Tuple{Symbol}, Tuple{Int64}, Tuple{typeof(Oceananigans.Operators.identity1)}}, Oceananigans.Forcings.ContinuousForcing{Center, Center, Center, NamedTuple{(:Nx, :Ny, :Nz, :Lx, :Ly, :Lz, :constant_wind_stress, :surface_buoyancy_flux, :amplitude, :wavelength, :wavenumber, :frequency, :Uˢ, :N², :N2, :Qᵘ, :u_star, :La_t, :Qʰ, :Qᵇ, :initial_mixed_layer_depth, :w_star, :Λ, :sponge_depth, :sponge_bottom, :sponge_top, :σ), Tuple{Int64, Int64, Int64, Int64, Int64, Int64, Int64, Int64, Float64, Int64, Float64, Float64, Float64, Float64, Float64, Float64, Float64, Float64, Int64, Int64, Int64, Vararg{Float64, 6}}}, typeof(sponge_b), Tuple{Symbol}, Tuple{Int64}, Tuple{typeof(Oceananigans.Operators.identity2)}}}}, WENO{3, Float64, Nothing, Nothing, Nothing, true, Nothing, WENO{2, Float64, Nothing, Nothing, Nothing, true, Nothing, UpwindBiased{1, Float64, Nothing, Nothing, Nothing, Nothing, Centered{1, Float64, Nothing, Nothing, Nothing, Nothing}}, Centered{1, Float64, Nothing, Nothing, Nothing, Nothing}}, Centered{2, Float64, Nothing, Nothing, Nothing, Centered{1, Float64, Nothing, Nothing, Nothing, Nothing}}}, Oceananigans.Solvers.FFTBasedPoissonSolver{RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, NamedTuple{(:λx, :λy, :λz), Tuple{CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}}}, CuArray{ComplexF64, 3, CUDA.Mem.DeviceBuffer}, CuArray{ComplexF64, 3, CUDA.Mem.DeviceBuffer}, NamedTuple{(:forward, :backward), Tuple{Tuple{Oceananigans.Solvers.DiscreteTransform{CUDA.CUFFT.cCuFFTPlan{ComplexF64, -1, true, 3}, Oceananigans.Solvers.Forward, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Int64, Bounded, Int64, NamedTuple{(:forward, :backward), Tuple{CuArray{ComplexF64, 3, CUDA.Mem.DeviceBuffer}, CuArray{ComplexF64, 3, CUDA.Mem.DeviceBuffer}}}, Nothing}, Oceananigans.Solvers.DiscreteTransform{CUDA.CUFFT.cCuFFTPlan{ComplexF64, -1, true, 3}, Oceananigans.Solvers.Forward, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Vector{Int64}, Vector{Periodic}, Int64, Nothing, Nothing}}, Tuple{Oceananigans.Solvers.DiscreteTransform{AbstractFFTs.ScaledPlan{ComplexF64, CUDA.CUFFT.cCuFFTPlan{ComplexF64, 1, true, 3}, Float64}, Oceananigans.Solvers.Backward, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Vector{Int64}, Vector{Periodic}, Int64, Nothing, Nothing}, Oceananigans.Solvers.DiscreteTransform{AbstractFFTs.ScaledPlan{ComplexF64, CUDA.CUFFT.cCuFFTPlan{ComplexF64, 1, true, 3}, Float64}, Oceananigans.Solvers.Backward, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Int64, Bounded, Int64, NamedTuple{(:forward, :backward), Tuple{CuArray{ComplexF64, 3, CUDA.Mem.DeviceBuffer}, CuArray{ComplexF64, 3, CUDA.Mem.DeviceBuffer}}}, Nothing}}}}}, NamedTuple{(:νₑ, :κₑ), Tuple{Field{Center, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}, NamedTuple{(:b,), Tuple{Field{Center, Center, Center, Nothing, RectilinearGrid{Float64, Periodic, Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, GPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}}, Float64, FieldBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Periodic, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}}, Nothing, Oceananigans.Fields.FieldBoundaryBuffers{Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}}}}}, NamedTuple{(:velocities, :tracers), Tuple{NamedTuple{(:u, :v, :w), Tuple{Oceananigans.Fields.ZeroField{Int64, 3}, Oceananigans.Fields.ZeroField{Int64, 3}, Oceananigans.Fields.ZeroField{Int64, 3}}}, NamedTuple{(:b,), Tuple{Oceananigans.Fields.ZeroField{Int64, 3}}}}}, Nothing, Nothing, Nothing, NamedTuple{(), Tuple{}}}, outputs::NamedTuple{(:u, :v, :w), Tuple{typeof(subsample_u), typeof(subsample_v), typeof(subsample_w)}}; filename::String, schedule::TimeInterval, dir::String, array_type::Type{Array{Float32}}, indices::Tuple{Colon, Colon, Colon}, with_halos::Bool, global_attributes::Dict{Any, Any}, output_attributes::Dict{Any, Any}, dimensions::Dict{Any, Any}, overwrite_existing::Nothing, deflatelevel::Int64, verbose::Bool)
   @ Oceananigans.OutputWriters ~/.julia/packages/Oceananigans/3LHMs/src/OutputWriters/netcdf_output_writer.jl:422
 [4] top-level scope
   @ In[8]:1

Apr 02 '24 16:04 iuryt

I think if you're passing anything to NetCDFWriter that isn't a Field it can't infer dimensions from it. So you need to pass dimensions manually.

This example is pretty illustrative, but basically you need to pass a dimensions flag that's a Dict (or possibly a NamedTuple) with dimensions.

Apr 02 '24 17:04 tomchor

Heh. I would never imply you can't take my code and adapt it to do nice things. You can of course do that.

For example if you have outputs already assembled, you can do

function create_subsampled_output(field)
    loc = location(field)
    subsampled_field = Field(loc, subsampled_grid)
    function output(model)
        interpolate!(subsampled_field, field)
        return subsampled_field
    end
    return output
end

subsampled_outputs = NamedTuple(name => create_subsampled_output(outputs[name]) for name in keys(outputs))

And there's an infinity of other ways you can write your code.

Apr 02 '24 20:04 glwagner

It wouldn't be crazy to add this kind of feature to the output writers directly --- for convenience, since obviously it's not strictly necessary. Certainly, we'll need such conveniences for global simulations. All in due time naturally, I feel this "manual" way is not all that hard either and subsampling output sort of implies you're doing something slightly beyond the standard.

Apr 02 '24 20:04 glwagner

I'll let @glwagner comment on the possibility of an easier way, but worst case scenario, there's always the possibility of using metaprogramming to automate the creation of these functions. For example as in https://docs.julialang.org/en/v1/manual/metaprogramming/#Code-Generation

Oceananigans uses that approach in a few places too.

My suggestion above uses a closure over subsampled_field. Possibly, one could also use a macro to achieve the same thing.

Apr 02 '24 21:04 glwagner

Thanks, @glwagner!! I will test that and see how it goes. I really think it would be nice for us to develop this functionality in for the writers. Saving a subset of the simulation is pretty common when running high-res simulations.

Apr 02 '24 21:04 iuryt

Ok, then we need to think about the user interface. We could add a keyword argument output_grid to the output writer, or perhaps a positional argument as in

JLD2OutputWriter(model, outputs, grid; kw...)

I think a feature that might be useful for implementing this feature is an InterpolatedField, which looks something like

struct InterpolatedField
    grid
    data
    interpoland
end

function compute!(interp::InterpolatedField)
    compute!(interp.interpoland)
    interpolate!(interp, interp.interpoland)
    return interp
end

Then we could use this feature in construct_output:

https://github.com/CliMA/Oceananigans.jl/blob/main/src/OutputWriters/output_construction.jl

and maybe elsewhere.

This is a chunk of work of course. I don't need it right now personally.

Apr 02 '24 21:04 glwagner

Just for reference, I created the MWE example below of a column model with a sheared u that evolves in time, where I write both the full u and an interpolated u to half the original resolution. What I'm doing here is creating a coarse grid, coarse model, and coarse field that takes the interpolated u. I use a Callback to keep interpolating u into coarse_u and then it's only a matter of setting up a coarse NetCDFWriter with coarse_model.

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_model = NonhydrostaticModel(; grid=coarse_grid, closure = ScalarDiffusivity(ν=1e-2))
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(coarse_model, (; coarse_u,),
                                                              filename="coarsefields.nc",
                                                              schedule=TimeInterval(5),
                                                              overwrite_existing=true,)

run!(simulation)

This seems to be working. In the figure below the lines and triangles come from the full res output, while the cirlces come from the coarse (interpolated) output:

You can definitely take nicer/fancier approaches that scale much more easily (@glwagner mentioned a few options), but I feel like this is very simple to understand and can be scaled up to a list of outputs reasonably easily; you just gotta create the list of coarsened outputs based on the original list of outputs and interpolate the list in the Callback.

@iuryt is this good enough for your purposes?

Also a quick note: if you want outputs at levels that aren't equally spaced from each other, you can set-up coarse_grid using an array of depths as the vertical coordinate.

May 02 '24 15:05 tomchor

Ah interesting, why do you create coarse_model? Is that so coarse_model.grid gets saved down? I think we can make it so that isn't needed probably...

May 02 '24 16:05 glwagner

Ah interesting, why do you create coarse_model? Is that so coarse_model.grid gets saved down? I think we can make it so that isn't needed probably...

Right now it's needed solely because NetCDFWriter needs a model to get the grid from (and maybe other stuff). I think it would be pretty easy to modify NetCDFWriter so that it only needs a grid passed, but I haven't looked into that.

May 02 '24 16:05 tomchor

You can definitely take nicer/fancier approaches that scale much more easily (@glwagner mentioned a few options)

Note that I was also suggesting source code changes that could be contributed to stream line this (not just fancy things that are hard to understand 😆)! A source code contribution is of course the best option, we must all agree 😄

May 02 '24 16:05 glwagner

Ah interesting, why do you create coarse_model? Is that so coarse_model.grid gets saved down? I think we can make it so that isn't needed probably...

Right now it's needed solely because NetCDFWriter needs a model to get the grid from (and maybe other stuff). I think it would be pretty easy to modify NetCDFWriter so that it only needs a grid passed, but I haven't looked into that.

I think just modifying the constructor for the output writer to

NetCDFOutputWriter(model, outputs, grid=model.grid; kw...)

would do it.

The next level would be to look at the grid for the outputs and automagically build interpolated outputs if the grids are different (but that could be added separately from the above change)

May 02 '24 16:05 glwagner

Thanks, @tomchor ! Is that a reason why you define the closure for the coarse_model?

Until now, what I was doing was something like this

for i = 1:10:model.Nz
    key = Symbol("xy", i)  # Create the dictionary key dynamically
    fname = "vxy_z$(@sprintf("%05d", i)).nc"
    simulation.output_writers[key] = NetCDFOutputWriter(model, output_fields,
                                                        schedule = TimeInterval(output_interval),
                                                        filename = fname,
                                                        indices = (:,:,i),
                                                        with_halos = false,
                                                        overwrite_existing = overwrite_existing,
                                                        array_type = Array{Float32})
end

Which creates a file for each subset level.

While I think that @tomchor solution is the best because it is more general and can be for any arbitrary new grid, I still think we should also be able to simply pass the indices to NetCDFOutputWriter. For example, this should work

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 coarse output
simulation.output_writers[:coarsefields] = NetCDFOutputWriter(model, (; model.velocities.u),
                                                              filename="coarsefields.nc",
                                                              schedule=TimeInterval(5),
                                                              indices = (:,:,1:10:model.Nz),
                                                              overwrite_existing=true,)

run!(simulation)

May 02 '24 17:05 iuryt

Is that a reason why you define the closure for the coarse_model?

There reason is that at the moment NetCDFWriter needs it to get some info on the grid. But as @glwagner and I pointed out, it's probably pretty easy to change NetCDFWriter to avoid that. I might try a PR soon that makes the simplest change possible and see it tests pass.

May 02 '24 17:05 tomchor

Created a PR in https://github.com/CliMA/Oceananigans.jl/pull/3576 to see if tests pass (i.e. if NetCDFWriter still works with that trivial change).

May 02 '24 17:05 tomchor

For example, this should work

You can probably make this work pretty easily for precomputed fields. It's no different from simply calling Array(view(field, indices...)). So you just have to add some stuff to view(f::Field, i, j, k) to special-case StepRange indices.

It's more work / requires more ingenuity and cleverness to get that working with AbstractOperations and allocating the right amount of memory, etc, plus index gymnastics. That said I don't see any issue. Just some coding and arithmetic.

May 02 '24 17:05 glwagner

Oceananigans.jl Oceananigans.jl copied to clipboard

Can we calculate a `Field` at every `N` grid points?

Oceananigans.jl
Oceananigans.jl copied to clipboard