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 new Fields are then used to compute the BPE per unit volume. Calling Integral(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$... ?

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

[tomchor]

This line doesn't read well. But I also don't have immediate suggestions on how to improve it... irrc the whole nomenclature of BuoyancyModel is kinda confusing and we discussed changing it over at Oceananigans a few times.

Maybe changing it from buoyancy_model to buoyancy_treatment might help...?

[tomchor]

This is the core of the PR, right? It looks correct to me, but we should make sure to test that the area/volume calculations are correct with a stretched grid.

It seems to me that this works for horizontally-stretched grids as well, correct?

[tomchor]

When comparing b_sorted and sorted_buoiyancy, it's always gonna return true by construction, right? Since b_sorted is calculated with the permutations that are used to calculate sorted_buoyancy? I think we should still keep this test (since it tests that the code is working and all), but maybe we need another one that's a little more rigorous.

[tomchor]

Would starting with linear buoyancy distribution in the x direction (where we know db/dx) and sorting it, making sure that db/dz matches, work? We could try that for grids that are stretched in the horizontal and vertical directions separately. If we could test something along those lines it would be great.

[tomchor]

The comment above extends to these other model.buoyancy types. A separate comment here is that you're writing a lot of code 3 times. You can probably reuse code, so that you write most things only once.

[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!

[jbisits]

Sorry to have fallen asleep on this PR! The way I was implementing this was to define a OneDReferenceField which has the total number of elements from the domain reshaped into a vector over the vertical extent of the domain,

function OneDReferenceField(f::Union{Field, AbstractOperation}; rev = false)

    area = sum(AreaField(f.grid))
    volume_field = VolumeField(f.grid)
    v = reshape(interior(volume_field), :)
    field_data = reshape(interior(f), :)

    p = sortperm(field_data; rev)
    sorted_field_data = field_data[p]
    z✶ = -cumsum(v[p]) / area # divide by area at surface

    grid_arch = f.grid.architecture
    grid_size = prod(size(f.grid))
    new_grid = RectilinearGrid(grid_arch, size = grid_size, z = (-f.grid.Lz, 0), topology=(Flat, Flat, Bounded))

    sorted_field = CenterField(new_grid)
    set!(sorted_field, reshape(sorted_field_data, size(new_grid)))
    z✶_field = CenterField(new_grid)
    set!(z✶_field, reshape(z✶, size(new_grid)))

    return sorted_field, z✶_field
end

It looks like CumulativeIntegral could simplify some of the calculation in what I have above. I don't really like the implementation above as there is a lot going on at each time step so I tried for some time to move this to a KernelFunctionOperation, or look for parts that could be simplified/optimized, but could not figure out how to and quite a bit of time went by very quickly (again my apologies!). I know there is quite a bit to be done for this calculation but returning new fields computed from modified model fields with new grids even though it works seems like a lot of overhead.

Further, for this to work on grids that do not have uniform horizontal area I think the area at each depth of the water column still needs to be calculated (as pointed out here). With that in mind is there a way that @hdrake's method for non-uniform horizontal area to be put into a GPU kernel? I am happy to look into this if that is something that will give the best results.

@tomchor depending on what you think about the method above it may also be worth closing this PR and implementing any new changes/different methods in a new PR (maybe with this PR converted to an issue though have not looked into if that can be done). Or I can try to improve on what is already in place in this PR --- maybe the hiatus from this PR will have helped! I will leave everything as is until a decision is made.

Again sorry to have let this go!

[tomchor]

@jbisits don't worry about it! We all have lots to do we have to prioritize :)

It looks like CumulativeIntegral could simplify some of the calculation in what I have above. I don't really like the implementation above as there is a lot going on at each time step so I tried for some time to move this to a KernelFunctionOperation, or look for parts that could be simplified/optimized, but could not figure out how to and quite a bit of time went by very quickly (again my apologies!). I know there is quite a bit to be done for this calculation but returning new fields computed from modified model fields with new grids even though it works seems like a lot of overhead.

I don't think this can be implemented using KernelFunctionOperations, since you're rearranging the grid. But I'd be happy to be proven wrong.

Further, for this to work on grids that do not have uniform horizontal area I think the area at each depth of the water column still needs to be calculated (as pointed out here). With that in mind is there a way that @hdrake's method for non-uniform horizontal area to be put into a GPU kernel? I am happy to look into this if that is something that will give the best results.

@tomchor depending on what you think about the method above it may also be worth closing this PR and implementing any new changes/different methods in a new PR (maybe with this PR converted to an issue though have not looked into if that can be done). Or I can try to improve on what is already in place in this PR --- maybe the hiatus from this PR will have helped! I will leave everything as is until a decision is made.

Again sorry to have let this go!

It does look like @hdrake's method would give the best results. I haven't looked at that method in a while (so I'll let @hdrake chime in here) but if this PR has something that already works (even though it's limited to regular grids and you're not 100% happy with the main function) I'd say it's okay to merge. Then we can improve this design later or even replace it if we think it's necessary. Personally, I think having something you can already use helps to visualize better ways of doing things.

[jbisits]

...if this PR has something that already works (even though it's limited to regular grids and you're not 100% happy with the main function) I'd say it's okay to merge. Then we can improve this design later or even replace it if we think it's necessary. Personally, I think having something you can already use helps to visualize better ways of doing things.

Ok sounds good. I will try to get this cleaned up and checked over the next week or so!