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]

@jbisits thanks this is great!

I'll only have time to look at this more carefully next week, but a couple of comments that I'll say right now are:

1. You're right, I think it would be best to avoid KernelFuncionOperation nesting. I think this can be done quite easily, no?
2. The sorting procedure is only an approximation to compute the BPE, so I think this should be mentioned somewhere in the docstring. (Or maybe even reflected in the function name? ApproximateBackgroundPotentialEnergy?)
3. Even then, sorting is only valid when the grid is regular (i.e. all the grid cells have equal volume). So we should test for that and spit out a useful error otherwise. (Or if you're up for it, implement the appropriate integral-form calculation of z*, as it would be useful to have something that works on stretched grids.)

A quick note is that I'm surprised that a KFO of reshape(sort(reshape(f, :)), size(f)) works out of the box! I'd have thought that, because it's a nonlocal calculation (sorting), we would need to some "massaging" here. Glad I was wrong :)

[glwagner]

I'm a little confused about the sorting. Don't you want to sort only in the vertical? In the example you posted, the sorted field not only has a vertical buoyancy gradient but also a left-right buoyancy gradient which seems arbitrary...

Also, you may just want to extend sort! and sort for Field, rather than defining a new function sorted_field or whatever.

[glwagner]

A quick note is that I'm surprised that a KFO of reshape(sort(reshape(f, :)), size(f)) works out of the box! I'd have thought that, because it's a nonlocal calculation (sorting), we would need to some "massaging" here. Glad I was wrong :)

Probably sort just uses getindex and length.

Personally I think that sort of a Field should return Field, not Array. (For example, this will almost certainly throw a scalar indexing error on GPU since we can't use getindex naively there.)

I think we should be able to support this without too much difficulty, provided we only support dimension-by-dimension sorting. If we do actually need to support sorting of Fields that are reshaped, we have a bit more work to do with ReshapedArray. Nothing too hard but it might be nice to know if we really need to support that before implementing it...

[glwagner]

A little bit more about that... consider

julia> grid = RectilinearGrid(size=(1, 2, 3), x=(0, 1), y=(0, 1), z=(0, 1))
1×2×3 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
├── Periodic x ∈ [0.0, 1.0) regularly spaced with Δx=1.0
├── Periodic y ∈ [0.0, 1.0) regularly spaced with Δy=0.5
└── Bounded  z ∈ [0.0, 1.0] regularly spaced with Δz=0.333333

julia> c = CenterField(grid)
1×2×3 Field{Center, Center, Center} on RectilinearGrid on CPU
├── grid: 1×2×3 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Periodic, east: Periodic, south: Periodic, north: Periodic, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 7×8×9 OffsetArray(::Array{Float64, 3}, -2:4, -2:5, -2:6) with eltype Float64 with indices -2:4×-2:5×-2:6
    └── max=0.0, min=0.0, mean=0.0

julia> rc = reshape(c, :)
6-element reshape(::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}}, CPU}, Tuple{Colon, Colon, Colon}, OffsetArrays.OffsetArray{Float64, 3, Array{Float64, 3}}, 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}}, 6) with eltype Float64:
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0

julia> @which sort(rc)
sort(v::AbstractVector; kws...)
     @ Base.Sort sort.jl:1489

This calls

https://github.com/JuliaLang/julia/blob/9107bd3a2db63865b8818597feda0728807da613/base/sort.jl#L1730

But copymutable(rc) is just

julia> Base.copymutable(rc)
6-element Vector{Float64}:
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0

because we haven't defined similar for ReshapedArray{<:Any, <:Any, <:Field}:

help?> Base.ReshapedArray
  No documentation found.

  Summary
  ≡≡≡≡≡≡≡

  struct Base.ReshapedArray{T, N, P<:AbstractArray, MI<:Tuple{Vararg{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64}}}}

Supporting sort! seems rather easy, possibly as easy as

sort!(f::Field; kw...) = sort!(interior(f); kw...)

Note this also works on GPU. For sort we can take the approach that CUDA takes:

sort(f::Field; kw...) = sort!(copy(f); kw...)

