diff --git a/Project.toml b/Project.toml
index 5256cde..b9d0474 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "TwoLayerDirectNumericalShenanigans"
 uuid = "40aaee9f-3595-48be-b36c-f1067009652f"
 authors = ["Josef Bisits <jbisits@gmail.com>"]
-version = "0.6.1"
+version = "0.6.2"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
diff --git a/src/TwoLayerDirectNumericalShenanigans.jl b/src/TwoLayerDirectNumericalShenanigans.jl
index c54e49c..fe09fe7 100644
--- a/src/TwoLayerDirectNumericalShenanigans.jl
+++ b/src/TwoLayerDirectNumericalShenanigans.jl
@@ -3,12 +3,11 @@ module TwoLayerDirectNumericalShenanigans
 using Oceananigans, Printf, Reexport, JLD2, NCDatasets, GibbsSeaWater
 using Oceananigans: AbstractModel, Operators.ℑzᵃᵃᶜ
 using Oceananigans: Models.seawater_density
-using SeawaterPolynomials: BoussinesqEquationOfState
+using SeawaterPolynomials: BoussinesqEquationOfState, ρ
 using SeawaterPolynomials.TEOS10
 using SeawaterPolynomials.SecondOrderSeawaterPolynomials
-import SeawaterPolynomials.ρ
 using SpecialFunctions: erf
-using Oceanostics: KineticEnergyDissipationRate
+using Oceanostics: KineticEnergyDissipationRate, KineticEnergy
 using CUDA: allowscalar, CuArray
 import Base: show, iterate
 
diff --git a/src/kernelfunctions.jl b/src/kernelfunctions.jl
index 0682f35..ed648fc 100644
--- a/src/kernelfunctions.jl
+++ b/src/kernelfunctions.jl
@@ -1,6 +1,10 @@
 @inline tracer_perturbationᶜᶜᶜ(i, j, k, grid, tracer, tracer_mean) =
     tracer[i, j, k] - tracer_mean[i, j, k]
 # Only needed for testing
+"""
+    function tracer_perturbation(model, tracer)
+Calculate the perturbation of `tracer` from the spatial `tracer` mean.
+"""
 @inline function tracer_perturbation(model, tracer)
 
     grid = model.grid
@@ -11,6 +15,10 @@ end
 
 @inline vtf(i, j, k, grid, tracer, tracer_mean, w) =
         -ℑzᵃᵃᶜ(i, j, k, w) * tracer_perturbationᶜᶜᶜ(i, j, k, grid, tracer, tracer_mean)
+"""
+    function vertical_tracer_flux(model, tracer)
+Calculate the vertical tracer flux of `tracer` from `model`.
+"""
 @inline function vertical_tracer_flux(model, tracer)
 
     grid = model.grid
@@ -19,3 +27,21 @@ end
 
     return KernelFunctionOperation{Center, Center, Center}(vtf, grid, tracer, tracer_mean, w)
 end
