Prescribing the buoyancy field and using thermal wind for the velocities

[iuryt]

It might not be the great idea, but what I am trying to do is to use the architecture of Oceananigans to run the following simulation:

I have a submesoscale resolving simulation with fronts getting unstable and I want to use the 10km averaged buoyancy from this simulation to run another simulation where we only advect biogeochemical with the thermal-wind balanced fields.
Link to input file with the averaged fields.

My idea is that at each timestep I prescribe the buoyancy field and update the velocity field from the vertical integration of the horizontal buoyancy gradients. In other words, I don't want to evolve the buoyancy and velocity field, but use the architecture of the model to advect and force my tracers with the prescribed fields. Does that make sense?

It seems to work well when I am just prescribing the buoyancy field, but it returns NaN if I try to add a function to compute and update velocity fields using thermal wind (assuming zero velocities at the bottom).

The code below has a simplified version that reproduces the error. I removed all the forcing for the tracers and mixed-layer depth to make is simpler and reproductible.

using Oceananigans
using Oceananigans.Units
using KernelAbstractions.Extras.LoopInfo: @unroll

#--------------- Grid

using NCDatasets
ds = Dataset("input_coarse.nc")

const initial_time = ds["time"][1]
const H = 1000 # maximum depth


const Nx, Ny, Nz = size(ds["b"][:,:,:,1])

grid = RectilinearGrid(CPU(),
    size=(Nx, Ny, Nz),
    x=(minimum(ds["xF"]),maximum(ds["xF"])+diff(ds["xF"])[end]),
    y=(minimum(ds["yF"]),maximum(ds["yF"])+diff(ds["yF"])[end]),
    z=ds["zF"][:],
    topology=(Periodic, Bounded, Bounded)
)


#--------------- Coriolis

coriolis = FPlane(latitude=60)

#--------------- Sponges

horizontal_closure = HorizontalScalarDiffusivity(ν=10, κ=10)
vertical_closure = ScalarDiffusivity(ν=1e-5, κ=1e-5)

#--------------- Instantiate Model


model = NonhydrostaticModel(grid = grid,
                            advection = WENO5(),
                            coriolis = coriolis,
                            closure=(horizontal_closure,vertical_closure),
                            tracers = (:b, :P, :N, :Nr))

bᵢ = Float64.(ds["b"][:,:,:,1])
Pᵢ = Float64.(ds["P"][:,:,:,1])
Nᵢ = Float64.(ds["N"][:,:,:,1])
Nrᵢ = Float64.(ds["Nr"][:,:,:,1])

set!(model; b=bᵢ, P=Pᵢ, N=Nᵢ, Nr=Nrᵢ)


#--------------- Initial Geostrophic Velocities

b = model.tracers.b
f = model.coriolis.f

# shear operations
uz_op = @at((Face, Center, Center), -∂y(b) / f );
vz_op = @at((Center, Face, Center),  ∂x(b) / f );
# compute shear
uz = compute!(Field(uz_op))
vz = compute!(Field(vz_op))

# include function for cumulative integration
include("cumulative_vertical_integration.jl")

# compute geostrophic velocities
uᵢ = cumulative_vertical_integration!(uz)
vᵢ = cumulative_vertical_integration!(vz)

# prescribe geostrophic velocities for initial condition
set!(model; u = uᵢ, v = vᵢ)

#--------------- Simulation

simulation = Simulation(model, Δt = 3hour, stop_iteration = 19)


#--------------- Prescribe buoyancy


function update_b!(sim)
    ti = sim.model.clock.time
    
    set!(model,
        b = Float64.(ds["b"][:,:,:,Int(round(ti/3hours))+1])
    )
    return nothing
end
simulation.callbacks[:update_b] = Callback(update_b!)


#--------------- Compute geostrophic velocities

b = model.tracers.b
f = model.coriolis.f

# shear operations
uz_op = Field(@at((Face, Center, Center), -∂y(b) / f ));
vz_op = Field(@at((Center, Face, Center),  ∂x(b) / f ));