(we need to check that copy works correctly, and I'm not sure if we want to use deepcopy or something more complicated).

Then, if we don't need to support ReshapedArray, we are done with these two lines.

[jbisits]

Thanks for the comments and suggestions @tomchor and @glwagner!

@glwagner that does look like a cleaner (and better) implementation for sorting a Field. I tried quite a few ways (including your suggestion) but I was often making things a little more difficult than they needed to be --- so I ended up using the simplest implementation I came up with. I will use your suggestions and come up with a better (and seemingly even simpler) method to sort a Field.

I'm a little confused about the sorting. Don't you want to sort only in the vertical? In the example you posted, the sorted field not only has a vertical buoyancy gradient but also a left-right buoyancy gradient which seems arbitrary...

Yes the horizontal buoyancy gradient that appears in the figures is arbitrary but part of that arises from sorting a field which is just random noise. In the case of a stratified fluid I think we would see approximately uniform density over each horizontal layer. For reference here is the excerpt from section 6 in Winters et al. (1995) which I have used when approximating the background potential energy

[jbisits]

@jbisits thanks this is great!

I'll only have time to look at this more carefully next week, but a couple of comments that I'll say right now are:

1. You're right, I think it would be best to avoid KernelFuncionOperation nesting. I think this can be done quite easily, no?
2. The sorting procedure is only an approximation to compute the BPE, so I think this should be mentioned somewhere in the docstring. (Or maybe even reflected in the function name? ApproximateBackgroundPotentialEnergy?)
3. Even then, sorting is only valid when the grid is regular (i.e. all the grid cells have equal volume). So we should test for that and spit out a useful error otherwise. (Or if you're up for it, implement the appropriate integral-form calculation of z*, as it would be useful to have something that works on stretched grids.)

A quick note is that I'm surprised that a KFO of reshape(sort(reshape(f, :)), size(f)) works out of the box! I'd have thought that, because it's a nonlocal calculation (sorting), we would need to some "massaging" here. Glad I was wrong :)

No problem! In response to the questions above:

1. Yep I think so especially given @glwagner's comments above
2. Definitely! Certainly in the docstring and probably no harm in putting it in the function name as well.
3. Again definitely! In this PR I will add a validation check with an error message for non-uniform grids. Then, if/when I get to it I can try and code up a calculation of z*.

[glwagner]

@jbisits did you try

sort!(f::Field; kw...) = sort!(interior(f); kw...)
sort(f::Field; kw...) = sort!(copy(f); kw...)

?

In the case of a stratified fluid I think we would see approximately uniform density over each horizontal layer.

But wouldn't any horizontal structure be destroyed if we resort in the horizontal? Consider doing this for the global ocean. It's cold at the poles and warm on the equator. Resorting in the horizontal would put cold water at the north pole and warm water at the south pole. That doesn't make sense. We must be assuming horizontal homogeneity. Is this necessary?

[jbisits]

@jbisits did you try

sort!(f::Field; kw...) = sort!(interior(f); kw...)
sort(f::Field; kw...) = sort!(copy(f); kw...)

?

Not yet I will get to it soon.

In the case of a stratified fluid I think we would see approximately uniform density over each horizontal layer.

But wouldn't any horizontal structure be destroyed if we resort in the horizontal? Consider doing this for the global ocean. It's cold at the poles and warm on the equator. Resorting in the horizontal would put cold water at the north pole and warm water at the south pole. That doesn't make sense. We must be assuming horizontal homogeneity. Is this necessary?

I am not sure if it matters that horizontal structure is not maintained after resorting (sorry if I am not understanding something here) because using the sorted array we want to attain an approximation to the minimal potential energy state which can be attained through adiabatic rearrangement of fluid parcels (or volume elements in Boussinesq case).

For your example above, if we first reshape a 3D density field for the global ocean into a 1D array and then sort it we would see the lighter, warmer water around the equator at the start of the sorted array and the deep, dense bottom water at the end end of the sorted array (with other water masses in between). Redistributing this sorted array throughout the original grid would then have the lightest waters (e.g. from the around the equator) at the top of the grid $(z=0)$ and the densest water at the bottom of the grid (with the other water masses distributed in between). The result should be a stably stratified density field which can the be used to compute an approximation to the minimal potential energy state after adiabatic rearrangement which is the background potential energy.

As you point out this could potentially result in warm water from the equator at the poles which does not make sense but I am not sure if that matters as we just want the lightest water on top to give a minimal potential energy state.

In the case of some experiments I have been running (more along direct numerical simulation lines) I would like to be able first reshape a 3D density field into a 1D array, sort this, then redistribute throughout the original grid. We could add a keyword argument that only sorts over the vertical dimension of the domain though? So we would have two cases one that sorts the reshaped array and redistributes over the grid

function sort!(f::Field; vertical_only = false, kw...)
    sorted_1D_field = reshape(interior(f), :)
    sort!(sorted_1D_field)
    f.data = reshape(sorted_1D_field, size(interior(f)) # I know this wont work because data has the halos is different size is immutable but this is the idea
    return f
end

and another that just sorts over the vertical dimension (assuming vertical dimension is 3 in example below) of the domain

sort!(f::Field; vertical_only = true, kw...) = sort!(interior(f), dims = 3; kw...)

Again sorry if I have misunderstood something but this is how I have been thinking about the BPE and its implementation for experiments I have been running.

[glwagner]

I am not sure if it matters that horizontal structure is not maintained after resorting (sorry if I am not understanding something here) because using the sorted array we want to attain an approximation to the minimal potential energy state which can be attained through adiabatic rearrangement of fluid parcels (or volume elements in Boussinesq case).

For your example above, if we first reshape a 3D density field for the global ocean into a 1D array and then sort it we would see the lighter, warmer water around the equator at the start of the sorted array and the deep, dense bottom water at the end end of the sorted array (with other water masses in between). Redistributing this sorted array throughout the original grid would then have the lightest waters (e.g. from the around the equator) at the top of the grid and the densest water at the bottom of the grid (with the other water masses distributed in between). The result should be a stably stratified density field which can the be used to compute an approximation to the minimal potential energy state after adiabatic rearrangement which is the background potential energy.

It's easy to construct a counter example. If we simply flip the definition of the y coordinate (so that increasing j-index implies decreasing latitude), we end up with a different background energy state. But there's only one physical background energy state --- right?

More severely I think the background energy state cannot be a function of horizontal position. It can only be a function of z. At least, that's the case unless we introduce some additional concepts like a horizontal filter or something like that. See Winters et al 1995.

Possibly the solution is as simple as averaging the sorted state in the horizontal directions. I'm not quite sure though. I think we are making additional assumptions. For example, it seems like we are assuming a regular grid (both vertical and horizontal). I think we need to account for the differing grid volumes to define the sorted background state in a way that conserves mass...

It's also important to understand that this only works because of the indexing convention that we use. In other words there is an implicit assumption in the code that the z-index is the last index. Somehow, it might be nicer to write the code so that this assumption is either more clear or so that the code would still work if the z-index were 1st or 2nd.

[tomchor]

As far as I understand, the physical background state is indeed unique. It necessarily has no horizontal buoyancy gradients. But we're only approximating that numerically, and the "numerical" background state doesn't need to be unique. In fact, like @glwagner pointed out, just flipping one of the horizontal directions would produce a different state.

I think that's the best we can do numerically and that's what winters et al suggest in their papers. I think averaging the horizontal directions would go against the premise of the BPE, which is an adiabatic resorting of the buoyancy field by definition.

Usually though the result of the sorting is better than what's pictured here and the horizontal gradients are quite small (the original paper gives estimates for when this is a good approximation). But since @jbisits is using white noise here as an example, the approximation isn't very good and the horizontal gradients are more pronounced than they'd normally be.

[glwagner]

I think that's the best we can do numerically and that's what winters et al suggest in their papers.

In the Winters papers, the background state is only a function of z and not of horizontal directions. Shouldn't our background state also be 1D then?

The solution is fairly easy in the uniform grid case --- we just average the sorted field in the horizontal directions.

For stretched grids we need to sort more carefully, taking into account the cell volumes. I don't think this is too difficult though and we have GridMetricOperation to help if we need it.

[tomchor]

I think that's the best we can do numerically and that's what winters et al suggest in their papers.

In the Winters papers, the background state is only a function of z and not of horizontal directions. Shouldn't our background state also be 1D then?

In the physical world the background state is a function of $z$ because water parcels can be infinitesimally small). In our case each "parcel" is a grid cell that we're not willing to further subdivide, so it's natural that a resorting won't get rid of all horizontal gradients.

Your point of forcing the numerical, discrete background state to also be a function of $z$ only is interesting, but I'm not sure modifying the sorting procedure is the way to go. Here are some thoughts in no particular order:

1. This sorting procedure has become standard in the literature and experience seems to point to the errors produced by this simple sorting generally being small. My own experience corroborates that. The example concocted here by @jbisits is in no way physical so it overestimates this error in relation to realistic buoyancy field distribution.
2. I'd be reluctant to add an averaging procedure whose effects on the calculations we haven't investigated before. So I think if we do decide that we want to average in the horizontal, we need to do some investigating first with "realistic" simulations (whatever we think that means).
3. The sorting computation was always just meant to be an approximation to the background state. I think if we want to be more precise with the calculation, the way to go is to use the actual definition of the background state, rather than modifying the sorting procedure:

I feel like the equation above will be less performant (I think it'll always scale with O($n^2$)), but we can understand that as price to pay for precision.

@glwagner thoughts?

[glwagner]

I'm not suggesting to change the sorting. I'm saying the background state should be derived by

1. Sorting as in the present example
2. Averaging the sorted field horizontally to obtain a background state independent of z

This algorithm produces identical background states even if the horizontal coordinate direciton is reversed for example.

I'm hazy on the relationship between the sorted field and $z\star$. Isn't the result of this computation to actually obtain the background field $\rho(z\star)$? The additional step of computing $z_\star$ is actually what one needs to do in order to relate the computed background to physical space $z$... ?

[glwagner]

@tomchor can you please share the literature on this sorting-based calculation?

I'm also missing how the sorted field is used to compute the available potential energy. Do you compute the difference between the current state and the sorted state then?

[glwagner]

Here's a way to think about it... the variable $z_\star$ is the reference / resting height for the water parcel that is currently at $x, t$,

The sorting that we are doing is precisely that adiabatic rearrangement being referred to. Somehow in fact, the information we would need to compute $z_\star$ must be contained in the permutator (ie output of sortperm) if we were to compute that.

When we average the sorted field in the horizontal direction, we obtain the density profile $\rho\star(z\star)$. This is a finite volume expression of what's written in Winters et al. If we did not want to sort in the horizontal direction, we could perhaps further imagine that each parcel is stretched to fill the whole horizontal domain. This could generate a sorted field on a very fine vertical grid with dimension Nx * Ny * Nz.

I think once we get the sorted field we can compute $z\star$ more easily than equation 11 in Winters et al. It seems like for each grid point / density value we have to find the vertical level in the reference profile associated with it. That gives us a 3D field $z\star$? It's not quite N^2, I think that's Nx * Ny * Nz^2 because we have to search in z for every grid point). Maybe with an extra log(Nz). Probably doable...

[glwagner]

There might be even cleverer ways to do this...

[glwagner]

Here's an idea that generalizes to a (horizontally) stretched grid...

1. Define the total cell density $R = \rho \Delta V$. Also define a field of cell volumes, $\Delta V$.
2. Compute the sorting permutation for $R$, and use it sort $R$ and $\Delta V$.
3. Average both $R$ and $\Delta V$ over the "horizontal dimensions". The purpose of this is to produce a background field with the same "resolution" as the primitive fields. However, they will be on different grids, so the next task is to compute the "grid" of the background field...
4. Compute $\rho_\star(z_k) = \bar R / \overline{\Delta V}$, and the height of each increment $\Delta z_k$ by $\Delta z_k = \overline{\Delta V}_k / A$ where $A$ is the area of the domain.

Would this work? Hmm, as I wrote this I realized this will not work for irregular domains. Anyways, it's a start perhaps.

[tomchor]

I'm not suggesting to change the sorting. I'm saying the background state should be derived by

1. Sorting as in the present example
2. Averaging the sorted field horizontally to obtain a background state independent of z

I understand that. Sorry, I guess I'm expressing myself poorly. That's what I meant by "modifying the sorting procedure", which is admittedly a misnomer.

This algorithm produces identical background states even if the horizontal coordinate direciton is reversed for example.

I'm hazy on the relationship between the sorted field and z⋆. Isn't the result of this computation to actually obtain the background field ρ(z⋆)? The additional step of computing z⋆ is actually what one needs to do in order to relate the computed background to physical space z... ?

The variable $z*$ is the height at which a parcel would be if all parcels were adiabatically rearranged to the state of minimal PE. If the parcels can be made infinitesimal this is state is unique and there are no horizontal buoyancy gradients. Which is why the equation presented to calculate it (Eq (11) in their paper, which I reproduced here) has a unique answer. That equation should be the preferred way to calculate $z\star$, but afaik people rarely use it because it's computationally expensive to calculate. The calculation to get the BPE goes like this:

1. From a field $b$ (or $\rho$) you use Equation (11) to get a field $z_\star$ (the height of each parcel if all parcels were adiabatically resorted to a state of total minimum PE).
2. You then multiply $b$ by $z_\star$ to get the PE of each parcel in this state of minimum total PE. Note that in this calculation never really resorted anything, as really we only need the height at which each parcel would sit, and Equation (11) gives us that.
3. Finally you integrate everything and then you have the total PE of the domain in a minimum PE state: the BPE. Note that, in the paper, most (all?) the calculations are done with domain-integrated quantities since the BPE is necessarily a nonlocal quantity; that's these calculation are always a bit of a hassle.

Now, if you wanna avoid calculating Equation (11), instead of manually calculating $z_\star$ we can, in the words of Winters et al., "let the fluid elements correspond to volume elements of a numerical grid" and actually sort buoyancy values. Then the vertical coordinate (originally $z$) approximately becomes $z_\star$ by definition, since each parcel is approximately where it would be if the values were resorted to a statement of minimum PE. I'm emphasizing "approximately" here because as we discussed, grid cells aren't infinitesimal, and their sorting only approximates the state of minimum PE, and their sorted state not only isn't unique, but in general has some small horizontal gradients.

• The advantage of this latter method is that is faster since we never explicitly calculate $z_*$; it implicitly appears after the sorting is complete. Basically we exchange a $O(n^2)$ operation (Equatoin (11)) for a $\approx O(n log(n))$ operation.
• The disadvantage is that it's only an approximation.

Then the calculation the get the BPE through this approach is:

1. You 3D sort a field of $b$.
2. After resorting the original vertical coordinate $z$ becomes $z\star$. We ignore horizontal coordinates. Then we can just multiply the sorted $b$ by its coordinate (which now approximates $z\star$) to obtain the PE of each parcel in the minimum PE state.
3. Integrate everything to get the (approximated) BPE.

My take is that the latter method (i.e. sorting) will always be an approximation to the former method (i.e. direct calculation of $z\star$ with Eq. (11)). Thus we can use the sorting approach as it is until an example/simulation comes along where that approximation fails and we need something better. Then we can either implement the first method (direct calculation of $z\star$) or try to improve on the approximated one.

@jbisits @glwagner thoughts?

[tomchor]

@tomchor can you please share the literature on this sorting-based calculation?

I don't think there's more literature specifically on the sorting-based calculation other than the original paper. Winter and D'Asaro touch on it again on their follow-up 1996 paper but I think that's it.

My impression that the sorting-based approach works well comes from reading papers that employed it without complications, plus talking to a few people that corroborate that impression. I'll try to remember some of the papers that I read which use that calculation and post here later (I don't remember any of them off the top of my head).

[glwagner]

Then the vertical coordinate (originally ) approximately becomes by definition, since each parcel is approximately where it would be if the values were resorted to a statement of minimum PE

I agree that the vertical coordinate of the sorted field is $z_\star$. But this does not give you the function $f(x, t) = z_\star$ in terms of the unsorted coordinates $(x, t)$. Something is funky here...

[tomchor]

Then the vertical coordinate (originally ) approximately becomes by definition, since each parcel is approximately where it would be if the values were resorted to a statement of minimum PE

I agree that the vertical coordinate of the sorted field is z⋆. But this does not give you the function f(x,t)=z⋆ in terms of the unsorted coordinates (x,t). Something is funky here...

I think you're right. It doesn't. But I think we don't need that function. Just $z_\star$.

[glwagner]

Can you be more explicit? What exactly is the calculation that you do after you obtain the sorted field (which is the resting density field, not a coordinate mapping)?

For example, equation 19 in Winters et al. 1995 is

For this we need the function $f(x, t) = z\star$. How do we compute equation 19 without explicitly obtaining $z\star$?

[jbisits]

This algorithm produces identical background states even if the horizontal coordinate direciton is reversed for example. I'm hazy on the relationship between the sorted field and z⋆. Isn't the result of this computation to actually obtain the background field ρ(z⋆)? The additional step of computing z⋆ is actually what one needs to do in order to relate the computed background to physical space z... ?

The variable z∗ is the height at which a parcel would be if all parcels were adiabatically rearranged to the state of minimal PE. If the parcels can be made infinitesimal this is state is unique and there are no horizontal buoyancy gradients. Which is why the equation presented to calculate it (Eq (11) in their paper, which I reproduced here) has a unique answer. That equation should be the preferred way to calculate z⋆, but afaik people rarely use it because it's computationally expensive to calculate. The calculation to get the BPE goes like this:

1. From a field b (or ρ) you use Equation (11) to get a field z⋆ (the height of each parcel if all parcels were adiabatically resorted to a state of total minimum PE).
2. You then multiply b by z⋆ to get the PE of each parcel in this state of minimum total PE. Note that in this calculation never really resorted anything, as really we only need the height at which each parcel would sit, and Equation (11) gives us that.
3. Finally you integrate everything and then you have the total PE of the domain in a minimum PE state: the BPE. Note that, in the paper, most (all?) the calculations are done with domain-integrated quantities since the BPE is necessarily a nonlocal quantity; that's these calculation are always a bit of a hassle.

Now, if you wanna avoid calculating Equation (11), instead of manually calculating z⋆ we can, in the words of Winters et al., "let the fluid elements correspond to volume elements of a numerical grid" and actually sort buoyancy values. Then the vertical coordinate (originally z) approximately becomes z⋆ by definition, since each parcel is approximately where it would be if the values were resorted to a statement of minimum PE. I'm emphasizing "approximately" here because as we discussed, grid cells aren't infinitesimal, and their sorting only approximates the state of minimum PE, and their sorted state not only isn't unique, but in general has some small horizontal gradients.

• The advantage of this latter method is that is faster since we never explicitly calculate z∗; it implicitly appears after the sorting is complete. Basically we exchange a O(n2) operation (Equatoin (11)) for a ≈O(nlog(n)) operation.
• The disadvantage is that it's only an approximation.

Then the calculation the get the BPE through this approach is:

1. You 3D sort a field of b.
2. After resorting the original vertical coordinate z becomes z⋆. We ignore horizontal coordinates. Then we can just multiply the sorted b by its coordinate (which now approximates z⋆) to obtain the PE of each parcel in the minimum PE state.
3. Integrate everything to get the (approximated) BPE.

My take is that the latter method (i.e. sorting) will always be an approximation to the former method (i.e. direct calculation of z⋆ with Eq. (11)). Thus we can use the sorting approach as it is until an example/simulation comes along where that approximation fails and we need something better. Then we can either implement the first method (direct calculation of z⋆) or try to improve on the approximated one.

@jbisits @glwagner thoughts?

Yes this is my understanding of how to obtain an approximation for the volume integrated BPE and is what I have tried to implement in this PR.

[jbisits]

@tomchor can you please share the literature on this sorting-based calculation?

I don't think there's more literature specifically on the sorting-based calculation other than the original paper. Winter and D'Asaro touch on it again on their follow-up 1996 paper but I think that's it.

My impression that the sorting-based approach works well comes from reading papers that employed it without complications, plus talking to a few people that corroborate that impression. I'll try to remember some of the papers that I read which use that calculation and post here later (I don't remember any of them off the top of my head).

