Merge pull request #13 from mncrowe/main

merge energy functions, tests and examples into joss_draft branch

docs/src/Functions.md

@@ -52,6 +52,15 @@ Base.summary
 Base.show
 ```
 
+## `energetics.jl`
+
+```@docs
+EnergyLQG
+EnstrophyLQG
+EnergySQG
+AreaInteg2
+```
+
 ## `monopoles.jl`
 
 ```@docs

docs/src/index.md

@@ -244,13 +244,14 @@ Full descriptions can be found on the [List of Functions](https://mncrowe.github
 | `Calc_∇` | `src/create_modon.jl` | Function | Calculates the x and y derivatives of a given field |
 | `CartesianGrid` | `src/create_modon.jl` | Function | Outputs 2D (x, y) Arrays from `grid` |
 | `PolarGrid` | `src/create_modon.jl` | Function | Outputs 2D (r, θ) Arrays created from `CartesianGrid` outputs |
+| `EnergyLQG` | `src/energetics.jl` | Function | Calculates the kinetic and potential energy for the LQG model |
+| `EnstrophyLQG` | `src/energetics.jl` | Function | Calculates the enstrophy for the LQG model |
+| `EnergySQG` | `src/energetics.jl` | Function | Calculates domain-integrated energy and surface potential energy for the SQG model |
+| `AreaInteg2` | `src/energetics.jl` | Function | Performs area integration in Fourier space |
 | `CreateRankine` | `src/monopoles.jl` | Function | High level wrapper function for calculating ``\psi``, ``q``, ``u`` and ``v`` for the Rankine vortex using given parameters |
 | `CreateMonopole` | `src/monopoles.jl` | Function | High level wrapper function for calculating ``\psi``, ``q``, ``u`` and ``v`` for a monopolar QG vortex using given parameters |
 | `InvertVorticity1LQG` | `src/monopoles.jl` | Function | Calculates ``\psi`` from ``q`` for the 1-layer QG model on the f-plane |
 
-
-
-
 [^1]: [Johnson, E. R., and M. N. Crowe, 2023, Oceanic dipoles in a surface quasigeostrophic model, J. Fluid Mech., 958, R2](https://doi.org/10.1017/jfm.2023.87).
 [^2]: [Crowe, M. N., and E. R. Johnson, 2023, The evolution of surface quasi-geostrophic modons on sloping topography, J. Fluid. Mech., 970, A10](https://doi.org/10.1017/jfm.2023.607).
 [^3]: [Crowe, M. N., and E. R. Johnson, 2024, Modon solutions in an N-layer quasi-geostrophic model, J. Fluid. Mech., 994, R1](https://doi.org/10.1017/jfm.2024.619).

examples/Energy_LQG.jl

@@ -0,0 +1,33 @@
+# Example script for calculating the energy and enstrophy for 2-Layer QG modon
+
+using QGDipoles
+
+# Set problem parameters
+
+U, ℓ = 1, 1	# vortex speed and radius
+R = [1, 1]	# Rossby radius in each layer
+β = [0, 1]	# background PV gradient in each layer
+
+M = 8		# number of coefficients in Zernike expansion
+tol = 1e-8	# maximum error in solution evaluation
+cuda = false	# use CuArrays for grid
+
+# Set grid parameters
+
+Nx, Ny = 512, 512
+Lx, Ly = 10, 10
+
+# Create modon solution and grid
+
+grid = CreateGrid(Nx, Ny, Lx, Ly; cuda)
+ψ, q, K, a = CreateModonLQG(grid, M, U, ℓ, R, β; tol)
+
+# Calculate kinetic and potential energy, we've used layer depth 
+# H = R (assuming g'/f^2 = 1 in nondimensional units)
+
+KE, PE = EnergyLQG(grid, ψ, R, R.^2)
+
+# Calculate the enstrophy
+
+Q = EnstrophyLQG(grid, q, R.^2)
+

examples/Energy_SQG.jl

@@ -0,0 +1,27 @@
+# Example script for calculating the energy of an SQG modon
+
+using QGDipoles
+
+# Set problem parameters
+
+U, ℓ = 1, 1	# vortex speed and radius
+R = [Inf, Inf]	# Baroclinic and Barotropic Rossby radii
+β = 0		# background PV gradient in the interior
+
+M = 20		# number of coefficients in Zernike expansion
+tol = 1e-6	# maximum error in solution evaluation
+cuda = false	# use CuArrays for grid
+
+# Set grid parameters
+
+Nx, Ny = 512, 512
+Lx, Ly = 10, 10
+
+# create modon solution and grid
+
+grid = CreateGrid(Nx, Ny, Lx, Ly; cuda)
+ψ, b, _, _ = CreateModonSQG(grid, M, U, ℓ, R, β; tol)
+
+# calculate the domain integrated energy and surface potential energy
+
+E, SPE = EnergySQG(grid, ψ, b, R[2])

examples/readme.md

@@ -8,6 +8,10 @@ The available example scripts are:
 
 * 2_Layer_Modon.jl: Constructs a dipolar vortex solution in a 2-layer QG model.
 
+* Energy_LQG.jl: Constructs a 2-layer dipolar vortex solution and calculates the kinetic energy, potential energy and enstrophy
+
+* Energy_SQG.jl: Constructs an SQG dipolar vortex solution and calculates the domain integrated energy and the surface potential energy
+
 * 3_Layer_Modon.jl: Constructs a dipolar vortex solution in a 3-layer QG model (corresponds to Example 2 [here](https://mncrowe.github.io/QGDipoles.jl/dev/Examples/)).
 
 * GeophysicalFlows_Example.jl: Demostrates how to used `QGDipoles.jl` to generate initial conditions for `GeophysicalFlows.jl` (corresponds to Example 5 [here](https://mncrowe.github.io/QGDipoles.jl/dev/Examples/)).

src/QGDipoles.jl

@@ -52,6 +52,12 @@ export
 	Calc_∇,
 	CartesianGrid,
 	PolarGrid,
+
+	# energetics.jl
+	EnergyLQG,
+	EnstrophyLQG,
+	EnergySQG,
+	AreaInteg2,
 
 	# monopoles.jl
 	CreateRankine,
@@ -63,6 +69,7 @@ export
 include("JJ_integ.jl")
 include("lin_sys.jl")
 include("create_modon.jl")
+include("energetics.jl")
 include("monopoles.jl")
 
 end

src/energetics.jl

@@ -0,0 +1,175 @@
+"""
+This file contains functions to calculate the energy and enstrophy for dipolar vortex
+solutions.
+
+For the (multi-layer) LQG model we have:
+
+``KE = \\frac{H_i}{2H} \\int_A |\\nabla\\psi_i|^2 dx dy,``
+
+``PE = \\frac{H_i}{2H R_i^2} \\int_A |\\psi_i - \\psi_{i+1}|^2 dx dy,``
+
+For 1 layer (equivalent barotropic) QG, the PE is given by
+
+``PE = \\frac{1}{2 R^2} \\int_A |\\psi|^2 dx dy,``
+
+and the KE is the same as the multi-layer case.
+
+For SQG, the two quantities of interest are the total domain averaged energy
+
+``E = -\\frac{1}{2} \\int_A \\psi b dx dy,``
+
+and the surface potential energy
+
+``SPE = \\frac{1}{2} \\int_A |b + \\psi / R^\\prime|^2 dx dy.``
+
+"""
+
+
+"""
+Function: `EnergyLQG(grid, ψ, R, H=[1])`
+
+Calculates the kinetic and potential energy for the LQG system
+
+Arguments:
+ - `grid`: grid structure containing Krsq
+ - `ψ`: streamfunction in each layer, Array or CuArray
+ - `R`: Rossby radius in each layer, Number or Vector
+ - `H`: Thickness of each layer, Number or Vector
+"""
+function EnergyLQG(grid, ψ::Union{CuArray,Array}, R::Union{Number,Vector}, H::Union{Number,Vector}=[1])
+
+	Nx, Ny = size(ψ)[1:2]
+	N = Int(length(ψ) / (Nx * Ny))
+
+	if N == 1
+
+		ψh = rfft(ψ, [1, 2])
+		eh = sqrt.(grid.Krsq) .* ψh
+
+		KE = [AreaInteg2(eh, grid) / 2]
+		PE = [AreaInteg2(ψh, grid) ./ R .^ 2 / 2]
+
+	else
+
+		D = sum(H)
+		KE, PE = zeros(N), zeros(N-1)
+
+		ψh = rfft(ψ, [1, 2])
+		eh = sqrt.(grid.Krsq) .* ψh
+
+		for i in 1:N
+			KE[i] = H[i] / (2 * D) * AreaInteg2(eh[:, :, i], grid)
+		end
+
+		for i = 1:N-1
+			PE[i] = H[i] / (2 * D * R[i]^2) * AreaInteg2(ψh[:, :, i] - ψh[:, :, i+1], grid)
+		end		
+
+	end
+
+	return KE, PE
+
+end
+
+"""
+Function: `EnstrophyLQG(grid, q, H=[1])`
+
+Calculates the enstrophy for the LQG system
+
+Arguments:
+ - `grid`: grid structure containing Krsq
+ - `q`: potential vorticity anomaly in each layer, Array or CuArray
+ - `H`: Thickness of each layer, Number or Vector
+"""
+function EnstrophyLQG(grid, q::Union{CuArray,Array}, H::Union{Number,Vector}=[1])
+
+	Nx, Ny = size(q)[1:2]
+	N = Int(length(q) / (Nx * Ny))
+
+	D = sum(H)
+
+	EN = zeros(N)
+
+	for i = 1:N
+		EN[i] = H[i] / (2 * D) * AreaInteg2(q[:, :, i], grid)
+	end
+
+	return EN
+
+end
+
+"""
+Function: `EnergySQG(grid, ψ, b, R′)`
+
+Calculates the energies for the SQG system; the total domain integrated energy
+and the surface potential energy
+
+Arguments:
+ - `grid`: grid structure containing Krsq
+ - `ψ`: surface streamfunction, Array or CuArray
+ - `b`: surface buoyancy, , Array or CuArray
+ - `R′`: reduced barotropic Rossby radius, Number (default: `Inf`)
+
+Note: the surface potential energy is sometimes referred to as the generalised
+enstrophy or the buoyancy variance.
+"""
+function EnergySQG(grid, ψ::Union{CuArray,Array}, b::Union{CuArray,Array}, R′::Number=Inf)
+
+	ψh = rfft(ψ)
+	bh = rfft(b)
+
+	eh = sqrt.(ψh .* bh)
+	sh = bh + ψh / R′
+
+	E   = [AreaInteg2(eh, grid) / 2]
+	SPE = [AreaInteg2(sh, grid) / 2]
+
+	return E, SPE
+
+end
+
+"""
+Function: `AreaInteg2(f, grid)`
+
+Calculates the integral ``I = \\int_A f^2 \\mathrm{d}\\A`` where ``A``
+is the 2D domain described by `grid`.
+
+Arguments:
+ - `f`: input Array in real or Fourier space
+ - `grid`: grid structure
+
+Note: f can be entered in real space or Fourier space, we use the rfft function
+to calculate the Fourier transform so array sizes can distinguish the two.
+"""
+function AreaInteg2(f::Union{CuArray,Array}, grid::GridStruct, exponent::Int=1)
+
+	# Get grid parameters
+
+	Nkr = length(grid.kr)
+	Nx, Ny = length(grid.x), length(grid.y)
+	Δx, Δy = grid.x[2] - grid.x[1], grid.y[2] - grid.y[1]
+
+	# If input array is in real space, apply Fourier transform
+
+	if size(f)[1:2] == (Nx, Ny)
+
+		fh = rfft(f, [1, 2])
+
+	else
+
+		fh = f
+
+	end
+
+	# Add up components in Fourier space, non-edge elements are counted twice for an rfft
+
+	I =   sum(abs2, fh[1, :]) +             # kr = 0 (edge)
+	      sum(abs2, fh[Nkr , :]) +         	# kr = Nx/2 (edge)
+	    2*sum(abs2, fh[2:Nkr-1, :])  	# sum twice for non-edge modes (as using rfft)
+
+  	I = I * (Δx * Δy) / (Nx * Ny) 		# normalization factor for fft
+
+  return I
+
+end
+

src/vortex_types.jl

@@ -0,0 +1,35 @@
+"""
+This file contains functions which build vortices as a `VortexSolution` type.
+
+
+Plan:
+
+Single type, `VortexSolution` which can be SQG or LQG (maybe Monopole?)
+
+show and summary functions to display properties
+
+Maybe types for SQG and LQG parameter sets?
+
+
+
+"""
+
+
+"""
+Function: `...`
+
+...
+
+Arguments:
+ - `n`: order, Integer
+ - `x`: evaluation point, Number or Array
+
+Note: ...
+"""
+function ...
+
+
+
+	return ...
+end
+