Reorganize module and notations (#204)

Author: tomchor

.github/workflows/ci.yml

@@ -65,7 +65,7 @@ jobs:
           TEST_GROUP: "tracer_diagnostics"
 
   ke_diagnostics:
-    name: KE and TKE diagnostics - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+    name: KE diagnostics - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
     runs-on: ${{ matrix.os }}
     timeout-minutes: 120
     strategy:
@@ -89,6 +89,31 @@ jobs:
         env:
           TEST_GROUP: "ke_diagnostics"
 
+  tke_diagnostics:
+    name: TKE diagnostics - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+    runs-on: ${{ matrix.os }}
+    timeout-minutes: 120
+    strategy:
+      fail-fast: false
+      matrix:
+        version:
+          - '1.10'
+        os:
+          - ubuntu-latest
+        arch:
+          - x64
+    steps:
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: ${{ matrix.version }}
+          arch: ${{ matrix.arch }}
+      - uses: julia-actions/cache@v2
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1
+        env:
+          TEST_GROUP: "tke_diagnostics"
+
   pe_diagnostics:
     name: PE diagnostics - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
     runs-on: ${{ matrix.os }}

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceanostics"
 uuid = "d0ccf422-c8fb-49b5-a76d-74acdde946ac"
 authors = ["tomchor <[email protected]>"]
-version = "0.15.3"
+version = "0.16.0"
 
 [deps]
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
@@ -11,7 +11,7 @@ SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40"
 
 [compat]
 DocStringExtensions = "0.9"
-Oceananigans = "^0.96.30"
+Oceananigans = "^0.96.30, 0.97"
 SeawaterPolynomials = "0.3"
 julia = "1.9"
 

docs/examples/kelvin_helmholtz.jl

@@ -78,8 +78,8 @@ Q = QVelocityGradientTensorInvariant(model)
 # variance dissipation rate and diffusive term. When volume-integrated, these two quantities should
 # be equal.
 
-∫χᴰ = Integral(TracerVarianceDissipationRate(model, :b))
-∫χ = Integral(TracerVarianceDiffusiveTerm(model, :b))
+∫χᴰ = Integral(TracerVarianceEquation.DissipationRate(model, :b))
+∫χ = Integral(TracerVarianceEquation.Diffusion(model, :b))
 
 
 # Now we write these quantities, along with `b`, to a NetCDF:
@@ -163,6 +163,6 @@ end
 # ![](kelvin_helmholtz.mp4)
 #
 # Similarly to the kinetic energy dissipation rate (see the [Two-dimensional turbulence example](@ref two_d_turbulence_example)),
-# `TracerVarianceDissipationRate` and `TracerVarianceDiffusiveTerm` are implemented
+# `TracerVarianceDissipationRate` and `TracerVarianceDiffusion` are implemented
 # with a energy-conserving formulation, which means that (for `NoFlux` boundary conditions) their
 # volume-integral should be exactly (up to machine precision) the same.

docs/examples/tilted_bottom_boundary_layer.jl

@@ -177,9 +177,7 @@ run!(simulation)
 
 # Now we'll read the results and plot an animation
 
-using Rasters
-
-ds = RasterStack(simulation.output_writers[:nc].filepath)
+ds = NCDataset(simulation.output_writers[:nc].filepath)
 
 # We now use Makie to create the figure and its axes
 
@@ -198,31 +196,31 @@ ax3 = Axis(fig[2, 3]; title = "PV", kwargs...);
 
 n = Observable(1)
 
-x_caa = Array(dims(ds, :x_caa))
-x_faa = Array(dims(ds, :x_faa))
-z_aac = Array(dims(ds, :z_aac))
-z_aaf = Array(dims(ds, :z_aaf))
+x_caa = ds["x_caa"][:]
+x_faa = ds["x_faa"][:]
+z_aac = ds["z_aac"][:]
+z_aaf = ds["z_aaf"][:]
 
-bₙ = @lift Array(ds.b[Ti=$n, y_aca=Near(0)])
+bₙ = @lift ds["b"][:, :, $n]
 
-Riₙ = @lift Array(ds.Ri[Ti=$n, y_aca=Near(0)])
+Riₙ = @lift ds["Ri"][:, :, $n]
 hm1 = heatmap!(ax1, x_caa, z_aaf, Riₙ; colormap = :coolwarm, colorrange = (-1, +1))
 contour!(ax1, x_caa, z_aac, bₙ; levels=10, color=:white, linestyle=:dash, linewidth=0.5)
 Colorbar(fig[3, 1], hm1, vertical=false, height=8, ticklabelsize=14)
 
-Roₙ = @lift Array(ds.Ro[Ti=$n, y_afa=Near(0)])
+Roₙ = @lift ds["Ro"][:, :, $n]
 hm2 = heatmap!(ax2, x_faa, z_aaf, Roₙ; colormap = :balance, colorrange = (-10, +10))
 contour!(ax2, x_caa, z_aac, bₙ; levels=10, color=:black, linestyle=:dash, linewidth=0.5)
 Colorbar(fig[3, 2], hm2, vertical=false, height=8, ticklabelsize=14)
 
-PVₙ = @lift Array(ds.PV[Ti=$n, y_afa=Near(0)])
+PVₙ = @lift ds["PV"][:, :, $n]
 hm3 = heatmap!(ax3, x_faa, z_aaf, PVₙ; colormap = :coolwarm, colorrange = N²*f₀.*(-1.5, +1.5))
 contour!(ax3, x_caa, z_aac, bₙ; levels=10, color=:white, linestyle=:dash, linewidth=0.5)
 Colorbar(fig[3, 3], hm3, vertical=false, height=8, ticklabelsize=14);
 
 # Now we mark the time by placing a vertical line in the bottom panel and adding a helpful title
 
-times = dims(ds, :Ti)
+times = ds["time"][:]
 title = @lift "Time = " * string(prettytime(times[$n]))
 fig[1, 1:3] = Label(fig, title, fontsize=24, tellwidth=false);
 
@@ -235,6 +233,8 @@ record(fig, filename * ".mp4", 1:length(times), framerate=10) do i
        n[] = i
 end
 
+close(ds)
+
 # ![](tilted_bottom_boundary_layer.mp4)
 #
 # The animation shows negative PV being produced at the bottom due to drag, which leads to the

docs/examples/two_dimensional_turbulence.jl

@@ -61,24 +61,24 @@ simulation.callbacks[:progress] = Callback(progress, IterationInterval(100))
 # Using Oceanostics we can easily calculate two important diagnostics, the kinetic energy KE and
 # its dissipation rate ε
 
-KE = KineticEnergy(model)
-ε = KineticEnergyDissipationRate(model)
+KE = KineticEnergyEquation.KineticEnergy(model)
+ε = KineticEnergyEquation.DissipationRate(model)
 
 # And we can define their volume-integrals
 
 ∫KE = Integral(KE)
 ∫ε = Integral(ε)
 
 # We also create another integrated quantity that appears in the TKE evolution equation: the
-# `KineticEnergyStressTerm`, which in our case is
+# `KineticEnergyStress`, which in our case is
 #
 # ```math
 # \varepsilon^D = u_i \partial_j \tau_{ij}
 # ```
 # where ``\tau_{ij}`` is the diffusive flux of ``i`` momentum in the ``j``-th direction.
 #
 
-∫εᴰ = Integral(KineticEnergyStressTerm(model))
+∫εᴰ = Integral(KineticEnergyEquation.Stress(model))
 
 # The idea in calculating this term is that, in integrated form, all transport contributions in it
 # should equal zero and `∫εᴰ` should equal `∫ε`.
@@ -190,6 +190,6 @@ nothing #hide
 # Second, again as expected, the volume-integrated KE dissipation rate is the same as the
 # volume-integrated KE diffusion term (since all the non-dissipation parts of the term
 # volume-integrate to zero). In fact, both `KineticEnergyDissipationRate` and
-# `KineticEnergyStressTerm` in Oceanostics are implemented in an energy-conserving form (i.e.,
+# `KineticEnergyStress` in Oceanostics are implemented in an energy-conserving form (i.e.,
 # they use the exact same discretization scheme and interpolations as used in Oceananigans), so they
 # agree to machine-precision, and are great for closing budgets.

docs/make.jl

@@ -49,7 +49,7 @@ makedocs(sitename = "Oceanostics.jl",
          doctest = true,
          clean = true,
          format = format,
-         checkdocs = :exports
+         checkdocs = :none
          )
 #---
 

docs/src/index.md

@@ -27,13 +27,13 @@ julia> simulation = Simulation(model, Δt=1, stop_time=10);
 
 julia> simulation.callbacks[:progress] = Callback(ProgressMessengers.TimedMessenger(), IterationInterval(5));
 
-julia> ke = KineticEnergy(model)
+julia> ke = KineticEnergyEquation.KineticEnergy(model)
 KernelFunctionOperation at (Center, Center, Center)
 ├── grid: 4×5×6 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
-├── kernel_function: turbulent_kinetic_energy_ccc (generic function with 1 method)
-└── arguments: ("Field", "Field", "Field", "Int64", "Int64", "Int64")
+├── kernel_function: kinetic_energy_ccc (generic function with 1 method)
+└── arguments: ("Field", "Field", "Field")
 
-julia> ε = KineticEnergyDissipationRate(model)
+julia> ε = KineticEnergyEquation.KineticEnergyDissipationRate(model)
 KernelFunctionOperation at (Center, Center, Center)
 ├── grid: 4×5×6 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
 ├── kernel_function: viscous_dissipation_rate_ccc (generic function with 1 method)

docs/src/library.md

@@ -1,45 +1,20 @@
-# Library
-
-Documentation for the public user interface.
-
-## Oceanostics.jl
-
-```@autodocs
-Modules = [Oceanostics]
-Private = false
+```@contents
+Pages = ["library.md"]
 ```
 
-## Oceanostics.TKEBudgetTerms
-
-```@autodocs
-Modules = [Oceanostics.TKEBudgetTerms]
-Private = false
-```
-
-## Oceanostics.TracerBudgetTerms
-
-```@autodocs
-Modules = [Oceanostics.TracerBudgetTerms]
-Private = false
-```
-
-## Oceanostics.TracerVarianceBudgetTerms
-
-```@autodocs
-Modules = [Oceanostics.TracerVarianceBudgetTerms]
-Private = false
-```
-
-## Oceanostics.FlowDiagnostics
-
-```@autodocs
-Modules = [Oceanostics.FlowDiagnostics]
-Private = false
-```
+# Library
 
-## Oceanostics.PotentialEnergyEquationTerms
+Documentation for the public user interface.
 
 ```@autodocs
-Modules = [Oceanostics.PotentialEnergyEquationTerms]
-Private = false
+Modules = [Oceanostics,
+           Oceanostics.TracerEquation,
+           Oceanostics.TracerVarianceEquation,
+           Oceanostics.TurbulentKineticEnergyEquation,
+           Oceanostics.KineticEnergyEquation,
+           Oceanostics.FlowDiagnostics,
+           Oceanostics.PotentialEnergyEquation,
+           Oceanostics.ProgressMessengers]
+Order = [:function, :type, :macro]
+Filter = t -> true
 ```

src/FlowDiagnostics.jl

@@ -10,7 +10,8 @@ export BottomCellValue
 using Oceanostics: validate_location,
                    validate_dissipative_closure,
                    add_background_fields,
-                   get_coriolis_frequency_components
+                   get_coriolis_frequency_components,
+                   CustomKFO
 
 using Oceananigans: NonhydrostaticModel, FPlane, ConstantCartesianCoriolis, BuoyancyField, BuoyancyTracer, location
 using Oceananigans.BuoyancyFormulations: get_temperature_and_salinity, SeawaterBuoyancy, g_Earth, buoyancy_perturbationᶜᶜᶜ
@@ -63,6 +64,8 @@ end
     return dbdz / duₕdz^2
 end
 
+const RichardsonNumber = CustomKFO{<:typeof(richardson_number_ccf)}
+
 """
     $(SIGNATURES)
 
@@ -116,6 +119,8 @@ end
     return (ω_x*params.fx + ω_y*params.fy + ω_z*params.fz)/(params.fx^2 + params.fy^2 + params.fz^2)
 end
 
+const RossbyNumber = CustomKFO{<:typeof(rossby_number_fff)}
+
 """
     $(SIGNATURES)
 
@@ -175,6 +180,8 @@ end
     return pv_barot + pv_baroc
 end
 
+const ThermalWindPotentialVorticity = CustomKFO{<:typeof(potential_vorticity_in_thermal_wind_fff)}
+
 """
     $(SIGNATURES)
 
@@ -220,6 +227,8 @@ end
     return pv_x + pv_y + pv_z
 end
 
+const ErtelPotentialVorticity = CustomKFO{<:typeof(ertel_potential_vorticity_fff)}
+
 """
     $(SIGNATURES)
 
@@ -255,7 +264,7 @@ KernelFunctionOperation at (Face, Face, Face)
 ├── kernel_function: ertel_potential_vorticity_fff (generic function with 1 method)
 └── arguments: ("Field", "Field", "Field", "Field", "Int64", "Int64", "Float64")
 
-julia> interior(compute!(Field(EPV)))
+julia> interior(Field(EPV))
 1×1×5 view(::Array{Float64, 3}, 1:1, 1:1, 4:8) with eltype Float64:
 [:, :, 1] =
  0.0
@@ -313,6 +322,7 @@ end
     return (params.f_dir + ω_dir) * dbddir
 end
 
+const DirectionalErtelPotentialVorticity = CustomKFO{<:typeof(directional_ertel_potential_vorticity_fff)}
 
 """
     $(SIGNATURES)
@@ -358,6 +368,8 @@ function strain_rate_tensor_modulus_ccc(i, j, k, grid, u, v, w)
     return √(Sˣˣ² + Sʸʸ² + Sᶻᶻ² + 2 * (Sˣʸ² + Sˣᶻ² + Sʸᶻ²))
 end
 
+const StrainRateTensorModulus = CustomKFO{<:typeof(strain_rate_tensor_modulus_ccc)}
+
 """
     $(SIGNATURES)
 
@@ -392,6 +404,8 @@ function vorticity_tensor_modulus_ccc(i, j, k, grid, u, v, w)
     return √(Ωˣʸ² + Ωˣᶻ² + Ωʸᶻ² + Ωʸˣ² + Ωᶻˣ² + Ωᶻʸ²)
 end
 
+const VorticityTensorModulus = CustomKFO{<:typeof(vorticity_tensor_modulus_ccc)}
+
 """
     $(SIGNATURES)
 
@@ -419,6 +433,8 @@ end
     Ω² = vorticity_tensor_modulus_ccc(i, j, k, grid, u, v, w)^2
     return (Ω² - S²) / 2
 end
+
+const QVelocityGradientTensorInvariant = CustomKFO{<:typeof(Q_velocity_gradient_tensor_invariant_ccc)}
 #---
 
 #+++ Mixed layer depth
@@ -475,7 +491,7 @@ function MixedLayerDepth(grid::AbstractGrid, args...; criterion = BuoyancyAnomal
     return KernelFunctionOperation{Center, Center, Nothing}(MLD, grid, args...)
 end
 
-function (MLD::MixedLayerDepthKernel)(i, j, k, grid, args...)
+@inline function (MLD::MixedLayerDepthKernel)(i, j, k, grid, args...)
     kₘₗ = -1
 
     for k in grid.Nz-1:-1:1