# include function for cumulative integration
include("cumulative_vertical_integration.jl")


function compute_uv!(sim)
    u,v,w = sim.model.velocities
    u = cumulative_vertical_integration!(compute!(uz))
    v = cumulative_vertical_integration!(compute!(vz))
    return nothing
end
simulation.callbacks[:compute_uv] = Callback(compute_uv!)


#--------------- Writing Outputs

bi, Pi, Ni, Nri = simulation.model.tracers
u, v, w = simulation.model.velocities


extra_outputs = (; 
    u=@at((Center, Center, Center), u),
    v=@at((Center, Center, Center), v),
    w=@at((Center, Center, Center), w),
    N2=@at((Center, Center, Center), ∂z(bi)),
    ∇b=@at((Center, Center, Center), sqrt(∂x(bi)^2 + ∂y(bi)^2)),
    Ro=@at((Center, Center, Center), (∂x(v)-∂y(u))/f),
)

simulation.output_writers[:fields] =
    NetCDFOutputWriter(model, merge(model.tracers, extra_outputs), 
                    filename = "output_coarse.nc",
                    overwrite_existing=true,
                    schedule=IterationInterval(1))

#--------------- Printing Progress

using Printf

function print_progress(simulation)
    u, v, w = simulation.model.velocities

    # Print a progress message
    msg = @sprintf("i: %04d, t: %s, Δt: %s, umax = (%.1e, %.1e, %.1e) ms⁻¹, wall time: %s\n",
                   iteration(simulation),
                   prettytime(time(simulation)),
                   prettytime(simulation.Δt),
                   maximum(abs, u), maximum(abs, v), maximum(abs, w),
                   prettytime(simulation.run_wall_time))

    @info msg

    return nothing
end

simulation.callbacks[:progress] = Callback(print_progress, IterationInterval(1))

run!(simulation)

The cumulative_vertical_integration.jl file is:

using KernelAbstractions: @index, @kernel
using KernelAbstractions.Extras.LoopInfo: @unroll
using Oceananigans.Architectures: device_event, architecture
using Oceananigans.Utils: launch!
using Oceananigans.Grids
using Oceananigans.Operators: Δzᶜᶜᶜ

@kernel function _cumulative_vertical_integration!(field, grid)
    i, j = @index(Global, NTuple)

    integ = 0
    @unroll for k in 1 : grid.Nz
        integ = integ + @inbounds field[i, j, k] * Δzᶜᶜᶜ(i, j, k, grid)
        @inbounds field[i, j, k] = integ
    end
end

function cumulative_vertical_integration!(input_field)
    field = Field(input_field)
    grid = field.grid
    arch = architecture(grid)

    event = launch!(arch, grid, :xy,
                    _cumulative_vertical_integration!, field, grid,
                    dependencies = device_event(arch))

    wait(device_event(arch), event)

    return field
end

The simulation run returns the following output:

