From 5efa6a6056488193b63f5637daafc142628d8163 Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Fri, 29 Mar 2024 09:10:46 +1100
Subject: [PATCH 01/37] Try some sorting idea for background state

---
 src/PotentialEnergyEquationTerms.jl | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index fb208222..1c91383c 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -3,6 +3,7 @@ module PotentialEnergyEquationTerms
 using DocStringExtensions
 
 export PotentialEnergy
+export sort_field, sorted_field
 
 using Oceananigans.AbstractOperations: KernelFunctionOperation
 using Oceananigans.Models: seawater_density
@@ -12,6 +13,7 @@ using Oceananigans.Grids: NegativeZDirection
 using Oceananigans.BuoyancyModels: Buoyancy, BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState
 using Oceananigans.BuoyancyModels: buoyancy_perturbationᶜᶜᶜ, Zᶜᶜᶜ
 using Oceananigans.Models: ShallowWaterModel
+using Oceananigans.Fields: Field
 using Oceanostics: validate_location
 using SeawaterPolynomials: BoussinesqEquationOfState
 
@@ -184,4 +186,23 @@ end
 
 @inline g′z_ccc(i, j, k, grid, ρ, p) = (p.g / p.ρ₀) * ρ[i, j, k] * Zᶜᶜᶜ(i, j, k, grid)
 
+## Background Potential Energy drafts
+
+resorted_field(i, j, k, grid, sf) = sf[i, j, k]
+function sorted_field(f::Field)
+    grid = f.grid
+    sorted_f = reshape(sort(reshape(f, :)), size(f))
+    return KernelFunctionOperation{Center, Center, Center}(resorted_field, grid, sorted_f)
+end
+
+function background_potential_energy(model)
+
+    return nothing
+end
+
+## Background potential energy would look like
+# Eb = gρ✶z/ρ₀ where ρ✶ is the sorted density field
+# or Eb = gρz✶/ρ₀ where z✶ = z[sortperm(ρ)]
+# so could have g/ρ₀ * ρ✶[i, j, k] * zᶜᶜᶜ(i, j, k, grid)
+# how to minimise the computation here? ie can I just return one kernel function operation?
 end # module

From 711369774813e023d2b068841fde241a62ec4bab Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Tue, 2 Apr 2024 09:37:48 +1100
Subject: [PATCH 02/37] First attempt at bpe implementation

---
 src/PotentialEnergyEquationTerms.jl | 129 ++++++++++++++++++++++++----
 1 file changed, 113 insertions(+), 16 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 1c91383c..818c63d1 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -2,8 +2,7 @@ module PotentialEnergyEquationTerms
 
 using DocStringExtensions
 
-export PotentialEnergy
-export sort_field, sorted_field
+export PotentialEnergy, BackgroundPotentialEnergy
 
 using Oceananigans.AbstractOperations: KernelFunctionOperation
 using Oceananigans.Models: seawater_density
@@ -13,7 +12,7 @@ using Oceananigans.Grids: NegativeZDirection
 using Oceananigans.BuoyancyModels: Buoyancy, BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState
 using Oceananigans.BuoyancyModels: buoyancy_perturbationᶜᶜᶜ, Zᶜᶜᶜ
 using Oceananigans.Models: ShallowWaterModel
-using Oceananigans.Fields: Field
+using Oceananigans.Fields: Field, compute!
 using Oceanostics: validate_location
 using SeawaterPolynomials: BoussinesqEquationOfState
 
@@ -186,23 +185,121 @@ end
 
 @inline g′z_ccc(i, j, k, grid, ρ, p) = (p.g / p.ρ₀) * ρ[i, j, k] * Zᶜᶜᶜ(i, j, k, grid)
 
-## Background Potential Energy drafts
-
-resorted_field(i, j, k, grid, sf) = sf[i, j, k]
-function sorted_field(f::Field)
+# For testing `sort`ed fields.
+@inline resorted_field(i, j, k, grid, sf) = sf[i, j, k]
+@inline function sort_field(f::Field)
     grid = f.grid
-    sorted_f = reshape(sort(reshape(f, :)), size(f))
-    return KernelFunctionOperation{Center, Center, Center}(resorted_field, grid, sorted_f)
+    sorted_field = reshape(sort(reshape(f, :)), size(f))
+    return KernelFunctionOperation{Center, Center, Center}(resorted_field, grid, sorted_field)
 end
 
-function background_potential_energy(model)
+"""
+    $(SIGNATURES)
+Return a `kernelFunctionOperation` to compute the `BackgroundPotentialEnergy` per unit
+volume,
+```math
+Eₚ = \\frac{gρ✶z}{ρ₀}.
+```
+The `BackgroundPotentialEnergy` is the potential energy computed after adiabatically resorting
+the buoyancy or density field into a reference state of minimal potential energy.
+The reference state is computed by reshaping the gridded buoyancy or density field and
+`sort`ing into a monotonically increasing `Vector`. This `sort`ed vector is then reshaped into
+the `size(model.grid)`.
+
+Examples
+========
+
+```jldoctest
+julia> using Oceananigans
+
+julia> using Oceanostics.PotentialEnergyEquationTerms: PotentialEnergy
+
+julia> grid = RectilinearGrid(size=5, z=(-5, 0), topology=(Flat, Flat, Bounded))
+1×1×5 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
+├── Flat x
+├── Flat y
+└── Bounded  z ∈ [-5.0, 0.0]      regularly spaced with Δz=1.0
+
+julia> model = NonhydrostaticModel(; grid, buoyancy=BuoyancyTracer(), tracers=(:b,))
+NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
+├── grid: 1×1×5 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
+├── timestepper: QuasiAdamsBashforth2TimeStepper
+├── tracers: b
+├── closure: Nothing
+├── buoyancy: BuoyancyTracer with ĝ = NegativeZDirection()
+└── coriolis: Nothing
+
+julia> set!(model, b = randn(size(model.grid)));
+
+julia> interior(model.tracers.b, 1, 1, :)
+5-element view(::Array{Float64, 3}, 1, 1, 4:8) with eltype Float64:
+  1.6132888274716937
+ -0.7877032120116347
+ -0.9357749288520266
+  1.310757334669075
+  0.09514543268446216
+
+julia> bpe = BackgroundPotentialEnergy(model)
+KernelFunctionOperation at (Center, Center, Center)
+├── grid: 1×1×5 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
+├── kernel_function: bz_ccc (generic function with 2 methods)
+└── arguments: ("1×1×5 Array{Float64, 3}",)
+
+julia> bpe_field = Field(bpe);
+
+julia> compute!(bpe_field);
+
+julia> interior(bpe_field, 1, 1, :)
+5-element view(::Array{Float64, 3}, 1, 1, 4:8) with eltype Float64:
+  4.210987179834119
+  2.7569612420407217
+ -0.2378635817111554
+ -1.9661360020036125
+ -0.8066444137358468
+```
+"""
+@inline function BackgroundPotentialEnergy(model; location = (Center, Center, Center),
+                                                  geopotential_height = model_geopotential_height(model))
+
+    validate_location(location, "BackgroundPotentialEnergy")
+    isnothing(model.buoyancy) ? nothing : validate_gravity_unit_vector(model.buoyancy.gravity_unit_vector)
+
+    return BackgroundPotentialEnergy(model, model.buoyancy, geopotential_height)
+end
+
+@inline function BackgroundPotentialEnergy(model, buoyancy_model::BuoyancyTracerModel, geopotential_height)
+
+    grid = model.grid
+    b✶ = reshape(sort(reshape(model.tracers.b, :)), size(grid))
+
+    return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b✶)
+end
+
+linear_eos_buoyancy(grid, buoyancy, tracers) = KernelFunctionOperation{Center, Center, Center}(buoyancy_perturbationᶜᶜᶜ, grid, buoyancy, tracers)
+
+@inline function BackgroundPotentialEnergy(model, buoyancy_model::BuoyancyLinearEOSModel, geopotential_height)
+
+    grid = model.grid
+    buoyancy = model.buoyancy.model
+    tracers = model.tracers
+    b = linear_eos_buoyancy(grid, buoyancy, tracers)
+    b = Field(b)
+    compute!(b)
+    b✶ = reshape(sort(reshape(b, :)), size(grid))
+
+    return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b✶)
+end
+
+@inline function BackgroundPotentialEnergy(model, buoyancy_model::BuoyancyBoussinesqEOSModel, geopotential_height)
+
+    grid = model.grid
+    ρ = Field(seawater_density(model; geopotential_height))
+    compute!(ρ)
+    ρ✶ = reshape(sort(reshape(ρ, :)), size(grid))
+    parameters = (g = model.buoyancy.model.gravitational_acceleration,
+                  ρ₀ = model.buoyancy.model.equation_of_state.reference_density)
 
-    return nothing
+    return KernelFunctionOperation{Center, Center, Center}(g′z_ccc, grid, ρ✶, parameters)
 end
 
-## Background potential energy would look like
-# Eb = gρ✶z/ρ₀ where ρ✶ is the sorted density field
-# or Eb = gρz✶/ρ₀ where z✶ = z[sortperm(ρ)]
-# so could have g/ρ₀ * ρ✶[i, j, k] * zᶜᶜᶜ(i, j, k, grid)
-# how to minimise the computation here? ie can I just return one kernel function operation?
 end # module

From 1590116fce72598001bfb2d886e3800631491c6d Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Tue, 2 Apr 2024 12:10:39 +1100
Subject: [PATCH 03/37] Change docstring example

---
 src/PotentialEnergyEquationTerms.jl | 34 +++++------------------------
 1 file changed, 6 insertions(+), 28 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 818c63d1..a32e2943 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -214,48 +214,26 @@ julia> using Oceananigans
 
 julia> using Oceanostics.PotentialEnergyEquationTerms: PotentialEnergy
 
-julia> grid = RectilinearGrid(size=5, z=(-5, 0), topology=(Flat, Flat, Bounded))
-1×1×5 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
+julia> grid = RectilinearGrid(size=100, z=(-1000, 0), topology=(Flat, Flat, Bounded))
+1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
 ├── Flat x
 ├── Flat y
-└── Bounded  z ∈ [-5.0, 0.0]      regularly spaced with Δz=1.0
+└── Bounded  z ∈ [-1000.0, 0.0]   regularly spaced with Δz=10.0
 
 julia> model = NonhydrostaticModel(; grid, buoyancy=BuoyancyTracer(), tracers=(:b,))
 NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
-├── grid: 1×1×5 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
+├── grid: 1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
 ├── timestepper: QuasiAdamsBashforth2TimeStepper
 ├── tracers: b
 ├── closure: Nothing
 ├── buoyancy: BuoyancyTracer with ĝ = NegativeZDirection()
 └── coriolis: Nothing
 
-julia> set!(model, b = randn(size(model.grid)));
-
-julia> interior(model.tracers.b, 1, 1, :)
-5-element view(::Array{Float64, 3}, 1, 1, 4:8) with eltype Float64:
-  1.6132888274716937
- -0.7877032120116347
- -0.9357749288520266
-  1.310757334669075
-  0.09514543268446216
-
 julia> bpe = BackgroundPotentialEnergy(model)
 KernelFunctionOperation at (Center, Center, Center)
-├── grid: 1×1×5 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
+├── grid: 1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
 ├── kernel_function: bz_ccc (generic function with 2 methods)
-└── arguments: ("1×1×5 Array{Float64, 3}",)
-
-julia> bpe_field = Field(bpe);
-
-julia> compute!(bpe_field);
-
-julia> interior(bpe_field, 1, 1, :)
-5-element view(::Array{Float64, 3}, 1, 1, 4:8) with eltype Float64:
-  4.210987179834119
-  2.7569612420407217
- -0.2378635817111554
- -1.9661360020036125
- -0.8066444137358468
+└── arguments: ("1×1×100 Array{Float64, 3}",)
 ```
 """
 @inline function BackgroundPotentialEnergy(model; location = (Center, Center, Center),

From 8ace5a80cbde3d2ee474232b7f0bf198f0fb68cf Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Tue, 2 Apr 2024 12:48:38 +1100
Subject: [PATCH 04/37] Add space

---
 src/PotentialEnergyEquationTerms.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index a32e2943..ce11299d 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -195,6 +195,7 @@ end
 
 """
     $(SIGNATURES)
+
 Return a `kernelFunctionOperation` to compute the `BackgroundPotentialEnergy` per unit
 volume,
 ```math

