CliMA / Oceananigans.jl

🌊 Julia software for fast, friendly, flexible, ocean-flavored fluid dynamics on CPUs and GPUs
https://clima.github.io/OceananigansDocumentation/stable
MIT License
991 stars 194 forks source link

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

Open tomchor opened 9 months ago

tomchor commented 9 months ago

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

glwagner commented 9 months ago

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?

iuryt commented 8 months ago

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?

glwagner commented 8 months ago

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?

iuryt commented 8 months ago

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.

glwagner commented 8 months ago

"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.

iuryt commented 7 months ago

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?

glwagner commented 7 months ago

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)
iuryt commented 7 months ago
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?

tomchor commented 7 months ago

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.

iuryt commented 7 months ago

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
tomchor commented 7 months ago

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.

glwagner commented 7 months ago

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.

glwagner commented 7 months ago

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.

glwagner commented 7 months ago

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.

iuryt commented 7 months ago

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.

glwagner commented 7 months ago

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.

tomchor commented 6 months ago

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:

image

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.

glwagner commented 6 months ago

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...

tomchor commented 6 months ago

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.

glwagner commented 6 months ago

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 😄

glwagner commented 6 months ago

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)

iuryt commented 6 months ago

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)
tomchor commented 6 months ago

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.

tomchor commented 6 months ago

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).

glwagner commented 6 months ago

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.