Caulfield and Peltier (2000) and Carpenter et al. (2012) compute the BPE of a 3D scalar field using a sorted 1D array. The excerpt from Carpenter et al. (2012) where they show how they do this is

In this calculation they have computed an "equivalent" (not sure if this is useful terminology) $z$ coordinate, $z_{b}$, for the sorted 1D scalar array. I have also used this method but I thought that the most simple way to have the BPE implemented in Oceanostics was reshape the sorted scalar field back on to the original grid so that Integral could be called (as pointed out by @tomchor it is the integrated quantity we almost always, or even always, want). I think this is also similar to what @glwagner wrote above?

One way to calculate this "equivalent" $z$ coordinate is to:

1. compute the volume, $\Delta V$, of each grid cell over a 3D domain
2. reshape the gridded values of $\Delta V$ into a 1D array and cumulatively sum to get the amount of volume up to a given value of iso-scalar (i.e. a given value in the sorted 1D array of scalar values)
3. divide the cumulatively summed volume array by the horizontal surface area (if it is constant, otherwise divide each volume element by each horizontal area element , $\Delta v / \Delta A$) and you have an equivalent $z$ array.

This is similar to the methods developed in Sohail et al (2021) where they reshape a 3D scalar field into a 1D array and sort to look at the amount of a scalar quantity up to a fixed volume (or equivalent depth after dividing by surface area).