From 5d2eee14a543667add22f008f43d7752ad136476 Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Wed, 3 Apr 2024 14:54:18 +1100
Subject: [PATCH 05/37] `sort` reverse so that matches `z`

---
 src/PotentialEnergyEquationTerms.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index ce11299d..3db7a8e2 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -189,7 +189,7 @@ end
 @inline resorted_field(i, j, k, grid, sf) = sf[i, j, k]
 @inline function sort_field(f::Field)
     grid = f.grid
-    sorted_field = reshape(sort(reshape(f, :)), size(f))
+    sorted_field = reshape(sort(reshape(f, :), rev = true), size(f))
     return KernelFunctionOperation{Center, Center, Center}(resorted_field, grid, sorted_field)
 end
 
@@ -249,7 +249,7 @@ end
 @inline function BackgroundPotentialEnergy(model, buoyancy_model::BuoyancyTracerModel, geopotential_height)
 
     grid = model.grid
-    b✶ = reshape(sort(reshape(model.tracers.b, :)), size(grid))
+    b✶ = reshape(sort(reshape(model.tracers.b, :), rev = true), size(grid))
 
     return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b✶)
 end
@@ -264,7 +264,7 @@ linear_eos_buoyancy(grid, buoyancy, tracers) = KernelFunctionOperation{Center, C
     b = linear_eos_buoyancy(grid, buoyancy, tracers)
     b = Field(b)
     compute!(b)
-    b✶ = reshape(sort(reshape(b, :)), size(grid))
+    b✶ = reshape(sort(reshape(b, :), rev = true), size(grid))
 
     return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b✶)
 end
@@ -274,7 +274,7 @@ end
     grid = model.grid
     ρ = Field(seawater_density(model; geopotential_height))
     compute!(ρ)
-    ρ✶ = reshape(sort(reshape(ρ, :)), size(grid))
+    ρ✶ = reshape(sort(reshape(ρ, :), rev = true), size(grid))
     parameters = (g = model.buoyancy.model.gravitational_acceleration,
                   ρ₀ = model.buoyancy.model.equation_of_state.reference_density)
 

From 84265a8cb8ac7caed9086c5e4555649c09e938c1 Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Wed, 3 Apr 2024 15:37:36 +1100
Subject: [PATCH 06/37] `sort(rev=true)` for density only (not buoyancy)

---
 src/PotentialEnergyEquationTerms.jl | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 3db7a8e2..573e0f94 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -188,6 +188,11 @@ end
 # For testing `sort`ed fields.
 @inline resorted_field(i, j, k, grid, sf) = sf[i, j, k]
 @inline function sort_field(f::Field)
+    grid = f.grid
+    sorted_field = reshape(sort(reshape(f, :)), size(f))
+    return KernelFunctionOperation{Center, Center, Center}(resorted_field, grid, sorted_field)
+end
+@inline function sort_field_rev(f::Field)
     grid = f.grid
     sorted_field = reshape(sort(reshape(f, :), rev = true), size(f))
     return KernelFunctionOperation{Center, Center, Center}(resorted_field, grid, sorted_field)
@@ -249,7 +254,7 @@ end
 @inline function BackgroundPotentialEnergy(model, buoyancy_model::BuoyancyTracerModel, geopotential_height)
 
     grid = model.grid
-    b✶ = reshape(sort(reshape(model.tracers.b, :), rev = true), size(grid))
+    b✶ = reshape(sort(reshape(model.tracers.b, :)), size(grid))
 
     return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b✶)
 end
@@ -264,7 +269,7 @@ linear_eos_buoyancy(grid, buoyancy, tracers) = KernelFunctionOperation{Center, C
     b = linear_eos_buoyancy(grid, buoyancy, tracers)
     b = Field(b)
     compute!(b)
-    b✶ = reshape(sort(reshape(b, :), rev = true), size(grid))
+    b✶ = reshape(sort(reshape(b, :), rev = true))
 
     return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b✶)
 end

From 44d81c5031de7ee90e1cdfe891b9cd2718c8c7be Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Wed, 3 Apr 2024 15:45:29 +1100
Subject: [PATCH 07/37] Removed wrong thing

---
 src/PotentialEnergyEquationTerms.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 573e0f94..a93a64d5 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -269,7 +269,7 @@ linear_eos_buoyancy(grid, buoyancy, tracers) = KernelFunctionOperation{Center, C
     b = linear_eos_buoyancy(grid, buoyancy, tracers)
     b = Field(b)
     compute!(b)
-    b✶ = reshape(sort(reshape(b, :), rev = true))
+    b✶ = reshape(sort(reshape(b, :)), size(grid))
 
     return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b✶)
 end

From 9362a271bef83622670dd400a12578c96e480910 Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Wed, 3 Apr 2024 16:00:22 +1100
Subject: [PATCH 08/37] Correction in docstring

---
 src/PotentialEnergyEquationTerms.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index a93a64d5..714d3a0f 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -204,7 +204,7 @@ end
 Return a `kernelFunctionOperation` to compute the `BackgroundPotentialEnergy` per unit
 volume,
 ```math
-Eₚ = \\frac{gρ✶z}{ρ₀}.
+E_{b} = \\frac{gρ✶z}{ρ₀}.
 ```
 The `BackgroundPotentialEnergy` is the potential energy computed after adiabatically resorting
 the buoyancy or density field into a reference state of minimal potential energy.

From fbe8c0219081b2943935df4dd23d7ea7937851e9 Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Mon, 22 Apr 2024 10:49:08 +1000
Subject: [PATCH 09/37] Cleanup function definitions

---
 src/PotentialEnergyEquationTerms.jl | 54 +++++++++++++++++++++++------
 1 file changed, 43 insertions(+), 11 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 714d3a0f..565a912f 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -12,7 +12,7 @@ using Oceananigans.Grids: NegativeZDirection
 using Oceananigans.BuoyancyModels: Buoyancy, BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState
 using Oceananigans.BuoyancyModels: buoyancy_perturbationᶜᶜᶜ, Zᶜᶜᶜ
 using Oceananigans.Models: ShallowWaterModel
-using Oceananigans.Fields: Field, compute!
+using Oceananigans.Fields: Field
 using Oceanostics: validate_location
 using SeawaterPolynomials: BoussinesqEquationOfState
 
@@ -185,19 +185,35 @@ end
 
 @inline g′z_ccc(i, j, k, grid, ρ, p) = (p.g / p.ρ₀) * ρ[i, j, k] * Zᶜᶜᶜ(i, j, k, grid)
 
-# For testing `sort`ed fields.
-@inline resorted_field(i, j, k, grid, sf) = sf[i, j, k]
+"""
+    function sort_field(f)
+Reshape a `Field` `f` into an array (that is the lenght of the interior grid dimensions),
+`sort` this array (default behaviour of `sort` is ascending) then distribute this `sort`ed
+array throughout the grid such that the _largest_ values are at the top of the grid and
+the smallest values are at the _bottom_ of the grid.
+"""
+@inline sort_field(f) = reshape(sort(reshape(f, :)), size(f))
+
 @inline function sort_field(f::Field)
     grid = f.grid
-    sorted_field = reshape(sort(reshape(f, :)), size(f))
+    sorted_field = sort_field(f)
     return KernelFunctionOperation{Center, Center, Center}(resorted_field, grid, sorted_field)
 end
-@inline function sort_field_rev(f::Field)
+"""
+    function sort_field_revesre(f)
+Same method as [`sort_field`](@ref) but the 1D array is sorted in _descending_ order.
+"""
+@inline sort_field_reverse(f) = reshape(sort(reshape(f, :), rev = true), size(f))
+
+@inline function sort_field_reverse(f::Field)
     grid = f.grid
-    sorted_field = reshape(sort(reshape(f, :), rev = true), size(f))
+    sorted_field = sort_field_reverse(f)
     return KernelFunctionOperation{Center, Center, Center}(resorted_field, grid, sorted_field)
 end
 
+# For testing and validating `sort`ed fields.
+@inline resorted_field(i, j, k, grid, sf) = sf[i, j, k]
+
 """
     $(SIGNATURES)
 
@@ -254,7 +270,7 @@ end
 @inline function BackgroundPotentialEnergy(model, buoyancy_model::BuoyancyTracerModel, geopotential_height)
 
     grid = model.grid
-    b✶ = reshape(sort(reshape(model.tracers.b, :)), size(grid))
+    b✶ = sort_field(model.tracers.b)
 
     return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b✶)
 end
@@ -266,10 +282,9 @@ linear_eos_buoyancy(grid, buoyancy, tracers) = KernelFunctionOperation{Center, C
     grid = model.grid
     buoyancy = model.buoyancy.model
     tracers = model.tracers
-    b = linear_eos_buoyancy(grid, buoyancy, tracers)
-    b = Field(b)
+    b = Field(linear_eos_buoyancy(grid, buoyancy, tracers))
     compute!(b)
-    b✶ = reshape(sort(reshape(b, :)), size(grid))
+    b✶ = sort_field(b)
 
     return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b✶)
 end
@@ -279,11 +294,28 @@ end
     grid = model.grid
     ρ = Field(seawater_density(model; geopotential_height))
     compute!(ρ)
-    ρ✶ = reshape(sort(reshape(ρ, :), rev = true), size(grid))
+    ρ✶ = sort_field_reverse(ρ)
     parameters = (g = model.buoyancy.model.gravitational_acceleration,
                   ρ₀ = model.buoyancy.model.equation_of_state.reference_density)
 
     return KernelFunctionOperation{Center, Center, Center}(g′z_ccc, grid, ρ✶, parameters)
 end
 
+## By defining custom field
+export ReferenceStateOperand, ReferenceStateField
+
+struct ReferenceStateOperand{O}
+    sorting_function :: O
+end
+const ReferenceStateField = Field{Center, Center, Center, <:ReferenceStateOperand}
+
+import Oceananigans.Fields: compute!
+
+function compute!(rs::ReferenceStateField, time=nothing)
+
+    rs.data = rs.operand(rs.data)
+
+    return rs
+end
+
 end # module

From e59fbadfc0c24be8926cb9adccc7a3e6e6fc249f Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Mon, 22 Apr 2024 10:58:30 +1000
Subject: [PATCH 10/37] Remove bad Type design

---
 src/PotentialEnergyEquationTerms.jl | 15 ---------------
 1 file changed, 15 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 565a912f..85781810 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -193,27 +193,12 @@ array throughout the grid such that the _largest_ values are at the top of the g
 the smallest values are at the _bottom_ of the grid.
 """
 @inline sort_field(f) = reshape(sort(reshape(f, :)), size(f))
-
-@inline function sort_field(f::Field)
-    grid = f.grid
-    sorted_field = sort_field(f)
-    return KernelFunctionOperation{Center, Center, Center}(resorted_field, grid, sorted_field)
-end
 """
     function sort_field_revesre(f)
 Same method as [`sort_field`](@ref) but the 1D array is sorted in _descending_ order.
 """
 @inline sort_field_reverse(f) = reshape(sort(reshape(f, :), rev = true), size(f))
 
