What if I want to start with geostrophically-balanced currents from initial buoyancy field?

[iuryt]

I have a function for the buoyancy that is basically three fronts constructed from a background buoyancy:

const L = 10kilometers
const amp = 1kilometers
const g = 9.82
const ρₒ = 1026

# background density profile based on Argo data
@inline bg(z) = 0.25*tanh(0.0027*(-653.3-z))-6.8*z/1e5+1027.56

# decay function for fronts
@inline decay(z) = (tanh((z+500)/300)+1)/2

# front function
@inline front(x, y, z, cy) = tanh((y-(cy+sin(2pi*x/L)*amp))/12kilometers)

# initial density function
@inline D(x, y, z) = bg(z) + 0.8*decay(z)*((front(x,y,z,-120kilometers)+front(x,y,z,0)+front(x,y,z,120kilometers))-3)/6

# initial buoyancy function
@inline B(x, y, z) = -(g/ρₒ)*D(x,y,z)

set!(model;b=B)

For now, I am starting with zero velocity field, but I would like to start with a geostrophic balanced current from B(x,y,z). In other words, I would like to prescribe the initial velocity field using thermal wind and integrating in the vertical the horizontal gradient of buoyancy (assuming zero velocity at the bottom).

What is the most julian and easy way to do that?

[glwagner]

The best way I can think of now would be to first define the geostrophic streamfunction, and then calculate the geostrophic buoyancy and geostrophic velocity from that. You can do that using Field and AbstractOperations or analytically. We've used that method for other applications.

If you prefer to start from the buoyancy field, it is in principle possible to compute a geostrophic streamfunction given the buoyancy field. If we had CumulativeIntegralField (as discussed at #2385) that would be fairly easy. Without that you'll have to use cumsum! to compute the cumulative integral for the streamfunction.

But I think it will also help others understand your code and science if you derive a buoyancy and velocity field from a geostrophic streamfunction.

[glwagner]

This also works, also less ideal

julia> vz_op_fcc = -∂x(b) / f
BinaryOperation at (Face, Center, Center)
├── grid: 1×1×1 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×1 halo
└── tree:
    / at (Face, Center, Center)
    ├── - at (Face, Center, Center) via identity
    │   └── ∂xᶠᶜᶜ at (Face, Center, Center) via identity
        │   └── 1×1×1 Field{Center, Center, Center} on RectilinearGrid on CPU
    └── 1

julia> vz_op = @at (Center, Face, Face) vz_op_fcc
UnaryOperation at (Center, Face, Face)
├── grid: 1×1×1 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×1 halo
└── tree:
    interpolate_identity at (Center, Face, Face) via ℑxyzᶜᶠᶠ
    └── / at (Face, Center, Center)
        ├── - at (Face, Center, Center) via identity
        │   └── ∂xᶠᶜᶜ at (Face, Center, Center) via identity
            │   └── 1×1×1 Field{Center, Center, Center} on RectilinearGrid on CPU
        └── 1

[iuryt]

I tried to do like this:

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

uz = compute!(Field(uz_op))
vz = compute!(Field(vz_op))

but it gives error because of -+ 😕 . So I had to do:

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

uz = compute!(Field(uz_op))
vz = compute!(Field(vz_op))

For now, I have the code below. I am still doing it centered in z.
I will think about doing it for Face in z. Do you think this is worth implementing to Oceananigans' source as GeostrophicU and GeostrophicV or something like that?
I think this is going to be more complicated with ImmersedBoundaries and a 3D Coriolis, but this is a starting point right?

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

uz = compute!(Field(uz_op))
vz = compute!(Field(vz_op))

# ----- kernel

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

# -----

U = cumulative_vertical_integration!(uz)
V = cumulative_vertical_integration!(vz)

set!(model; u = U, v = V)

Also, about the barotropic velocity, should I simply start the integration with it? How does that work for NonhydrostraticModel?

[iuryt]

Noticed you avoid making operations inplace. Good to know that as well. Learning good practices!

[glwagner]

If we implement accumulated fields (eg Accumulation to pair with Reduction) then we might be able to support syntax like

dudz = @at (Face, Center, Face) (+∂y(b) / f)
dvdz = @at (Center, Face, Face) (-∂x(b) / f)

u = Field(CumulativeIntegral(dudz, dims=3))
v = Field(CumulativeIntegral(dvdz, dims=3))

What do you think about that?

I think generic code like the above might be preferred over specific functions just for the geostrophic velocity field. It's also just about the same amount of work to implement CumulativeIntegral I think.

[iuryt]

That would be awesome! I think that this is a better and more general solution (matches Oceananigans way of thinking).
So, I will leave this discussion with no answer until then.