Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add finite depth functionality to SurfaceQG #375

Open
wants to merge 15 commits into
base: main
Choose a base branch
from
Open

Add finite depth functionality to SurfaceQG #375

Show file tree
Hide file tree
Changes from 11 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
b494935
Update surfaceqg.jl with finite depth
mncrowe Jan 21, 2025
d850f33
fudge types in Params
mncrowe Jan 21, 2025
22e2bd4
add helper function to calculate streamfunction from b
mncrowe Jan 21, 2025
e34c677
add multiple dispatch for finite and infinite depth SQG cases
mncrowe Jan 22, 2025
b061230
fix typos with setting parameter structures
mncrowe Jan 22, 2025
ee89d90
fix typo in defining Params and delete commented out code
mncrowe Jan 22, 2025
61f965c
add streamfunctionfromb!
navidcy Jan 23, 2025
166e3a9
simplify finite-depth sqg
navidcy Jan 24, 2025
32d3876
remove unused param type
navidcy Jan 24, 2025
7de8a99
fix typo
navidcy Jan 24, 2025
6da9ee1
fix constructor + tests
navidcy Jan 24, 2025
974d139
fix docstring
navidcy Jan 24, 2025
a833b5e
fix docstring
navidcy Jan 24, 2025
1648288
Add total energy diagnostic, currently preliminary
mncrowe Jan 24, 2025
9b00fe3
Merge branch 'main' into main
navidcy Jan 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/src/lib/functions.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@ GeophysicalFlows.MultiLayerQG.addforcing!

```@docs
GeophysicalFlows.SurfaceQG.Problem
GeophysicalFlows.streamfunctionfromb!
GeophysicalFlows.SurfaceQG.buoyancy_dissipation
GeophysicalFlows.SurfaceQG.buoyancy_work
```
Expand Down
4 changes: 2 additions & 2 deletions src/multilayerqg.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -406,7 +406,7 @@ end
LinearEquation(params, grid)

Return the equation for a multi-layer quasi-geostrophic problem with `params` and `grid`.
The linear opeartor ``L`` includes only (hyper)-viscosity and is computed via
The linear operator ``L`` includes only (hyper)-viscosity and is computed via
`hyperviscosity(params, grid)`.

The nonlinear term is computed via function `calcNlinear!`.
Expand All @@ -421,7 +421,7 @@ end
Equation(params, grid)

Return the equation for a multi-layer quasi-geostrophic problem with `params` and `grid`.
The linear opeartor ``L`` includes only (hyper)-viscosity and is computed via
The linear operator ``L`` includes only (hyper)-viscosity and is computed via
`hyperviscosity(params, grid)`.

The nonlinear term is computed via function `calcN!`.
Expand Down
115 changes: 73 additions & 42 deletions src/surfaceqg.jl
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
Problem,
set_b!,
updatevars!,
streamfunctionfromb!,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the end this is not used, right? Delete?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's not used. However, it might still be useful as the streamfunction is not available in vars and often it's very useful to know. Removing it would be ok as we have the inversion relation saved, though it's perhaps more user-friendly to have a function that calculates it rather than expecting users to do the Fourier transforms themselves?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note that "exporting" is really about choosing a stable user-API (not necessarily defining all functions that users need). Its ok to require users to qualify their names or write using GeophysicalFlows: streamfunctionfromb!. Not expressing an opinion about exporting or not here, just want to make sure that the goals are clear.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, this reply was intended to answer this and the question below in one go so I didn't separate removing the function and removing the export. I'll leave the choice of exporting vs not to you but I think including it (or perhaps a non-! version) would be useful.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ah I was confused about @navidcy's suggestion

Copy link
Member

@navidcy navidcy Jan 24, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok, sorry if I created confusion.
I think I understand now the point of the method -- thanks @mncrowe

I agree with @glwagner, we can include the method (if we also test it) but we don't necessarily need to export it.


kinetic_energy,
buoyancy_variance,
Expand All @@ -28,6 +29,7 @@
Lx = 2π,
ny = nx,
Ly = Lx,
H = Inf,
ν = 0,
nν = 1,
dt = 0.01,
Expand All @@ -49,6 +51,7 @@
- `ny`: Number of grid points in ``y``-domain.
- `Lx`: Extent of the ``x``-domain.
- `Ly`: Extent of the ``y``-domain.
- `H`: Layer depth, set Inf for standard SQG.
- `ν`: Small-scale (hyper)-viscosity coefficient.
- `nν`: (Hyper)-viscosity order, `nν```≥ 1``.
- `dt`: Time-step.
Expand All @@ -64,6 +67,7 @@
ny = nx,
Lx = 2π,
Ly = Lx,
H = Inf,
# Hyper-viscosity parameters
ν = 0,
nν = 1,
Expand All @@ -78,8 +82,8 @@

grid = TwoDGrid(dev; nx, Lx, ny, Ly, aliased_fraction, T)

params = Params{T}(ν, nν, calcF)