-@inline function sort_field_reverse(f::Field)
-    grid = f.grid
-    sorted_field = sort_field_reverse(f)
-    return KernelFunctionOperation{Center, Center, Center}(resorted_field, grid, sorted_field)
-end
-
-# For testing and validating `sort`ed fields.
-@inline resorted_field(i, j, k, grid, sf) = sf[i, j, k]
-
 """
     $(SIGNATURES)
 

From 1ac15f32f6e65204a52309887e6503c6c4792e43 Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Mon, 22 Apr 2024 12:09:45 +1000
Subject: [PATCH 11/37] Using more kfo's

---
 src/PotentialEnergyEquationTerms.jl | 49 +++++++++++++++--------------
 1 file changed, 25 insertions(+), 24 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 85781810..c8003072 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -187,17 +187,37 @@ end
 
 """
     function sort_field(f)
-Reshape a `Field` `f` into an array (that is the lenght of the interior grid dimensions),
+Reshape a `Field` `f` into an array (that is the length of the product of the grid dimensions),
 `sort` this array (default behaviour of `sort` is ascending) then distribute this `sort`ed
 array throughout the grid such that the _largest_ values are at the top of the grid and
 the smallest values are at the _bottom_ of the grid.
 """
-@inline sort_field(f) = reshape(sort(reshape(f, :)), size(f))
+@inline sorted_field(f::Field) = reshape(sort(reshape(f, :)), size(f))
+
+@inline function sort_field(f::Field)
+
+    compute!(f)
+    grid = f.grid
+    sf = sorted_field(f)
+
+    return KernelFunctionOperation{Center, Center, Center}(sorted_field, grid, sf)
+end
 """
     function sort_field_revesre(f)
 Same method as [`sort_field`](@ref) but the 1D array is sorted in _descending_ order.
 """
-@inline sort_field_reverse(f) = reshape(sort(reshape(f, :), rev = true), size(f))
+@inline sorted_field_reverse(f::Field) = reshape(sort(reshape(f, :), rev = true), size(f))
+
+@inline function sort_field_reverse(f::Field)
+
+    compute!(f)
+    grid = f.grid
+    sf = sorted_field_reverse(f)
+
+    return KernelFunctionOperation{Center, Center, Center}(sorted_field, grid, sf)
+end
+
+@inline sorted_field(i, j, k, grid, sf) = sf[i, j, k]
 
 """
     $(SIGNATURES)
@@ -267,8 +287,7 @@ linear_eos_buoyancy(grid, buoyancy, tracers) = KernelFunctionOperation{Center, C
     grid = model.grid
     buoyancy = model.buoyancy.model
     tracers = model.tracers
-    b = Field(linear_eos_buoyancy(grid, buoyancy, tracers))
-    compute!(b)
+    b = linear_eos_buoyancy(grid, buoyancy, tracers)
     b✶ = sort_field(b)
 
     return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b✶)
@@ -277,8 +296,7 @@ end
 @inline function BackgroundPotentialEnergy(model, buoyancy_model::BuoyancyBoussinesqEOSModel, geopotential_height)
 
     grid = model.grid
-    ρ = Field(seawater_density(model; geopotential_height))
-    compute!(ρ)
+    ρ = seawater_density(model; geopotential_height)
     ρ✶ = sort_field_reverse(ρ)
     parameters = (g = model.buoyancy.model.gravitational_acceleration,
                   ρ₀ = model.buoyancy.model.equation_of_state.reference_density)
@@ -286,21 +304,4 @@ end
     return KernelFunctionOperation{Center, Center, Center}(g′z_ccc, grid, ρ✶, parameters)
 end
 
-## By defining custom field
-export ReferenceStateOperand, ReferenceStateField
-
-struct ReferenceStateOperand{O}
-    sorting_function :: O
-end
-const ReferenceStateField = Field{Center, Center, Center, <:ReferenceStateOperand}
-
-import Oceananigans.Fields: compute!
-
-function compute!(rs::ReferenceStateField, time=nothing)
-
-    rs.data = rs.operand(rs.data)
-
-    return rs
-end
-
 end # module

From 6d3d7da50944c507d69f699f003602bb451496fe Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Mon, 22 Apr 2024 12:12:49 +1000
Subject: [PATCH 12/37] Update PotentialEnergyEquationTerms.jl

---
 src/PotentialEnergyEquationTerms.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index c8003072..b490918e 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -12,7 +12,7 @@ using Oceananigans.Grids: NegativeZDirection
 using Oceananigans.BuoyancyModels: Buoyancy, BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState
 using Oceananigans.BuoyancyModels: buoyancy_perturbationᶜᶜᶜ, Zᶜᶜᶜ
 using Oceananigans.Models: ShallowWaterModel
-using Oceananigans.Fields: Field
+using Oceananigans.Fields: Field, compute!
 using Oceanostics: validate_location
 using SeawaterPolynomials: BoussinesqEquationOfState
 

From 62b0a47c16c2aa95f8b9f5bf12fd0a39c03c0a55 Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Mon, 22 Apr 2024 13:30:56 +1000
Subject: [PATCH 13/37] Update BPE docstring

---
 src/PotentialEnergyEquationTerms.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index b490918e..968752e8 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -239,7 +239,7 @@ Examples
 ```jldoctest
 julia> using Oceananigans
 
-julia> using Oceanostics.PotentialEnergyEquationTerms: PotentialEnergy
+julia> using Oceanostics.PotentialEnergyEquationTerms: BackgroundPotentialEnergy
 
 julia> grid = RectilinearGrid(size=100, z=(-1000, 0), topology=(Flat, Flat, Bounded))
 1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
@@ -260,7 +260,7 @@ julia> bpe = BackgroundPotentialEnergy(model)
 KernelFunctionOperation at (Center, Center, Center)
 ├── grid: 1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
 ├── kernel_function: bz_ccc (generic function with 2 methods)
-└── arguments: ("1×1×100 Array{Float64, 3}",)
+└── arguments: ("KernelFunctionOperation at (Center, Center, Center)",)
 ```
 """
 @inline function BackgroundPotentialEnergy(model; location = (Center, Center, Center),

From 1c730777d443add43ebca99bec6001b9ef1e1b14 Mon Sep 17 00:00:00 2001
From: Josef Bisits <jbisits@gmail.com>
Date: Wed, 24 Apr 2024 09:48:40 +1000
Subject: [PATCH 14/37] Update src/PotentialEnergyEquationTerms.jl

Co-authored-by: Tomas Chor <tomaschor@gmail.com>
---
 src/PotentialEnergyEquationTerms.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 968752e8..e34ab231 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -202,6 +202,7 @@ the smallest values are at the _bottom_ of the grid.
 
     return KernelFunctionOperation{Center, Center, Center}(sorted_field, grid, sf)
 end
+
 """
     function sort_field_revesre(f)
 Same method as [`sort_field`](@ref) but the 1D array is sorted in _descending_ order.

From 2c4faad4f1cf97f653fc1a2f234820291cc0682c Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Wed, 8 May 2024 13:32:43 +1000
Subject: [PATCH 15/37] Try `OneDReferenceField`

---
 src/PotentialEnergyEquationTerms.jl | 68 +++++++++++++++++------------
 1 file changed, 41 insertions(+), 27 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index e34ab231..228f8729 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -4,15 +4,15 @@ using DocStringExtensions
 
 export PotentialEnergy, BackgroundPotentialEnergy
 
-using Oceananigans.AbstractOperations: KernelFunctionOperation
+using Oceananigans.AbstractOperations: KernelFunctionOperation, volume, Az, GridMetricOperation
 using Oceananigans.Models: seawater_density
 using Oceananigans.Models: model_geopotential_height
-using Oceananigans.Grids: Center, Face
-using Oceananigans.Grids: NegativeZDirection
+using Oceananigans.Grids
+using Oceananigans.Grids: Center, Face, NegativeZDirection, interior, CenterField, regrid!
 using Oceananigans.BuoyancyModels: Buoyancy, BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState
 using Oceananigans.BuoyancyModels: buoyancy_perturbationᶜᶜᶜ, Zᶜᶜᶜ
 using Oceananigans.Models: ShallowWaterModel
-using Oceananigans.Fields: Field, compute!
+using Oceananigans.Fields: Field, compute!, field, set!
 using Oceanostics: validate_location
 using SeawaterPolynomials: BoussinesqEquationOfState
 
@@ -185,40 +185,54 @@ end
 
 @inline g′z_ccc(i, j, k, grid, ρ, p) = (p.g / p.ρ₀) * ρ[i, j, k] * Zᶜᶜᶜ(i, j, k, grid)
 
-"""
-    function sort_field(f)
-Reshape a `Field` `f` into an array (that is the length of the product of the grid dimensions),
-`sort` this array (default behaviour of `sort` is ascending) then distribute this `sort`ed
-array throughout the grid such that the _largest_ values are at the top of the grid and
-the smallest values are at the _bottom_ of the grid.
-"""
-@inline sorted_field(f::Field) = reshape(sort(reshape(f, :)), size(f))
+## Grid metrics from https://github.com/tomchor/Oceanostics.jl/issues/163#issuecomment-2012623824
 
-@inline function sort_field(f::Field)
+function MetricField(loc, grid, metric)
 
-    compute!(f)
-    grid = f.grid
-    sf = sorted_field(f)
+    metric_operation = GridMetricOperation(loc, metric, grid)
+    metric_field = Field(metric_operation)
 
-    return KernelFunctionOperation{Center, Center, Center}(sorted_field, grid, sf)
+    return compute!(metric_field)
 end
 
+VolumeField(grid, loc=(Center, Center, Center)) = MetricField(loc, grid, volume)
+  AreaField(grid, loc=(Center, Center, Nothing)) = MetricField(loc, grid, Az)
+
 """
-    function sort_field_revesre(f)
-Same method as [`sort_field`](@ref) but the 1D array is sorted in _descending_ order.
+    function OneDReferenceField(f::Field)
+Return a `OneDReferenceField` of the gridded data from the `Field` `f` and the `z✶` for the `Field`.
+The gridded data is first reshaped into a 1D `Array` then sorted. Returned is a new `Field` of this sorted data
+on a z✶ `grid`. The z✶ `grid` is defined as
+```math
+\\frac{1}{A}\\int_{f\\mathrm{min}}^{f\\mathrm{max}} \\mathrm{d}V.
+```
+and is computed by cumulatively summing the 1D `Array` of grid volumes `ΔV`.
+**Note:** the `OneDReferenceField` is only appropriate for grids that have uniform horizontal
+area.
 """
-@inline sorted_field_reverse(f::Field) = reshape(sort(reshape(f, :), rev = true), size(f))
+function OneDReferenceField(f::Field)
 
-@inline function sort_field_reverse(f::Field)
+    area = sum(AreaField(f.grid))
+    volume_field = VolumeField(f.grid)
+    v = reshape(interior(volume_field), :)
+    field_data = reshape(interior(f), :)
 
-    compute!(f)
-    grid = f.grid
-    sf = sorted_field_reverse(f)
+    p = sortperm(field_data)
+    sorted_field_data = field_data[p]
+    z✶ = cumsum(v[p]) / area
 
-    return KernelFunctionOperation{Center, Center, Center}(sorted_field, grid, sf)
-end
+    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)))
 
-@inline sorted_field(i, j, k, grid, sf) = sf[i, j, k]
+    return sorted_field, z✶_field
+
+end
 
 """
     $(SIGNATURES)

From 0a87fbbd0d74d017123aad179638379e31d421cd Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Wed, 8 May 2024 13:34:45 +1000
Subject: [PATCH 16/37] Export the function

---
 src/PotentialEnergyEquationTerms.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 228f8729..a80cebac 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -2,7 +2,7 @@ module PotentialEnergyEquationTerms
 
 using DocStringExtensions
 
-export PotentialEnergy, BackgroundPotentialEnergy
+export PotentialEnergy, BackgroundPotentialEnergy, OneDReferenceField
 
 using Oceananigans.AbstractOperations: KernelFunctionOperation, volume, Az, GridMetricOperation
 using Oceananigans.Models: seawater_density

From 73479d5b82b9ab9916f14b76d52d61bba21000ca Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Wed, 8 May 2024 15:50:48 +1000
Subject: [PATCH 17/37] Correct BPE function, still need to check sorting

---
 src/PotentialEnergyEquationTerms.jl | 31 +++++++++++++++++------------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index a80cebac..dac6f4e0 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -210,14 +210,14 @@ and is computed by cumulatively summing the 1D `Array` of grid volumes `ΔV`.
 **Note:** the `OneDReferenceField` is only appropriate for grids that have uniform horizontal
 area.
 """