I have used this method to obtain an equivalent $z$ and then compute the integrated BPE. Comparing the results (for the very small number of cases I have looked at) using this method vs the method implemented in the PR there was negligible difference in the two values.

[jbisits]

Can you be more explicit? What exactly is the calculation that you do after you obtain the sorted field (which is the resting density field, not a coordinate mapping)?

For example, equation 19 in Winters et al. 1995 is

For this we need the function f(x,t)=z⋆. How do we compute equation 19 without explicitly obtaining z⋆?

Equation (19) is for the available potential energy

$$ E{a} = g \int{V}\rho (z - z^{*}) \mathrm{d}V $$

which in Winters et al. (1995) is defined as the difference between the potential energy

$$ E{p} = g\int{V} \rho z \mathrm{d}V $$

and background potential energy

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

In our case, we can approximate $E_{b}$ by

$$ E{b}' \approx g \int{V}\rho^{*} z \mathrm{d}V, $$

where $\rho^{*}$ is the sorted 1D array of the density field that has been reshaped back onto the original grid. The function to compute the potential energy is already in Oceanostics so we can then approximate the available potential energy equation by

$$ E{a}' \approx E{p} - E_{b}' $$

and take a finite difference approximation of $E_{a}'$ to get an approximation of equation (19).

I believe this is the method suggested by Winters et al. (1995) in the section titled numerical implementation.