params = Params(T(H), T(ν), nν, calcF, grid)
vars = calcF == nothingfunction ? DecayingVars(grid) : (stochastic ? StochasticForcedVars(grid) : ForcedVars(grid))

equation = Equation(params, grid)
Expand All @@ -92,24 +96,34 @@
# Parameters
# ----------

abstract type SurfaceQGParams <: AbstractParams end

"""
Params{T}(ν, nν, calcF!)
struct Params{T, Atrans}(H, ν, nν, calcF!, DirNeu) <: SurfaceQGParams

A struct containing the parameters for Surface QG dynamics. Included are:

$(TYPEDFIELDS)
"""
struct Params{T} <: AbstractParams
struct Params{T, Atrans <: AbstractArray} <: SurfaceQGParams
"layer depth"
H :: T
"buoyancy (hyper)-viscosity coefficient"
ν :: T
ν :: T
"buoyancy (hyper)-viscosity order"
nν :: Int
nν :: Int
"function that calculates the Fourier transform of the forcing, ``F̂``"
calcF! :: Function
calcF! :: Function
"array containing Dirichlet-to-Neumann operator for buoyancy-streamfunction relation"
ψhfrombh :: Atrans
end

Params(ν, nν) = Params(ν, nν, nothingfunction)
function Params(H, ν, nν, calcF!, grid::AbstractGrid)
ψhfrombh = @. sqrt(grid.invKrsq) * coth(H / sqrt(grid.invKrsq))
return Params(H, ν, nν, calcF!, ψhfrombh)
end

Params(ν, nν, grid) = Params(Inf, ν, nν, nothingfunction, grid)

# ---------
# Equations
Expand All @@ -119,7 +133,7 @@
Equation(params, grid)

Return the `equation` for surface QG dynamics with `params` and `grid`. The linear
opeartor ``L`` includes (hyper)-viscosity of order ``n_ν`` with coefficient ``ν``,
operator ``L`` includes (hyper)-viscosity of order ``n_ν`` with coefficient ``ν``,