-function OneDReferenceField(f::Field)
+function OneDReferenceField(f::Field; 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)
+    p = sortperm(field_data; rev)
     sorted_field_data = field_data[p]
     z✶ = cumsum(v[p]) / area
 
@@ -240,13 +240,12 @@ end
 Return a `kernelFunctionOperation` to compute the `BackgroundPotentialEnergy` per unit
 volume,
 ```math
-E_{b} = \\frac{gρ✶z}{ρ₀}.
+E_{b} = \\frac{gρz✶}{ρ₀}.
 ```
 The `BackgroundPotentialEnergy` is the potential energy computed after adiabatically resorting
 the buoyancy or density field into a reference state of minimal potential energy.
 The reference state is computed by reshaping the gridded buoyancy or density field and
-`sort`ing into a monotonically increasing `Vector`. This `sort`ed vector is then reshaped into
-the `size(model.grid)`.
+`sort`ing into a monotonically increasing `Vector`.
 
 Examples
 ========
@@ -274,8 +273,8 @@ NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
 julia> bpe = BackgroundPotentialEnergy(model)
 KernelFunctionOperation at (Center, Center, Center)
 ├── grid: 1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
-├── kernel_function: bz_ccc (generic function with 2 methods)
-└── arguments: ("KernelFunctionOperation at (Center, Center, Center)",)
+├── kernel_function: bz✶_ccc (generic function with 1 method)
+└── arguments: ("1×1×100 Field{Center, Center, Center} on RectilinearGrid on CPU", "1×1×100 Field{Center, Center, Center} on RectilinearGrid on CPU")
 ```
 """
 @inline function BackgroundPotentialEnergy(model; location = (Center, Center, Center),
@@ -290,11 +289,13 @@ end
 @inline function BackgroundPotentialEnergy(model, buoyancy_model::BuoyancyTracerModel, geopotential_height)
 
     grid = model.grid
-    b✶ = sort_field(model.tracers.b)
+    sorted_buoyancy, z✶ = OneDReferenceField(model.tracers.b)
 
-    return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b✶)
+    return KernelFunctionOperation{Center, Center, Center}(bz✶_ccc, grid, sorted_buoyancy, z✶)
 end
 
+@inline bz✶_ccc(i, j, k, grid, b, z✶) = b[i, j, k] * z✶[i, j, k]
+
 linear_eos_buoyancy(grid, buoyancy, tracers) = KernelFunctionOperation{Center, Center, Center}(buoyancy_perturbationᶜᶜᶜ, grid, buoyancy, tracers)
 
 @inline function BackgroundPotentialEnergy(model, buoyancy_model::BuoyancyLinearEOSModel, geopotential_height)
@@ -303,20 +304,24 @@ linear_eos_buoyancy(grid, buoyancy, tracers) = KernelFunctionOperation{Center, C
     buoyancy = model.buoyancy.model
     tracers = model.tracers
     b = linear_eos_buoyancy(grid, buoyancy, tracers)
-    b✶ = sort_field(b)
+    compute!(b)
+    sorted_buoyancy, z✶ = OneDReferenceField(b)
 
-    return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b✶)
+    return KernelFunctionOperation{Center, Center, Center}(bz✶_ccc, grid, sorted_buoyancy, z✶)
 end
 
 @inline function BackgroundPotentialEnergy(model, buoyancy_model::BuoyancyBoussinesqEOSModel, geopotential_height)
 
     grid = model.grid
     ρ = seawater_density(model; geopotential_height)
-    ρ✶ = sort_field_reverse(ρ)
+    compute!(ρ)
+    sorted_density, z✶ = OneDReferenceField(ρ)
     parameters = (g = model.buoyancy.model.gravitational_acceleration,
                   ρ₀ = model.buoyancy.model.equation_of_state.reference_density)
 
-    return KernelFunctionOperation{Center, Center, Center}(g′z_ccc, grid, ρ✶, parameters)
+    return KernelFunctionOperation{Center, Center, Center}(g′z_ccc, grid, sorted_density, z✶, parameters)
 end
 
+@inline g′z✶_ccc(i, j, k, grid, ρ, z✶, p) = (p.g / p.ρ₀) * ρ[i, j, k] * z✶[i, j, k]
+
 end # module

From f28a26ae1b628d4f751128d84c2f8eaeeeafa28e Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Fri, 10 May 2024 09:59:33 +1000
Subject: [PATCH 18/37] Suppress some docstring output

---
 src/PotentialEnergyEquationTerms.jl | 62 +++++------------------------
 1 file changed, 10 insertions(+), 52 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 00969dfb..965cec27 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -50,21 +50,9 @@ julia> using Oceananigans
 
 julia> using Oceanostics.PotentialEnergyEquationTerms: PotentialEnergy
 
-julia> grid = RectilinearGrid(size=100, z=(-1000, 0), topology=(Flat, Flat, Bounded))
-1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
-├── Flat x
-├── Flat y
-└── Bounded  z ∈ [-1000.0, 0.0]   regularly spaced with Δz=10.0
-
-julia> model = NonhydrostaticModel(; grid, buoyancy=BuoyancyTracer(), tracers=(:b,))
-NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
-├── grid: 1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
-├── timestepper: QuasiAdamsBashforth2TimeStepper
-├── advection scheme: Centered reconstruction order 2
-├── tracers: b
-├── closure: Nothing
-├── buoyancy: BuoyancyTracer with ĝ = NegativeZDirection()
-└── coriolis: Nothing
+julia> grid = RectilinearGrid(size=100, z=(-1000, 0), topology=(Flat, Flat, Bounded));
+
+julia> model = NonhydrostaticModel(; grid, buoyancy=BuoyancyTracer(), tracers=(:b,));
 
 julia> PotentialEnergy(model)
 KernelFunctionOperation at (Center, Center, Center)
@@ -80,34 +68,15 @@ julia> using Oceananigans, SeawaterPolynomials.TEOS10
 
 julia> using Oceanostics.PotentialEnergyEquationTerms: PotentialEnergy
 
-julia> grid = RectilinearGrid(size=100, z=(-1000, 0), topology=(Flat, Flat, Bounded))
-1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
-├── Flat x
-├── Flat y
-└── Bounded  z ∈ [-1000.0, 0.0]   regularly spaced with Δz=10.0
+julia> grid = RectilinearGrid(size=100, z=(-1000, 0), topology=(Flat, Flat, Bounded));
 
-julia> tracers = (:T, :S)
-(:T, :S)
+julia> tracers = (:T, :S);
 
-julia> eos = TEOS10EquationOfState()
-BoussinesqEquationOfState{Float64}:
-    ├── seawater_polynomial: TEOS10SeawaterPolynomial{Float64}
-    └── reference_density: 1020.0
+julia> eos = TEOS10EquationOfState();
 
-julia> buoyancy = SeawaterBuoyancy(equation_of_state=eos)
-SeawaterBuoyancy{Float64}:
-├── gravitational_acceleration: 9.80665
-└── equation_of_state: BoussinesqEquationOfState{Float64}
+julia> buoyancy = SeawaterBuoyancy(equation_of_state=eos);
 
-julia> model = NonhydrostaticModel(; grid, buoyancy, tracers)
-NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
-├── grid: 1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
-├── timestepper: QuasiAdamsBashforth2TimeStepper
-├── advection scheme: Centered reconstruction order 2
-├── tracers: (T, S)
-├── closure: Nothing
-├── buoyancy: SeawaterBuoyancy with g=9.80665 and BoussinesqEquationOfState{Float64} with ĝ = NegativeZDirection()
-└── coriolis: Nothing
+julia> model = NonhydrostaticModel(; grid, buoyancy, tracers);
 
 julia> PotentialEnergy(model)
 KernelFunctionOperation at (Center, Center, Center)
@@ -257,20 +226,9 @@ julia> using Oceananigans
 
 julia> using Oceanostics.PotentialEnergyEquationTerms: BackgroundPotentialEnergy
 
-julia> grid = RectilinearGrid(size=100, z=(-1000, 0), topology=(Flat, Flat, Bounded))
-1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
-├── Flat x
-├── Flat y
-└── Bounded  z ∈ [-1000.0, 0.0]   regularly spaced with Δz=10.0
+julia> grid = RectilinearGrid(size=100, z=(-1000, 0), topology=(Flat, Flat, Bounded));
 
-julia> model = NonhydrostaticModel(; grid, buoyancy=BuoyancyTracer(), tracers=(:b,))
-NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
-├── grid: 1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
-├── timestepper: QuasiAdamsBashforth2TimeStepper
-├── tracers: b
-├── closure: Nothing
-├── buoyancy: BuoyancyTracer with ĝ = NegativeZDirection()
-└── coriolis: Nothing
+julia> model = NonhydrostaticModel(; grid, buoyancy=BuoyancyTracer(), tracers=(:b,));
 
 julia> bpe = BackgroundPotentialEnergy(model)
 KernelFunctionOperation at (Center, Center, Center)

From aad59a95244c0dd65a8d800d06025cf70b3ddda3 Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Fri, 10 May 2024 11:53:23 +1000
Subject: [PATCH 19/37] Docstring updates

---
 src/PotentialEnergyEquationTerms.jl | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 965cec27..59a3dff2 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -171,15 +171,17 @@ VolumeField(grid, loc=(Center, Center, Center)) = MetricField(loc, grid, volume)
 
 """
     function OneDReferenceField(f::Field)
-Return a `OneDReferenceField` of the gridded data from the `Field` `f` and the `z✶` for the `Field`.
-The gridded data is first reshaped into a 1D `Array` then sorted. Returned is a new `Field` of this sorted data
-on a z✶ `grid`. The z✶ `grid` is defined as
+Return a `OneDReferenceField` of the data in `f` and the corresponding `z✶` coordinate.
+The gridded data is first reshaped into a 1D `Array` then sorted. Returned is a new `Field`
+of this sorted data as well as the `z✶` `Field` that are both on a grid that has the same vertical extent
+as the original field (`f.grid.Lz`) but with the resolution of the whole domain (`Nx * Ny * Nz`).
+The `z✶` `Field` is defined as
 ```math
 \\frac{1}{A}\\int_{f\\mathrm{min}}^{f\\mathrm{max}} \\mathrm{d}V.
 ```
 and is computed by cumulatively summing the 1D `Array` of grid volumes `ΔV`.
-**Note:** the `OneDReferenceField` is only appropriate for grids that have uniform horizontal
-area.
+**Note:** `OneDReferenceField` is currently only appropriate for grids that have uniform
+horizontal area.
 """
 function OneDReferenceField(f::Field; rev = false)
 

From 6251dc8accac3275e0d8b3a70310e71f6038a8a3 Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Fri, 10 May 2024 12:40:24 +1000
Subject: [PATCH 20/37] Validate model grid + docstring updates

---
 src/PotentialEnergyEquationTerms.jl | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 59a3dff2..cf02c3a6 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -7,8 +7,9 @@ export PotentialEnergy, BackgroundPotentialEnergy, OneDReferenceField
 using Oceananigans.AbstractOperations: KernelFunctionOperation, volume, Az, GridMetricOperation
 using Oceananigans.Models: seawater_density
 using Oceananigans.Models: model_geopotential_height
+using Oceananigans.ImmersedBoundaryGrid: ImmersedBoundaryGrid
 using Oceananigans.Grids
-using Oceananigans.Grids: Center, Face, NegativeZDirection, interior, CenterField, regrid!
+using Oceananigans.Grids: Center, NegativeZDirection, interior, CenterField
 using Oceananigans.BuoyancyModels: Buoyancy, BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState
 using Oceananigans.BuoyancyModels: buoyancy_perturbationᶜᶜᶜ, Zᶜᶜᶜ
 using Oceananigans.Models: ShallowWaterModel
@@ -23,7 +24,10 @@ const BuoyancyBoussinesqEOSModel = Buoyancy{<:SeawaterBuoyancy{FT, <:BoussinesqE
 
 validate_gravity_unit_vector(gravity_unit_vector::NegativeZDirection) = nothing
 validate_gravity_unit_vector(gravity_unit_vector) =
-    throw(ArgumentError("`PotentialEnergy` is curently only defined for models that have a `NegativeZDirection` gravity unit vector."))
+    throw(ArgumentError("`PotentialEnergy` and `BackgroundPotentialEnergy` are curently only defined for models that have a `NegativeZDirection` gravity unit vector."))
+validate_grid(grid) = nothing
+validate_grid(grid::ImmersedBoundaryGrid) =
+    throw(ArgumentError("`PotentialEnergy` and `BackgroundPotentialEnergy` are not currently defined for $(grid)."))
 
 """
     $(SIGNATURES)
@@ -116,6 +120,7 @@ KernelFunctionOperation at (Center, Center, Center)
 
     validate_location(location, "PotentialEnergy")
     isnothing(model.buoyancy) ? nothing : validate_gravity_unit_vector(model.buoyancy.gravity_unit_vector)
+    validate_grid(model.grid)
 
     return PotentialEnergy(model, model.buoyancy, geopotential_height)
 end
@@ -210,15 +215,17 @@ end
 """
     $(SIGNATURES)
 
-Return a `kernelFunctionOperation` to compute the `BackgroundPotentialEnergy` per unit
-volume,
+Return a `kernelFunctionOperation` to compute an approximation to the `BackgroundPotentialEnergy`
+per unit volume,
 ```math
 E_{b} = \\frac{gρz✶}{ρ₀}.
 ```
-The `BackgroundPotentialEnergy` is the potential energy computed after adiabatically resorting
+The `BackgroundPotentialEnergy` is the potential energy computed by adiabatically resorting
 the buoyancy or density field into a reference state of minimal potential energy.
-The reference state is computed by reshaping the gridded buoyancy or density field and
-`sort`ing into a monotonically increasing `Vector`.
+The reference state, and corresponding `z✶` cooridinate, is approximated using [OneDReferenceField](@ref).
+For more informaion about the background potential energy state and its implementation see
+[Winters et al. (1995)](https://www.cambridge.org/core/journals/journal-of-fluid-mechanics/article/available-potential-energy-and-mixing-in-densitystratified-fluids/A45F1A40521FF0A0DC82BC705AD398DA)
+and [Carpenter et al. (2012)](https://www.cambridge.org/core/journals/journal-of-fluid-mechanics/article/simulations-of-a-doublediffusive-interface-in-the-diffusive-convection-regime/63D2ECE2AA41439E01A01F9A0D76F2E2).
 
 Examples
 ========
@@ -244,6 +251,7 @@ KernelFunctionOperation at (Center, Center, Center)
 
     validate_location(location, "BackgroundPotentialEnergy")
     isnothing(model.buoyancy) ? nothing : validate_gravity_unit_vector(model.buoyancy.gravity_unit_vector)
+    validate_grid(model.grid)
 
     return BackgroundPotentialEnergy(model, model.buoyancy, geopotential_height)
 end

From 2e98f34265d95e72e273b8d9ec3f9f8e92b9d7dd Mon Sep 17 00:00:00 2001
From: Josef Bisits <jbisits@gmail.com>
Date: Fri, 10 May 2024 12:46:28 +1000
Subject: [PATCH 21/37] Update src/PotentialEnergyEquationTerms.jl

Co-authored-by: Gregory L. Wagner <wagner.greg@gmail.com>
---
 src/PotentialEnergyEquationTerms.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index cf02c3a6..213aab05 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -56,7 +56,7 @@ julia> using Oceanostics.PotentialEnergyEquationTerms: PotentialEnergy
 
 julia> grid = RectilinearGrid(size=100, z=(-1000, 0), topology=(Flat, Flat, Bounded));
 
-julia> model = NonhydrostaticModel(; grid, buoyancy=BuoyancyTracer(), tracers=(:b,));
+julia> model = NonhydrostaticModel(; grid, buoyancy=BuoyancyTracer(), tracers=:b);
 
 julia> PotentialEnergy(model)
 KernelFunctionOperation at (Center, Center, Center)

From 38b9a836708d3ca92ed294a009184594eb83b85d Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Fri, 10 May 2024 12:46:41 +1000
Subject: [PATCH 22/37] Update PotentialEnergyEquationTerms.jl

---
 src/PotentialEnergyEquationTerms.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index cf02c3a6..29babfff 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -7,7 +7,7 @@ export PotentialEnergy, BackgroundPotentialEnergy, OneDReferenceField
 using Oceananigans.AbstractOperations: KernelFunctionOperation, volume, Az, GridMetricOperation
 using Oceananigans.Models: seawater_density
 using Oceananigans.Models: model_geopotential_height
-using Oceananigans.ImmersedBoundaryGrid: ImmersedBoundaryGrid
+using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
 using Oceananigans.Grids
 using Oceananigans.Grids: Center, NegativeZDirection, interior, CenterField
 using Oceananigans.BuoyancyModels: Buoyancy, BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState

From 3107ec61997059875b4a99e7fb88c32fec8f1f1d Mon Sep 17 00:00:00 2001
From: Josef Bisits <jbisits@gmail.com>
Date: Sat, 11 May 2024 19:43:03 +1000
Subject: [PATCH 23/37] Update src/PotentialEnergyEquationTerms.jl

Co-authored-by: Tomas Chor <tomaschor@gmail.com>
---
 src/PotentialEnergyEquationTerms.jl | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index fad27a0b..4235f75a 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -164,10 +164,8 @@ end
 ## Grid metrics from https://github.com/tomchor/Oceanostics.jl/issues/163#issuecomment-2012623824
 
 function MetricField(loc, grid, metric)
-
     metric_operation = GridMetricOperation(loc, metric, grid)
     metric_field = Field(metric_operation)
-
     return compute!(metric_field)
 end
 

From 23d095d069e9c6b81bd0f5428b1ac48a2a0b4187 Mon Sep 17 00:00:00 2001
From: Josef Bisits <jbisits@gmail.com>
Date: Sat, 11 May 2024 19:43:40 +1000
Subject: [PATCH 24/37] Update src/PotentialEnergyEquationTerms.jl

Co-authored-by: Tomas Chor <tomaschor@gmail.com>
---
 src/PotentialEnergyEquationTerms.jl | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 4235f75a..bca0a1cd 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -176,8 +176,9 @@ VolumeField(grid, loc=(Center, Center, Center)) = MetricField(loc, grid, volume)
     function OneDReferenceField(f::Field)
 Return a `OneDReferenceField` of the data in `f` and the corresponding `z✶` coordinate.
 The gridded data is first reshaped into a 1D `Array` then sorted. Returned is a new `Field`
-of this sorted data as well as the `z✶` `Field` that are both on a grid that has the same vertical extent
-as the original field (`f.grid.Lz`) but with the resolution of the whole domain (`Nx * Ny * Nz`).
+of this sorted data as well as the `z✶` `Field` that are both on a one-dimensional grid that
+has the same vertical extent as the original field (`f.grid.Lz`) but with the same _total_
+number of points of the original domain (`Nx * Ny * Nz`).
 The `z✶` `Field` is defined as
 ```math
 \\frac{1}{A}\\int_{f\\mathrm{min}}^{f\\mathrm{max}} \\mathrm{d}V.