Is this similar to what you have used in the past @tomchor?

[tomchor]

Is this similar to what you have used in the past @tomchor?

Yes, that's pretty much it!

[tomchor]

Can you be more explicit? What exactly is the calculation that you do after you obtain the sorted field (which is the resting density field, not a coordinate mapping)?

For example, equation 19 in Winters et al. 1995 is

For this we need the function f(x,t)=z⋆. How do we compute equation 19 without explicitly obtaining z⋆?

For the calculation of all the terms in that equation, yes, we need to explicitly calculate $z\star$ (only due to the last term I think). But we're not trying to calculate that whole equation in this PR. For what we're trying to calculate (which was summarized neatly by @jbisits in https://github.com/tomchor/Oceanostics.jl/pull/172#issuecomment-2087722148 and amounts to approximating the BPE so that we can calculate the APE) an explicit calculation of $z\star$ isn't needed (as long as you're okay with the approximation and its (hopefully small) errors).

[glwagner]

Is this similar to what you have used in the past @tomchor?

Yes, that's pretty much it!

But how do you compute $z - z_\star$? The math is not the problem, it's the code that is missing. I don't see how you are doing it after obtaining the sorted density field.

[glwagner]

Caulfield and Peltier (2000) and Carpenter et al. (2012) compute the BPE of a 3D scalar field using a sorted 1D array.

This clarifies things a lot and makes much more sense to me than reshaping the sorted field back to the dimensions of the 3D field (which doesn't make sense to me, or rather I think its wrong based on the arguments I gave above).

Comparing the results (for the very small number of cases I have looked at) using this method vs the method implemented in the PR there was negligible difference in the two values.

Isn't it easy to come up with an example that breaks it?

It's dangerous to write code we know is wrong. Sometimes it ends up being more effort than the time saved by making the approximation, because we have to laboriously explain when the code is wrong and give people advice about when / when not to use it, not to mention losing sleep at night...

[jbisits]

Is this similar to what you have used in the past @tomchor?

Yes, that's pretty much it!

But how do you compute z−z⋆? The math is not the problem, it's the code that is missing. I don't see how you are doing it after obtaining the sorted density field.

You do not have to compute $z - z$ when using this method. We compute $E{p}$ and approximate $E{b}$ and take a difference which gives an approximation to the APE without calculating $z - z$ (this is method proposed by Winters et al. (1995) which we are following). As @tomchor pointed out the point of sorting the density field is

After resorting the original vertical coordinate $z$ becomes $z^*$

So using the sorted density field (I have called this $\rho$ above) and the original vertical coordinate $z$ we approximate $\rho z$ with $\rho^ z$ and then use $\rho^ z$ in the approximation for $E{b}$. This means we can approximate the two quantities $E{p}$ and $E{b}$ and take the difference to approximate $E{a}$ all without calculating $z - z*$.

[jbisits]

Caulfield and Peltier (2000) and Carpenter et al. (2012) compute the BPE of a 3D scalar field using a sorted 1D array.

This clarifies things a lot and makes much more sense to me than reshaping the sorted field back to the dimensions of the 3D field (which doesn't make sense to me, or rather I think its wrong based on the arguments I gave above).

It is just a different method to approximate the background potential energy. It can be implemented like this it just means we need to define and compute $z{b}$ (to use Carpenter et al.'s terminology) and there are probably a few ways to do this.

Comparing the results (for the very small number of cases I have looked at) using this method vs the method implemented in the PR there was negligible difference in the two values.

Isn't it easy to come up with an example that breaks it?

It's dangerous to write code we know is wrong. Sometimes it ends up being more effort than the time saved by making the approximation, because we have to laboriously explain when the code is wrong and give people advice about when / when not to use it, not to mention losing sleep at night...

I am not quite sure this indicates that I have written code here that I know is wrong. I was saying that so far I have used several methods (e.g. the method presented in this PR and something similar to Carpenter et al. (2012) with two different versions of the $z_{b}$) to obtain an estimate for the volume integrated background potential energy (BPE). The estimates were all in agreement to around 9 decimal places. I know this is not perfect and perhaps, or even likely, there is an example that breaks it. The point I was trying to get across is there is more than one method to approximate the BPE.

The method I chose to implement in this PR is based off the original suggestion by Winters et al. (1995) and is the method that, after some experiments and discussions with my supervisors, we chose to use. For stratified fluids on uniform grids it should produce a reasonable approximation to the BPE.

[glwagner]

Is this similar to what you have used in the past @tomchor?

Yes, that's pretty much it!

But how do you compute z−z⋆? The math is not the problem, it's the code that is missing. I don't see how you are doing it after obtaining the sorted density field.

You do not have to compute z−z∗ when using this method. We compute Ep and approximate Eb and take a difference which gives an approximation to the APE without calculating z−z∗ (this is method proposed by Winters et al. (1995) which we are following). As @tomchor pointed out the point of sorting the density field is

After resorting the original vertical coordinate z becomes z∗

So using the sorted density field (I have called this ρ∗ above) and the original vertical coordinate z we approximate ρz∗ with ρ∗z and then use ρ∗z in the approximation for Eb. This means we can approximate the two quantities Ep and Eb and take the difference to approximate Ea all without calculating z−z∗.

Ok, thank you for the extended implementation (I got confused in your original explanation and didn't see what you wrote below that first line). And pointing to Winters et al. 1995 is helpful --- I agree they make exactly the same argument you have been making.

The notation is a little confusing, but I don't think there is any approximation between what you define as $E_b$ and $E'_b$. It's just a change of coordinates from physical coordinate system to the "resorted coordinate system". Since $dV$ are the same between the two, you can integrate the density in the resorted coordinate system just fine.

Winters et al. 1995 state that the approximate part is the reshaping, since the reference profile should be 1D but they (like you) are using a 3D state that isn't the true minimum potential energy state (because fluid elements with different densities are side by side rather than vertically sorted).

I'm sort of scratching my head though --- isn't it quite trivial to compute the true reference density, at least on a uniform grid? After going to all of the huge amount of effort to write a CFD code and do all of these computations, we really stop short of the endzone at this point? It doesn't make sense to me. We simply have to sum the 1D sorted field in a straightforward manner... ?

If the grid is stretched things are indeed more complicated, because we have to take into account the different volumes of the elements that contribute to the reference profile. I think what I proposed aabove is a solution though. Not that hard, at least in a domain that doesn't have immersed boundaries... or am I missing something?

[jbisits]

I'm sort of scratching my head though --- isn't it quite trivial to compute the true reference density, at least on a uniform grid? After going to all of the huge amount of effort to write a CFD code and do all of these computations, we really stop short of the endzone at this point? It doesn't make sense to me. We simply have to sum the 1D sorted field in a straightforward manner... ?

Yep that it is totally true!

If the grid is stretched things are indeed more complicated, because we have to take into account the different volumes of the elements that contribute to the reference profile. I think what I proposed above is a solution though. Not that hard, at least in a domain that doesn't have immersed boundaries... or am I missing something?

Nope, I do not think it should be too hard. I will implement a form of this over the next days.

[glwagner]

I'm sort of scratching my head though --- isn't it quite trivial to compute the true reference density, at least on a uniform grid? After going to all of the huge amount of effort to write a CFD code and do all of these computations, we really stop short of the endzone at this point? It doesn't make sense to me. We simply have to sum the 1D sorted field in a straightforward manner... ?

Yep that it is totally true!

If the grid is stretched things are indeed more complicated, because we have to take into account the different volumes of the elements that contribute to the reference profile. I think what I proposed above is a solution though. Not that hard, at least in a domain that doesn't have immersed boundaries... or am I missing something?

Nope, I do not think it should be too hard. I will implement a form of this over the next days.

Nice! Doing this right seems like a lot of payoff for the effort (I guess Julia helps).

Extending to irregular domains can remain an open challenge, I have an inkling that somebody, someday will want to tackle that.

[tomchor]

Isn't it easy to come up with an example that breaks it? It's dangerous to write code we know is wrong. Sometimes it ends up being more effort than the time saved by making the approximation, because we have to laboriously explain when the code is wrong and give people advice about when / when not to use it, not to mention losing sleep at night...

I am not quite sure this indicates that I have written code here that I know is wrong. I was saying that so far I have used several methods (e.g. the method presented in this PR and something similar to Carpenter et al. (2012) with two different versions of the zb) to obtain an estimate for the volume integrated background potential energy (BPE). The estimates were all in agreement to around 9 decimal places. I know this is not perfect and perhaps, or even likely, there is an example that breaks it. The point I was trying to get across is there is more than one method to approximate the BPE.

I agree with @jbisits here. The code isn't wrong. Both sorting procedures (1D and 3D) are approximations to D'Asaro's Equation (11).

Winters et al. 1995 state that the approximate part is the reshaping, since the reference profile should be 1D but they (like you) are using a 3D state that isn't the true minimum potential energy state (because fluid elements with different densities are side by side rather than vertically sorted).

Can you point to (or post) the statement where they say that here? The way I understand it, resorting everything with finite, grid-size, indivisible volumes is a source of error from the start.

I'm sort of scratching my head though --- isn't it quite trivial to compute the true reference density, at least on a uniform grid? After going to all of the huge amount of effort to write a CFD code and do all of these computations, we really stop short of the endzone at this point? It doesn't make sense to me. We simply have to sum the 1D sorted field in a straightforward manner... ?

Yep that it is totally true!

If the grid is stretched things are indeed more complicated, because we have to take into account the different volumes of the elements that contribute to the reference profile. I think what I proposed above is a solution though. Not that hard, at least in a domain that doesn't have immersed boundaries... or am I missing something?

Nope, I do not think it should be too hard. I will implement a form of this over the next days.

Nice! Doing this right seems like a lot of payoff for the effort (I guess Julia helps).

Extending to irregular domains can remain an open challenge, I have an inkling that somebody, someday will want to tackle that.

After all this back and forth I'm a bit confused. What are you calling the "right" approach here? IMO the "right" way to do things (i.e. the one with the fewest approximations) is implementing their Equation (11) and directly calculating $z_\star$. That approach can be used for both regular and stretched grids, as well as immersed boundary grids.

[jbisits]

After all this back and forth I'm a bit confused. What are you calling the "right" approach here? IMO the "right" way to do things (i.e. the one with the fewest approximations) is implementing their Equation (11) and directly calculating z⋆. That approach can be used for both regular and stretched grids, as well as immersed boundary grids.

Yes I think that is correct --- the most accurate method is implementing $z$. Perhaps we could go with your initial suggestion of putting approximate into the function name (e.g. ApproximateBackgroudPotentialEnergy) then look to implement $z$ in another PR?

[glwagner]

I agree with @jbisits here. The code isn't wrong. Both sorting procedures (1D and 3D) are approximations to D'Asaro's Equation (11).

It's just semantics. An approximation is "wrong", by definition. That's what I meant.

Can you point to (or post) the statement where they say that here?

The statement is here:

Doing an exact computation merely means generating a sorted density field that is parallel to gravity (and doing that correctly). It's not very hard...

IMO the "right" way to do things (i.e. the one with the fewest approximations) is implementing their Equation (11)

The only approximation is the reshaping. If we generate the reference state correctly, then we can evaluate the integrated background potential energy. This does not require explicitly computing $z_\star$.

However, if one wishes to compute $z_\star$ for other reasons (if we need more than just the integrated background potential energy), then more work is required. Equation (11) is one way, but I think there are more efficient ways as I have outlined, that would introduce an additional fact only in the vertical direction (rather than computing a 3D integral at every point, as equation 11 does).

[tomchor]

Okay I think I finally understand what @glwagner is proposing. Correct me if I'm wrong but what you want to do is to deform (flatten, really) each grid cell when sorting, so that they'd go from having a volume $\Delta v = \Delta x \Delta y \Delta z_i$ to having volume $\Delta v = Lx Ly \Delta z_f$. That way they can be stacked on top of each other (via a 1D sorting) and the final result will be a sorted domain with the same dimensions we started with. I guess that's what Carpenter et al. (2012) were describing here, but it definitely wasn't clear to me when reading that passage. Pretty ingenious! And I agree that's better.

A note is that while this can be generalized to stretched grids if you're careful with with spacings as @glwagner mentioned, it doesn't generalize to ImmersedGrids or (I think) cases where gravity isn't aligned with the grid's z direction. For that (plus computing other terms in the APE equation) we'd need a direct calculation of $z_\star$, which we probably should implement in another PR.

I think we're finally all on the same page now! :rocket:

[glwagner]

A note is that while this can be generalized to stretched grids if you're careful with with spacings as @glwagner mentioned, it doesn't generalize to ImmersedGrids or (I think) cases where gravity isn't aligned with the grid's z direction. For that (plus computing other terms in the APE equation) we'd need a direct calculation of z⋆, which we probably should implement in another PR.

I think we're finally all on the same page now! 🚀

Agree! Getting this right for complex domains requires some extra work. But it's ok to leave a few things for the next generation.

It might good to throw error when gravity direction is not z or when using immersed boundary grid, also to help people who may be interested in contributing that extension.

[tomchor]

It might good to throw error when gravity direction is not z or when using immersed boundary grid, also to help people who may be interested in contributing that extension.

For sure!

[jbisits]

I am still in the process of getting this finished but do keep comments as will likely speed things up. Once it is ready I will write up everything in the PR intro box (does this have a better name?).

[hdrake]

Just a note that @ikeshwani has been independently computing BPE, for now just by brute-force offline for the non-immersed case, but eventually including the ImmersedBoundary case. @liuchihl is probably interested in computing BPE with tilted gravity vector and/or on domains that are periodic in the direction of the gravity vector (e.g. where there is a stratified background buoyancy field).

We'll be keeping a close eye on these developments and hope to contribute to generalizations in these more complicated contexts.

[hdrake]

@jbisits, have you tested the new diagnostics on anything more complicated than a 1D profile and a random 2D field? @ikeshwani implemented a nice notebook that qualitatively replicates the Winters et al. (1995) shear instability example, which you might find useful.

[jbisits]

@jbisits, have you tested the new diagnostics on anything more complicated than a 1D profile and a random 2D field? @ikeshwani implemented a nice notebook that qualitatively replicates the Winters et al. (1995) shear instability example, which you might find useful.

Not the diagnostics in this PR (I have used PotentialEnergy in a DNS) but it would be great to use the notebook you have shared and this branch as an example. I will get to it soon - Thanks very much for sharing this!!!

It could be a good addition to the documentation once this PR is finished if @ikeshwani is interested?

[tomchor]

It could be a good addition to the documentation once this PR is finished if @ikeshwani is interested?

Agreed! To me it sounds like it'd be great to merge @ikeshwani's notebook with the Kelvin-Helmholtz example that's already in the docs. That example performs some very simple diagnostics, and it would be great to use that to reproduce results from Winters et al. (1995) in addition to what's already there.

[hdrake]

I've finally caught up on this whole thread and want to clarify a few things:

Regarding approximations

The sorting procedure is only an approximation to compute the BPE, so I think this should be mentioned somewhere in the docstring. (Or maybe even reflected in the function name? ApproximateBackgroundPotentialEnergy?)

@tomchor, could you clarify what the approximation being made is? My understanding is that the method is only approximate for equations of state that depend on pressure, and are exact for linear EOS or a buoyancy tracer. Some negligible error might be introduced if the discretization method does not align perfectly with the model's grid discretization, but that is a fundamentally different thing than saying we are making a physical approximation to the Boussinesq BPE. (Or maybe you are saying that Boussinesq BPE is already an approximate to non-Boussinesq BPE?) But the physics of our diagnostics should be consistent with our model physics.

Regarding irregular grids

Even then, sorting is only valid when the grid is regular (i.e. all the grid cells have equal volume). So we should test for that and spit out a useful error otherwise. (Or if you're up for it, implement the appropriate integral-form calculation of z*, as it would be useful to have something that works on stretched grids.)

@tomchor, the sorting method is easily generalized to account for unequal grid cell volumes; you just have to account for the actual volume of a parcel when you stack up the elevation increments to compute $z^{\star}$. Maybe this is what you had in mind, but you don't actually need to evaluate the integrals of the heaviside directly, which is an prohibitively expensive $\mathcal{O}(N^{2})$ operation.

Regarding global sorting, the sorted buoyancy field $b(z^\star)$, and horizontal averaging

But wouldn't any horizontal structure be destroyed if we resort in the horizontal? Consider doing this for the global ocean. It's cold at the poles and warm on the equator. Resorting in the horizontal would put cold water at the north pole and warm water at the south pole. That doesn't make sense. We must be assuming horizontal homogeneity. Is this necessary?

@glwagner and @jbisits: there is no need to sort horizontally, but the initial sort does need to be global. [Aside: Kial Stewart et al (2014) argue that it's justifiable to ignore topographic barriers in the adiabatic adjustment (or sorting) process in the computation of APE tendencies but not the actual pool of APE, which is very sensitive to the sorting method.]

The final sorted density field should indeed only be a function of depth. In practice, @glwagner is correct in pointing out that the sorting method results in . This can be very problematic if the horizontal-mean vertical density difference between two layers is smaller or comparable to the density differences between the $N{xy} = N{x} N_{y}$ sorted grid cells that make up a layer of the sorted buoyancy field array.

However, I want to emphasize that we don't actually need to use the sorted buoyancy field for the calculation of BPE; just $z^{\star}$, which is unaffected by the ordering of the initial array indices. While BPE and APE themselves will not be affected by the incorrect inclusion of horizontal buoyancy gradients in $b(z^{\star})$, some other terms in the APE budget do require $b(z^{\star})$ and so we will need to think more critically about something like a horizontal average.

It's easy to construct a counter example. If we simply flip the definition of the y coordinate (so that increasing j-index implies decreasing latitude), we end up with a different background energy state. But there's only one physical background energy state --- right? This is true, but want to highlight again that this would not impact the calculation of $z^{\star}$, only the reconstructed $b(z^{\star})$ array.

Possibly the solution is as simple as averaging the sorted state in the horizontal directions.

@glwagner, this would give a unique solution regardless of flipping the horizontal indices, right? I think this is the best approach.

I think averaging the horizontal directions would go against the premise of the BPE, which is an adiabatic resorting of the buoyancy field by definition.

@tomchor, I disagree; if you are only applying the average to the sorted buoyancy field then there is no conflict with the concept of adiabatic adjustment. You are merely correcting for the fact that we are trying to shove the adiabatically-sorted grid cells back into a vertically-coarse discrete grid.

Your point of forcing the numerical, discrete background state to also be a function of only is interesting, but I'm not sure modifying the sorting procedure is the way to go.

@tomchor, I think what @glwagner is proposing (and which I agree with) is that we only do the horizontal averaging as a very last step, which does not actually affect the sorting process itself. The $z^{\star}$ array that really encodes the adiabatic sorting information would be left intact and not averaged.

I'd be reluctant to add an averaging procedure whose effects on the calculations we haven't investigated before. So I think if we do decide that we want to average in the horizontal, we need to do some investigating first with "realistic" simulations (whatever we think that means).

@tomchor, I agree on this point. Here is a proof of concept implementation of the sorting-based method to compute $z^{\star}$ in @ikeshwani's horizontal convection configuration (with immersed boundaries, because why not)! My main takeaway is that, at least in this configuration, horizontally-averaging the sorted buoyancy field has a huge difference on the near-surface reference buoyancy $b(z^{\star}=0)$.

I feel like the equation above will be less performant (I think it'll always scale with $O(N^{2}$)), but we can understand that as price to pay for precision.

@tomchor, I think a correctly implemented sorting method will give identical results to the brute-force $O(N^{2})$ approach. @ikeshwani and I will investigate for a simple shear instability case in which we can actually afford the cost of $O(N^{2})$ and we have now implemented both methods (for offline buoyancy field snapshots).

[hdrake]

I've implemented a very efficient offline algorithm for computing $z^{\star}$ with the sorting method [in this notebook], that in principle handles both variable grid spacing and immersed boundaries. I don't yet understand the kernel techniques required to implement this on GPUs well enough myself, but I think it should be obvious from the notebook how to extend @jbisits' current diagnostic to handle these more general configurations.

Two modifications to the algorithm are required to generalize the method to support immersed boundaries: 1) Ocean area changes with depth: The vertical cumsum to get zstar needs to be done via a for-loop because the area needs to be updated as you go up in the water column and the horizontal ocean area at the sorted depth changes. I do this by creating an interpolation function for $A(z)$ and then interpolating to every $z^{star}$, but that may be overkill and a nearest-neighbor look-up could suffice. 2) Skipping land cells when iterating through grid cells: By default, sort and the more useful sortperm functions throw NaN values to the end of the sorted array, so I used this to essentially skip over the land cells. Also, when you map the sorted buoyancy field (as a 1D vector) back into the 3D array, you need to be careful to skip the land cells!

[hdrake]

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^{\star}(x,y,z,t)$ (as a 4D array) and $\frac{\text{d}z^{\star}}{\text{d}b}\Big|_{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^{\star})$ and $z^{\star}(b)$ as a sorted array, it is straight-forward to estimate the derivative $\frac{\text{d}z^{\star}}{\text{d}b}$ 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.)

[tomchor]

I've finally caught up on this whole thread and want to clarify a few things:

Regarding approximations

The sorting procedure is only an approximation to compute the BPE, so I think this should be mentioned somewhere in the docstring. (Or maybe even reflected in the function name? ApproximateBackgroundPotentialEnergy?)

@tomchor, could you clarify what the approximation being made is? My understanding is that the method is only approximate for equations of state that depend on pressure, and are exact for linear EOS or a buoyancy tracer. Some negligible error might be introduced if the discretization method does not align perfectly with the model's grid discretization, but that is a fundamentally different thing than saying we are making a physical approximation to the Boussinesq BPE. (Or maybe you are saying that Boussinesq BPE is already an approximate to non-Boussinesq BPE?) But the physics of our diagnostics should be consistent with our model physics.

At the time we were planning to do a simple 1D sorting, and the reshape everything back to 3D, exactly like described in the original paper. This would incur horizontal gradients, which is what I was talking about. Since then, after discussion, we decided to do a 1D sorting only the vertical and flatten the cell volumes, without any 3D reshaping (you can follow the previous comments for more details). With this, there's no further approximating other than the ones you mention.

Regarding irregular grids

Even then, sorting is only valid when the grid is regular (i.e. all the grid cells have equal volume). So we should test for that and spit out a useful error otherwise. (Or if you're up for it, implement the appropriate integral-form calculation of z*, as it would be useful to have something that works on stretched grids.)

@tomchor, the sorting method is easily generalized to account for unequal grid cell volumes; you just have to account for the actual volume of a parcel when you stack up the elevation increments to compute z⋆. Maybe this is what you had in mind, but you don't actually need to evaluate the integrals of the heaviside directly, which is an prohibitively expensive O(N2) operation.

Agreed. Again, this comment was before we decided to do a 1D sorting along with cell flattening. Note that we still probably want to implement the expensive $z_\star$ calculation at some point since it's still needed for tilted domains or domains with immersed boundaries. I wouldn't say it's prohibitively expensive. But yes, definitely expensive since it scales as $N^2$. (Although at some point you mentioned that it can scale as $nlog(n)$? I don't see how but hopefully you're right!)

[hdrake]

Note that we still probably want to implement the expensive $z^{\star}$ 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.

Although at some point you mentioned that it can scale as $n\log(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 $n\log(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 $n^{2}$ algorithm.