┌ Warning: defaulting to uniform WENO scheme with Float64 precision, use WENO5(grid = grid) if this was not intended
└ @ Oceananigans.Advection /home/isimoesdesousa/.julia/packages/Oceananigans/DpyWt/src/Advection/weno_fifth_order.jl:154
┌ Warning: Overwriting existing ./output_coarse.nc.
└ @ Oceananigans.OutputWriters /home/isimoesdesousa/.julia/packages/Oceananigans/DpyWt/src/OutputWriters/netcdf_output_writer.jl:316
┌ Info: Initializing simulation...
└ @ Oceananigans.Simulations /home/isimoesdesousa/.julia/packages/Oceananigans/DpyWt/src/Simulations/run.jl:167
┌ Info: i: 0000, t: 0 seconds, Δt: 3 hours, umax = (3.0e-01, 1.2e-04, 1.7e-18) ms⁻¹, wall time: 0 seconds
└ @ Main In[3]:150
┌ Info:     ... simulation initialization complete (5.478 seconds)
└ @ Oceananigans.Simulations /home/isimoesdesousa/.julia/packages/Oceananigans/DpyWt/src/Simulations/run.jl:202
┌ Info: Executing initial time step...
└ @ Oceananigans.Simulations /home/isimoesdesousa/.julia/packages/Oceananigans/DpyWt/src/Simulations/run.jl:112
┌ Info:     ... initial time step complete (4.264 ms).
└ @ Oceananigans.Simulations /home/isimoesdesousa/.julia/packages/Oceananigans/DpyWt/src/Simulations/run.jl:119
┌ Info: i: 0001, t: 3 hours, Δt: 3 hours, umax = (3.0e-01, 2.8e-01, 2.3e-03) ms⁻¹, wall time: 0 seconds
└ @ Main In[3]:150
┌ Info: i: 0002, t: 6 hours, Δt: 3 hours, umax = (2.9e-01, 5.5e-01, 4.8e-03) ms⁻¹, wall time: 5.524 seconds
└ @ Main In[3]:150
┌ Info: i: 0003, t: 9 hours, Δt: 3 hours, umax = (1.1e+00, 5.0e-01, 3.5e-03) ms⁻¹, wall time: 5.544 seconds
└ @ Main In[3]:150
┌ Info: i: 0004, t: 12 hours, Δt: 3 hours, umax = (2.5e-01, 2.6e+00, 1.9e-02) ms⁻¹, wall time: 5.563 seconds
└ @ Main In[3]:150
┌ Info: i: 0005, t: 15 hours, Δt: 3 hours, umax = (4.7e+00, 3.2e+00, 2.4e-02) ms⁻¹, wall time: 5.582 seconds
└ @ Main In[3]:150
┌ Info: i: 0006, t: 18 hours, Δt: 3 hours, umax = (1.5e+01, 1.2e+01, 1.2e-01) ms⁻¹, wall time: 5.603 seconds
└ @ Main In[3]:150
┌ Info: i: 0007, t: 21 hours, Δt: 3 hours, umax = (1.2e+02, 5.1e+01, 9.3e-01) ms⁻¹, wall time: 5.623 seconds
└ @ Main In[3]:150
┌ Info: i: 0008, t: 1 day, Δt: 3 hours, umax = (6.0e+04, 5.4e+03, 7.8e+01) ms⁻¹, wall time: 5.642 seconds
└ @ Main In[3]:150
┌ Info: i: 0009, t: 1.125 days, Δt: 3 hours, umax = (4.2e+09, 2.0e+08, 1.2e+07) ms⁻¹, wall time: 5.662 seconds
└ @ Main In[3]:150
┌ Info: i: 0010, t: 1.250 days, Δt: 3 hours, umax = (1.6e+19, 8.5e+17, 4.7e+16) ms⁻¹, wall time: 5.682 seconds
└ @ Main In[3]:150
┌ Info: i: 0011, t: 1.375 days, Δt: 3 hours, umax = (2.9e+38, 2.6e+37, 9.9e+35) ms⁻¹, wall time: 5.703 seconds
└ @ Main In[3]:150
┌ Info: i: 0012, t: 1.500 days, Δt: 3 hours, umax = (1.4e+77, 6.4e+75, 3.8e+74) ms⁻¹, wall time: 5.757 seconds
└ @ Main In[3]:150
┌ Info: i: 0013, t: 1.625 days, Δt: 3 hours, umax = (2.1e+154, 8.0e+152, 5.5e+151) ms⁻¹, wall time: 5.776 seconds
└ @ Main In[3]:150
┌ Info: i: 0014, t: 1.750 days, Δt: 3 hours, umax = (NaN, 0.0e+00, 0.0e+00) ms⁻¹, wall time: 5.795 seconds
└ @ Main In[3]:150
┌ Info: i: 0015, t: 1.875 days, Δt: 3 hours, umax = (NaN, 0.0e+00, 0.0e+00) ms⁻¹, wall time: 5.815 seconds
└ @ Main In[3]:150
┌ Info: i: 0016, t: 2 days, Δt: 3 hours, umax = (NaN, 0.0e+00, 0.0e+00) ms⁻¹, wall time: 5.834 seconds
└ @ Main In[3]:150
┌ Info: i: 0017, t: 2.125 days, Δt: 3 hours, umax = (NaN, 0.0e+00, 0.0e+00) ms⁻¹, wall time: 5.853 seconds
└ @ Main In[3]:150
┌ Info: i: 0018, t: 2.250 days, Δt: 3 hours, umax = (NaN, 0.0e+00, 0.0e+00) ms⁻¹, wall time: 5.872 seconds
└ @ Main In[3]:150
┌ Info: Simulation is stopping. Model iteration 19 has hit or exceeded simulation stop iteration 19.
└ @ Oceananigans.Simulations /home/isimoesdesousa/.julia/packages/Oceananigans/DpyWt/src/Simulations/simulation.jl:158
┌ Info: i: 0019, t: 2.375 days, Δt: 3 hours, umax = (NaN, 0.0e+00, 0.0e+00) ms⁻¹, wall time: 5.891 seconds
└ @ Main In[3]:150