From 769ecc0be7285812c968674f31603854a0d7f659 Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Sat, 11 May 2024 20:10:42 +1000
Subject: [PATCH 25/37] Wrong function calls

---
 src/PotentialEnergyEquationTerms.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index bca0a1cd..b31caea7 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -288,7 +288,7 @@ end
     parameters = (g = model.buoyancy.model.gravitational_acceleration,
                   ρ₀ = model.buoyancy.model.equation_of_state.reference_density)
 
-    return KernelFunctionOperation{Center, Center, Center}(g′z_ccc, grid, sorted_density, z✶, parameters)
+    return KernelFunctionOperation{Center, Center, Center}(g′z✶_ccc, grid, sorted_density, z✶, parameters)
 end
 
 @inline g′z✶_ccc(i, j, k, grid, ρ, z✶, p) = (p.g / p.ρ₀) * ρ[i, j, k] * z✶[i, j, k]

From 44c02065e1a3a5b4c26bf74de7fccec97d0c81f5 Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Tue, 14 May 2024 11:19:42 +1000
Subject: [PATCH 26/37] Reduce number of methods + update packages

Can remove the updated packages if need be.
---
 Manifest.toml                       | 28 ++++++++++++++--------------
 src/PotentialEnergyEquationTerms.jl | 24 +++++++-----------------
 2 files changed, 21 insertions(+), 31 deletions(-)

diff --git a/Manifest.toml b/Manifest.toml
index f4fabf13..1619d54d 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -1,6 +1,6 @@
 # This file is machine-generated - editing it directly is not advised
 
-julia_version = "1.10.2"
+julia_version = "1.10.3"
 manifest_format = "2.0"
 project_hash = "e19f51c46fc9161a1a3899ab99209f39adde0868"
 
@@ -141,9 +141,9 @@ version = "0.11.5"
 
 [[deps.Colors]]
 deps = ["ColorTypes", "FixedPointNumbers", "Reexport"]
-git-tree-sha1 = "fc08e5930ee9a4e03f84bfb5211cb54e7769758a"
+git-tree-sha1 = "362a287c3aa50601b0bc359053d5c2468f0e7ce0"
 uuid = "5ae59095-9a9b-59fe-a467-6f913c188581"
-version = "0.12.10"
+version = "0.12.11"
 
 [[deps.CommonDataModel]]
 deps = ["CFTime", "DataStructures", "Dates", "Preferences", "Printf", "Statistics"]
@@ -164,7 +164,7 @@ weakdeps = ["Dates", "LinearAlgebra"]
 [[deps.CompilerSupportLibraries_jll]]
 deps = ["Artifacts", "Libdl"]
 uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
-version = "1.1.0+0"
+version = "1.1.1+0"
 
 [[deps.ConstructionBase]]
 deps = ["LinearAlgebra"]
@@ -285,9 +285,9 @@ uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee"
 
 [[deps.FixedPointNumbers]]
 deps = ["Statistics"]
-git-tree-sha1 = "335bfdceacc84c5cdf16aadc768aa5ddfc5383cc"
+git-tree-sha1 = "05882d6995ae5c12bb5f36dd2ed3f61c98cbb172"
 uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
-version = "0.8.4"
+version = "0.8.5"
 
 [[deps.Future]]
 deps = ["Random"]
@@ -389,9 +389,9 @@ version = "1.0.0"
 
 [[deps.JLD2]]
 deps = ["FileIO", "MacroTools", "Mmap", "OrderedCollections", "Pkg", "PrecompileTools", "Printf", "Reexport", "Requires", "TranscodingStreams", "UUIDs"]
-git-tree-sha1 = "5ea6acdd53a51d897672edb694e3cc2912f3f8a7"
+git-tree-sha1 = "dca9ff5abdf5fab4456876bc93f80c59a37b81df"
 uuid = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
-version = "0.4.46"
+version = "0.4.47"
 
 [[deps.JLLWrappers]]
 deps = ["Artifacts", "Preferences"]
@@ -419,9 +419,9 @@ version = "0.2.1+0"
 
 [[deps.KernelAbstractions]]
 deps = ["Adapt", "Atomix", "InteractiveUtils", "LinearAlgebra", "MacroTools", "PrecompileTools", "Requires", "SparseArrays", "StaticArrays", "UUIDs", "UnsafeAtomics", "UnsafeAtomicsLLVM"]