```math
L = - ν |𝐤|^{2 n_ν} .
Expand All @@ -129,10 +143,10 @@

The nonlinear term is computed via function `calcN!()`.
"""
function Equation(params::Params, grid::AbstractGrid)
function Equation(params::SurfaceQGParams, grid::AbstractGrid)
L = @. - params.ν * grid.Krsq^params.nν
CUDA.@allowscalar L[1, 1] = 0

return FourierFlows.Equation(L, calcN!, grid)
end

Expand Down Expand Up @@ -199,7 +213,7 @@

@devzeros Dev T (grid.nx, grid.ny) b u v
@devzeros Dev Complex{T} (grid.nkr, grid.nl) bh uh vh Fh

return Vars(b, u, v, bh, uh, vh, Fh, nothing)
end

Expand All @@ -214,7 +228,7 @@

@devzeros Dev T (grid.nx, grid.ny) b u v
@devzeros Dev Complex{T} (grid.nkr, grid.nl) bh uh vh Fh prevsol

return Vars(b, u, v, bh, uh, vh, Fh, prevsol)
end

Expand All @@ -235,24 +249,24 @@
"""
function calcN_advection!(N, sol, t, clock, vars, params, grid)
@. vars.bh = sol
@. vars.uh = im * grid.l * sqrt(grid.invKrsq) * sol
@. vars.vh = - im * grid.kr * sqrt(grid.invKrsq) * sol
@. vars.uh = im * grid.l * params.ψhfrombh * sol
@. vars.vh = - im * grid.kr * params.ψhfrombh * sol

ldiv!(vars.u, grid.rfftplan, vars.uh)
ldiv!(vars.v, grid.rfftplan, vars.vh)
ldiv!(vars.b, grid.rfftplan, vars.bh)

ub, ubh = vars.u, vars.uh # use vars.u, vars.uh as scratch variables
vb, vbh = vars.v, vars.vh # use vars.v, vars.vh as scratch variables

@. ub *= vars.b # u*b
@. vb *= vars.b # v*b

mul!(ubh, grid.rfftplan, ub) # \hat{u*b}
mul!(vbh, grid.rfftplan, vb) # \hat{v*b}

@. N = - im * grid.kr * ubh - im * grid.l * vbh

return nothing
end

Expand All @@ -267,11 +281,11 @@
"""
function calcN!(N, sol, t, clock, vars, params, grid)
dealias!(sol, grid)

calcN_advection!(N, sol, t, clock, vars, params, grid)

addforcing!(N, sol, t, clock, vars, params, grid)

return nothing
end

Expand All @@ -286,7 +300,7 @@
function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid)
params.calcF!(vars.Fh, sol, t, clock, vars, params, grid)
@. N += vars.Fh

return nothing
end

Expand All @@ -295,9 +309,9 @@
@. vars.prevsol = sol # sol at previous time-step is needed to compute budgets for stochastic forcing
params.calcF!(vars.Fh, sol, t, clock, vars, params, grid)
end

@. N += vars.Fh

return nothing
end

Expand All @@ -307,26 +321,28 @@
# ----------------

"""
updatevars!(sol, vars, params, grid)
updatevars!(prob)

Update variables in `vars` with solution in `sol`.
"""
function updatevars!(prob)
vars, grid, sol = prob.vars, prob.grid, prob.sol

function updatevars!(sol, vars, params, grid)
dealias!(sol, grid)

@. vars.bh = sol
@. vars.uh = im * grid.l * sqrt(grid.invKrsq) * sol
@. vars.vh = - im * grid.kr * sqrt(grid.invKrsq) * sol
@. vars.uh = im * grid.l * params.ψhfrombh * sol
@. vars.vh = - im * grid.kr * params.ψhfrombh * sol

ldiv!(vars.b, grid.rfftplan, deepcopy(vars.bh))
ldiv!(vars.u, grid.rfftplan, deepcopy(vars.uh))
ldiv!(vars.v, grid.rfftplan, deepcopy(vars.vh))

return nothing
end

updatevars!(prob) = updatevars!(prob.sol, prob.vars, prob.params, prob.grid)


"""
set_b!(prob, b)

Expand All @@ -335,33 +351,48 @@
function set_b!(prob, b)
mul!(prob.sol, prob.grid.rfftplan, b)
CUDA.@allowscalar prob.sol[1, 1] = 0 # zero domain average

updatevars!(prob)


return nothing
end

"""
streamfunctionfromb!(ψ, bh, params, grid)
streamfunctionfromb!(ψ, prob)

Compute the streamfunction `ψ` from the buoyancy `bh` in Fourier space.
(Note that `bh = prob.sol`.)
"""
function streamfunctionfromb!(ψ, bh, params, grid)
ldiv!(ψ, grid.rfftplan, @. params.ψhfrombh * bh)

Check warning on line 368 in src/surfaceqg.jl

View check run for this annotation

Codecov / codecov/patch

src/surfaceqg.jl#L367-L368 

Added lines #L367 - L368 were not covered by tests
return nothing
end
Copy link
Member

@navidcy navidcy Jan 24, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Delete since it's not used?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same as above

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since the method anyway allocates memory, wouldn't it be simpler if it was something like:

function get_streamfunction(prob)
  ψh = @. prob.params.ψhfrombh * prob.sol
  return irfft(ψh, prob.grid.nx)
end

Copy link
Contributor Author

@mncrowe mncrowe Jan 24, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, looks good, this is what I was getting at with 'non-! version' above.

Also, is there a reason (other than the !) why rfft and irfft are used in some places and mul! and ldiv! in others?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We use ldiv! and mul! to avoid allocations when already the arrays we’ll put the values in are pre-existing.


streamfunctionfromb!(ψ, prob) = streamfunctionfromb!(ψ, prob.sol, prob.params, prob.grid)

Check warning on line 372 in src/surfaceqg.jl

View check run for this annotation

Codecov / codecov/patch

src/surfaceqg.jl#L372 

Added line #L372 was not covered by tests

"""
kinetic_energy(prob)
kinetic_energy(sol, vars, params, grid)

Return the domain-averaged surface kinetic energy. Since ``u² + v² = |{\\bf ∇} ψ|²``, we get
```math
\\int \\frac1{2} |{\\bf ∇} ψ|² \\frac{𝖽x 𝖽y}{L_x L_y} = \\sum_{𝐤} \\frac1{2} |𝐤|² |ψ̂|² .
```
In SQG, this is identical to half the domain-averaged surface buoyancy variance.
In SQG with infinite depth, this is identical to half the domain-averaged surface buoyancy variance.
"""
@inline function kinetic_energy(prob)
sol, vars, grid = prob.sol, prob.vars, prob.grid

@inline function kinetic_energy(sol, vars, params, grid)
ψh = vars.uh # use vars.uh as scratch variable
kinetic_energyh = vars.bh # use vars.bh as scratch variable
@. ψh = sqrt(grid.invKrsq) * sol

@. ψh = params.ψhfrombh * sol
@. kinetic_energyh = 1 / 2 * grid.Krsq * abs2(ψh)

return 1 / (grid.Lx * grid.Ly) * parsevalsum(kinetic_energyh, grid)
end

@inline kinetic_energy(prob) = kinetic_energy(prob.sol, prob.vars, prob.params, prob.grid)

"""
buoyancy_variance(prob)

Expand All @@ -370,11 +401,11 @@
\\int b² \\frac{𝖽x 𝖽y}{L_x L_y} = \\sum_{𝐤} |b̂|² .
```
In SQG, this is identical to the velocity variance (i.e., twice the domain-averaged kinetic
energy).
energy in the infinite-depth case).
"""
@inline function buoyancy_variance(prob)
sol, grid = prob.sol, prob.grid

return 1 / (grid.Lx * grid.Ly) * parsevalsum(abs2.(sol), grid)
end

Expand Down
Loading
Loading