Anyone has any idea on how I could fix this?

[glwagner]

I think you're looking for PrescribedVelocityFields:

https://clima.github.io/OceananigansDocumentation/stable/appendix/library/#Oceananigans.Models.HydrostaticFreeSurfaceModels.PrescribedVelocityFields-Tuple{}

which only works with HydrostaticFreeSurfaceModel.

Docs are sparse, but there is this validation test (which may be out of date):

Oceananigans.jl/validation/solid_body_rotation/solid_body_tracer_advection.jl

Line 79 in 3469e93

velocities = PrescribedVelocityFields(u=uᵢ),

[iuryt]

Do you think it works if I give an array/field u(x,y,z,t) instead of a function?

[glwagner]

Doesn't hurt to try...

[glwagner]

It looks like Field is supported:

Oceananigans.jl/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl

Lines 47 to 62 in 3ae3e6c

	PrescribedField(X, Y, Z, f::Function, grid; kwargs...) = FunctionField{X, Y, Z}(f, grid; kwargs...)
	PrescribedField(X, Y, Z, f::AbstractField, grid; kwargs...) = f
	function PrescribedField(X, Y, Z, f::Field, grid; kwargs...)
	fill_halo_regions!(f)
	return f
	end
	function HydrostaticFreeSurfaceVelocityFields(velocities::PrescribedVelocityFields, grid, clock, bcs)
	u = PrescribedField(Face, Center, Center, velocities.u, grid; clock=clock, parameters=velocities.parameters)
	v = PrescribedField(Center, Face, Center, velocities.v, grid; clock=clock, parameters=velocities.parameters)
	w = PrescribedField(Center, Center, Face, velocities.w, grid; clock=clock, parameters=velocities.parameters)
	return PrescribedVelocityFields(u, v, w, velocities.parameters)
	end

but the docstring doesn't mention it.

[glwagner]

It looks like Field is supported:

Oceananigans.jl/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl

Lines 47 to 62 in 3ae3e6c

	PrescribedField(X, Y, Z, f::Function, grid; kwargs...) = FunctionField{X, Y, Z}(f, grid; kwargs...)
	PrescribedField(X, Y, Z, f::AbstractField, grid; kwargs...) = f
	function PrescribedField(X, Y, Z, f::Field, grid; kwargs...)
	fill_halo_regions!(f)
	return f
	end
	function HydrostaticFreeSurfaceVelocityFields(velocities::PrescribedVelocityFields, grid, clock, bcs)
	u = PrescribedField(Face, Center, Center, velocities.u, grid; clock=clock, parameters=velocities.parameters)
	v = PrescribedField(Center, Face, Center, velocities.v, grid; clock=clock, parameters=velocities.parameters)
	w = PrescribedField(Center, Center, Face, velocities.w, grid; clock=clock, parameters=velocities.parameters)
	return PrescribedVelocityFields(u, v, w, velocities.parameters)
	end

but the docstring doesn't mention it.

[iuryt]

I am using it now and it is working. I will turn the repository public next week and post here as an example.