-git-tree-sha1 = "ed7167240f40e62d97c1f5f7735dea6de3cc5c49"
+git-tree-sha1 = "db02395e4c374030c53dc28f3c1d33dec35f7272"
 uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
-version = "0.9.18"
+version = "0.9.19"
 
     [deps.KernelAbstractions.extensions]
     EnzymeExt = "EnzymeCore"
@@ -861,9 +861,9 @@ version = "0.3.4"
 
 [[deps.SentinelArrays]]
 deps = ["Dates", "Random"]
-git-tree-sha1 = "0e7508ff27ba32f26cd459474ca2ede1bc10991f"
+git-tree-sha1 = "363c4e82b66be7b9f7c7c7da7478fdae07de44b9"
 uuid = "91c51154-3ec4-41a3-a24f-3f23e20d615c"
-version = "1.4.1"
+version = "1.4.2"
 
 [[deps.Serialization]]
 uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
@@ -884,9 +884,9 @@ version = "1.10.0"
 
 [[deps.SpecialFunctions]]
 deps = ["IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"]
-git-tree-sha1 = "e2cfc4012a19088254b3950b85c3c1d8882d864d"
+git-tree-sha1 = "2f5d4697f21388cbe1ff299430dd169ef97d7e14"
 uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
-version = "2.3.1"
+version = "2.4.0"
 
     [deps.SpecialFunctions.extensions]
     SpecialFunctionsChainRulesCoreExt = "ChainRulesCore"
diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index b31caea7..f4b13004 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -142,9 +142,9 @@ end
 
     grid = model.grid
     C = model.tracers
-    b = buoyancy_model.model
+    b_model = buoyancy_model.model
 
-    return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b, C)
+    return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b_model, C)
 end
 
 @inline bz_ccc(i, j, k, grid, b, C) = buoyancy_perturbationᶜᶜᶜ(i, j, k, grid, b, C) * Zᶜᶜᶜ(i, j, k, grid)
@@ -255,30 +255,20 @@ KernelFunctionOperation at (Center, Center, Center)
     return BackgroundPotentialEnergy(model, model.buoyancy, geopotential_height)
 end
 
-@inline function BackgroundPotentialEnergy(model, buoyancy_model::BuoyancyTracerModel, geopotential_height)
-
-    grid = model.grid
-    sorted_buoyancy, z✶ = OneDReferenceField(model.tracers.b)
-
-    return KernelFunctionOperation{Center, Center, Center}(bz✶_ccc, grid, sorted_buoyancy, z✶)
-end
-
-@inline bz✶_ccc(i, j, k, grid, b, z✶) = b[i, j, k] * z✶[i, j, k]
-
 linear_eos_buoyancy(grid, buoyancy, tracers) = KernelFunctionOperation{Center, Center, Center}(buoyancy_perturbationᶜᶜᶜ, grid, buoyancy, tracers)
 
-@inline function BackgroundPotentialEnergy(model, buoyancy_model::BuoyancyLinearEOSModel, geopotential_height)
+@inline function BackgroundPotentialEnergy(model, buoyancy_model::Union{BuoyancyTracerModel, BuoyancyLinearEOSModel}, geopotential_height)
 
     grid = model.grid
-    buoyancy = model.buoyancy.model
-    tracers = model.tracers
-    b = linear_eos_buoyancy(grid, buoyancy, tracers)
-    compute!(b)
+    b = buoyancy_model isa BuoyancyTracerModel ? model.tracers.b :
+                                                 compute!(Field(linear_eos_buoyancy(grid, buoyancy_model.model, model.tracers)))
     sorted_buoyancy, z✶ = OneDReferenceField(b)
 
     return KernelFunctionOperation{Center, Center, Center}(bz✶_ccc, grid, sorted_buoyancy, z✶)
 end
 
+@inline bz✶_ccc(i, j, k, grid, b, z✶) = b[i, j, k] * z✶[i, j, k]
+
 @inline function BackgroundPotentialEnergy(model, buoyancy_model::BuoyancyBoussinesqEOSModel, geopotential_height)
 
     grid = model.grid

From ac0ed37919daea14de42b72b808efdf010bd960e Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Thu, 16 May 2024 14:05:55 +1000
Subject: [PATCH 27/37] Default to `geopotential_height = 0`

---
 src/PotentialEnergyEquationTerms.jl | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index f4b13004..f5c77b43 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -6,7 +6,6 @@ export PotentialEnergy, BackgroundPotentialEnergy, OneDReferenceField
 
 using Oceananigans.AbstractOperations: KernelFunctionOperation, volume, Az, GridMetricOperation
 using Oceananigans.Models: seawater_density
-using Oceananigans.Models: model_geopotential_height
 using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
 using Oceananigans.Grids
 using Oceananigans.Grids: Center, NegativeZDirection, interior, CenterField
