Add background potential energy computation

[jbisits]

This PR adds functions to compute the BackgroundPotentialEnergy, BPE, (see #168 for discussion on function naming) during a simulation. It does so using the definition given in Winters et al. (1995),

$$ E{b} = g\int{V}\rho z^{*}\mathrm{d}V. $$

To compute $E_{b}$ the buoyancy/density Field from a model is first reshaped into a 1D Array and then sorted. The $z^{*}$ (for a density Field) is defined as,

$$ z^{*} = \frac{1}{A}\int{\rho{\mathrm{min}}}^{\rho_{\mathrm{max}}}\mathrm{d}V. $$

To compute $z^{*}$, each volume element from the model.grid is computed, reshaped into a 1D Array, sorted using the sortperm from the buoyancy/density sorting, cumulatively summed then divided by the horizontal area of the model.grid.

These computations are done in OneDReferenceField which returns two new fields, the sorted buoyancy/density Field and the $z^{}$ Field. These returned Fields are on a grid with a vertical extent equal to the original model.grid but with total number of elements equal to the number of points in the domain (i.e. `model.grid.Nx model.grid.Ny * model.grid.Nz). These two newFields are then used to compute the BPE per unit volume. CallingIntegral(Field(BackgroundPotentialEnergy))` will give the volume integrated $E_{b}$ defined above.

I have handled the sorting of buoyancy and density separately because in the case of buoyancy, the buoyancy in the sorted profile should decrease as $z*$ increases,

$$ z^{*} = \frac{1}{A}\int{b{\mathrm{max}}}^{b_{\mathrm{min}}}\mathrm{d}V, $$

whereas the density should increase as $z*$ increases. If this is not clear (or wrong!) please let me know!

I have not tried this using a GPU but I think that these functions (reshape and sort) work with CuArrays but I might have some scalar indexing that will cause problems --- when I get a chance to I will try on a GPU.

The tests I have handled are really only validations that OneDReferenceField is computing things and returning things in the correct way. If you have any other suggestions for tests let me know.

Examples of the sorting for buoyancy and density using random noise as input data are below (they should be able to be run provided this branch of Oceanostics.jl is being used).

using Oceananigans, GLMakie
using Oceananigans.Models: seawater_density
using SeawaterPolynomials: TEOS10EquationOfState
using Oceanostics.PotentialEnergyEquationTerms

grid = RectilinearGrid(size=(10, 10, 50), x=(-10, 10), y=(-10, 10), z=(-100, 0),
                       topology=(Bounded, Bounded, Bounded))

## Using buoyancy
model = NonhydrostaticModel(; grid, buoyancy=BuoyancyTracer(), tracers=(:b,))
set!(model, b=randn(size(grid)))

buoyancy_reference_profile, z✶ = OneDReferenceField(model.tracers.b, rev = true)

fig1 = Figure(size = (1000, 500))
ax = Axis(fig1[1, 1], title = "x-z slice of buoyancy field",
            xlabel = "x",
            ylabel = "z")
x, z = xnodes(model.grid, Center()), znodes(model.grid, Center())
hm = heatmap!(ax, x, z, interior(model.tracers.b, :, 1, :), colormap = :balance)
Colorbar(fig1[1, 2], hm, label = "buoyancy")
ax2 = Axis(fig1[1, 3], title = "Buoyancy reference profile",
            xlabel = "buoyancy",
            ylabel = "z✶")
lines!(ax2, interior(buoyancy_reference_profile, 1, 1, :), interior(z✶, 1, 1, :))

## Using density
eos = TEOS10EquationOfState()
model = NonhydrostaticModel(; grid, buoyancy=SeawaterBuoyancy(equation_of_state=eos), tracers=(:S, :T))
set!(model, S=randn(size(grid)), T=randn(size(grid)))
ρ = Field(seawater_density(model))
compute!(ρ)

density_reference_profile, z✶ = OneDReferenceField(ρ)

fig2 = Figure(size = (1000, 500))
ax = Axis(fig2[1, 1], title = "x-z slice of density field",
            xlabel = "x",
            ylabel = "z")
x, z = xnodes(model.grid, Center()), znodes(model.grid, Center())
hm = heatmap!(ax, x, z, interior(ρ, :, 1, :), colormap = :dense)
Colorbar(fig2[1, 2], hm, label = "density")
ax = Axis(fig2[1, 3], title = "Density reference profile",
            xlabel = "Density",
            ylabel = "z✶")
lines!(ax, interior(density_reference_profile, 1, 1, :), interior(z✶, 1, 1, :))

cc @janzika

[tomchor]

Taking a step back @tomchor, @glwagner, @jbisits: Why do we actually care about the sorted buoyancy field or the "background state" as a function of physical space (i.e. as a 3D/4D array)? We never actually use that to calculate any terms in the BPE or APE budgets, as far as I can tell. All we ever need is z⋆(x,y,z,t) (as a 4D array) and dz⋆db|x,y,z,t (also as a 4D array), both of which are accurately and uniquely provided by the sorting methods we're discussing (aside from limit cases where the density/buoyancy of two grid cells are equal within machine precision, but that is a negligible effect). Since we have b(z⋆) and z⋆(b) as a sorted array, it is straight-forward to estimate the derivative dz⋆db and then map this back to each grid cell. There is never any need to average those fields and it doesn't make any sense to do so.

Am I missing something? (Feel free to suggest any additional calculations to make this discussion more concrete.)

I think indeed if we can get what we need without explicitly computing $b\star$ that would be fine. I also think that once you do the hard work of computing the permutations (which if I understand correctly you need to get $z\star$), then getting $b_\star$ is trivial anyway... right?

@jbisits you're the one implementing this; what do you think?

[tomchor]

Note that we still probably want to implement the expensive z⋆ calculation at some point since it's still needed for tilted domains or domains with immersed boundaries.

I don't think that's true. I've already implemented a sorting-based method for immersed boundaries (see my notebook linked above) and I think the method could be generalized to tilted domains but that one is more complicated because of the background buoyancy field and boundary conditions.

Ah, yes, I just now saw that comment. I'll look into it later.

Although at some point you mentioned that it can scale as nlog⁡(n)? I don't see how but hopefully you're right!

I meant that the sorting-based method scales like whatever sort scales like, which I think is nlog⁡(n). If we can get the sorting-based method to work for immersed boundaries then I don't see why we would want the brute-force n2 algorithm.

Agreed, if we can get the sorting method to work with immersed boundaries and tilted domains we don't need/want the brute-force approach of a double-loop heaviside iteration.

[jbisits]

I am sorry I have been slow in replying to the recent comments. I will be able to get back to this around the middle of this coming week - my apologies!