+"""
+    function potential_energy(model)
+Calculate the potential energy
+```math
+Eₚ = g∫ᵥρzdV
+```
+for `model` during a simulation.
+"""
+@inline function potential_energy(model)
+
+    grid = model.grid
+    ρ = seawater_density(model)
+    Z = Oceananigans.Models.model_geopotential_height(model)
+    g = tuple(model.buoyancy.model.gravitational_acceleration)
+
+    return KernelFunctionOperation{Center, Center, Center}(Ep, grid, ρ, Z, g)
+end
+@inline Ep(i, j, k, grid, ρ, Z, g) = g[1] * ρ[i, j, k] * Z[i, j, k]
diff --git a/src/twolayerdns.jl b/src/twolayerdns.jl
index f7e5874..ededa11 100644
--- a/src/twolayerdns.jl
+++ b/src/twolayerdns.jl
@@ -292,6 +292,10 @@ function save_computed_output!(simulation, tldns, save_schedule, save_file, outp
     ϵ = KineticEnergyDissipationRate(model)
     ∫ϵ = Integral(ϵ)
     ϵ_maximum = Reduction(maximum!, ϵ, dims = (1, 2, 3))
+    Eₖ = KineticEnergy(model)
+    ∫Eₖ = Integral(Eₖ)
+    Eₚ = potential_energy(model)
+    ∫Eₚ = Integral(Eₚ)
 
     # Horizontally integrated vertical temperature flux and vertical temperature gradient
     T = model.tracers.T
@@ -300,13 +304,16 @@ function save_computed_output!(simulation, tldns, save_schedule, save_file, outp
     ∂T∂z = ∂z(T)
     ∫ₐ∂T∂z = Integral(∂T∂z, dims = (1, 2))
 
-    computed_outputs = Dict("σ" => σ, "∫ϵ" => ∫ϵ, "ϵ_maximum" => ϵ_maximum, "∫ₐ∂T∂z" => ∫ₐ∂T∂z,
-                            "∫ₐw′T′" => ∫ₐw′T′)
+    computed_outputs = Dict("σ" => σ, "∫ϵ" => ∫ϵ, "ϵ_maximum" => ϵ_maximum, "∫Eₖ" => ∫Eₖ,
+                            "∫Eₚ" => ∫Eₚ, "∫ₐ∂T∂z" => ∫ₐ∂T∂z, "∫ₐw′T′" => ∫ₐw′T′)
+
     oa = Dict(
         "σ" => Dict("longname" => "Seawater potential density calculated using TEOS-10 at $(reference_gp_height)dbar",
                     "units" => "kgm⁻³"),
         "ϵ_maximum" => Dict("longname" => "Maximum (in space) TKE dissipation"),
         "∫ϵ" => Dict("longname" => "Volume integrated turbulent kintetic energy dissipation"),
+        "∫Eₖ" => Dict("longname" => "Volume integrated turbulent kinetic energy"),
+        "∫Eₚ" => Dict("longname" => "Volume integrated potential energy (g∫ᵥρzdV)"),
         "∫ₐw′T′" => Dict("longname" => "Horizontally integrated vertical temeprature flux"),
         "∫ₐ∂T∂z" => Dict("longname" => "Horizontally integrated vertical temperature gradient")
         )
diff --git a/test/kernelfunction_test.jl b/test/kernelfunction_test.jl
index 3b46ee5..ac23d8e 100644
--- a/test/kernelfunction_test.jl
+++ b/test/kernelfunction_test.jl
@@ -13,3 +13,12 @@ profile_function = StepChange(transition_depth)
 tldns = TwoLayerDNS(model, profile_function, initial_conditions)
 
 set_two_layer_initial_conditions!(tldns)
+
+Eₚ_model = DNSModel(architecture, DOMAIN_EXTENT, resolution, diffusivities)
+set!(Eₚ_model, T = -1.5, S = 34.568)
+ρ_model = Field(seawater_density(Eₚ_model))
+compute!(ρ_model)
+z_grid = Field(Oceananigans.Models.model_geopotential_height(Eₚ_model))
+compute!(z_grid)
+dV = xspacings(model.grid, Center()) * yspacings(model.grid, Center()) * zspacings(model.grid, Center())
+computed_Eₚ = model.buoyancy.model.gravitational_acceleration * ρ_model.data .* z_grid.data
diff --git a/test/runtests.jl b/test/runtests.jl
index 4df577e..a6bfd72 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,6 +1,7 @@
 using TwoLayerDirectNumericalShenanigans, Test
 using TwoLayerDirectNumericalShenanigans: perturb_tracer
 using Oceananigans.Fields
+using Oceananigans: Models.seawater_density
 using SeawaterPolynomials
 using NCDatasets, JLD2
 import SeawaterPolynomials.ρ
@@ -189,3 +190,11 @@ include("kernelfunction_test.jl")
     @test all(vtflux.data .≈ 0 )
 
 end
+
+@testset "Potential energy" begin
+
+    Eₚ = Field(TLDNS.potential_energy(Eₚ_model))
+    compute!(Eₚ)
+    @test all(Eₚ.data == computed_Eₚ)
+
+end