@@ -115,8 +114,7 @@ KernelFunctionOperation at (Center, Center, Center)
 └── arguments: ("KernelFunctionOperation at (Center, Center, Center)", "(g=9.80665, ρ₀=1020.0)")
 ```
 """
-@inline function PotentialEnergy(model; location = (Center, Center, Center),
-                                 geopotential_height = model_geopotential_height(model))
+@inline function PotentialEnergy(model; location = (Center, Center, Center), geopotential_height = 0)
 
     validate_location(location, "PotentialEnergy")
     isnothing(model.buoyancy) ? nothing : validate_gravity_unit_vector(model.buoyancy.gravity_unit_vector)
@@ -196,7 +194,7 @@ function OneDReferenceField(f::Field; rev = false)
 
     p = sortperm(field_data; rev)
     sorted_field_data = field_data[p]
-    z✶ = cumsum(v[p]) / area
+    z✶ = cumsum(v[p]) / area # divide by area at surface
 
     grid_arch = f.grid.architecture
     grid_size = prod(size(f.grid))
@@ -246,7 +244,7 @@ KernelFunctionOperation at (Center, Center, Center)
 ```
 """
 @inline function BackgroundPotentialEnergy(model; location = (Center, Center, Center),
-                                                  geopotential_height = model_geopotential_height(model))
+                                                  geopotential_height = 0)
 
     validate_location(location, "BackgroundPotentialEnergy")
     isnothing(model.buoyancy) ? nothing : validate_gravity_unit_vector(model.buoyancy.gravity_unit_vector)
@@ -262,7 +260,7 @@ linear_eos_buoyancy(grid, buoyancy, tracers) = KernelFunctionOperation{Center, C
     grid = model.grid
     b = buoyancy_model isa BuoyancyTracerModel ? model.tracers.b :
                                                  compute!(Field(linear_eos_buoyancy(grid, buoyancy_model.model, model.tracers)))
-    sorted_buoyancy, z✶ = OneDReferenceField(b)
+    sorted_buoyancy, z✶ = OneDReferenceField(b, rev = true)
 
     return KernelFunctionOperation{Center, Center, Center}(bz✶_ccc, grid, sorted_buoyancy, z✶)
 end
@@ -272,7 +270,7 @@ end
 @inline function BackgroundPotentialEnergy(model, buoyancy_model::BuoyancyBoussinesqEOSModel, geopotential_height)
 
     grid = model.grid
-    ρ = seawater_density(model; geopotential_height)
+    ρ = seawater_density(model; geopotential_height) # default to σ₀, potential density referenced to sea surface height
     compute!(ρ)
     sorted_density, z✶ = OneDReferenceField(ρ)
     parameters = (g = model.buoyancy.model.gravitational_acceleration,

From 8689673ffba8148b3b70f6b3e2ac688f03bf4068 Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Thu, 16 May 2024 14:56:00 +1000
Subject: [PATCH 28/37] Simplify testing

This should also fix the test error
---
 test/runtests.jl | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 3d9066fa..b70ad55f 100755
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -364,10 +364,9 @@ function test_potential_energy_equation_terms_errors(model)
     return nothing
 end
 
-function test_potential_energy_equation_terms(model; geopotential_height = nothing)
+function test_potential_energy_equation_terms(model; geopotential_height = 0)
 
-    Eₚ = isnothing(geopotential_height) ? PotentialEnergy(model) :
-                                          PotentialEnergy(model; geopotential_height)
+    Eₚ = PotentialEnergy(model; geopotential_height)
 
     Eₚ_field = Field(Eₚ)
     @test Eₚ isa AbstractOperation
@@ -375,8 +374,7 @@ function test_potential_energy_equation_terms(model; geopotential_height = nothi
     compute!(Eₚ_field)
 
     if model.buoyancy isa BuoyancyBoussinesqEOSModel
-        ρ = isnothing(geopotential_height) ? Field(seawater_density(model)) :
-                                             Field(seawater_density(model; geopotential_height))
+        ρ = Field(seawater_density(model); geopotential_height)
 
         compute!(ρ)
         Z = Field(model_geopotential_height(model))
@@ -577,7 +575,7 @@ model_types = (NonhydrostaticModel, HydrostaticFreeSurfaceModel)
                     test_potential_energy_equation_terms_errors(model)
                 else
                     test_potential_energy_equation_terms(model)
-                    test_potential_energy_equation_terms(model, geopotential_height = 0)
+                    test_potential_energy_equation_terms(model, geopotential_height = 1000)
                 end
 
             end

From 8219b78169c8ca684bf389e58a8ebe6f74035b1b Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Thu, 16 May 2024 15:12:44 +1000
Subject: [PATCH 29/37] Correct misaligned brackets

---
 test/runtests.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index b70ad55f..e70515cb 100755
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -374,7 +374,7 @@ function test_potential_energy_equation_terms(model; geopotential_height = 0)
     compute!(Eₚ_field)
 
     if model.buoyancy isa BuoyancyBoussinesqEOSModel
-        ρ = Field(seawater_density(model); geopotential_height)
+        ρ = Field(seawater_density(model; geopotential_height))
 
         compute!(ρ)
         Z = Field(model_geopotential_height(model))

From 71b55371407bd27afe0aae050f028a0c8720a1e3 Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Fri, 17 May 2024 12:49:13 +1000
Subject: [PATCH 30/37] Add tests for `OneDReferenceField` + bpe

Still need a test for checking z*
---
 src/PotentialEnergyEquationTerms.jl | 11 +++---
 test/runtests.jl                    | 54 +++++++++++++++++++++++++++--
 2 files changed, 58 insertions(+), 7 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index f5c77b43..a916f0b4 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -255,12 +255,15 @@ end
 
 linear_eos_buoyancy(grid, buoyancy, tracers) = KernelFunctionOperation{Center, Center, Center}(buoyancy_perturbationᶜᶜᶜ, grid, buoyancy, tracers)
 
+@inline BackgroundPotentialEnergy(model, buoyancy_model::NoBuoyancyModel, geopotential_height) =
+    throw(ArgumentError("Cannot calculate gravitational potential energy without a Buoyancy model."))
+
 @inline function BackgroundPotentialEnergy(model, buoyancy_model::Union{BuoyancyTracerModel, BuoyancyLinearEOSModel}, geopotential_height)
 
-    grid = model.grid
     b = buoyancy_model isa BuoyancyTracerModel ? model.tracers.b :
-                                                 compute!(Field(linear_eos_buoyancy(grid, buoyancy_model.model, model.tracers)))
+                                                 compute!(Field(linear_eos_buoyancy(model.grid, buoyancy_model.model, model.tracers)))
     sorted_buoyancy, z✶ = OneDReferenceField(b, rev = true)
+    grid = sorted_buoyancy.grid
 
     return KernelFunctionOperation{Center, Center, Center}(bz✶_ccc, grid, sorted_buoyancy, z✶)
 end
@@ -269,10 +272,10 @@ end
 
 @inline function BackgroundPotentialEnergy(model, buoyancy_model::BuoyancyBoussinesqEOSModel, geopotential_height)
 
-    grid = model.grid
-    ρ = seawater_density(model; geopotential_height) # default to σ₀, potential density referenced to sea surface height
+    ρ = Field(seawater_density(model; geopotential_height)) # default to σ₀, potential density referenced to sea surface height
     compute!(ρ)
     sorted_density, z✶ = OneDReferenceField(ρ)
+    grid = sorted_density.grid
     parameters = (g = model.buoyancy.model.gravitational_acceleration,
                   ρ₀ = model.buoyancy.model.equation_of_state.reference_density)
 
diff --git a/test/runtests.jl b/test/runtests.jl
index e70515cb..6d97a26e 100755
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -11,7 +11,9 @@ using SeawaterPolynomials: RoquetEquationOfState, TEOS10EquationOfState
 using Oceanostics
 using Oceanostics: TKEBudgetTerms, TracerVarianceBudgetTerms, FlowDiagnostics, PressureRedistributionTerm, BuoyancyProductionTerm, AdvectionTerm
 using Oceanostics.TKEBudgetTerms: AdvectionTerm
-using Oceanostics: PotentialEnergy, PotentialEnergyEquationTerms.BuoyancyBoussinesqEOSModel
+using Oceanostics: PotentialEnergy, BackgroundPotentialEnergy
+using Oceanostics.PotentialEnergyEquationTerms: OneDReferenceField, linear_eos_buoyancy
+using Oceanostics.PotentialEnergyEquationTerms: BuoyancyTracerModel, BuoyancyLinearEOSModel, BuoyancyBoussinesqEOSModel
 using Oceanostics.ProgressMessengers
 using Oceanostics: perturbation_fields
 
@@ -360,6 +362,8 @@ function test_potential_energy_equation_terms_errors(model)
 
     @test_throws ArgumentError PotentialEnergy(model)
     @test_throws ArgumentError PotentialEnergy(model, geopotential_height = 0)
+    @test_throws ArgumentError BackgroundPotentialEnergy(model)
+    @test_throws ArgumentError BackgroundPotentialEnergy(model, geopotential_height = 0)
 
     return nothing
 end
@@ -375,10 +379,11 @@ function test_potential_energy_equation_terms(model; geopotential_height = 0)
 
     if model.buoyancy isa BuoyancyBoussinesqEOSModel
         ρ = Field(seawater_density(model; geopotential_height))
-
         compute!(ρ)
+
         Z = Field(model_geopotential_height(model))
         compute!(Z)
+
         ρ₀ = model.buoyancy.model.equation_of_state.reference_density
         g = model.buoyancy.model.gravitational_acceleration
 
@@ -388,6 +393,45 @@ function test_potential_energy_equation_terms(model; geopotential_height = 0)
         end
     end
 
+    Eb = BackgroundPotentialEnergy(model; geopotential_height)
+
+    Eb_field = Field(Eb)
+    @test Eb isa AbstractOperation
+    @test Eb_field isa Field
+    compute!(Eb_field)
+
+    return nothing
+end
+function test_1D_reference_field(model; geopotential_height = 0)
+
+    if model.buoyancy isa BuoyancyTracerModel
+
+        b_sorted = sort(reshape(interior(model.tracers.b), :), rev = true)
+
+        sorted_buoyancy, z✶ = OneDReferenceField(model.tracers.b, rev = true)
+
+        @test isequal(reshape(interior(sorted_buoyancy), :), b_sorted)
+
+    elseif model.buoyancy isa BuoyancyLinearEOSModel
+
+        b = Field(linear_eos_buoyancy(model.grid, model.buoyancy.model, model.tracers))
+        compute!(b)
+        b_sorted = sort(reshape(interior(b), :), rev = true)
+
+        sorted_buoyancy, z✶ = OneDReferenceField(b, rev = true)
+
+        @test isequal(reshape(interior(sorted_buoyancy), :), b_sorted)
+
+    elseif model.buoyancy isa BuoyancyBoussinesqEOSModel
+        ρ = Field(seawater_density(model; geopotential_height))
+        compute!(ρ)
+        ρ_data_sorted = sort(reshape(interior(ρ), :))
+
+        sorted_density, z✶ = OneDReferenceField(ρ)
+
+        @test isequal(reshape(interior(sorted_density), :), ρ_data_sorted)
+    end
+
     return nothing
 end
 #---
@@ -565,7 +609,7 @@ model_types = (NonhydrostaticModel, HydrostaticFreeSurfaceModel)
 
             end
 
-            @info "Testing `PotentialEnergy`"
+            @info "Testing `PotentialEnergy`, `BackgroundPotentialEnergy` and `OneDReferenceField`"
             for buoyancy in buoyancy_models
 
                 tracers = buoyancy isa BuoyancyTracer ? :b : (:S, :T)
@@ -576,6 +620,10 @@ model_types = (NonhydrostaticModel, HydrostaticFreeSurfaceModel)
                 else
                     test_potential_energy_equation_terms(model)
                     test_potential_energy_equation_terms(model, geopotential_height = 1000)
+                    buoyancy isa BuoyancyTracer ? set!(model, b = randn(size(model.grid))) :
+                                                  set!(model, S = randn(size(model.grid)), T = randn(size(model.grid)))
+                    test_1D_reference_field(model)
+                    test_1D_reference_field(model, geopotential_height = 1000)
                 end
 
             end

From 82c24212e5d7f290a0109e184345f7e7fee7dcda Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Fri, 17 May 2024 14:39:39 +1000
Subject: [PATCH 31/37] Test z*

More validation that the `OneDReferenceField` is behaving as it should than test.
---
 test/runtests.jl | 32 ++++++++++++++++++++++++++++----
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 6d97a26e..a4b856cc 100755
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -14,6 +14,7 @@ using Oceanostics.TKEBudgetTerms: AdvectionTerm
 using Oceanostics: PotentialEnergy, BackgroundPotentialEnergy
 using Oceanostics.PotentialEnergyEquationTerms: OneDReferenceField, linear_eos_buoyancy
 using Oceanostics.PotentialEnergyEquationTerms: BuoyancyTracerModel, BuoyancyLinearEOSModel, BuoyancyBoussinesqEOSModel
+using Oceanostics.PotentialEnergyEquationTerms: VolumeField, AreaField
 using Oceanostics.ProgressMessengers
 using Oceanostics: perturbation_fields
 
@@ -406,30 +407,53 @@ function test_1D_reference_field(model; geopotential_height = 0)
 
     if model.buoyancy isa BuoyancyTracerModel
 
-        b_sorted = sort(reshape(interior(model.tracers.b), :), rev = true)
+        b = reshape(interior(model.tracers.b), :)
+        p = sortperm(b, rev = true)
+        b_sorted = b[p]
+
+        v_grid = VolumeField(model.grid)
+        A = sum(AreaField(model.grid))
+        z✶_check = cumsum(reshape(v_grid, :)[p]) / A
 
         sorted_buoyancy, z✶ = OneDReferenceField(model.tracers.b, rev = true)
 
         @test isequal(reshape(interior(sorted_buoyancy), :), b_sorted)
+        @test isequal(reshape(interior(z✶), :), z✶_check)
 
     elseif model.buoyancy isa BuoyancyLinearEOSModel
 
         b = Field(linear_eos_buoyancy(model.grid, model.buoyancy.model, model.tracers))
         compute!(b)
-        b_sorted = sort(reshape(interior(b), :), rev = true)
+        b_flat = reshape(interior(b), :)
+        p = sortperm(b_flat, rev = true)
+        b_sorted = b_flat[p]
+
+        v_grid = VolumeField(model.grid)
+        A = sum(AreaField(model.grid))
+        z✶_check = cumsum(reshape(v_grid, :)[p]) / A
 
         sorted_buoyancy, z✶ = OneDReferenceField(b, rev = true)
 
         @test isequal(reshape(interior(sorted_buoyancy), :), b_sorted)
+        @test isequal(reshape(interior(z✶), :), z✶_check)
 
     elseif model.buoyancy isa BuoyancyBoussinesqEOSModel
+
         ρ = Field(seawater_density(model; geopotential_height))
         compute!(ρ)
-        ρ_data_sorted = sort(reshape(interior(ρ), :))
+        ρ_flat = reshape(interior(ρ), :)
+        p = sortperm(ρ_flat)
+        ρ_sorted = ρ_flat[p]
+
+        v_grid = VolumeField(model.grid)
+        A = sum(AreaField(model.grid))
+        z✶_check = cumsum(reshape(v_grid, :)[p]) / A
 
         sorted_density, z✶ = OneDReferenceField(ρ)
 
-        @test isequal(reshape(interior(sorted_density), :), ρ_data_sorted)
+        @test isequal(reshape(interior(sorted_density), :), ρ_sorted)
+        @test isequal(reshape(interior(z✶), :), z✶_check)
+
     end
 
     return nothing

From d4d396ea92aabd7887d800f0bd7486f6ea09740b Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Mon, 20 May 2024 10:16:27 +1000
Subject: [PATCH 32/37] Docstring updates

---
 src/PotentialEnergyEquationTerms.jl | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index a916f0b4..468b1352 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -40,9 +40,9 @@ and `SeawaterBuoyancy`. See the relevant Oceananigans.jl documentation on
 [buoyancy models](https://clima.github.io/OceananigansDocumentation/dev/model_setup/buoyancy_and_equation_of_state/)
 for more information about available options.
 
-The optional keyword argument `geopotential_height` is only used
-if ones wishes to calculate `Eₚ` with a potential density referenced to `geopotential_height`,
-rather than in-situ density, when using a `BoussinesqEquationOfState`.
+The keyword argument `geopotential_height` allows users to choose a set a different potential
+density for use in the calulation of `Eₚ` when using a `BoussinesqEquationOfState`. See the
+example below for how to do this.
 
 Example
 =======
@@ -64,8 +64,9 @@ KernelFunctionOperation at (Center, Center, Center)
 └── arguments: ("1×1×100 Field{Center, Center, Center} on RectilinearGrid on CPU",)
 ```
 
-The default behaviour of `PotentialEnergy` uses the *in-situ density* in the calculation
-when the equation of state is a `BoussinesqEquationOfState`:
+The default behaviour of `PotentialEnergy` uses *potential density* with reference
+`geopotential_height = 0` (i.e. σ₀ is the density variable) in the calculation when the
+equation of state is a `BoussinesqEquationOfState`:
 ```jldoctest
 julia> using Oceananigans, SeawaterPolynomials.TEOS10
 
@@ -88,8 +89,8 @@ KernelFunctionOperation at (Center, Center, Center)
 └── arguments: ("KernelFunctionOperation at (Center, Center, Center)", "(g=9.80665, ρ₀=1020.0)")
 ```
 
-To use a reference density set a constant value for the keyword argument `geopotential_height`
-and pass this the function. For example,
+To use a different potential density set a constant value for the keyword argument
+`geopotential_height` and pass this the function. For example,
 ```jldoctest
 julia> using Oceananigans, SeawaterPolynomials.TEOS10;
 
@@ -105,9 +106,9 @@ julia> buoyancy = SeawaterBuoyancy(equation_of_state=eos);
 
 julia> model = NonhydrostaticModel(; grid, buoyancy, tracers);
 
-julia> geopotential_height = 0; # density variable will be σ₀
+julia> geopotential_height = 1000; # density variable will be σ₁
 
-julia> PotentialEnergy(model)
+julia> PotentialEnergy(model; geopotential_height)
 KernelFunctionOperation at (Center, Center, Center)
 ├── grid: 1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
 ├── kernel_function: g′z_ccc (generic function with 1 method)
@@ -150,7 +151,7 @@ end
 @inline function PotentialEnergy(model, buoyancy_model::BuoyancyBoussinesqEOSModel, geopotential_height)
 
     grid = model.grid
-    ρ = seawater_density(model; geopotential_height)
+    ρ = seawater_density(model; geopotential_height) # default to σ₀, potential density referenced to sea surface height
     parameters = (g = model.buoyancy.model.gravitational_acceleration,
                   ρ₀ = model.buoyancy.model.equation_of_state.reference_density)
 
@@ -212,7 +213,7 @@ end
 """
     $(SIGNATURES)
 
-Return a `kernelFunctionOperation` to compute an approximation to the `BackgroundPotentialEnergy`
+Return a `KernelFunctionOperation` to compute the `BackgroundPotentialEnergy`
 per unit volume,
 ```math
 E_{b} = \\frac{gρz✶}{ρ₀}.

From 63a58b297513968441e678ead240a9842152f330 Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Wed, 22 May 2024 11:56:41 +1000
Subject: [PATCH 33/37] =?UTF-8?q?Align=20`z`=20and=20`z=E2=9C=B6`?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

So that `z✶` increases in the same direction as `z`.
---
 src/PotentialEnergyEquationTerms.jl | 9 +++++----
 test/runtests.jl                    | 6 +++---
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 468b1352..76b82098 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -180,9 +180,10 @@ has the same vertical extent as the original field (`f.grid.Lz`) but with the sa
 number of points of the original domain (`Nx * Ny * Nz`).
 The `z✶` `Field` is defined as
 ```math
-\\frac{1}{A}\\int_{f\\mathrm{min}}^{f\\mathrm{max}} \\mathrm{d}V.
+-\\frac{1}{A}\\int_{f\\mathrm{min}}^{f\\mathrm{max}} \\mathrm{d}V.
 ```
-and is computed by cumulatively summing the 1D `Array` of grid volumes `ΔV`.
+and is computed by cumulatively summing the 1D `Array` of grid volumes `ΔV`. The negative sign
+ensures that `z✶` is increasing over the same range and same direction as `f.grid`.
 **Note:** `OneDReferenceField` is currently only appropriate for grids that have uniform
 horizontal area.
 """
@@ -195,7 +196,7 @@ function OneDReferenceField(f::Field; rev = false)
 
     p = sortperm(field_data; rev)
     sorted_field_data = field_data[p]
-    z✶ = cumsum(v[p]) / area # divide by area at surface
+    z✶ = -cumsum(v[p]) / area # divide by area at surface
 
     grid_arch = f.grid.architecture
     grid_size = prod(size(f.grid))
@@ -216,7 +217,7 @@ end
 Return a `KernelFunctionOperation` to compute the `BackgroundPotentialEnergy`
 per unit volume,
 ```math
-E_{b} = \\frac{gρz✶}{ρ₀}.
+E_{b} = \\frac{gρz✶}{ρ₀} = -bz.
 ```
 The `BackgroundPotentialEnergy` is the potential energy computed by adiabatically resorting
 the buoyancy or density field into a reference state of minimal potential energy.
diff --git a/test/runtests.jl b/test/runtests.jl
index a4b856cc..75c90c07 100755
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -413,7 +413,7 @@ function test_1D_reference_field(model; geopotential_height = 0)
 
         v_grid = VolumeField(model.grid)
         A = sum(AreaField(model.grid))
-        z✶_check = cumsum(reshape(v_grid, :)[p]) / A
+        z✶_check = -cumsum(reshape(v_grid, :)[p]) / A
 
         sorted_buoyancy, z✶ = OneDReferenceField(model.tracers.b, rev = true)
 
@@ -430,7 +430,7 @@ function test_1D_reference_field(model; geopotential_height = 0)
 
         v_grid = VolumeField(model.grid)
         A = sum(AreaField(model.grid))
-        z✶_check = cumsum(reshape(v_grid, :)[p]) / A
+        z✶_check = -cumsum(reshape(v_grid, :)[p]) / A
 
         sorted_buoyancy, z✶ = OneDReferenceField(b, rev = true)
 
@@ -447,7 +447,7 @@ function test_1D_reference_field(model; geopotential_height = 0)
 
         v_grid = VolumeField(model.grid)
         A = sum(AreaField(model.grid))
-        z✶_check = cumsum(reshape(v_grid, :)[p]) / A
+        z✶_check = -cumsum(reshape(v_grid, :)[p]) / A
 
         sorted_density, z✶ = OneDReferenceField(ρ)
 

From 3f752bf8acfdff968a32e9dd643a51978910c23c Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Wed, 22 May 2024 11:58:47 +1000
Subject: [PATCH 34/37] =?UTF-8?q?Missed=20the=20=E2=9C=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/PotentialEnergyEquationTerms.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 76b82098..39173836 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -217,7 +217,7 @@ end
 Return a `KernelFunctionOperation` to compute the `BackgroundPotentialEnergy`
 per unit volume,
 ```math
-E_{b} = \\frac{gρz✶}{ρ₀} = -bz.
+E_{b} = \\frac{gρz✶}{ρ₀} = -bz✶.
 ```
 The `BackgroundPotentialEnergy` is the potential energy computed by adiabatically resorting
 the buoyancy or density field into a reference state of minimal potential energy.

From 1fd9d64ccb93f2b822b42e82d88c895b489f82a2 Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Wed, 22 May 2024 20:09:04 +1000
Subject: [PATCH 35/37] Correct `BuoyancyTracer()` and linear eos sorting
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Also rename functions to `minus_bz✶_ccc` to try and make clearer what they are doing
---
 src/PotentialEnergyEquationTerms.jl | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 39173836..846971ce 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -208,7 +208,6 @@ function OneDReferenceField(f::Field; rev = false)
     set!(z✶_field, reshape(z✶, size(new_grid)))
 
     return sorted_field, z✶_field
-
 end
 
 """
@@ -217,7 +216,7 @@ end
 Return a `KernelFunctionOperation` to compute the `BackgroundPotentialEnergy`
 per unit volume,
 ```math
-E_{b} = \\frac{gρz✶}{ρ₀} = -bz✶.
+E_{b} = \\frac{g}{ρ₀}ρz✶ = -bz✶.
 ```
 The `BackgroundPotentialEnergy` is the potential energy computed by adiabatically resorting
 the buoyancy or density field into a reference state of minimal potential energy.
@@ -241,7 +240,7 @@ julia> model = NonhydrostaticModel(; grid, buoyancy=BuoyancyTracer(), tracers=(:
 julia> bpe = BackgroundPotentialEnergy(model)
 KernelFunctionOperation at (Center, Center, Center)
 ├── grid: 1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
-├── kernel_function: bz✶_ccc (generic function with 1 method)
+├── kernel_function: minus_bz✶_ccc (generic function with 2 methods)
 └── arguments: ("1×1×100 Field{Center, Center, Center} on RectilinearGrid on CPU", "1×1×100 Field{Center, Center, Center} on RectilinearGrid on CPU")
 ```
 """
@@ -255,22 +254,24 @@ KernelFunctionOperation at (Center, Center, Center)
     return BackgroundPotentialEnergy(model, model.buoyancy, geopotential_height)
 end
 
-linear_eos_buoyancy(grid, buoyancy, tracers) = KernelFunctionOperation{Center, Center, Center}(buoyancy_perturbationᶜᶜᶜ, grid, buoyancy, tracers)
+linear_eos_buoyancy(grid, buoyancy, tracers) =
+    KernelFunctionOperation{Center, Center, Center}(buoyancy_perturbationᶜᶜᶜ, grid, buoyancy, tracers)
 
 @inline BackgroundPotentialEnergy(model, buoyancy_model::NoBuoyancyModel, geopotential_height) =
     throw(ArgumentError("Cannot calculate gravitational potential energy without a Buoyancy model."))
 
 @inline function BackgroundPotentialEnergy(model, buoyancy_model::Union{BuoyancyTracerModel, BuoyancyLinearEOSModel}, geopotential_height)
 
-    b = buoyancy_model isa BuoyancyTracerModel ? model.tracers.b :
-                                                 compute!(Field(linear_eos_buoyancy(model.grid, buoyancy_model.model, model.tracers)))
-    sorted_buoyancy, z✶ = OneDReferenceField(b, rev = true)
-    grid = sorted_buoyancy.grid
+    minus_b = buoyancy_model isa BuoyancyTracerModel ? -model.tracers.b :
+                                                       compute!(Field(-linear_eos_buoyancy(model.grid, buoyancy_model.model, model.tracers)))
+
+    sorted_minus_b, z✶ = OneDReferenceField(minus_b)
+    grid = sorted_minus_b.grid
 
-    return KernelFunctionOperation{Center, Center, Center}(bz✶_ccc, grid, sorted_buoyancy, z✶)
+    return KernelFunctionOperation{Center, Center, Center}(minus_bz✶_ccc, grid, sorted_minus_b, z✶)
 end
 
-@inline bz✶_ccc(i, j, k, grid, b, z✶) = b[i, j, k] * z✶[i, j, k]
+@inline minus_bz✶_ccc(i, j, k, grid, sorted_minus_b::Field, z✶::Field) = sorted_minus_b[i, j, k] * z✶[i, j, k]
 
 @inline function BackgroundPotentialEnergy(model, buoyancy_model::BuoyancyBoussinesqEOSModel, geopotential_height)
 
@@ -284,6 +285,6 @@ end
     return KernelFunctionOperation{Center, Center, Center}(g′z✶_ccc, grid, sorted_density, z✶, parameters)
 end
 
-@inline g′z✶_ccc(i, j, k, grid, ρ, z✶, p) = (p.g / p.ρ₀) * ρ[i, j, k] * z✶[i, j, k]
+@inline minus_bz✶_ccc(i, j, k, grid, sorted_ρ::Field, z✶::Field, p::NamedTuple) = (p.g / p.ρ₀) * sorted_ρ[i, j, k] * z✶[i, j, k]
 
 end # module

From 7514adf689cf44e3ed021664c393cbde9ac1a463 Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Thu, 23 May 2024 10:07:16 +1000
Subject: [PATCH 36/37] Update `Type` input for `OneDReferenceField`

---
 src/PotentialEnergyEquationTerms.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 846971ce..0f789047 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -4,7 +4,7 @@ using DocStringExtensions
 
 export PotentialEnergy, BackgroundPotentialEnergy, OneDReferenceField
 
-using Oceananigans.AbstractOperations: KernelFunctionOperation, volume, Az, GridMetricOperation
+using Oceananigans.AbstractOperations: AbstractOperation, KernelFunctionOperation, volume, Az, GridMetricOperation
 using Oceananigans.Models: seawater_density
 using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
 using Oceananigans.Grids
@@ -187,7 +187,7 @@ ensures that `z✶` is increasing over the same range and same direction as `f.g
 **Note:** `OneDReferenceField` is currently only appropriate for grids that have uniform
 horizontal area.
 """
-function OneDReferenceField(f::Field; rev = false)
+function OneDReferenceField(f::Union{Field, AbstractOperation}; rev = false)
 
     area = sum(AreaField(f.grid))
     volume_field = VolumeField(f.grid)

From 7cbc0527b9ef087ba4138f1dd5ec479d723429cd Mon Sep 17 00:00:00 2001
From: jbisits <jbisits@gmail.com>
Date: Thu, 23 May 2024 10:20:50 +1000
Subject: [PATCH 37/37] Correct function call

---
 src/PotentialEnergyEquationTerms.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/PotentialEnergyEquationTerms.jl b/src/PotentialEnergyEquationTerms.jl
index 0f789047..17999e2e 100644
--- a/src/PotentialEnergyEquationTerms.jl
+++ b/src/PotentialEnergyEquationTerms.jl
@@ -282,7 +282,7 @@ end
     parameters = (g = model.buoyancy.model.gravitational_acceleration,
                   ρ₀ = model.buoyancy.model.equation_of_state.reference_density)
 
-    return KernelFunctionOperation{Center, Center, Center}(g′z✶_ccc, grid, sorted_density, z✶, parameters)
+    return KernelFunctionOperation{Center, Center, Center}(minus_bz✶_ccc, grid, sorted_density, z✶, parameters)
 end
 
 @inline minus_bz✶_ccc(i, j, k, grid, sorted_ρ::Field, z✶::Field, p::NamedTuple) = (p.g / p.ρ₀) * sorted_ρ[i, j, k] * z✶[i, j, k]