From 641628632ea1c510112ee9ac7ccc6e5edef28770 Mon Sep 17 00:00:00 2001 From: Matt Lobo Date: Fri, 9 Jun 2023 10:16:07 -0400 Subject: [PATCH 01/16] changed reduced gravity to have const. denominator --- src/multilayerqg.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/multilayerqg.jl b/src/multilayerqg.jl index 31924cd3..0ab95d83 100644 --- a/src/multilayerqg.jl +++ b/src/multilayerqg.jl @@ -347,7 +347,7 @@ function Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ ρ = reshape(T.(ρ), (1, 1, nlayers)) H = reshape(T.(H), (1, 1, nlayers)) - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface + g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[1] # reduced gravity at each interface Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) From 04ee8fe4067365572794217f0770baa2b2726d85 Mon Sep 17 00:00:00 2001 From: Matt Lobo Date: Fri, 9 Jun 2023 10:42:16 -0400 Subject: [PATCH 02/16] changed sign of Fp in multi layer Qy calc, line 367 --- src/multilayerqg.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/multilayerqg.jl b/src/multilayerqg.jl index 0ab95d83..e50206ad 100644 --- a/src/multilayerqg.jl +++ b/src/multilayerqg.jl @@ -364,7 +364,7 @@ function Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ CUDA.@allowscalar @views Qy[:, :, 1] = @. Qy[:, :, 1] - Fp[1] * (U[:, :, 2] - U[:, :, 1]) for j = 2:nlayers-1 - CUDA.@allowscalar @views Qy[:, :, j] = @. Qy[:, :, j] - Fp[j] * (U[:, :, j+1] - U[:, :, j]) + Fm[j-1] * (U[:, :, j-1] - U[:, :, j]) + CUDA.@allowscalar @views Qy[:, :, j] = @. Qy[:, :, j] + Fp[j] * (U[:, :, j+1] - U[:, :, j]) + Fm[j-1] * (U[:, :, j-1] - U[:, :, j]) end CUDA.@allowscalar @views Qy[:, :, nlayers] = @. Qy[:, :, nlayers] - Fm[nlayers-1] * (U[:, :, nlayers-1] - U[:, :, nlayers]) From 427cdff2e8a0bb92412c8ecfb6bd9b429a39fcc6 Mon Sep 17 00:00:00 2001 From: Matt Lobo Date: Fri, 9 Jun 2023 10:56:54 -0400 Subject: [PATCH 03/16] reverting to old definition of reduced grav to test bug fix with sign in Qy def --- src/multilayerqg.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/multilayerqg.jl b/src/multilayerqg.jl index e50206ad..bd49c747 100644 --- a/src/multilayerqg.jl +++ b/src/multilayerqg.jl @@ -347,7 +347,7 @@ function Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ ρ = reshape(T.(ρ), (1, 1, nlayers)) H = reshape(T.(H), (1, 1, nlayers)) - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[1] # reduced gravity at each interface + g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) From 3ce77d3d5628553fcdea68b4c3f8e224d56cfbe9 Mon Sep 17 00:00:00 2001 From: Matt Lobo Date: Fri, 9 Jun 2023 11:05:43 -0400 Subject: [PATCH 04/16] alternate (but still incorrect) definition of reduced gravity (shifted denominator by 1 index) to verify that you must have that denominator as constant for zero interior PV --- src/multilayerqg.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/multilayerqg.jl b/src/multilayerqg.jl index bd49c747..f21597b5 100644 --- a/src/multilayerqg.jl +++ b/src/multilayerqg.jl @@ -347,7 +347,7 @@ function Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ ρ = reshape(T.(ρ), (1, 1, nlayers)) H = reshape(T.(H), (1, 1, nlayers)) - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface + g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[1:nlayers-1] # reduced gravity at each interface Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) From b9098cec7965525d5d73e259d3d5f7448aec4275 Mon Sep 17 00:00:00 2001 From: Matt Lobo Date: Fri, 9 Jun 2023 11:29:34 -0400 Subject: [PATCH 05/16] changed sign error, but leaving reduced grav; just to create pull request from this fork to main repo --- src/multilayerqg.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/multilayerqg.jl b/src/multilayerqg.jl index f21597b5..e50206ad 100644 --- a/src/multilayerqg.jl +++ b/src/multilayerqg.jl @@ -347,7 +347,7 @@ function Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ ρ = reshape(T.(ρ), (1, 1, nlayers)) H = reshape(T.(H), (1, 1, nlayers)) - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[1:nlayers-1] # reduced gravity at each interface + g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[1] # reduced gravity at each interface Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) From a7509bd54beb489ddb1fe587b44f6baa0834dfca Mon Sep 17 00:00:00 2001 From: Matt Lobo Date: Fri, 9 Jun 2023 11:37:26 -0400 Subject: [PATCH 06/16] changing things so that I only have the sign error fixed, to create pull request to main GFjl repo --- src/multilayerqg.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/multilayerqg.jl b/src/multilayerqg.jl index f21597b5..bd49c747 100644 --- a/src/multilayerqg.jl +++ b/src/multilayerqg.jl @@ -347,7 +347,7 @@ function Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ ρ = reshape(T.(ρ), (1, 1, nlayers)) H = reshape(T.(H), (1, 1, nlayers)) - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[1:nlayers-1] # reduced gravity at each interface + g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) From 96add480d8a202c2716258b332450dc62e9ef7bd Mon Sep 17 00:00:00 2001 From: Matt Lobo Date: Fri, 9 Jun 2023 11:41:48 -0400 Subject: [PATCH 07/16] changing reduced grav definition to match current one in main repo --- src/multilayerqg.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/multilayerqg.jl b/src/multilayerqg.jl index e50206ad..bd49c747 100644 --- a/src/multilayerqg.jl +++ b/src/multilayerqg.jl @@ -347,7 +347,7 @@ function Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ ρ = reshape(T.(ρ), (1, 1, nlayers)) H = reshape(T.(H), (1, 1, nlayers)) - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[1] # reduced gravity at each interface + g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) From cd6f29ef6e5046f9136bbad82a23645a41456a22 Mon Sep 17 00:00:00 2001 From: Matt Lobo Date: Fri, 9 Jun 2023 12:02:12 -0400 Subject: [PATCH 08/16] changing reduced gravity definition...hopefully last time --- src/multilayerqg.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/multilayerqg.jl b/src/multilayerqg.jl index bd49c747..e50206ad 100644 --- a/src/multilayerqg.jl +++ b/src/multilayerqg.jl @@ -347,7 +347,7 @@ function Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ ρ = reshape(T.(ρ), (1, 1, nlayers)) H = reshape(T.(H), (1, 1, nlayers)) - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface + g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[1] # reduced gravity at each interface Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) From 25a1d735705f1cf72add550c1aaac6c2bff3d25b Mon Sep 17 00:00:00 2001 From: Matt Lobo Date: Fri, 9 Jun 2023 12:16:08 -0400 Subject: [PATCH 09/16] trying to fix reduced gravity merge conflict :') --- src/multilayerqg.jl | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/multilayerqg.jl b/src/multilayerqg.jl index e780d0b8..e50206ad 100644 --- a/src/multilayerqg.jl +++ b/src/multilayerqg.jl @@ -347,11 +347,7 @@ function Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ ρ = reshape(T.(ρ), (1, 1, nlayers)) H = reshape(T.(H), (1, 1, nlayers)) -<<<<<<< HEAD g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[1] # reduced gravity at each interface -======= - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface ->>>>>>> matts_dev Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) From 0ec43c680502666f20581008ab2b11366059c73b Mon Sep 17 00:00:00 2001 From: Matt Lobo Date: Fri, 9 Jun 2023 12:18:57 -0400 Subject: [PATCH 10/16] remove backup files --- src/multilayerqg_BACKUP_444028.jl | 1111 ----------------------------- src/multilayerqg_BASE_444028.jl | 1107 ---------------------------- src/multilayerqg_LOCAL_444028.jl | 1107 ---------------------------- src/multilayerqg_REMOTE_444028.jl | 1107 ---------------------------- 4 files changed, 4432 deletions(-) delete mode 100644 src/multilayerqg_BACKUP_444028.jl delete mode 100644 src/multilayerqg_BASE_444028.jl delete mode 100644 src/multilayerqg_LOCAL_444028.jl delete mode 100644 src/multilayerqg_REMOTE_444028.jl diff --git a/src/multilayerqg_BACKUP_444028.jl b/src/multilayerqg_BACKUP_444028.jl deleted file mode 100644 index e780d0b8..00000000 --- a/src/multilayerqg_BACKUP_444028.jl +++ /dev/null @@ -1,1111 +0,0 @@ -module MultiLayerQG - -export - fwdtransform!, - invtransform!, - streamfunctionfrompv!, - pvfromstreamfunction!, - updatevars!, - - set_q!, - set_ψ!, - energies, - fluxes - -using - FFTW, - CUDA, - LinearAlgebra, - StaticArrays, - Reexport, - DocStringExtensions - -@reexport using FourierFlows - -using FourierFlows: parsevalsum, parsevalsum2, superzeros, plan_flows_rfft - -nothingfunction(args...) = nothing - -""" - Problem(nlayers :: Int, - dev = CPU(); - nx = 128, - ny = nx, - Lx = 2π, - Ly = Lx, - f₀ = 1.0, - β = 0.0, - g = 1.0, - U = zeros(nlayers), - H = 1/nlayers * ones(nlayers), - ρ = Array{Float64}(1:nlayers), - eta = nothing, - topographic_pv_gradient = (0, 0), - μ = 0, - ν = 0, - nν = 1, - dt = 0.01, - stepper = "RK4", - calcFq = nothingfunction, - stochastic = false, - linear = false, - aliased_fraction = 1/3, - T = Float64) - -Construct a multi-layer quasi-geostrophic problem with `nlayers` fluid layers on device `dev`. - -Arguments -========= -- `nlayers`: (required) Number of fluid layers. -- `dev`: (required) `CPU()` (default) or `GPU()`; computer architecture used to time-step `problem`. - -Keyword arguments -================= - - `nx`: Number of grid points in ``x``-domain. - - `ny`: Number of grid points in ``y``-domain. - - `Lx`: Extent of the ``x``-domain. - - `Ly`: Extent of the ``y``-domain. - - `f₀`: Constant planetary vorticity. - - `β`: Planetary vorticity ``y``-gradient. - - `g`: Gravitational acceleration constant. - - `U`: The imposed constant zonal flow ``U(y)`` in each fluid layer. - - `H`: Rest height of each fluid layer. - - `ρ`: Density of each fluid layer. - - `eta`: Periodic component of the topographic potential vorticity. - - `topographic_pv_gradient`: The ``(x, y)`` components of the topographic PV large-scale gradient. - - `μ`: Linear bottom drag coefficient. - - `ν`: Small-scale (hyper)-viscosity coefficient. - - `nν`: (Hyper)-viscosity order, `nν```≥ 1``. - - `dt`: Time-step. - - `stepper`: Time-stepping method. - - `calcF`: Function that calculates the Fourier transform of the forcing, ``F̂``. - - `stochastic`: `true` or `false` (default); boolean denoting whether `calcF` is temporally stochastic. - - `linear`: `true` or `false` (default); boolean denoting whether the linearized equations of motions are used. - - `aliased_fraction`: the fraction of high-wavenumbers that are zero-ed out by `dealias!()`. - - `T`: `Float32` or `Float64` (default); floating point type used for `problem` data. -""" -function Problem(nlayers::Int, # number of fluid layers - dev = CPU(); - # Numerical parameters - nx = 128, - ny = nx, - Lx = 2π, - Ly = Lx, - # Physical parameters - f₀ = 1.0, # Coriolis parameter - β = 0.0, # y-gradient of Coriolis parameter - g = 1.0, # gravitational constant - U = zeros(nlayers), # imposed zonal flow U(y) in each layer - H = 1/nlayers * ones(nlayers), # rest fluid height of each layer - ρ = Array{Float64}(1:nlayers), # density of each layer - eta = nothing, # periodic component of the topographic PV - topographic_pv_gradient = (0, 0), # tuple with the ``(x, y)`` components of topographic PV large-scale gradient - # Bottom Drag and/or (hyper)-viscosity - μ = 0, - ν = 0, - nν = 1, - # Timestepper and equation options - dt = 0.01, - stepper = "RK4", - calcFq = nothingfunction, - stochastic = false, - linear = false, - # Float type and dealiasing - aliased_fraction = 1/3, - T = Float64) - - if dev == GPU() && nlayers > 2 - @warn """MultiLayerQG module is not optimized on the GPU yet for configurations with - 3 fluid layers or more! - - See issues on Github at https://github.com/FourierFlows/GeophysicalFlows.jl/issues/112 - and https://github.com/FourierFlows/GeophysicalFlows.jl/issues/267. - - To use MultiLayerQG with 3 fluid layers or more we suggest, for now, to restrict running - on CPU.""" - end - - if nlayers == 1 - @warn """MultiLayerQG module does work for single-layer configuration but may not be as - optimized. We suggest using SingleLayerQG module for single-layer QG simulation unless - you have reasons to use MultiLayerQG in a single-layer configuration, e.g., you want to - compare solutions with varying number of fluid layers.""" - end - - # topographic PV - eta === nothing && (eta = zeros(dev, T, (nx, ny))) - - grid = TwoDGrid(dev; nx, Lx, ny, Ly, aliased_fraction, T) - - params = Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ, ν, nν, grid; calcFq) - - vars = calcFq == nothingfunction ? DecayingVars(grid, params) : (stochastic ? StochasticForcedVars(grid, params) : ForcedVars(grid, params)) - - equation = linear ? LinearEquation(params, grid) : Equation(params, grid) - - FourierFlows.Problem(equation, stepper, dt, grid, vars, params) -end - -""" - struct Params{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams - -The parameters for the `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct Params{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams - # prescribed params - "number of fluid layers" - nlayers :: Int - "gravitational constant" - g :: T - "constant planetary vorticity" - f₀ :: T - "planetary vorticity ``y``-gradient" - β :: T - "array with density of each fluid layer" - ρ :: Aphys3D - "array with rest height of each fluid layer" - H :: Aphys3D - "array with imposed constant zonal flow ``U(y)`` in each fluid layer" - U :: Aphys3D - "array containing the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear bottom drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "array with the reduced gravity constants for each fluid interface" - g′ :: Aphys1D - "array containing ``x``-gradient of PV due to topographic PV in each fluid layer" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV in each fluid layer" - Qy :: Aphys3D - "array containing coeffients for getting PV from streamfunction" - S :: Atrans4D - "array containing coeffients for inverting PV to streamfunction" - S⁻¹ :: Atrans4D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -""" - struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - -The parameters for the a single-layer `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - # prescribed params - "planetary vorticity ``y``-gradient" - β :: T - "array with imposed constant zonal flow ``U(y)``" - U :: Aphys3D - "array containing the periodic component of the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "array containing ``x``-gradient of PV due to topographic PV" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV" - Qy :: Aphys3D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -""" - struct TwoLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - -The parameters for the a two-layer `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct TwoLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - # prescribed params - "gravitational constant" - g :: T - "constant planetary vorticity" - f₀ :: T - "planetary vorticity ``y``-gradient" - β :: T - "array with density of each fluid layer" - ρ :: Aphys3D - "tuple with rest height of each fluid layer" - H :: Tuple - "array with imposed constant zonal flow ``U(y)`` in each fluid layer" - U :: Aphys3D - "array containing periodic component of the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear bottom drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "the reduced gravity constants for the fluid interface" - g′ :: T - "array containing ``x``-gradient of PV due to topographic PV in each fluid layer" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV in each fluid layer" - Qy :: Aphys3D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 1}) where TU - T = eltype(grid) - if length(U) == nlayers - U_2D = zeros(dev, T, (1, nlayers)) - U_2D[:] = U - U_2D = repeat(U_2D, outer=(grid.ny, 1)) - else - U_2D = zeros(dev, T, (grid.ny, 1)) - U_2D[:] = U - end - U_3D = zeros(dev, T, (1, grid.ny, nlayers)) - @views U_3D[1, :, :] = U_2D - return U_3D -end - -function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 2}) where TU - T = eltype(grid) - U_3D = zeros(dev, T, (1, grid.ny, nlayers)) - @views U_3D[1, :, :] = U - return U_3D -end - -function convert_U_to_U3D(dev, nlayers, grid, U::Number) - T = eltype(grid) - A = device_array(dev) - U_3D = reshape(repeat([T(U)], outer=(grid.ny, 1)), (1, grid.ny, nlayers)) - return A(U_3D) -end - -function Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ, ν, nν, grid; calcFq=nothingfunction, effort=FFTW.MEASURE) - dev = grid.device - T = eltype(grid) - A = device_array(dev) - - ny, nx = grid.ny , grid.nx - nkr, nl = grid.nkr, grid.nl - kr, l = grid.kr , grid.l - - U = convert_U_to_U3D(dev, nlayers, grid, U) - - Uyy = real.(ifft(-l.^2 .* fft(U))) - Uyy = CUDA.@allowscalar repeat(Uyy, outer=(nx, 1, 1)) - - # Calculate periodic components of the topographic PV gradients. - etah = rfft(A(eta)) - etax = irfft(im * kr .* etah, nx) # ∂η/∂x - etay = irfft(im * l .* etah, nx) # ∂η/∂y - - # Add topographic PV large-scale gradient - topographic_pv_gradient = T.(topographic_pv_gradient) - @. etax += topographic_pv_gradient[1] - @. etay += topographic_pv_gradient[2] - - Qx = zeros(dev, T, (nx, ny, nlayers)) - @views @. Qx[:, :, nlayers] += etax - - Qy = zeros(dev, T, (nx, ny, nlayers)) - Qy = T(β) .- Uyy # T(β) is needed to ensure that Qy remains same type as U - @views @. Qy[:, :, nlayers] += etay - - rfftplanlayered = plan_flows_rfft(A{T, 3}(undef, grid.nx, grid.ny, nlayers), [1, 2]; flags=effort) - - if nlayers==1 - return SingleLayerParams(T(β), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, Qx, Qy, rfftplanlayered) - - else # if nlayers≥2 - - ρ = reshape(T.(ρ), (1, 1, nlayers)) - H = reshape(T.(H), (1, 1, nlayers)) - -<<<<<<< HEAD - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[1] # reduced gravity at each interface -======= - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface ->>>>>>> matts_dev - - Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) - Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) - - typeofSkl = SArray{Tuple{nlayers, nlayers}, T, 2, nlayers^2} # StaticArrays of type T and dims = (nlayers x nlayers) - - S = Array{typeofSkl, 2}(undef, (nkr, nl)) - calcS!(S, Fp, Fm, nlayers, grid) - - S⁻¹ = Array{typeofSkl, 2}(undef, (nkr, nl)) - calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) - - S, S⁻¹, Fp, Fm = A(S), A(S⁻¹), A(Fp), A(Fm) # convert to appropriate ArrayType - - CUDA.@allowscalar @views Qy[:, :, 1] = @. Qy[:, :, 1] - Fp[1] * (U[:, :, 2] - U[:, :, 1]) - for j = 2:nlayers-1 - CUDA.@allowscalar @views Qy[:, :, j] = @. Qy[:, :, j] + Fp[j] * (U[:, :, j+1] - U[:, :, j]) + Fm[j-1] * (U[:, :, j-1] - U[:, :, j]) - end - CUDA.@allowscalar @views Qy[:, :, nlayers] = @. Qy[:, :, nlayers] - Fm[nlayers-1] * (U[:, :, nlayers-1] - U[:, :, nlayers]) - - if nlayers==2 - return TwoLayerParams(T(g), T(f₀), T(β), A(ρ), (T(H[1]), T(H[2])), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, T(g′[1]), Qx, Qy, rfftplanlayered) - else # if nlayers>2 - return Params(nlayers, T(g), T(f₀), T(β), A(ρ), A(H), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, A(g′), Qx, Qy, S, S⁻¹, rfftplanlayered) - end - end -end - -numberoflayers(params) = params.nlayers -numberoflayers(::SingleLayerParams) = 1 -numberoflayers(::TwoLayerParams) = 2 - -# --------- -# Equations -# --------- - -""" - hyperviscosity(params, grid) - -Return the linear operator `L` that corresponds to (hyper)-viscosity of order ``n_ν`` with -coefficient ``ν`` for ``n`` fluid layers. -```math -L_j = - ν |𝐤|^{2 n_ν}, \\ j = 1, ...,n . -``` -""" -function hyperviscosity(params, grid) - dev = grid.device - T = eltype(grid) - - L = device_array(dev){T}(undef, (grid.nkr, grid.nl, numberoflayers(params))) - @. L = - params.ν * grid.Krsq^params.nν - @views @. L[1, 1, :] = 0 - - return L -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 -`hyperviscosity(params, grid)`. - -The nonlinear term is computed via function `calcNlinear!`. -""" -function LinearEquation(params, grid) - L = hyperviscosity(params, grid) - - return FourierFlows.Equation(L, calcNlinear!, grid) -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 -`hyperviscosity(params, grid)`. - -The nonlinear term is computed via function `calcN!`. -""" -function Equation(params, grid) - L = hyperviscosity(params, grid) - - return FourierFlows.Equation(L, calcN!, grid) -end - - -# ---- -# Vars -# ---- - -""" - struct Vars{Aphys, Atrans, F, P} <: AbstractVars - -The variables for multi-layer QG problem. - -$(FIELDS) -""" -struct Vars{Aphys, Atrans, F, P} <: AbstractVars - "relative vorticity + vortex stretching" - q :: Aphys - "streamfunction" - ψ :: Aphys - "x-component of velocity" - u :: Aphys - "y-component of velocity" - v :: Aphys - "Fourier transform of relative vorticity + vortex stretching" - qh :: Atrans - "Fourier transform of streamfunction" - ψh :: Atrans - "Fourier transform of ``x``-component of velocity" - uh :: Atrans - "Fourier transform of ``y``-component of velocity" - vh :: Atrans - "Fourier transform of forcing" - Fqh :: F - "`sol` at previous time-step" - prevsol :: P -end - -const DecayingVars = Vars{<:AbstractArray, <:AbstractArray, Nothing, Nothing} -const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing} -const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray} - -""" - DecayingVars(grid, params) - -Return the variables for an unforced multi-layer QG problem with `grid` and `params`. -""" -function DecayingVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, nothing, nothing) -end - -""" - ForcedVars(grid, params) - -Return the variables for a forced multi-layer QG problem with `grid` and `params`. -""" -function ForcedVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, nothing) -end - -""" - StochasticForcedVars(grid, params) - -Return the variables for a forced multi-layer QG problem with `grid` and `params`. -""" -function StochasticForcedVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh prevsol - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, prevsol) -end - -""" - fwdtransform!(varh, var, params) - -Compute the Fourier transform of `var` and store it in `varh`. -""" -fwdtransform!(varh, var, params::AbstractParams) = mul!(varh, params.rfftplan, var) - -""" - invtransform!(var, varh, params) - -Compute the inverse Fourier transform of `varh` and store it in `var`. -""" -invtransform!(var, varh, params::AbstractParams) = ldiv!(var, params.rfftplan, varh) - -""" - pvfromstreamfunction!(qh, ψh, params, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` in each layer using -`qh = params.S * ψh`. -""" -function pvfromstreamfunction!(qh, ψh, params, grid) - for j=1:grid.nl, i=1:grid.nkr - CUDA.@allowscalar @views qh[i, j, :] .= params.S[i, j] * ψh[i, j, :] - end - - return nothing -end - -""" - pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` for the special -case of a single fluid layer configuration. In this case, ``q̂ = - k² ψ̂``. -""" -function pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) - @. qh = -grid.Krsq * ψh - - return nothing -end - -""" - pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` for the special -case of a two fluid layer configuration. In this case we have, - -```math -q̂₁ = - k² ψ̂₁ + f₀² / (g′ H₁) * (ψ̂₂ - ψ̂₁) , -``` - -```math -q̂₂ = - k² ψ̂₂ + f₀² / (g′ H₂) * (ψ̂₁ - ψ̂₂) . -``` - -(Here, the PV-streamfunction relationship is hard-coded to avoid scalar operations -on the GPU.) -""" -function pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) - f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] - - ψ1h, ψ2h = view(ψh, :, :, 1), view(ψh, :, :, 2) - - @views @. qh[:, :, 1] = - grid.Krsq * ψ1h + f₀^2 / (g′ * H₁) * (ψ2h - ψ1h) - @views @. qh[:, :, 2] = - grid.Krsq * ψ2h + f₀^2 / (g′ * H₂) * (ψ1h - ψ2h) - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` in each layer from -`qh` using `ψh = params.S⁻¹ qh`. -""" -function streamfunctionfrompv!(ψh, qh, params, grid) - for j=1:grid.nl, i=1:grid.nkr - CUDA.@allowscalar @views ψh[i, j, :] .= params.S⁻¹[i, j] * qh[i, j, :] - end - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` for the special -case of a single fluid layer configuration. In this case, ``ψ̂ = - k⁻² q̂``. -""" -function streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) - @. ψh = -grid.invKrsq * qh - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` for the special -case of a two fluid layer configuration. In this case we have, - -```math -ψ̂₁ = - [k² q̂₁ + (f₀² / g′) (q̂₁ / H₂ + q̂₂ / H₁)] / Δ , -``` - -```math -ψ̂₂ = - [k² q̂₂ + (f₀² / g′) (q̂₁ / H₂ + q̂₂ / H₁)] / Δ , -``` - -where ``Δ = k² [k² + f₀² (H₁ + H₂) / (g′ H₁ H₂)]``. - -(Here, the PV-streamfunction relationship is hard-coded to avoid scalar operations -on the GPU.) -""" -function streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) - f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] - - q1h, q2h = view(qh, :, :, 1), view(qh, :, :, 2) - - @views @. ψh[:, :, 1] = - grid.Krsq * q1h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) - @views @. ψh[:, :, 2] = - grid.Krsq * q2h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) - - for j in 1:2 - @views @. ψh[:, :, j] *= grid.invKrsq / (grid.Krsq + f₀^2 / g′ * (H₁ + H₂) / (H₁ * H₂)) - end - - return nothing -end - -""" - calcS!(S, Fp, Fm, nlayers, grid) - -Construct the array ``𝕊``, which consists of `nlayer` x `nlayer` static arrays ``𝕊_𝐤`` that -relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``q̂_𝐤 = 𝕊_𝐤 ψ̂_𝐤``. -""" -function calcS!(S, Fp, Fm, nlayers, grid) - F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) - - for n=1:grid.nl, m=1:grid.nkr - k² = CUDA.@allowscalar grid.Krsq[m, n] - Skl = SMatrix{nlayers, nlayers}(- k² * I + F) - S[m, n] = Skl - end - - return nothing -end - -""" - calcS⁻¹!(S, Fp, Fm, nlayers, grid) - -Construct the array ``𝕊⁻¹``, which consists of `nlayer` x `nlayer` static arrays ``(𝕊_𝐤)⁻¹`` -that relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``ψ̂_𝐤 = (𝕊_𝐤)⁻¹ q̂_𝐤``. -""" -function calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) - F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) - - for n=1:grid.nl, m=1:grid.nkr - k² = CUDA.@allowscalar grid.Krsq[m, n] == 0 ? 1 : grid.Krsq[m, n] - Skl = - k² * I + F - S⁻¹[m, n] = SMatrix{nlayers, nlayers}(I / Skl) - end - - T = eltype(grid) - S⁻¹[1, 1] = SMatrix{nlayers, nlayers}(zeros(T, (nlayers, nlayers))) - - return nothing -end - - -# ------- -# Solvers -# ------- - -""" - calcN!(N, sol, t, clock, vars, params, grid) - -Compute the nonlinear term, that is the advection term, the bottom drag, and the forcing: - -```math -N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . -``` -""" -function calcN!(N, sol, t, clock, vars, params, grid) - nlayers = numberoflayers(params) - - dealias!(sol, grid) - - calcN_advection!(N, sol, vars, params, grid) - - @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag - - addforcing!(N, sol, t, clock, vars, params, grid) - - return nothing -end - -""" - calcNlinear!(N, sol, t, clock, vars, params, grid) - -Compute the nonlinear term of the linearized equations: - -```math -N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} + \\widehat{(∂_y ψ_j)(∂_x Q_j)} -- \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . -``` -""" -function calcNlinear!(N, sol, t, clock, vars, params, grid) - nlayers = numberoflayers(params) - - calcN_linearadvection!(N, sol, vars, params, grid) - @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag - addforcing!(N, sol, t, clock, vars, params, grid) - - return nothing -end - -""" - calcN_advection!(N, sol, vars, params, grid) - -Compute the advection term and stores it in `N`: - -```math -N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} . -``` -""" -function calcN_advection!(N, sol, vars, params, grid) - @. vars.qh = sol - - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.u, vars.uh, params) - @. vars.u += params.U # add the imposed zonal flow U - - uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables - @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x - fwdtransform!(uQxh, uQx, params) - @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} - - invtransform!(vars.v, vars.vh, params) - - vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables - @. vQy = vars.v * params.Qy # v*∂Q/∂y - fwdtransform!(vQyh, vQy, params) - @. N -= vQyh # -\hat{v*∂Q/∂y} - - invtransform!(vars.q, vars.qh, params) - - uq , vq = vars.u , vars.v # use vars.u and vars.v as scratch variables - uqh, vqh = vars.uh, vars.vh # use vars.uh and vars.vh as scratch variables - @. uq *= vars.q # (U+u)*q - @. vq *= vars.q # v*q - - fwdtransform!(uqh, uq, params) - fwdtransform!(vqh, vq, params) - - @. N -= im * grid.kr * uqh + im * grid.l * vqh # -\hat{∂[(U+u)q]/∂x} - \hat{∂[vq]/∂y} - - return nothing -end - - -""" - calcN_linearadvection!(N, sol, vars, params, grid) - -Compute the advection term of the linearized equations and stores it in `N`: - -```math -N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} . -``` -""" -function calcN_linearadvection!(N, sol, vars, params, grid) - @. vars.qh = sol - - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.u, vars.uh, params) - @. vars.u += params.U # add the imposed zonal flow U - uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables - @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x - fwdtransform!(uQxh, uQx, params) - @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} - - invtransform!(vars.v, vars.vh, params) - - vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables - - @. vQy = vars.v * params.Qy # v*∂Q/∂y - fwdtransform!(vQyh, vQy, params) - @. N -= vQyh # -\hat{v*∂Q/∂y} - - invtransform!(vars.q, vars.qh, params) - - @. vars.u = params.U - Uq , Uqh = vars.u , vars.uh # use vars.u and vars.uh as scratch variables - @. Uq *= vars.q # U*q - - fwdtransform!(Uqh, Uq, params) - - @. N -= im * grid.kr * Uqh # -\hat{∂[U*q]/∂x} - - return nothing -end - - -""" - addforcing!(N, sol, t, clock, vars, params, grid) - -When the problem includes forcing, calculate the forcing term ``F̂`` for each layer and add -it to the nonlinear term ``N``. -""" -addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing - -function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) - params.calcFq!(vars.Fqh, sol, t, clock, vars, params, grid) - @. N += vars.Fqh - - return nothing -end - - -# ---------------- -# Helper functions -# ---------------- - -""" - updatevars!(vars, params, grid, sol) - updatevars!(prob) - -Update all problem variables using `sol`. -""" -function updatevars!(vars, params, grid, sol) - dealias!(sol, grid) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.q, deepcopy(vars.qh), params) - invtransform!(vars.ψ, deepcopy(vars.ψh), params) - invtransform!(vars.u, deepcopy(vars.uh), params) - invtransform!(vars.v, deepcopy(vars.vh), params) - - return nothing -end - -updatevars!(prob) = updatevars!(prob.vars, prob.params, prob.grid, prob.sol) - - -""" - set_q!(sol, params, vars, grid, q) - set_q!(prob, q) - -Set the solution `prob.sol` as the transform of `q` and update variables. -""" -function set_q!(sol, params, vars, grid, q) - A = typeof(vars.q) - fwdtransform!(vars.qh, A(q), params) - @. vars.qh[1, 1, :] = 0 - @. sol = vars.qh - updatevars!(vars, params, grid, sol) - - return nothing -end - -function set_q!(sol, params::SingleLayerParams, vars, grid, q::AbstractArray{T, 2}) where T - A = typeof(vars.q[:, :, 1]) - q_3D = vars.q - @views q_3D[:, :, 1] = A(q) - set_q!(sol, params, vars, grid, q_3D) - - return nothing -end - -set_q!(prob, q) = set_q!(prob.sol, prob.params, prob.vars, prob.grid, q) - - -""" - set_ψ!(params, vars, grid, sol, ψ) - set_ψ!(prob, ψ) - -Set the solution `prob.sol` to the transform `qh` that corresponds to streamfunction `ψ` -and update variables. -""" -function set_ψ!(sol, params, vars, grid, ψ) - A = typeof(vars.q) - fwdtransform!(vars.ψh, A(ψ), params) - pvfromstreamfunction!(vars.qh, vars.ψh, params, grid) - invtransform!(vars.q, vars.qh, params) - - set_q!(sol, params, vars, grid, vars.q) - - return nothing -end - -function set_ψ!(sol, params::SingleLayerParams, vars, grid, ψ::AbstractArray{T, 2}) where T - A = typeof(vars.ψ[:, :, 1]) - ψ_3D = vars.ψ - @views ψ_3D[:, :, 1] = A(ψ) - - set_ψ!(sol, params, vars, grid, ψ_3D) - - return nothing -end - -set_ψ!(prob, ψ) = set_ψ!(prob.sol, prob.params, prob.vars, prob.grid, ψ) - - -""" - energies(vars, params, grid, sol) - energies(prob) - -Return the kinetic energy of each fluid layer KE``_1, ...,`` KE``_{n}``, and the -potential energy of each fluid interface PE``_{3/2}, ...,`` PE``_{n-1/2}``, where ``n`` -is the number of layers in the fluid. (When ``n=1``, only the kinetic energy is returned.) - -The kinetic energy at the ``j``-th fluid layer is - -```math -𝖪𝖤_j = \\frac{H_j}{H} \\int \\frac1{2} |{\\bf ∇} ψ_j|^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{H_j}{H} \\sum_{𝐤} |𝐤|² |ψ̂_j|², \\ j = 1, ..., n , -``` - -while the potential energy that corresponds to the interface ``j+1/2`` (i.e., the interface -between the ``j``-th and ``(j+1)``-th fluid layer) is - -```math -𝖯𝖤_{j+1/2} = \\int \\frac1{2} \\frac{f₀^2}{g'_{j+1/2} H} (ψ_j - ψ_{j+1})^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{f₀^2}{g'_{j+1/2} H} \\sum_{𝐤} |ψ̂_j - ψ̂_{j+1}|², \\ j = 1, ..., n-1 . -``` -""" -function energies(vars, params, grid, sol) - nlayers = numberoflayers(params) - KE, PE = zeros(nlayers), zeros(nlayers-1) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - for j = 1:nlayers - CUDA.@allowscalar KE[j] = 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) - end - - for j = 1:nlayers-1 - CUDA.@allowscalar PE[j] = 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′[j] * parsevalsum(abs2.(vars.ψh[:, :, j+1] .- vars.ψh[:, :, j]), grid) - end - - return KE, PE -end - -function energies(vars, params::TwoLayerParams, grid, sol) - nlayers = numberoflayers(params) - KE, PE = zeros(nlayers), zeros(nlayers-1) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - ψ1h, ψ2h = view(vars.ψh, :, :, 1), view(vars.ψh, :, :, 2) - - for j = 1:nlayers - CUDA.@allowscalar KE[j] = @views 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) - end - - PE = @views 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′ * parsevalsum(abs2.(ψ2h .- ψ1h), grid) - - return KE, PE -end - -function energies(vars, params::SingleLayerParams, grid, sol) - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - return 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h, grid) -end - -energies(prob) = energies(prob.vars, prob.params, prob.grid, prob.sol) - -""" - fluxes(vars, params, grid, sol) - fluxes(prob) - -Return the lateral eddy fluxes within each fluid layer, lateralfluxes``_1,...,``lateralfluxes``_n`` -and also the vertical eddy fluxes at each fluid interface, -verticalfluxes``_{3/2},...,``verticalfluxes``_{n-1/2}``, where ``n`` is the total number of layers in the fluid. -(When ``n=1``, only the lateral fluxes are returned.) - -The lateral eddy fluxes within the ``j``-th fluid layer are - -```math -\\textrm{lateralfluxes}_j = \\frac{H_j}{H} \\int U_j v_j ∂_y u_j -\\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n , -``` - -while the vertical eddy fluxes at the ``j+1/2``-th fluid interface (i.e., interface between -the ``j``-th and ``(j+1)``-th fluid layer) are - -```math -\\textrm{verticalfluxes}_{j+1/2} = \\int \\frac{f₀²}{g'_{j+1/2} H} (U_j - U_{j+1}) \\, -v_{j+1} ψ_{j} \\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n-1. -``` -""" -function fluxes(vars, params, grid, sol) - nlayers = numberoflayers(params) - - lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) - - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.H * params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - for j = 1:nlayers-1 - CUDA.@allowscalar verticalfluxes[j] = sum(@views @. params.f₀^2 / params.g′[j] * (params.U[: ,:, j] - params.U[:, :, j+1]) * vars.v[:, :, j+1] * vars.ψ[:, :, j]; dims=(1, 2))[1] - CUDA.@allowscalar verticalfluxes[j] *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - end - - return lateralfluxes, verticalfluxes -end - -function fluxes(vars, params::TwoLayerParams, grid, sol) - nlayers = numberoflayers(params) - - lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) - - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - @. lateralfluxes *= params.H - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - U₁, U₂ = view(params.U, :, :, 1), view(params.U, :, :, 2) - ψ₁ = view(vars.ψ, :, :, 1) - v₂ = view(vars.v, :, :, 2) - - verticalfluxes = sum(@views @. params.f₀^2 / params.g′ * (U₁ - U₂) * v₂ * ψ₁; dims=(1, 2)) - verticalfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - return lateralfluxes, verticalfluxes -end - -function fluxes(vars, params::SingleLayerParams, grid, sol) - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly) - - return lateralfluxes -end - -fluxes(prob) = fluxes(prob.vars, prob.params, prob.grid, prob.sol) - -end # module diff --git a/src/multilayerqg_BASE_444028.jl b/src/multilayerqg_BASE_444028.jl deleted file mode 100644 index f21597b5..00000000 --- a/src/multilayerqg_BASE_444028.jl +++ /dev/null @@ -1,1107 +0,0 @@ -module MultiLayerQG - -export - fwdtransform!, - invtransform!, - streamfunctionfrompv!, - pvfromstreamfunction!, - updatevars!, - - set_q!, - set_ψ!, - energies, - fluxes - -using - FFTW, - CUDA, - LinearAlgebra, - StaticArrays, - Reexport, - DocStringExtensions - -@reexport using FourierFlows - -using FourierFlows: parsevalsum, parsevalsum2, superzeros, plan_flows_rfft - -nothingfunction(args...) = nothing - -""" - Problem(nlayers :: Int, - dev = CPU(); - nx = 128, - ny = nx, - Lx = 2π, - Ly = Lx, - f₀ = 1.0, - β = 0.0, - g = 1.0, - U = zeros(nlayers), - H = 1/nlayers * ones(nlayers), - ρ = Array{Float64}(1:nlayers), - eta = nothing, - topographic_pv_gradient = (0, 0), - μ = 0, - ν = 0, - nν = 1, - dt = 0.01, - stepper = "RK4", - calcFq = nothingfunction, - stochastic = false, - linear = false, - aliased_fraction = 1/3, - T = Float64) - -Construct a multi-layer quasi-geostrophic problem with `nlayers` fluid layers on device `dev`. - -Arguments -========= -- `nlayers`: (required) Number of fluid layers. -- `dev`: (required) `CPU()` (default) or `GPU()`; computer architecture used to time-step `problem`. - -Keyword arguments -================= - - `nx`: Number of grid points in ``x``-domain. - - `ny`: Number of grid points in ``y``-domain. - - `Lx`: Extent of the ``x``-domain. - - `Ly`: Extent of the ``y``-domain. - - `f₀`: Constant planetary vorticity. - - `β`: Planetary vorticity ``y``-gradient. - - `g`: Gravitational acceleration constant. - - `U`: The imposed constant zonal flow ``U(y)`` in each fluid layer. - - `H`: Rest height of each fluid layer. - - `ρ`: Density of each fluid layer. - - `eta`: Periodic component of the topographic potential vorticity. - - `topographic_pv_gradient`: The ``(x, y)`` components of the topographic PV large-scale gradient. - - `μ`: Linear bottom drag coefficient. - - `ν`: Small-scale (hyper)-viscosity coefficient. - - `nν`: (Hyper)-viscosity order, `nν```≥ 1``. - - `dt`: Time-step. - - `stepper`: Time-stepping method. - - `calcF`: Function that calculates the Fourier transform of the forcing, ``F̂``. - - `stochastic`: `true` or `false` (default); boolean denoting whether `calcF` is temporally stochastic. - - `linear`: `true` or `false` (default); boolean denoting whether the linearized equations of motions are used. - - `aliased_fraction`: the fraction of high-wavenumbers that are zero-ed out by `dealias!()`. - - `T`: `Float32` or `Float64` (default); floating point type used for `problem` data. -""" -function Problem(nlayers::Int, # number of fluid layers - dev = CPU(); - # Numerical parameters - nx = 128, - ny = nx, - Lx = 2π, - Ly = Lx, - # Physical parameters - f₀ = 1.0, # Coriolis parameter - β = 0.0, # y-gradient of Coriolis parameter - g = 1.0, # gravitational constant - U = zeros(nlayers), # imposed zonal flow U(y) in each layer - H = 1/nlayers * ones(nlayers), # rest fluid height of each layer - ρ = Array{Float64}(1:nlayers), # density of each layer - eta = nothing, # periodic component of the topographic PV - topographic_pv_gradient = (0, 0), # tuple with the ``(x, y)`` components of topographic PV large-scale gradient - # Bottom Drag and/or (hyper)-viscosity - μ = 0, - ν = 0, - nν = 1, - # Timestepper and equation options - dt = 0.01, - stepper = "RK4", - calcFq = nothingfunction, - stochastic = false, - linear = false, - # Float type and dealiasing - aliased_fraction = 1/3, - T = Float64) - - if dev == GPU() && nlayers > 2 - @warn """MultiLayerQG module is not optimized on the GPU yet for configurations with - 3 fluid layers or more! - - See issues on Github at https://github.com/FourierFlows/GeophysicalFlows.jl/issues/112 - and https://github.com/FourierFlows/GeophysicalFlows.jl/issues/267. - - To use MultiLayerQG with 3 fluid layers or more we suggest, for now, to restrict running - on CPU.""" - end - - if nlayers == 1 - @warn """MultiLayerQG module does work for single-layer configuration but may not be as - optimized. We suggest using SingleLayerQG module for single-layer QG simulation unless - you have reasons to use MultiLayerQG in a single-layer configuration, e.g., you want to - compare solutions with varying number of fluid layers.""" - end - - # topographic PV - eta === nothing && (eta = zeros(dev, T, (nx, ny))) - - grid = TwoDGrid(dev; nx, Lx, ny, Ly, aliased_fraction, T) - - params = Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ, ν, nν, grid; calcFq) - - vars = calcFq == nothingfunction ? DecayingVars(grid, params) : (stochastic ? StochasticForcedVars(grid, params) : ForcedVars(grid, params)) - - equation = linear ? LinearEquation(params, grid) : Equation(params, grid) - - FourierFlows.Problem(equation, stepper, dt, grid, vars, params) -end - -""" - struct Params{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams - -The parameters for the `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct Params{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams - # prescribed params - "number of fluid layers" - nlayers :: Int - "gravitational constant" - g :: T - "constant planetary vorticity" - f₀ :: T - "planetary vorticity ``y``-gradient" - β :: T - "array with density of each fluid layer" - ρ :: Aphys3D - "array with rest height of each fluid layer" - H :: Aphys3D - "array with imposed constant zonal flow ``U(y)`` in each fluid layer" - U :: Aphys3D - "array containing the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear bottom drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "array with the reduced gravity constants for each fluid interface" - g′ :: Aphys1D - "array containing ``x``-gradient of PV due to topographic PV in each fluid layer" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV in each fluid layer" - Qy :: Aphys3D - "array containing coeffients for getting PV from streamfunction" - S :: Atrans4D - "array containing coeffients for inverting PV to streamfunction" - S⁻¹ :: Atrans4D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -""" - struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - -The parameters for the a single-layer `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - # prescribed params - "planetary vorticity ``y``-gradient" - β :: T - "array with imposed constant zonal flow ``U(y)``" - U :: Aphys3D - "array containing the periodic component of the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "array containing ``x``-gradient of PV due to topographic PV" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV" - Qy :: Aphys3D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -""" - struct TwoLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - -The parameters for the a two-layer `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct TwoLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - # prescribed params - "gravitational constant" - g :: T - "constant planetary vorticity" - f₀ :: T - "planetary vorticity ``y``-gradient" - β :: T - "array with density of each fluid layer" - ρ :: Aphys3D - "tuple with rest height of each fluid layer" - H :: Tuple - "array with imposed constant zonal flow ``U(y)`` in each fluid layer" - U :: Aphys3D - "array containing periodic component of the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear bottom drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "the reduced gravity constants for the fluid interface" - g′ :: T - "array containing ``x``-gradient of PV due to topographic PV in each fluid layer" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV in each fluid layer" - Qy :: Aphys3D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 1}) where TU - T = eltype(grid) - if length(U) == nlayers - U_2D = zeros(dev, T, (1, nlayers)) - U_2D[:] = U - U_2D = repeat(U_2D, outer=(grid.ny, 1)) - else - U_2D = zeros(dev, T, (grid.ny, 1)) - U_2D[:] = U - end - U_3D = zeros(dev, T, (1, grid.ny, nlayers)) - @views U_3D[1, :, :] = U_2D - return U_3D -end - -function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 2}) where TU - T = eltype(grid) - U_3D = zeros(dev, T, (1, grid.ny, nlayers)) - @views U_3D[1, :, :] = U - return U_3D -end - -function convert_U_to_U3D(dev, nlayers, grid, U::Number) - T = eltype(grid) - A = device_array(dev) - U_3D = reshape(repeat([T(U)], outer=(grid.ny, 1)), (1, grid.ny, nlayers)) - return A(U_3D) -end - -function Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ, ν, nν, grid; calcFq=nothingfunction, effort=FFTW.MEASURE) - dev = grid.device - T = eltype(grid) - A = device_array(dev) - - ny, nx = grid.ny , grid.nx - nkr, nl = grid.nkr, grid.nl - kr, l = grid.kr , grid.l - - U = convert_U_to_U3D(dev, nlayers, grid, U) - - Uyy = real.(ifft(-l.^2 .* fft(U))) - Uyy = CUDA.@allowscalar repeat(Uyy, outer=(nx, 1, 1)) - - # Calculate periodic components of the topographic PV gradients. - etah = rfft(A(eta)) - etax = irfft(im * kr .* etah, nx) # ∂η/∂x - etay = irfft(im * l .* etah, nx) # ∂η/∂y - - # Add topographic PV large-scale gradient - topographic_pv_gradient = T.(topographic_pv_gradient) - @. etax += topographic_pv_gradient[1] - @. etay += topographic_pv_gradient[2] - - Qx = zeros(dev, T, (nx, ny, nlayers)) - @views @. Qx[:, :, nlayers] += etax - - Qy = zeros(dev, T, (nx, ny, nlayers)) - Qy = T(β) .- Uyy # T(β) is needed to ensure that Qy remains same type as U - @views @. Qy[:, :, nlayers] += etay - - rfftplanlayered = plan_flows_rfft(A{T, 3}(undef, grid.nx, grid.ny, nlayers), [1, 2]; flags=effort) - - if nlayers==1 - return SingleLayerParams(T(β), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, Qx, Qy, rfftplanlayered) - - else # if nlayers≥2 - - ρ = reshape(T.(ρ), (1, 1, nlayers)) - H = reshape(T.(H), (1, 1, nlayers)) - - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[1:nlayers-1] # reduced gravity at each interface - - Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) - Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) - - typeofSkl = SArray{Tuple{nlayers, nlayers}, T, 2, nlayers^2} # StaticArrays of type T and dims = (nlayers x nlayers) - - S = Array{typeofSkl, 2}(undef, (nkr, nl)) - calcS!(S, Fp, Fm, nlayers, grid) - - S⁻¹ = Array{typeofSkl, 2}(undef, (nkr, nl)) - calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) - - S, S⁻¹, Fp, Fm = A(S), A(S⁻¹), A(Fp), A(Fm) # convert to appropriate ArrayType - - CUDA.@allowscalar @views Qy[:, :, 1] = @. Qy[:, :, 1] - Fp[1] * (U[:, :, 2] - U[:, :, 1]) - for j = 2:nlayers-1 - CUDA.@allowscalar @views Qy[:, :, j] = @. Qy[:, :, j] + Fp[j] * (U[:, :, j+1] - U[:, :, j]) + Fm[j-1] * (U[:, :, j-1] - U[:, :, j]) - end - CUDA.@allowscalar @views Qy[:, :, nlayers] = @. Qy[:, :, nlayers] - Fm[nlayers-1] * (U[:, :, nlayers-1] - U[:, :, nlayers]) - - if nlayers==2 - return TwoLayerParams(T(g), T(f₀), T(β), A(ρ), (T(H[1]), T(H[2])), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, T(g′[1]), Qx, Qy, rfftplanlayered) - else # if nlayers>2 - return Params(nlayers, T(g), T(f₀), T(β), A(ρ), A(H), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, A(g′), Qx, Qy, S, S⁻¹, rfftplanlayered) - end - end -end - -numberoflayers(params) = params.nlayers -numberoflayers(::SingleLayerParams) = 1 -numberoflayers(::TwoLayerParams) = 2 - -# --------- -# Equations -# --------- - -""" - hyperviscosity(params, grid) - -Return the linear operator `L` that corresponds to (hyper)-viscosity of order ``n_ν`` with -coefficient ``ν`` for ``n`` fluid layers. -```math -L_j = - ν |𝐤|^{2 n_ν}, \\ j = 1, ...,n . -``` -""" -function hyperviscosity(params, grid) - dev = grid.device - T = eltype(grid) - - L = device_array(dev){T}(undef, (grid.nkr, grid.nl, numberoflayers(params))) - @. L = - params.ν * grid.Krsq^params.nν - @views @. L[1, 1, :] = 0 - - return L -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 -`hyperviscosity(params, grid)`. - -The nonlinear term is computed via function `calcNlinear!`. -""" -function LinearEquation(params, grid) - L = hyperviscosity(params, grid) - - return FourierFlows.Equation(L, calcNlinear!, grid) -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 -`hyperviscosity(params, grid)`. - -The nonlinear term is computed via function `calcN!`. -""" -function Equation(params, grid) - L = hyperviscosity(params, grid) - - return FourierFlows.Equation(L, calcN!, grid) -end - - -# ---- -# Vars -# ---- - -""" - struct Vars{Aphys, Atrans, F, P} <: AbstractVars - -The variables for multi-layer QG problem. - -$(FIELDS) -""" -struct Vars{Aphys, Atrans, F, P} <: AbstractVars - "relative vorticity + vortex stretching" - q :: Aphys - "streamfunction" - ψ :: Aphys - "x-component of velocity" - u :: Aphys - "y-component of velocity" - v :: Aphys - "Fourier transform of relative vorticity + vortex stretching" - qh :: Atrans - "Fourier transform of streamfunction" - ψh :: Atrans - "Fourier transform of ``x``-component of velocity" - uh :: Atrans - "Fourier transform of ``y``-component of velocity" - vh :: Atrans - "Fourier transform of forcing" - Fqh :: F - "`sol` at previous time-step" - prevsol :: P -end - -const DecayingVars = Vars{<:AbstractArray, <:AbstractArray, Nothing, Nothing} -const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing} -const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray} - -""" - DecayingVars(grid, params) - -Return the variables for an unforced multi-layer QG problem with `grid` and `params`. -""" -function DecayingVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, nothing, nothing) -end - -""" - ForcedVars(grid, params) - -Return the variables for a forced multi-layer QG problem with `grid` and `params`. -""" -function ForcedVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, nothing) -end - -""" - StochasticForcedVars(grid, params) - -Return the variables for a forced multi-layer QG problem with `grid` and `params`. -""" -function StochasticForcedVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh prevsol - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, prevsol) -end - -""" - fwdtransform!(varh, var, params) - -Compute the Fourier transform of `var` and store it in `varh`. -""" -fwdtransform!(varh, var, params::AbstractParams) = mul!(varh, params.rfftplan, var) - -""" - invtransform!(var, varh, params) - -Compute the inverse Fourier transform of `varh` and store it in `var`. -""" -invtransform!(var, varh, params::AbstractParams) = ldiv!(var, params.rfftplan, varh) - -""" - pvfromstreamfunction!(qh, ψh, params, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` in each layer using -`qh = params.S * ψh`. -""" -function pvfromstreamfunction!(qh, ψh, params, grid) - for j=1:grid.nl, i=1:grid.nkr - CUDA.@allowscalar @views qh[i, j, :] .= params.S[i, j] * ψh[i, j, :] - end - - return nothing -end - -""" - pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` for the special -case of a single fluid layer configuration. In this case, ``q̂ = - k² ψ̂``. -""" -function pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) - @. qh = -grid.Krsq * ψh - - return nothing -end - -""" - pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` for the special -case of a two fluid layer configuration. In this case we have, - -```math -q̂₁ = - k² ψ̂₁ + f₀² / (g′ H₁) * (ψ̂₂ - ψ̂₁) , -``` - -```math -q̂₂ = - k² ψ̂₂ + f₀² / (g′ H₂) * (ψ̂₁ - ψ̂₂) . -``` - -(Here, the PV-streamfunction relationship is hard-coded to avoid scalar operations -on the GPU.) -""" -function pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) - f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] - - ψ1h, ψ2h = view(ψh, :, :, 1), view(ψh, :, :, 2) - - @views @. qh[:, :, 1] = - grid.Krsq * ψ1h + f₀^2 / (g′ * H₁) * (ψ2h - ψ1h) - @views @. qh[:, :, 2] = - grid.Krsq * ψ2h + f₀^2 / (g′ * H₂) * (ψ1h - ψ2h) - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` in each layer from -`qh` using `ψh = params.S⁻¹ qh`. -""" -function streamfunctionfrompv!(ψh, qh, params, grid) - for j=1:grid.nl, i=1:grid.nkr - CUDA.@allowscalar @views ψh[i, j, :] .= params.S⁻¹[i, j] * qh[i, j, :] - end - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` for the special -case of a single fluid layer configuration. In this case, ``ψ̂ = - k⁻² q̂``. -""" -function streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) - @. ψh = -grid.invKrsq * qh - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` for the special -case of a two fluid layer configuration. In this case we have, - -```math -ψ̂₁ = - [k² q̂₁ + (f₀² / g′) (q̂₁ / H₂ + q̂₂ / H₁)] / Δ , -``` - -```math -ψ̂₂ = - [k² q̂₂ + (f₀² / g′) (q̂₁ / H₂ + q̂₂ / H₁)] / Δ , -``` - -where ``Δ = k² [k² + f₀² (H₁ + H₂) / (g′ H₁ H₂)]``. - -(Here, the PV-streamfunction relationship is hard-coded to avoid scalar operations -on the GPU.) -""" -function streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) - f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] - - q1h, q2h = view(qh, :, :, 1), view(qh, :, :, 2) - - @views @. ψh[:, :, 1] = - grid.Krsq * q1h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) - @views @. ψh[:, :, 2] = - grid.Krsq * q2h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) - - for j in 1:2 - @views @. ψh[:, :, j] *= grid.invKrsq / (grid.Krsq + f₀^2 / g′ * (H₁ + H₂) / (H₁ * H₂)) - end - - return nothing -end - -""" - calcS!(S, Fp, Fm, nlayers, grid) - -Construct the array ``𝕊``, which consists of `nlayer` x `nlayer` static arrays ``𝕊_𝐤`` that -relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``q̂_𝐤 = 𝕊_𝐤 ψ̂_𝐤``. -""" -function calcS!(S, Fp, Fm, nlayers, grid) - F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) - - for n=1:grid.nl, m=1:grid.nkr - k² = CUDA.@allowscalar grid.Krsq[m, n] - Skl = SMatrix{nlayers, nlayers}(- k² * I + F) - S[m, n] = Skl - end - - return nothing -end - -""" - calcS⁻¹!(S, Fp, Fm, nlayers, grid) - -Construct the array ``𝕊⁻¹``, which consists of `nlayer` x `nlayer` static arrays ``(𝕊_𝐤)⁻¹`` -that relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``ψ̂_𝐤 = (𝕊_𝐤)⁻¹ q̂_𝐤``. -""" -function calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) - F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) - - for n=1:grid.nl, m=1:grid.nkr - k² = CUDA.@allowscalar grid.Krsq[m, n] == 0 ? 1 : grid.Krsq[m, n] - Skl = - k² * I + F - S⁻¹[m, n] = SMatrix{nlayers, nlayers}(I / Skl) - end - - T = eltype(grid) - S⁻¹[1, 1] = SMatrix{nlayers, nlayers}(zeros(T, (nlayers, nlayers))) - - return nothing -end - - -# ------- -# Solvers -# ------- - -""" - calcN!(N, sol, t, clock, vars, params, grid) - -Compute the nonlinear term, that is the advection term, the bottom drag, and the forcing: - -```math -N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . -``` -""" -function calcN!(N, sol, t, clock, vars, params, grid) - nlayers = numberoflayers(params) - - dealias!(sol, grid) - - calcN_advection!(N, sol, vars, params, grid) - - @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag - - addforcing!(N, sol, t, clock, vars, params, grid) - - return nothing -end - -""" - calcNlinear!(N, sol, t, clock, vars, params, grid) - -Compute the nonlinear term of the linearized equations: - -```math -N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} + \\widehat{(∂_y ψ_j)(∂_x Q_j)} -- \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . -``` -""" -function calcNlinear!(N, sol, t, clock, vars, params, grid) - nlayers = numberoflayers(params) - - calcN_linearadvection!(N, sol, vars, params, grid) - @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag - addforcing!(N, sol, t, clock, vars, params, grid) - - return nothing -end - -""" - calcN_advection!(N, sol, vars, params, grid) - -Compute the advection term and stores it in `N`: - -```math -N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} . -``` -""" -function calcN_advection!(N, sol, vars, params, grid) - @. vars.qh = sol - - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.u, vars.uh, params) - @. vars.u += params.U # add the imposed zonal flow U - - uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables - @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x - fwdtransform!(uQxh, uQx, params) - @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} - - invtransform!(vars.v, vars.vh, params) - - vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables - @. vQy = vars.v * params.Qy # v*∂Q/∂y - fwdtransform!(vQyh, vQy, params) - @. N -= vQyh # -\hat{v*∂Q/∂y} - - invtransform!(vars.q, vars.qh, params) - - uq , vq = vars.u , vars.v # use vars.u and vars.v as scratch variables - uqh, vqh = vars.uh, vars.vh # use vars.uh and vars.vh as scratch variables - @. uq *= vars.q # (U+u)*q - @. vq *= vars.q # v*q - - fwdtransform!(uqh, uq, params) - fwdtransform!(vqh, vq, params) - - @. N -= im * grid.kr * uqh + im * grid.l * vqh # -\hat{∂[(U+u)q]/∂x} - \hat{∂[vq]/∂y} - - return nothing -end - - -""" - calcN_linearadvection!(N, sol, vars, params, grid) - -Compute the advection term of the linearized equations and stores it in `N`: - -```math -N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} . -``` -""" -function calcN_linearadvection!(N, sol, vars, params, grid) - @. vars.qh = sol - - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.u, vars.uh, params) - @. vars.u += params.U # add the imposed zonal flow U - uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables - @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x - fwdtransform!(uQxh, uQx, params) - @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} - - invtransform!(vars.v, vars.vh, params) - - vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables - - @. vQy = vars.v * params.Qy # v*∂Q/∂y - fwdtransform!(vQyh, vQy, params) - @. N -= vQyh # -\hat{v*∂Q/∂y} - - invtransform!(vars.q, vars.qh, params) - - @. vars.u = params.U - Uq , Uqh = vars.u , vars.uh # use vars.u and vars.uh as scratch variables - @. Uq *= vars.q # U*q - - fwdtransform!(Uqh, Uq, params) - - @. N -= im * grid.kr * Uqh # -\hat{∂[U*q]/∂x} - - return nothing -end - - -""" - addforcing!(N, sol, t, clock, vars, params, grid) - -When the problem includes forcing, calculate the forcing term ``F̂`` for each layer and add -it to the nonlinear term ``N``. -""" -addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing - -function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) - params.calcFq!(vars.Fqh, sol, t, clock, vars, params, grid) - @. N += vars.Fqh - - return nothing -end - - -# ---------------- -# Helper functions -# ---------------- - -""" - updatevars!(vars, params, grid, sol) - updatevars!(prob) - -Update all problem variables using `sol`. -""" -function updatevars!(vars, params, grid, sol) - dealias!(sol, grid) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.q, deepcopy(vars.qh), params) - invtransform!(vars.ψ, deepcopy(vars.ψh), params) - invtransform!(vars.u, deepcopy(vars.uh), params) - invtransform!(vars.v, deepcopy(vars.vh), params) - - return nothing -end - -updatevars!(prob) = updatevars!(prob.vars, prob.params, prob.grid, prob.sol) - - -""" - set_q!(sol, params, vars, grid, q) - set_q!(prob, q) - -Set the solution `prob.sol` as the transform of `q` and update variables. -""" -function set_q!(sol, params, vars, grid, q) - A = typeof(vars.q) - fwdtransform!(vars.qh, A(q), params) - @. vars.qh[1, 1, :] = 0 - @. sol = vars.qh - updatevars!(vars, params, grid, sol) - - return nothing -end - -function set_q!(sol, params::SingleLayerParams, vars, grid, q::AbstractArray{T, 2}) where T - A = typeof(vars.q[:, :, 1]) - q_3D = vars.q - @views q_3D[:, :, 1] = A(q) - set_q!(sol, params, vars, grid, q_3D) - - return nothing -end - -set_q!(prob, q) = set_q!(prob.sol, prob.params, prob.vars, prob.grid, q) - - -""" - set_ψ!(params, vars, grid, sol, ψ) - set_ψ!(prob, ψ) - -Set the solution `prob.sol` to the transform `qh` that corresponds to streamfunction `ψ` -and update variables. -""" -function set_ψ!(sol, params, vars, grid, ψ) - A = typeof(vars.q) - fwdtransform!(vars.ψh, A(ψ), params) - pvfromstreamfunction!(vars.qh, vars.ψh, params, grid) - invtransform!(vars.q, vars.qh, params) - - set_q!(sol, params, vars, grid, vars.q) - - return nothing -end - -function set_ψ!(sol, params::SingleLayerParams, vars, grid, ψ::AbstractArray{T, 2}) where T - A = typeof(vars.ψ[:, :, 1]) - ψ_3D = vars.ψ - @views ψ_3D[:, :, 1] = A(ψ) - - set_ψ!(sol, params, vars, grid, ψ_3D) - - return nothing -end - -set_ψ!(prob, ψ) = set_ψ!(prob.sol, prob.params, prob.vars, prob.grid, ψ) - - -""" - energies(vars, params, grid, sol) - energies(prob) - -Return the kinetic energy of each fluid layer KE``_1, ...,`` KE``_{n}``, and the -potential energy of each fluid interface PE``_{3/2}, ...,`` PE``_{n-1/2}``, where ``n`` -is the number of layers in the fluid. (When ``n=1``, only the kinetic energy is returned.) - -The kinetic energy at the ``j``-th fluid layer is - -```math -𝖪𝖤_j = \\frac{H_j}{H} \\int \\frac1{2} |{\\bf ∇} ψ_j|^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{H_j}{H} \\sum_{𝐤} |𝐤|² |ψ̂_j|², \\ j = 1, ..., n , -``` - -while the potential energy that corresponds to the interface ``j+1/2`` (i.e., the interface -between the ``j``-th and ``(j+1)``-th fluid layer) is - -```math -𝖯𝖤_{j+1/2} = \\int \\frac1{2} \\frac{f₀^2}{g'_{j+1/2} H} (ψ_j - ψ_{j+1})^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{f₀^2}{g'_{j+1/2} H} \\sum_{𝐤} |ψ̂_j - ψ̂_{j+1}|², \\ j = 1, ..., n-1 . -``` -""" -function energies(vars, params, grid, sol) - nlayers = numberoflayers(params) - KE, PE = zeros(nlayers), zeros(nlayers-1) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - for j = 1:nlayers - CUDA.@allowscalar KE[j] = 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) - end - - for j = 1:nlayers-1 - CUDA.@allowscalar PE[j] = 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′[j] * parsevalsum(abs2.(vars.ψh[:, :, j+1] .- vars.ψh[:, :, j]), grid) - end - - return KE, PE -end - -function energies(vars, params::TwoLayerParams, grid, sol) - nlayers = numberoflayers(params) - KE, PE = zeros(nlayers), zeros(nlayers-1) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - ψ1h, ψ2h = view(vars.ψh, :, :, 1), view(vars.ψh, :, :, 2) - - for j = 1:nlayers - CUDA.@allowscalar KE[j] = @views 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) - end - - PE = @views 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′ * parsevalsum(abs2.(ψ2h .- ψ1h), grid) - - return KE, PE -end - -function energies(vars, params::SingleLayerParams, grid, sol) - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - return 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h, grid) -end - -energies(prob) = energies(prob.vars, prob.params, prob.grid, prob.sol) - -""" - fluxes(vars, params, grid, sol) - fluxes(prob) - -Return the lateral eddy fluxes within each fluid layer, lateralfluxes``_1,...,``lateralfluxes``_n`` -and also the vertical eddy fluxes at each fluid interface, -verticalfluxes``_{3/2},...,``verticalfluxes``_{n-1/2}``, where ``n`` is the total number of layers in the fluid. -(When ``n=1``, only the lateral fluxes are returned.) - -The lateral eddy fluxes within the ``j``-th fluid layer are - -```math -\\textrm{lateralfluxes}_j = \\frac{H_j}{H} \\int U_j v_j ∂_y u_j -\\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n , -``` - -while the vertical eddy fluxes at the ``j+1/2``-th fluid interface (i.e., interface between -the ``j``-th and ``(j+1)``-th fluid layer) are - -```math -\\textrm{verticalfluxes}_{j+1/2} = \\int \\frac{f₀²}{g'_{j+1/2} H} (U_j - U_{j+1}) \\, -v_{j+1} ψ_{j} \\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n-1. -``` -""" -function fluxes(vars, params, grid, sol) - nlayers = numberoflayers(params) - - lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) - - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.H * params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - for j = 1:nlayers-1 - CUDA.@allowscalar verticalfluxes[j] = sum(@views @. params.f₀^2 / params.g′[j] * (params.U[: ,:, j] - params.U[:, :, j+1]) * vars.v[:, :, j+1] * vars.ψ[:, :, j]; dims=(1, 2))[1] - CUDA.@allowscalar verticalfluxes[j] *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - end - - return lateralfluxes, verticalfluxes -end - -function fluxes(vars, params::TwoLayerParams, grid, sol) - nlayers = numberoflayers(params) - - lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) - - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - @. lateralfluxes *= params.H - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - U₁, U₂ = view(params.U, :, :, 1), view(params.U, :, :, 2) - ψ₁ = view(vars.ψ, :, :, 1) - v₂ = view(vars.v, :, :, 2) - - verticalfluxes = sum(@views @. params.f₀^2 / params.g′ * (U₁ - U₂) * v₂ * ψ₁; dims=(1, 2)) - verticalfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - return lateralfluxes, verticalfluxes -end - -function fluxes(vars, params::SingleLayerParams, grid, sol) - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly) - - return lateralfluxes -end - -fluxes(prob) = fluxes(prob.vars, prob.params, prob.grid, prob.sol) - -end # module diff --git a/src/multilayerqg_LOCAL_444028.jl b/src/multilayerqg_LOCAL_444028.jl deleted file mode 100644 index e50206ad..00000000 --- a/src/multilayerqg_LOCAL_444028.jl +++ /dev/null @@ -1,1107 +0,0 @@ -module MultiLayerQG - -export - fwdtransform!, - invtransform!, - streamfunctionfrompv!, - pvfromstreamfunction!, - updatevars!, - - set_q!, - set_ψ!, - energies, - fluxes - -using - FFTW, - CUDA, - LinearAlgebra, - StaticArrays, - Reexport, - DocStringExtensions - -@reexport using FourierFlows - -using FourierFlows: parsevalsum, parsevalsum2, superzeros, plan_flows_rfft - -nothingfunction(args...) = nothing - -""" - Problem(nlayers :: Int, - dev = CPU(); - nx = 128, - ny = nx, - Lx = 2π, - Ly = Lx, - f₀ = 1.0, - β = 0.0, - g = 1.0, - U = zeros(nlayers), - H = 1/nlayers * ones(nlayers), - ρ = Array{Float64}(1:nlayers), - eta = nothing, - topographic_pv_gradient = (0, 0), - μ = 0, - ν = 0, - nν = 1, - dt = 0.01, - stepper = "RK4", - calcFq = nothingfunction, - stochastic = false, - linear = false, - aliased_fraction = 1/3, - T = Float64) - -Construct a multi-layer quasi-geostrophic problem with `nlayers` fluid layers on device `dev`. - -Arguments -========= -- `nlayers`: (required) Number of fluid layers. -- `dev`: (required) `CPU()` (default) or `GPU()`; computer architecture used to time-step `problem`. - -Keyword arguments -================= - - `nx`: Number of grid points in ``x``-domain. - - `ny`: Number of grid points in ``y``-domain. - - `Lx`: Extent of the ``x``-domain. - - `Ly`: Extent of the ``y``-domain. - - `f₀`: Constant planetary vorticity. - - `β`: Planetary vorticity ``y``-gradient. - - `g`: Gravitational acceleration constant. - - `U`: The imposed constant zonal flow ``U(y)`` in each fluid layer. - - `H`: Rest height of each fluid layer. - - `ρ`: Density of each fluid layer. - - `eta`: Periodic component of the topographic potential vorticity. - - `topographic_pv_gradient`: The ``(x, y)`` components of the topographic PV large-scale gradient. - - `μ`: Linear bottom drag coefficient. - - `ν`: Small-scale (hyper)-viscosity coefficient. - - `nν`: (Hyper)-viscosity order, `nν```≥ 1``. - - `dt`: Time-step. - - `stepper`: Time-stepping method. - - `calcF`: Function that calculates the Fourier transform of the forcing, ``F̂``. - - `stochastic`: `true` or `false` (default); boolean denoting whether `calcF` is temporally stochastic. - - `linear`: `true` or `false` (default); boolean denoting whether the linearized equations of motions are used. - - `aliased_fraction`: the fraction of high-wavenumbers that are zero-ed out by `dealias!()`. - - `T`: `Float32` or `Float64` (default); floating point type used for `problem` data. -""" -function Problem(nlayers::Int, # number of fluid layers - dev = CPU(); - # Numerical parameters - nx = 128, - ny = nx, - Lx = 2π, - Ly = Lx, - # Physical parameters - f₀ = 1.0, # Coriolis parameter - β = 0.0, # y-gradient of Coriolis parameter - g = 1.0, # gravitational constant - U = zeros(nlayers), # imposed zonal flow U(y) in each layer - H = 1/nlayers * ones(nlayers), # rest fluid height of each layer - ρ = Array{Float64}(1:nlayers), # density of each layer - eta = nothing, # periodic component of the topographic PV - topographic_pv_gradient = (0, 0), # tuple with the ``(x, y)`` components of topographic PV large-scale gradient - # Bottom Drag and/or (hyper)-viscosity - μ = 0, - ν = 0, - nν = 1, - # Timestepper and equation options - dt = 0.01, - stepper = "RK4", - calcFq = nothingfunction, - stochastic = false, - linear = false, - # Float type and dealiasing - aliased_fraction = 1/3, - T = Float64) - - if dev == GPU() && nlayers > 2 - @warn """MultiLayerQG module is not optimized on the GPU yet for configurations with - 3 fluid layers or more! - - See issues on Github at https://github.com/FourierFlows/GeophysicalFlows.jl/issues/112 - and https://github.com/FourierFlows/GeophysicalFlows.jl/issues/267. - - To use MultiLayerQG with 3 fluid layers or more we suggest, for now, to restrict running - on CPU.""" - end - - if nlayers == 1 - @warn """MultiLayerQG module does work for single-layer configuration but may not be as - optimized. We suggest using SingleLayerQG module for single-layer QG simulation unless - you have reasons to use MultiLayerQG in a single-layer configuration, e.g., you want to - compare solutions with varying number of fluid layers.""" - end - - # topographic PV - eta === nothing && (eta = zeros(dev, T, (nx, ny))) - - grid = TwoDGrid(dev; nx, Lx, ny, Ly, aliased_fraction, T) - - params = Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ, ν, nν, grid; calcFq) - - vars = calcFq == nothingfunction ? DecayingVars(grid, params) : (stochastic ? StochasticForcedVars(grid, params) : ForcedVars(grid, params)) - - equation = linear ? LinearEquation(params, grid) : Equation(params, grid) - - FourierFlows.Problem(equation, stepper, dt, grid, vars, params) -end - -""" - struct Params{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams - -The parameters for the `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct Params{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams - # prescribed params - "number of fluid layers" - nlayers :: Int - "gravitational constant" - g :: T - "constant planetary vorticity" - f₀ :: T - "planetary vorticity ``y``-gradient" - β :: T - "array with density of each fluid layer" - ρ :: Aphys3D - "array with rest height of each fluid layer" - H :: Aphys3D - "array with imposed constant zonal flow ``U(y)`` in each fluid layer" - U :: Aphys3D - "array containing the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear bottom drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "array with the reduced gravity constants for each fluid interface" - g′ :: Aphys1D - "array containing ``x``-gradient of PV due to topographic PV in each fluid layer" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV in each fluid layer" - Qy :: Aphys3D - "array containing coeffients for getting PV from streamfunction" - S :: Atrans4D - "array containing coeffients for inverting PV to streamfunction" - S⁻¹ :: Atrans4D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -""" - struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - -The parameters for the a single-layer `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - # prescribed params - "planetary vorticity ``y``-gradient" - β :: T - "array with imposed constant zonal flow ``U(y)``" - U :: Aphys3D - "array containing the periodic component of the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "array containing ``x``-gradient of PV due to topographic PV" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV" - Qy :: Aphys3D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -""" - struct TwoLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - -The parameters for the a two-layer `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct TwoLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - # prescribed params - "gravitational constant" - g :: T - "constant planetary vorticity" - f₀ :: T - "planetary vorticity ``y``-gradient" - β :: T - "array with density of each fluid layer" - ρ :: Aphys3D - "tuple with rest height of each fluid layer" - H :: Tuple - "array with imposed constant zonal flow ``U(y)`` in each fluid layer" - U :: Aphys3D - "array containing periodic component of the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear bottom drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "the reduced gravity constants for the fluid interface" - g′ :: T - "array containing ``x``-gradient of PV due to topographic PV in each fluid layer" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV in each fluid layer" - Qy :: Aphys3D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 1}) where TU - T = eltype(grid) - if length(U) == nlayers - U_2D = zeros(dev, T, (1, nlayers)) - U_2D[:] = U - U_2D = repeat(U_2D, outer=(grid.ny, 1)) - else - U_2D = zeros(dev, T, (grid.ny, 1)) - U_2D[:] = U - end - U_3D = zeros(dev, T, (1, grid.ny, nlayers)) - @views U_3D[1, :, :] = U_2D - return U_3D -end - -function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 2}) where TU - T = eltype(grid) - U_3D = zeros(dev, T, (1, grid.ny, nlayers)) - @views U_3D[1, :, :] = U - return U_3D -end - -function convert_U_to_U3D(dev, nlayers, grid, U::Number) - T = eltype(grid) - A = device_array(dev) - U_3D = reshape(repeat([T(U)], outer=(grid.ny, 1)), (1, grid.ny, nlayers)) - return A(U_3D) -end - -function Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ, ν, nν, grid; calcFq=nothingfunction, effort=FFTW.MEASURE) - dev = grid.device - T = eltype(grid) - A = device_array(dev) - - ny, nx = grid.ny , grid.nx - nkr, nl = grid.nkr, grid.nl - kr, l = grid.kr , grid.l - - U = convert_U_to_U3D(dev, nlayers, grid, U) - - Uyy = real.(ifft(-l.^2 .* fft(U))) - Uyy = CUDA.@allowscalar repeat(Uyy, outer=(nx, 1, 1)) - - # Calculate periodic components of the topographic PV gradients. - etah = rfft(A(eta)) - etax = irfft(im * kr .* etah, nx) # ∂η/∂x - etay = irfft(im * l .* etah, nx) # ∂η/∂y - - # Add topographic PV large-scale gradient - topographic_pv_gradient = T.(topographic_pv_gradient) - @. etax += topographic_pv_gradient[1] - @. etay += topographic_pv_gradient[2] - - Qx = zeros(dev, T, (nx, ny, nlayers)) - @views @. Qx[:, :, nlayers] += etax - - Qy = zeros(dev, T, (nx, ny, nlayers)) - Qy = T(β) .- Uyy # T(β) is needed to ensure that Qy remains same type as U - @views @. Qy[:, :, nlayers] += etay - - rfftplanlayered = plan_flows_rfft(A{T, 3}(undef, grid.nx, grid.ny, nlayers), [1, 2]; flags=effort) - - if nlayers==1 - return SingleLayerParams(T(β), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, Qx, Qy, rfftplanlayered) - - else # if nlayers≥2 - - ρ = reshape(T.(ρ), (1, 1, nlayers)) - H = reshape(T.(H), (1, 1, nlayers)) - - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[1] # reduced gravity at each interface - - Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) - Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) - - typeofSkl = SArray{Tuple{nlayers, nlayers}, T, 2, nlayers^2} # StaticArrays of type T and dims = (nlayers x nlayers) - - S = Array{typeofSkl, 2}(undef, (nkr, nl)) - calcS!(S, Fp, Fm, nlayers, grid) - - S⁻¹ = Array{typeofSkl, 2}(undef, (nkr, nl)) - calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) - - S, S⁻¹, Fp, Fm = A(S), A(S⁻¹), A(Fp), A(Fm) # convert to appropriate ArrayType - - CUDA.@allowscalar @views Qy[:, :, 1] = @. Qy[:, :, 1] - Fp[1] * (U[:, :, 2] - U[:, :, 1]) - for j = 2:nlayers-1 - CUDA.@allowscalar @views Qy[:, :, j] = @. Qy[:, :, j] + Fp[j] * (U[:, :, j+1] - U[:, :, j]) + Fm[j-1] * (U[:, :, j-1] - U[:, :, j]) - end - CUDA.@allowscalar @views Qy[:, :, nlayers] = @. Qy[:, :, nlayers] - Fm[nlayers-1] * (U[:, :, nlayers-1] - U[:, :, nlayers]) - - if nlayers==2 - return TwoLayerParams(T(g), T(f₀), T(β), A(ρ), (T(H[1]), T(H[2])), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, T(g′[1]), Qx, Qy, rfftplanlayered) - else # if nlayers>2 - return Params(nlayers, T(g), T(f₀), T(β), A(ρ), A(H), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, A(g′), Qx, Qy, S, S⁻¹, rfftplanlayered) - end - end -end - -numberoflayers(params) = params.nlayers -numberoflayers(::SingleLayerParams) = 1 -numberoflayers(::TwoLayerParams) = 2 - -# --------- -# Equations -# --------- - -""" - hyperviscosity(params, grid) - -Return the linear operator `L` that corresponds to (hyper)-viscosity of order ``n_ν`` with -coefficient ``ν`` for ``n`` fluid layers. -```math -L_j = - ν |𝐤|^{2 n_ν}, \\ j = 1, ...,n . -``` -""" -function hyperviscosity(params, grid) - dev = grid.device - T = eltype(grid) - - L = device_array(dev){T}(undef, (grid.nkr, grid.nl, numberoflayers(params))) - @. L = - params.ν * grid.Krsq^params.nν - @views @. L[1, 1, :] = 0 - - return L -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 -`hyperviscosity(params, grid)`. - -The nonlinear term is computed via function `calcNlinear!`. -""" -function LinearEquation(params, grid) - L = hyperviscosity(params, grid) - - return FourierFlows.Equation(L, calcNlinear!, grid) -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 -`hyperviscosity(params, grid)`. - -The nonlinear term is computed via function `calcN!`. -""" -function Equation(params, grid) - L = hyperviscosity(params, grid) - - return FourierFlows.Equation(L, calcN!, grid) -end - - -# ---- -# Vars -# ---- - -""" - struct Vars{Aphys, Atrans, F, P} <: AbstractVars - -The variables for multi-layer QG problem. - -$(FIELDS) -""" -struct Vars{Aphys, Atrans, F, P} <: AbstractVars - "relative vorticity + vortex stretching" - q :: Aphys - "streamfunction" - ψ :: Aphys - "x-component of velocity" - u :: Aphys - "y-component of velocity" - v :: Aphys - "Fourier transform of relative vorticity + vortex stretching" - qh :: Atrans - "Fourier transform of streamfunction" - ψh :: Atrans - "Fourier transform of ``x``-component of velocity" - uh :: Atrans - "Fourier transform of ``y``-component of velocity" - vh :: Atrans - "Fourier transform of forcing" - Fqh :: F - "`sol` at previous time-step" - prevsol :: P -end - -const DecayingVars = Vars{<:AbstractArray, <:AbstractArray, Nothing, Nothing} -const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing} -const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray} - -""" - DecayingVars(grid, params) - -Return the variables for an unforced multi-layer QG problem with `grid` and `params`. -""" -function DecayingVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, nothing, nothing) -end - -""" - ForcedVars(grid, params) - -Return the variables for a forced multi-layer QG problem with `grid` and `params`. -""" -function ForcedVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, nothing) -end - -""" - StochasticForcedVars(grid, params) - -Return the variables for a forced multi-layer QG problem with `grid` and `params`. -""" -function StochasticForcedVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh prevsol - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, prevsol) -end - -""" - fwdtransform!(varh, var, params) - -Compute the Fourier transform of `var` and store it in `varh`. -""" -fwdtransform!(varh, var, params::AbstractParams) = mul!(varh, params.rfftplan, var) - -""" - invtransform!(var, varh, params) - -Compute the inverse Fourier transform of `varh` and store it in `var`. -""" -invtransform!(var, varh, params::AbstractParams) = ldiv!(var, params.rfftplan, varh) - -""" - pvfromstreamfunction!(qh, ψh, params, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` in each layer using -`qh = params.S * ψh`. -""" -function pvfromstreamfunction!(qh, ψh, params, grid) - for j=1:grid.nl, i=1:grid.nkr - CUDA.@allowscalar @views qh[i, j, :] .= params.S[i, j] * ψh[i, j, :] - end - - return nothing -end - -""" - pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` for the special -case of a single fluid layer configuration. In this case, ``q̂ = - k² ψ̂``. -""" -function pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) - @. qh = -grid.Krsq * ψh - - return nothing -end - -""" - pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` for the special -case of a two fluid layer configuration. In this case we have, - -```math -q̂₁ = - k² ψ̂₁ + f₀² / (g′ H₁) * (ψ̂₂ - ψ̂₁) , -``` - -```math -q̂₂ = - k² ψ̂₂ + f₀² / (g′ H₂) * (ψ̂₁ - ψ̂₂) . -``` - -(Here, the PV-streamfunction relationship is hard-coded to avoid scalar operations -on the GPU.) -""" -function pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) - f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] - - ψ1h, ψ2h = view(ψh, :, :, 1), view(ψh, :, :, 2) - - @views @. qh[:, :, 1] = - grid.Krsq * ψ1h + f₀^2 / (g′ * H₁) * (ψ2h - ψ1h) - @views @. qh[:, :, 2] = - grid.Krsq * ψ2h + f₀^2 / (g′ * H₂) * (ψ1h - ψ2h) - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` in each layer from -`qh` using `ψh = params.S⁻¹ qh`. -""" -function streamfunctionfrompv!(ψh, qh, params, grid) - for j=1:grid.nl, i=1:grid.nkr - CUDA.@allowscalar @views ψh[i, j, :] .= params.S⁻¹[i, j] * qh[i, j, :] - end - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` for the special -case of a single fluid layer configuration. In this case, ``ψ̂ = - k⁻² q̂``. -""" -function streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) - @. ψh = -grid.invKrsq * qh - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` for the special -case of a two fluid layer configuration. In this case we have, - -```math -ψ̂₁ = - [k² q̂₁ + (f₀² / g′) (q̂₁ / H₂ + q̂₂ / H₁)] / Δ , -``` - -```math -ψ̂₂ = - [k² q̂₂ + (f₀² / g′) (q̂₁ / H₂ + q̂₂ / H₁)] / Δ , -``` - -where ``Δ = k² [k² + f₀² (H₁ + H₂) / (g′ H₁ H₂)]``. - -(Here, the PV-streamfunction relationship is hard-coded to avoid scalar operations -on the GPU.) -""" -function streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) - f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] - - q1h, q2h = view(qh, :, :, 1), view(qh, :, :, 2) - - @views @. ψh[:, :, 1] = - grid.Krsq * q1h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) - @views @. ψh[:, :, 2] = - grid.Krsq * q2h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) - - for j in 1:2 - @views @. ψh[:, :, j] *= grid.invKrsq / (grid.Krsq + f₀^2 / g′ * (H₁ + H₂) / (H₁ * H₂)) - end - - return nothing -end - -""" - calcS!(S, Fp, Fm, nlayers, grid) - -Construct the array ``𝕊``, which consists of `nlayer` x `nlayer` static arrays ``𝕊_𝐤`` that -relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``q̂_𝐤 = 𝕊_𝐤 ψ̂_𝐤``. -""" -function calcS!(S, Fp, Fm, nlayers, grid) - F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) - - for n=1:grid.nl, m=1:grid.nkr - k² = CUDA.@allowscalar grid.Krsq[m, n] - Skl = SMatrix{nlayers, nlayers}(- k² * I + F) - S[m, n] = Skl - end - - return nothing -end - -""" - calcS⁻¹!(S, Fp, Fm, nlayers, grid) - -Construct the array ``𝕊⁻¹``, which consists of `nlayer` x `nlayer` static arrays ``(𝕊_𝐤)⁻¹`` -that relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``ψ̂_𝐤 = (𝕊_𝐤)⁻¹ q̂_𝐤``. -""" -function calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) - F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) - - for n=1:grid.nl, m=1:grid.nkr - k² = CUDA.@allowscalar grid.Krsq[m, n] == 0 ? 1 : grid.Krsq[m, n] - Skl = - k² * I + F - S⁻¹[m, n] = SMatrix{nlayers, nlayers}(I / Skl) - end - - T = eltype(grid) - S⁻¹[1, 1] = SMatrix{nlayers, nlayers}(zeros(T, (nlayers, nlayers))) - - return nothing -end - - -# ------- -# Solvers -# ------- - -""" - calcN!(N, sol, t, clock, vars, params, grid) - -Compute the nonlinear term, that is the advection term, the bottom drag, and the forcing: - -```math -N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . -``` -""" -function calcN!(N, sol, t, clock, vars, params, grid) - nlayers = numberoflayers(params) - - dealias!(sol, grid) - - calcN_advection!(N, sol, vars, params, grid) - - @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag - - addforcing!(N, sol, t, clock, vars, params, grid) - - return nothing -end - -""" - calcNlinear!(N, sol, t, clock, vars, params, grid) - -Compute the nonlinear term of the linearized equations: - -```math -N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} + \\widehat{(∂_y ψ_j)(∂_x Q_j)} -- \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . -``` -""" -function calcNlinear!(N, sol, t, clock, vars, params, grid) - nlayers = numberoflayers(params) - - calcN_linearadvection!(N, sol, vars, params, grid) - @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag - addforcing!(N, sol, t, clock, vars, params, grid) - - return nothing -end - -""" - calcN_advection!(N, sol, vars, params, grid) - -Compute the advection term and stores it in `N`: - -```math -N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} . -``` -""" -function calcN_advection!(N, sol, vars, params, grid) - @. vars.qh = sol - - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.u, vars.uh, params) - @. vars.u += params.U # add the imposed zonal flow U - - uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables - @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x - fwdtransform!(uQxh, uQx, params) - @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} - - invtransform!(vars.v, vars.vh, params) - - vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables - @. vQy = vars.v * params.Qy # v*∂Q/∂y - fwdtransform!(vQyh, vQy, params) - @. N -= vQyh # -\hat{v*∂Q/∂y} - - invtransform!(vars.q, vars.qh, params) - - uq , vq = vars.u , vars.v # use vars.u and vars.v as scratch variables - uqh, vqh = vars.uh, vars.vh # use vars.uh and vars.vh as scratch variables - @. uq *= vars.q # (U+u)*q - @. vq *= vars.q # v*q - - fwdtransform!(uqh, uq, params) - fwdtransform!(vqh, vq, params) - - @. N -= im * grid.kr * uqh + im * grid.l * vqh # -\hat{∂[(U+u)q]/∂x} - \hat{∂[vq]/∂y} - - return nothing -end - - -""" - calcN_linearadvection!(N, sol, vars, params, grid) - -Compute the advection term of the linearized equations and stores it in `N`: - -```math -N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} . -``` -""" -function calcN_linearadvection!(N, sol, vars, params, grid) - @. vars.qh = sol - - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.u, vars.uh, params) - @. vars.u += params.U # add the imposed zonal flow U - uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables - @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x - fwdtransform!(uQxh, uQx, params) - @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} - - invtransform!(vars.v, vars.vh, params) - - vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables - - @. vQy = vars.v * params.Qy # v*∂Q/∂y - fwdtransform!(vQyh, vQy, params) - @. N -= vQyh # -\hat{v*∂Q/∂y} - - invtransform!(vars.q, vars.qh, params) - - @. vars.u = params.U - Uq , Uqh = vars.u , vars.uh # use vars.u and vars.uh as scratch variables - @. Uq *= vars.q # U*q - - fwdtransform!(Uqh, Uq, params) - - @. N -= im * grid.kr * Uqh # -\hat{∂[U*q]/∂x} - - return nothing -end - - -""" - addforcing!(N, sol, t, clock, vars, params, grid) - -When the problem includes forcing, calculate the forcing term ``F̂`` for each layer and add -it to the nonlinear term ``N``. -""" -addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing - -function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) - params.calcFq!(vars.Fqh, sol, t, clock, vars, params, grid) - @. N += vars.Fqh - - return nothing -end - - -# ---------------- -# Helper functions -# ---------------- - -""" - updatevars!(vars, params, grid, sol) - updatevars!(prob) - -Update all problem variables using `sol`. -""" -function updatevars!(vars, params, grid, sol) - dealias!(sol, grid) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.q, deepcopy(vars.qh), params) - invtransform!(vars.ψ, deepcopy(vars.ψh), params) - invtransform!(vars.u, deepcopy(vars.uh), params) - invtransform!(vars.v, deepcopy(vars.vh), params) - - return nothing -end - -updatevars!(prob) = updatevars!(prob.vars, prob.params, prob.grid, prob.sol) - - -""" - set_q!(sol, params, vars, grid, q) - set_q!(prob, q) - -Set the solution `prob.sol` as the transform of `q` and update variables. -""" -function set_q!(sol, params, vars, grid, q) - A = typeof(vars.q) - fwdtransform!(vars.qh, A(q), params) - @. vars.qh[1, 1, :] = 0 - @. sol = vars.qh - updatevars!(vars, params, grid, sol) - - return nothing -end - -function set_q!(sol, params::SingleLayerParams, vars, grid, q::AbstractArray{T, 2}) where T - A = typeof(vars.q[:, :, 1]) - q_3D = vars.q - @views q_3D[:, :, 1] = A(q) - set_q!(sol, params, vars, grid, q_3D) - - return nothing -end - -set_q!(prob, q) = set_q!(prob.sol, prob.params, prob.vars, prob.grid, q) - - -""" - set_ψ!(params, vars, grid, sol, ψ) - set_ψ!(prob, ψ) - -Set the solution `prob.sol` to the transform `qh` that corresponds to streamfunction `ψ` -and update variables. -""" -function set_ψ!(sol, params, vars, grid, ψ) - A = typeof(vars.q) - fwdtransform!(vars.ψh, A(ψ), params) - pvfromstreamfunction!(vars.qh, vars.ψh, params, grid) - invtransform!(vars.q, vars.qh, params) - - set_q!(sol, params, vars, grid, vars.q) - - return nothing -end - -function set_ψ!(sol, params::SingleLayerParams, vars, grid, ψ::AbstractArray{T, 2}) where T - A = typeof(vars.ψ[:, :, 1]) - ψ_3D = vars.ψ - @views ψ_3D[:, :, 1] = A(ψ) - - set_ψ!(sol, params, vars, grid, ψ_3D) - - return nothing -end - -set_ψ!(prob, ψ) = set_ψ!(prob.sol, prob.params, prob.vars, prob.grid, ψ) - - -""" - energies(vars, params, grid, sol) - energies(prob) - -Return the kinetic energy of each fluid layer KE``_1, ...,`` KE``_{n}``, and the -potential energy of each fluid interface PE``_{3/2}, ...,`` PE``_{n-1/2}``, where ``n`` -is the number of layers in the fluid. (When ``n=1``, only the kinetic energy is returned.) - -The kinetic energy at the ``j``-th fluid layer is - -```math -𝖪𝖤_j = \\frac{H_j}{H} \\int \\frac1{2} |{\\bf ∇} ψ_j|^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{H_j}{H} \\sum_{𝐤} |𝐤|² |ψ̂_j|², \\ j = 1, ..., n , -``` - -while the potential energy that corresponds to the interface ``j+1/2`` (i.e., the interface -between the ``j``-th and ``(j+1)``-th fluid layer) is - -```math -𝖯𝖤_{j+1/2} = \\int \\frac1{2} \\frac{f₀^2}{g'_{j+1/2} H} (ψ_j - ψ_{j+1})^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{f₀^2}{g'_{j+1/2} H} \\sum_{𝐤} |ψ̂_j - ψ̂_{j+1}|², \\ j = 1, ..., n-1 . -``` -""" -function energies(vars, params, grid, sol) - nlayers = numberoflayers(params) - KE, PE = zeros(nlayers), zeros(nlayers-1) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - for j = 1:nlayers - CUDA.@allowscalar KE[j] = 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) - end - - for j = 1:nlayers-1 - CUDA.@allowscalar PE[j] = 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′[j] * parsevalsum(abs2.(vars.ψh[:, :, j+1] .- vars.ψh[:, :, j]), grid) - end - - return KE, PE -end - -function energies(vars, params::TwoLayerParams, grid, sol) - nlayers = numberoflayers(params) - KE, PE = zeros(nlayers), zeros(nlayers-1) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - ψ1h, ψ2h = view(vars.ψh, :, :, 1), view(vars.ψh, :, :, 2) - - for j = 1:nlayers - CUDA.@allowscalar KE[j] = @views 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) - end - - PE = @views 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′ * parsevalsum(abs2.(ψ2h .- ψ1h), grid) - - return KE, PE -end - -function energies(vars, params::SingleLayerParams, grid, sol) - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - return 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h, grid) -end - -energies(prob) = energies(prob.vars, prob.params, prob.grid, prob.sol) - -""" - fluxes(vars, params, grid, sol) - fluxes(prob) - -Return the lateral eddy fluxes within each fluid layer, lateralfluxes``_1,...,``lateralfluxes``_n`` -and also the vertical eddy fluxes at each fluid interface, -verticalfluxes``_{3/2},...,``verticalfluxes``_{n-1/2}``, where ``n`` is the total number of layers in the fluid. -(When ``n=1``, only the lateral fluxes are returned.) - -The lateral eddy fluxes within the ``j``-th fluid layer are - -```math -\\textrm{lateralfluxes}_j = \\frac{H_j}{H} \\int U_j v_j ∂_y u_j -\\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n , -``` - -while the vertical eddy fluxes at the ``j+1/2``-th fluid interface (i.e., interface between -the ``j``-th and ``(j+1)``-th fluid layer) are - -```math -\\textrm{verticalfluxes}_{j+1/2} = \\int \\frac{f₀²}{g'_{j+1/2} H} (U_j - U_{j+1}) \\, -v_{j+1} ψ_{j} \\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n-1. -``` -""" -function fluxes(vars, params, grid, sol) - nlayers = numberoflayers(params) - - lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) - - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.H * params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - for j = 1:nlayers-1 - CUDA.@allowscalar verticalfluxes[j] = sum(@views @. params.f₀^2 / params.g′[j] * (params.U[: ,:, j] - params.U[:, :, j+1]) * vars.v[:, :, j+1] * vars.ψ[:, :, j]; dims=(1, 2))[1] - CUDA.@allowscalar verticalfluxes[j] *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - end - - return lateralfluxes, verticalfluxes -end - -function fluxes(vars, params::TwoLayerParams, grid, sol) - nlayers = numberoflayers(params) - - lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) - - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - @. lateralfluxes *= params.H - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - U₁, U₂ = view(params.U, :, :, 1), view(params.U, :, :, 2) - ψ₁ = view(vars.ψ, :, :, 1) - v₂ = view(vars.v, :, :, 2) - - verticalfluxes = sum(@views @. params.f₀^2 / params.g′ * (U₁ - U₂) * v₂ * ψ₁; dims=(1, 2)) - verticalfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - return lateralfluxes, verticalfluxes -end - -function fluxes(vars, params::SingleLayerParams, grid, sol) - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly) - - return lateralfluxes -end - -fluxes(prob) = fluxes(prob.vars, prob.params, prob.grid, prob.sol) - -end # module diff --git a/src/multilayerqg_REMOTE_444028.jl b/src/multilayerqg_REMOTE_444028.jl deleted file mode 100644 index bd49c747..00000000 --- a/src/multilayerqg_REMOTE_444028.jl +++ /dev/null @@ -1,1107 +0,0 @@ -module MultiLayerQG - -export - fwdtransform!, - invtransform!, - streamfunctionfrompv!, - pvfromstreamfunction!, - updatevars!, - - set_q!, - set_ψ!, - energies, - fluxes - -using - FFTW, - CUDA, - LinearAlgebra, - StaticArrays, - Reexport, - DocStringExtensions - -@reexport using FourierFlows - -using FourierFlows: parsevalsum, parsevalsum2, superzeros, plan_flows_rfft - -nothingfunction(args...) = nothing - -""" - Problem(nlayers :: Int, - dev = CPU(); - nx = 128, - ny = nx, - Lx = 2π, - Ly = Lx, - f₀ = 1.0, - β = 0.0, - g = 1.0, - U = zeros(nlayers), - H = 1/nlayers * ones(nlayers), - ρ = Array{Float64}(1:nlayers), - eta = nothing, - topographic_pv_gradient = (0, 0), - μ = 0, - ν = 0, - nν = 1, - dt = 0.01, - stepper = "RK4", - calcFq = nothingfunction, - stochastic = false, - linear = false, - aliased_fraction = 1/3, - T = Float64) - -Construct a multi-layer quasi-geostrophic problem with `nlayers` fluid layers on device `dev`. - -Arguments -========= -- `nlayers`: (required) Number of fluid layers. -- `dev`: (required) `CPU()` (default) or `GPU()`; computer architecture used to time-step `problem`. - -Keyword arguments -================= - - `nx`: Number of grid points in ``x``-domain. - - `ny`: Number of grid points in ``y``-domain. - - `Lx`: Extent of the ``x``-domain. - - `Ly`: Extent of the ``y``-domain. - - `f₀`: Constant planetary vorticity. - - `β`: Planetary vorticity ``y``-gradient. - - `g`: Gravitational acceleration constant. - - `U`: The imposed constant zonal flow ``U(y)`` in each fluid layer. - - `H`: Rest height of each fluid layer. - - `ρ`: Density of each fluid layer. - - `eta`: Periodic component of the topographic potential vorticity. - - `topographic_pv_gradient`: The ``(x, y)`` components of the topographic PV large-scale gradient. - - `μ`: Linear bottom drag coefficient. - - `ν`: Small-scale (hyper)-viscosity coefficient. - - `nν`: (Hyper)-viscosity order, `nν```≥ 1``. - - `dt`: Time-step. - - `stepper`: Time-stepping method. - - `calcF`: Function that calculates the Fourier transform of the forcing, ``F̂``. - - `stochastic`: `true` or `false` (default); boolean denoting whether `calcF` is temporally stochastic. - - `linear`: `true` or `false` (default); boolean denoting whether the linearized equations of motions are used. - - `aliased_fraction`: the fraction of high-wavenumbers that are zero-ed out by `dealias!()`. - - `T`: `Float32` or `Float64` (default); floating point type used for `problem` data. -""" -function Problem(nlayers::Int, # number of fluid layers - dev = CPU(); - # Numerical parameters - nx = 128, - ny = nx, - Lx = 2π, - Ly = Lx, - # Physical parameters - f₀ = 1.0, # Coriolis parameter - β = 0.0, # y-gradient of Coriolis parameter - g = 1.0, # gravitational constant - U = zeros(nlayers), # imposed zonal flow U(y) in each layer - H = 1/nlayers * ones(nlayers), # rest fluid height of each layer - ρ = Array{Float64}(1:nlayers), # density of each layer - eta = nothing, # periodic component of the topographic PV - topographic_pv_gradient = (0, 0), # tuple with the ``(x, y)`` components of topographic PV large-scale gradient - # Bottom Drag and/or (hyper)-viscosity - μ = 0, - ν = 0, - nν = 1, - # Timestepper and equation options - dt = 0.01, - stepper = "RK4", - calcFq = nothingfunction, - stochastic = false, - linear = false, - # Float type and dealiasing - aliased_fraction = 1/3, - T = Float64) - - if dev == GPU() && nlayers > 2 - @warn """MultiLayerQG module is not optimized on the GPU yet for configurations with - 3 fluid layers or more! - - See issues on Github at https://github.com/FourierFlows/GeophysicalFlows.jl/issues/112 - and https://github.com/FourierFlows/GeophysicalFlows.jl/issues/267. - - To use MultiLayerQG with 3 fluid layers or more we suggest, for now, to restrict running - on CPU.""" - end - - if nlayers == 1 - @warn """MultiLayerQG module does work for single-layer configuration but may not be as - optimized. We suggest using SingleLayerQG module for single-layer QG simulation unless - you have reasons to use MultiLayerQG in a single-layer configuration, e.g., you want to - compare solutions with varying number of fluid layers.""" - end - - # topographic PV - eta === nothing && (eta = zeros(dev, T, (nx, ny))) - - grid = TwoDGrid(dev; nx, Lx, ny, Ly, aliased_fraction, T) - - params = Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ, ν, nν, grid; calcFq) - - vars = calcFq == nothingfunction ? DecayingVars(grid, params) : (stochastic ? StochasticForcedVars(grid, params) : ForcedVars(grid, params)) - - equation = linear ? LinearEquation(params, grid) : Equation(params, grid) - - FourierFlows.Problem(equation, stepper, dt, grid, vars, params) -end - -""" - struct Params{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams - -The parameters for the `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct Params{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams - # prescribed params - "number of fluid layers" - nlayers :: Int - "gravitational constant" - g :: T - "constant planetary vorticity" - f₀ :: T - "planetary vorticity ``y``-gradient" - β :: T - "array with density of each fluid layer" - ρ :: Aphys3D - "array with rest height of each fluid layer" - H :: Aphys3D - "array with imposed constant zonal flow ``U(y)`` in each fluid layer" - U :: Aphys3D - "array containing the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear bottom drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "array with the reduced gravity constants for each fluid interface" - g′ :: Aphys1D - "array containing ``x``-gradient of PV due to topographic PV in each fluid layer" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV in each fluid layer" - Qy :: Aphys3D - "array containing coeffients for getting PV from streamfunction" - S :: Atrans4D - "array containing coeffients for inverting PV to streamfunction" - S⁻¹ :: Atrans4D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -""" - struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - -The parameters for the a single-layer `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - # prescribed params - "planetary vorticity ``y``-gradient" - β :: T - "array with imposed constant zonal flow ``U(y)``" - U :: Aphys3D - "array containing the periodic component of the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "array containing ``x``-gradient of PV due to topographic PV" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV" - Qy :: Aphys3D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -""" - struct TwoLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - -The parameters for the a two-layer `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct TwoLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - # prescribed params - "gravitational constant" - g :: T - "constant planetary vorticity" - f₀ :: T - "planetary vorticity ``y``-gradient" - β :: T - "array with density of each fluid layer" - ρ :: Aphys3D - "tuple with rest height of each fluid layer" - H :: Tuple - "array with imposed constant zonal flow ``U(y)`` in each fluid layer" - U :: Aphys3D - "array containing periodic component of the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear bottom drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "the reduced gravity constants for the fluid interface" - g′ :: T - "array containing ``x``-gradient of PV due to topographic PV in each fluid layer" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV in each fluid layer" - Qy :: Aphys3D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 1}) where TU - T = eltype(grid) - if length(U) == nlayers - U_2D = zeros(dev, T, (1, nlayers)) - U_2D[:] = U - U_2D = repeat(U_2D, outer=(grid.ny, 1)) - else - U_2D = zeros(dev, T, (grid.ny, 1)) - U_2D[:] = U - end - U_3D = zeros(dev, T, (1, grid.ny, nlayers)) - @views U_3D[1, :, :] = U_2D - return U_3D -end - -function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 2}) where TU - T = eltype(grid) - U_3D = zeros(dev, T, (1, grid.ny, nlayers)) - @views U_3D[1, :, :] = U - return U_3D -end - -function convert_U_to_U3D(dev, nlayers, grid, U::Number) - T = eltype(grid) - A = device_array(dev) - U_3D = reshape(repeat([T(U)], outer=(grid.ny, 1)), (1, grid.ny, nlayers)) - return A(U_3D) -end - -function Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ, ν, nν, grid; calcFq=nothingfunction, effort=FFTW.MEASURE) - dev = grid.device - T = eltype(grid) - A = device_array(dev) - - ny, nx = grid.ny , grid.nx - nkr, nl = grid.nkr, grid.nl - kr, l = grid.kr , grid.l - - U = convert_U_to_U3D(dev, nlayers, grid, U) - - Uyy = real.(ifft(-l.^2 .* fft(U))) - Uyy = CUDA.@allowscalar repeat(Uyy, outer=(nx, 1, 1)) - - # Calculate periodic components of the topographic PV gradients. - etah = rfft(A(eta)) - etax = irfft(im * kr .* etah, nx) # ∂η/∂x - etay = irfft(im * l .* etah, nx) # ∂η/∂y - - # Add topographic PV large-scale gradient - topographic_pv_gradient = T.(topographic_pv_gradient) - @. etax += topographic_pv_gradient[1] - @. etay += topographic_pv_gradient[2] - - Qx = zeros(dev, T, (nx, ny, nlayers)) - @views @. Qx[:, :, nlayers] += etax - - Qy = zeros(dev, T, (nx, ny, nlayers)) - Qy = T(β) .- Uyy # T(β) is needed to ensure that Qy remains same type as U - @views @. Qy[:, :, nlayers] += etay - - rfftplanlayered = plan_flows_rfft(A{T, 3}(undef, grid.nx, grid.ny, nlayers), [1, 2]; flags=effort) - - if nlayers==1 - return SingleLayerParams(T(β), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, Qx, Qy, rfftplanlayered) - - else # if nlayers≥2 - - ρ = reshape(T.(ρ), (1, 1, nlayers)) - H = reshape(T.(H), (1, 1, nlayers)) - - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface - - Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) - Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) - - typeofSkl = SArray{Tuple{nlayers, nlayers}, T, 2, nlayers^2} # StaticArrays of type T and dims = (nlayers x nlayers) - - S = Array{typeofSkl, 2}(undef, (nkr, nl)) - calcS!(S, Fp, Fm, nlayers, grid) - - S⁻¹ = Array{typeofSkl, 2}(undef, (nkr, nl)) - calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) - - S, S⁻¹, Fp, Fm = A(S), A(S⁻¹), A(Fp), A(Fm) # convert to appropriate ArrayType - - CUDA.@allowscalar @views Qy[:, :, 1] = @. Qy[:, :, 1] - Fp[1] * (U[:, :, 2] - U[:, :, 1]) - for j = 2:nlayers-1 - CUDA.@allowscalar @views Qy[:, :, j] = @. Qy[:, :, j] + Fp[j] * (U[:, :, j+1] - U[:, :, j]) + Fm[j-1] * (U[:, :, j-1] - U[:, :, j]) - end - CUDA.@allowscalar @views Qy[:, :, nlayers] = @. Qy[:, :, nlayers] - Fm[nlayers-1] * (U[:, :, nlayers-1] - U[:, :, nlayers]) - - if nlayers==2 - return TwoLayerParams(T(g), T(f₀), T(β), A(ρ), (T(H[1]), T(H[2])), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, T(g′[1]), Qx, Qy, rfftplanlayered) - else # if nlayers>2 - return Params(nlayers, T(g), T(f₀), T(β), A(ρ), A(H), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, A(g′), Qx, Qy, S, S⁻¹, rfftplanlayered) - end - end -end - -numberoflayers(params) = params.nlayers -numberoflayers(::SingleLayerParams) = 1 -numberoflayers(::TwoLayerParams) = 2 - -# --------- -# Equations -# --------- - -""" - hyperviscosity(params, grid) - -Return the linear operator `L` that corresponds to (hyper)-viscosity of order ``n_ν`` with -coefficient ``ν`` for ``n`` fluid layers. -```math -L_j = - ν |𝐤|^{2 n_ν}, \\ j = 1, ...,n . -``` -""" -function hyperviscosity(params, grid) - dev = grid.device - T = eltype(grid) - - L = device_array(dev){T}(undef, (grid.nkr, grid.nl, numberoflayers(params))) - @. L = - params.ν * grid.Krsq^params.nν - @views @. L[1, 1, :] = 0 - - return L -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 -`hyperviscosity(params, grid)`. - -The nonlinear term is computed via function `calcNlinear!`. -""" -function LinearEquation(params, grid) - L = hyperviscosity(params, grid) - - return FourierFlows.Equation(L, calcNlinear!, grid) -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 -`hyperviscosity(params, grid)`. - -The nonlinear term is computed via function `calcN!`. -""" -function Equation(params, grid) - L = hyperviscosity(params, grid) - - return FourierFlows.Equation(L, calcN!, grid) -end - - -# ---- -# Vars -# ---- - -""" - struct Vars{Aphys, Atrans, F, P} <: AbstractVars - -The variables for multi-layer QG problem. - -$(FIELDS) -""" -struct Vars{Aphys, Atrans, F, P} <: AbstractVars - "relative vorticity + vortex stretching" - q :: Aphys - "streamfunction" - ψ :: Aphys - "x-component of velocity" - u :: Aphys - "y-component of velocity" - v :: Aphys - "Fourier transform of relative vorticity + vortex stretching" - qh :: Atrans - "Fourier transform of streamfunction" - ψh :: Atrans - "Fourier transform of ``x``-component of velocity" - uh :: Atrans - "Fourier transform of ``y``-component of velocity" - vh :: Atrans - "Fourier transform of forcing" - Fqh :: F - "`sol` at previous time-step" - prevsol :: P -end - -const DecayingVars = Vars{<:AbstractArray, <:AbstractArray, Nothing, Nothing} -const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing} -const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray} - -""" - DecayingVars(grid, params) - -Return the variables for an unforced multi-layer QG problem with `grid` and `params`. -""" -function DecayingVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, nothing, nothing) -end - -""" - ForcedVars(grid, params) - -Return the variables for a forced multi-layer QG problem with `grid` and `params`. -""" -function ForcedVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, nothing) -end - -""" - StochasticForcedVars(grid, params) - -Return the variables for a forced multi-layer QG problem with `grid` and `params`. -""" -function StochasticForcedVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh prevsol - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, prevsol) -end - -""" - fwdtransform!(varh, var, params) - -Compute the Fourier transform of `var` and store it in `varh`. -""" -fwdtransform!(varh, var, params::AbstractParams) = mul!(varh, params.rfftplan, var) - -""" - invtransform!(var, varh, params) - -Compute the inverse Fourier transform of `varh` and store it in `var`. -""" -invtransform!(var, varh, params::AbstractParams) = ldiv!(var, params.rfftplan, varh) - -""" - pvfromstreamfunction!(qh, ψh, params, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` in each layer using -`qh = params.S * ψh`. -""" -function pvfromstreamfunction!(qh, ψh, params, grid) - for j=1:grid.nl, i=1:grid.nkr - CUDA.@allowscalar @views qh[i, j, :] .= params.S[i, j] * ψh[i, j, :] - end - - return nothing -end - -""" - pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` for the special -case of a single fluid layer configuration. In this case, ``q̂ = - k² ψ̂``. -""" -function pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) - @. qh = -grid.Krsq * ψh - - return nothing -end - -""" - pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` for the special -case of a two fluid layer configuration. In this case we have, - -```math -q̂₁ = - k² ψ̂₁ + f₀² / (g′ H₁) * (ψ̂₂ - ψ̂₁) , -``` - -```math -q̂₂ = - k² ψ̂₂ + f₀² / (g′ H₂) * (ψ̂₁ - ψ̂₂) . -``` - -(Here, the PV-streamfunction relationship is hard-coded to avoid scalar operations -on the GPU.) -""" -function pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) - f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] - - ψ1h, ψ2h = view(ψh, :, :, 1), view(ψh, :, :, 2) - - @views @. qh[:, :, 1] = - grid.Krsq * ψ1h + f₀^2 / (g′ * H₁) * (ψ2h - ψ1h) - @views @. qh[:, :, 2] = - grid.Krsq * ψ2h + f₀^2 / (g′ * H₂) * (ψ1h - ψ2h) - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` in each layer from -`qh` using `ψh = params.S⁻¹ qh`. -""" -function streamfunctionfrompv!(ψh, qh, params, grid) - for j=1:grid.nl, i=1:grid.nkr - CUDA.@allowscalar @views ψh[i, j, :] .= params.S⁻¹[i, j] * qh[i, j, :] - end - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` for the special -case of a single fluid layer configuration. In this case, ``ψ̂ = - k⁻² q̂``. -""" -function streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) - @. ψh = -grid.invKrsq * qh - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` for the special -case of a two fluid layer configuration. In this case we have, - -```math -ψ̂₁ = - [k² q̂₁ + (f₀² / g′) (q̂₁ / H₂ + q̂₂ / H₁)] / Δ , -``` - -```math -ψ̂₂ = - [k² q̂₂ + (f₀² / g′) (q̂₁ / H₂ + q̂₂ / H₁)] / Δ , -``` - -where ``Δ = k² [k² + f₀² (H₁ + H₂) / (g′ H₁ H₂)]``. - -(Here, the PV-streamfunction relationship is hard-coded to avoid scalar operations -on the GPU.) -""" -function streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) - f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] - - q1h, q2h = view(qh, :, :, 1), view(qh, :, :, 2) - - @views @. ψh[:, :, 1] = - grid.Krsq * q1h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) - @views @. ψh[:, :, 2] = - grid.Krsq * q2h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) - - for j in 1:2 - @views @. ψh[:, :, j] *= grid.invKrsq / (grid.Krsq + f₀^2 / g′ * (H₁ + H₂) / (H₁ * H₂)) - end - - return nothing -end - -""" - calcS!(S, Fp, Fm, nlayers, grid) - -Construct the array ``𝕊``, which consists of `nlayer` x `nlayer` static arrays ``𝕊_𝐤`` that -relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``q̂_𝐤 = 𝕊_𝐤 ψ̂_𝐤``. -""" -function calcS!(S, Fp, Fm, nlayers, grid) - F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) - - for n=1:grid.nl, m=1:grid.nkr - k² = CUDA.@allowscalar grid.Krsq[m, n] - Skl = SMatrix{nlayers, nlayers}(- k² * I + F) - S[m, n] = Skl - end - - return nothing -end - -""" - calcS⁻¹!(S, Fp, Fm, nlayers, grid) - -Construct the array ``𝕊⁻¹``, which consists of `nlayer` x `nlayer` static arrays ``(𝕊_𝐤)⁻¹`` -that relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``ψ̂_𝐤 = (𝕊_𝐤)⁻¹ q̂_𝐤``. -""" -function calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) - F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) - - for n=1:grid.nl, m=1:grid.nkr - k² = CUDA.@allowscalar grid.Krsq[m, n] == 0 ? 1 : grid.Krsq[m, n] - Skl = - k² * I + F - S⁻¹[m, n] = SMatrix{nlayers, nlayers}(I / Skl) - end - - T = eltype(grid) - S⁻¹[1, 1] = SMatrix{nlayers, nlayers}(zeros(T, (nlayers, nlayers))) - - return nothing -end - - -# ------- -# Solvers -# ------- - -""" - calcN!(N, sol, t, clock, vars, params, grid) - -Compute the nonlinear term, that is the advection term, the bottom drag, and the forcing: - -```math -N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . -``` -""" -function calcN!(N, sol, t, clock, vars, params, grid) - nlayers = numberoflayers(params) - - dealias!(sol, grid) - - calcN_advection!(N, sol, vars, params, grid) - - @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag - - addforcing!(N, sol, t, clock, vars, params, grid) - - return nothing -end - -""" - calcNlinear!(N, sol, t, clock, vars, params, grid) - -Compute the nonlinear term of the linearized equations: - -```math -N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} + \\widehat{(∂_y ψ_j)(∂_x Q_j)} -- \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . -``` -""" -function calcNlinear!(N, sol, t, clock, vars, params, grid) - nlayers = numberoflayers(params) - - calcN_linearadvection!(N, sol, vars, params, grid) - @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag - addforcing!(N, sol, t, clock, vars, params, grid) - - return nothing -end - -""" - calcN_advection!(N, sol, vars, params, grid) - -Compute the advection term and stores it in `N`: - -```math -N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} . -``` -""" -function calcN_advection!(N, sol, vars, params, grid) - @. vars.qh = sol - - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.u, vars.uh, params) - @. vars.u += params.U # add the imposed zonal flow U - - uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables - @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x - fwdtransform!(uQxh, uQx, params) - @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} - - invtransform!(vars.v, vars.vh, params) - - vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables - @. vQy = vars.v * params.Qy # v*∂Q/∂y - fwdtransform!(vQyh, vQy, params) - @. N -= vQyh # -\hat{v*∂Q/∂y} - - invtransform!(vars.q, vars.qh, params) - - uq , vq = vars.u , vars.v # use vars.u and vars.v as scratch variables - uqh, vqh = vars.uh, vars.vh # use vars.uh and vars.vh as scratch variables - @. uq *= vars.q # (U+u)*q - @. vq *= vars.q # v*q - - fwdtransform!(uqh, uq, params) - fwdtransform!(vqh, vq, params) - - @. N -= im * grid.kr * uqh + im * grid.l * vqh # -\hat{∂[(U+u)q]/∂x} - \hat{∂[vq]/∂y} - - return nothing -end - - -""" - calcN_linearadvection!(N, sol, vars, params, grid) - -Compute the advection term of the linearized equations and stores it in `N`: - -```math -N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} . -``` -""" -function calcN_linearadvection!(N, sol, vars, params, grid) - @. vars.qh = sol - - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.u, vars.uh, params) - @. vars.u += params.U # add the imposed zonal flow U - uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables - @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x - fwdtransform!(uQxh, uQx, params) - @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} - - invtransform!(vars.v, vars.vh, params) - - vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables - - @. vQy = vars.v * params.Qy # v*∂Q/∂y - fwdtransform!(vQyh, vQy, params) - @. N -= vQyh # -\hat{v*∂Q/∂y} - - invtransform!(vars.q, vars.qh, params) - - @. vars.u = params.U - Uq , Uqh = vars.u , vars.uh # use vars.u and vars.uh as scratch variables - @. Uq *= vars.q # U*q - - fwdtransform!(Uqh, Uq, params) - - @. N -= im * grid.kr * Uqh # -\hat{∂[U*q]/∂x} - - return nothing -end - - -""" - addforcing!(N, sol, t, clock, vars, params, grid) - -When the problem includes forcing, calculate the forcing term ``F̂`` for each layer and add -it to the nonlinear term ``N``. -""" -addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing - -function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) - params.calcFq!(vars.Fqh, sol, t, clock, vars, params, grid) - @. N += vars.Fqh - - return nothing -end - - -# ---------------- -# Helper functions -# ---------------- - -""" - updatevars!(vars, params, grid, sol) - updatevars!(prob) - -Update all problem variables using `sol`. -""" -function updatevars!(vars, params, grid, sol) - dealias!(sol, grid) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.q, deepcopy(vars.qh), params) - invtransform!(vars.ψ, deepcopy(vars.ψh), params) - invtransform!(vars.u, deepcopy(vars.uh), params) - invtransform!(vars.v, deepcopy(vars.vh), params) - - return nothing -end - -updatevars!(prob) = updatevars!(prob.vars, prob.params, prob.grid, prob.sol) - - -""" - set_q!(sol, params, vars, grid, q) - set_q!(prob, q) - -Set the solution `prob.sol` as the transform of `q` and update variables. -""" -function set_q!(sol, params, vars, grid, q) - A = typeof(vars.q) - fwdtransform!(vars.qh, A(q), params) - @. vars.qh[1, 1, :] = 0 - @. sol = vars.qh - updatevars!(vars, params, grid, sol) - - return nothing -end - -function set_q!(sol, params::SingleLayerParams, vars, grid, q::AbstractArray{T, 2}) where T - A = typeof(vars.q[:, :, 1]) - q_3D = vars.q - @views q_3D[:, :, 1] = A(q) - set_q!(sol, params, vars, grid, q_3D) - - return nothing -end - -set_q!(prob, q) = set_q!(prob.sol, prob.params, prob.vars, prob.grid, q) - - -""" - set_ψ!(params, vars, grid, sol, ψ) - set_ψ!(prob, ψ) - -Set the solution `prob.sol` to the transform `qh` that corresponds to streamfunction `ψ` -and update variables. -""" -function set_ψ!(sol, params, vars, grid, ψ) - A = typeof(vars.q) - fwdtransform!(vars.ψh, A(ψ), params) - pvfromstreamfunction!(vars.qh, vars.ψh, params, grid) - invtransform!(vars.q, vars.qh, params) - - set_q!(sol, params, vars, grid, vars.q) - - return nothing -end - -function set_ψ!(sol, params::SingleLayerParams, vars, grid, ψ::AbstractArray{T, 2}) where T - A = typeof(vars.ψ[:, :, 1]) - ψ_3D = vars.ψ - @views ψ_3D[:, :, 1] = A(ψ) - - set_ψ!(sol, params, vars, grid, ψ_3D) - - return nothing -end - -set_ψ!(prob, ψ) = set_ψ!(prob.sol, prob.params, prob.vars, prob.grid, ψ) - - -""" - energies(vars, params, grid, sol) - energies(prob) - -Return the kinetic energy of each fluid layer KE``_1, ...,`` KE``_{n}``, and the -potential energy of each fluid interface PE``_{3/2}, ...,`` PE``_{n-1/2}``, where ``n`` -is the number of layers in the fluid. (When ``n=1``, only the kinetic energy is returned.) - -The kinetic energy at the ``j``-th fluid layer is - -```math -𝖪𝖤_j = \\frac{H_j}{H} \\int \\frac1{2} |{\\bf ∇} ψ_j|^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{H_j}{H} \\sum_{𝐤} |𝐤|² |ψ̂_j|², \\ j = 1, ..., n , -``` - -while the potential energy that corresponds to the interface ``j+1/2`` (i.e., the interface -between the ``j``-th and ``(j+1)``-th fluid layer) is - -```math -𝖯𝖤_{j+1/2} = \\int \\frac1{2} \\frac{f₀^2}{g'_{j+1/2} H} (ψ_j - ψ_{j+1})^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{f₀^2}{g'_{j+1/2} H} \\sum_{𝐤} |ψ̂_j - ψ̂_{j+1}|², \\ j = 1, ..., n-1 . -``` -""" -function energies(vars, params, grid, sol) - nlayers = numberoflayers(params) - KE, PE = zeros(nlayers), zeros(nlayers-1) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - for j = 1:nlayers - CUDA.@allowscalar KE[j] = 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) - end - - for j = 1:nlayers-1 - CUDA.@allowscalar PE[j] = 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′[j] * parsevalsum(abs2.(vars.ψh[:, :, j+1] .- vars.ψh[:, :, j]), grid) - end - - return KE, PE -end - -function energies(vars, params::TwoLayerParams, grid, sol) - nlayers = numberoflayers(params) - KE, PE = zeros(nlayers), zeros(nlayers-1) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - ψ1h, ψ2h = view(vars.ψh, :, :, 1), view(vars.ψh, :, :, 2) - - for j = 1:nlayers - CUDA.@allowscalar KE[j] = @views 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) - end - - PE = @views 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′ * parsevalsum(abs2.(ψ2h .- ψ1h), grid) - - return KE, PE -end - -function energies(vars, params::SingleLayerParams, grid, sol) - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - return 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h, grid) -end - -energies(prob) = energies(prob.vars, prob.params, prob.grid, prob.sol) - -""" - fluxes(vars, params, grid, sol) - fluxes(prob) - -Return the lateral eddy fluxes within each fluid layer, lateralfluxes``_1,...,``lateralfluxes``_n`` -and also the vertical eddy fluxes at each fluid interface, -verticalfluxes``_{3/2},...,``verticalfluxes``_{n-1/2}``, where ``n`` is the total number of layers in the fluid. -(When ``n=1``, only the lateral fluxes are returned.) - -The lateral eddy fluxes within the ``j``-th fluid layer are - -```math -\\textrm{lateralfluxes}_j = \\frac{H_j}{H} \\int U_j v_j ∂_y u_j -\\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n , -``` - -while the vertical eddy fluxes at the ``j+1/2``-th fluid interface (i.e., interface between -the ``j``-th and ``(j+1)``-th fluid layer) are - -```math -\\textrm{verticalfluxes}_{j+1/2} = \\int \\frac{f₀²}{g'_{j+1/2} H} (U_j - U_{j+1}) \\, -v_{j+1} ψ_{j} \\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n-1. -``` -""" -function fluxes(vars, params, grid, sol) - nlayers = numberoflayers(params) - - lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) - - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.H * params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - for j = 1:nlayers-1 - CUDA.@allowscalar verticalfluxes[j] = sum(@views @. params.f₀^2 / params.g′[j] * (params.U[: ,:, j] - params.U[:, :, j+1]) * vars.v[:, :, j+1] * vars.ψ[:, :, j]; dims=(1, 2))[1] - CUDA.@allowscalar verticalfluxes[j] *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - end - - return lateralfluxes, verticalfluxes -end - -function fluxes(vars, params::TwoLayerParams, grid, sol) - nlayers = numberoflayers(params) - - lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) - - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - @. lateralfluxes *= params.H - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - U₁, U₂ = view(params.U, :, :, 1), view(params.U, :, :, 2) - ψ₁ = view(vars.ψ, :, :, 1) - v₂ = view(vars.v, :, :, 2) - - verticalfluxes = sum(@views @. params.f₀^2 / params.g′ * (U₁ - U₂) * v₂ * ψ₁; dims=(1, 2)) - verticalfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - return lateralfluxes, verticalfluxes -end - -function fluxes(vars, params::SingleLayerParams, grid, sol) - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly) - - return lateralfluxes -end - -fluxes(prob) = fluxes(prob.vars, prob.params, prob.grid, prob.sol) - -end # module From 40691ca68908a7e9b83c6606a2d277f79ba2f6c4 Mon Sep 17 00:00:00 2001 From: Matt Lobo Date: Fri, 9 Jun 2023 16:07:33 -0400 Subject: [PATCH 11/16] adding examples for GFjl issue --- examples/matts_examples/GFjl_bug_plots.jl | 102 ++ examples/matts_examples/mjcl_stab.jl | 250 +++++ examples/matts_examples/phillips_stab.jl | 316 ++++++ examples/matts_examples/prof_stab.jl | 89 ++ examples/matts_examples/sphere_ex.jl | 33 + src/multilayerqg_no_fixes.jl | 1107 +++++++++++++++++++++ src/multilayerqg_sign_fix_only.jl | 1107 +++++++++++++++++++++ 7 files changed, 3004 insertions(+) create mode 100644 examples/matts_examples/GFjl_bug_plots.jl create mode 100644 examples/matts_examples/mjcl_stab.jl create mode 100644 examples/matts_examples/phillips_stab.jl create mode 100644 examples/matts_examples/prof_stab.jl create mode 100644 examples/matts_examples/sphere_ex.jl create mode 100644 src/multilayerqg_no_fixes.jl create mode 100644 src/multilayerqg_sign_fix_only.jl diff --git a/examples/matts_examples/GFjl_bug_plots.jl b/examples/matts_examples/GFjl_bug_plots.jl new file mode 100644 index 00000000..1d84f861 --- /dev/null +++ b/examples/matts_examples/GFjl_bug_plots.jl @@ -0,0 +1,102 @@ +# showing how two different bug fixes in GeophysicalFlows.jl +# modify vertical profile of interior PV + +## load packages + +using GeophysicalFlows, CairoMakie, Printf, FFTW, LinearAlgebra, Statistics + +here = "/home/matt/Desktop/research/QG/julia_QG/" +include(here*"mjcl_stab.jl") + +## build basic model + +dev = CPU() # device (CPU) + +# numerical params +Ny = Nx = n = 128 # 2D resolution = n² +stepper = "FilteredRK4" # time stepping scheme +dt = 1e-2 # time step +nsteps = 60000 # total number of time-steps 20000 +nsubs = 50 # number of time-steps for plotting (nsteps must be multiple of nsubs) + +# physical params +L = 2π # domain size +μ = 5e-2 # bottom drag +beta = β = 0. # no gradient of planetary PV + +nlayers = 6 # number of layers +f0 = f₀ = 1. # 0.0001236812857687059 # coriolis param [s^-1] +g = 9.81 # gravity + +H = ones(nlayers) # even depths +rho = ρ = collect(range(4.0,5.0,nlayers)) # constant N^2 +U = collect(range(1.0,0.0,nlayers)) # constant shear +V = zeros(nlayers) # zonal flow only + +## setting up problems +# setting up the problen with bug fixes (denoted ``f'') +prob_f = MultiLayerQG.Problem(nlayers, dev; nx=n, Lx=L, f₀, g, H, ρ, U, μ, β, + dt, stepper, aliased_fraction=0) + +sol_f, clock_f, params_f, vars_f, grid_f = prob_f.sol, prob_f.clock, prob_f.params, prob_f.vars, prob_f.grid + +# setting up the problen with bug fixes (denoted ``sf'') +include(here*"../gfjl/GeophysicalFlows.jl/src/multilayerqg_sign_fix_only.jl") + +prob_sf = MultiLayerQG_sf.Problem_sf(nlayers, dev; nx=n, Lx=L, f₀, g, H, ρ, U, μ, β, + dt, stepper, aliased_fraction=0) + +sol_sf, clock_sf, params_sf, vars_sf, grid_sf = prob_sf.sol, prob_sf.clock, prob_sf.params, prob_sf.vars, prob_sf.grid + +# setting up the problen with no bug fixes (denoted ``nf'') +include(here*"../gfjl/GeophysicalFlows.jl/src/multilayerqg_no_fixes.jl") + +prob_nf = MultiLayerQG_nf.Problem_nf(nlayers, dev; nx=n, Lx=L, f₀, g, H, ρ, U, μ, β, + dt, stepper, aliased_fraction=0) + +sol_nf, clock_nf, params_nf, vars_nf, grid_nf = prob_nf.sol, prob_nf.clock, prob_nf.params, prob_nf.vars, prob_nf.grid + +# linear stability analysis +Lx, Ly = grid_f.Lx, grid_f.Ly + +eta = 0 + +eve,eva,max_eve_f,max_eva_f,k_x,k_y,qx,qy = lin_stab(U,V,beta,eta,Nx,Ny,rho,f0,Lx,Ly,params_f.Qy) +eve,eva,max_eve_sf,max_eva_sf,k_x,k_y,qx,qy = lin_stab(U,V,beta,eta,Nx,Ny,rho,f0,Lx,Ly,params_sf.Qy) +eve,eva,max_eve_nf,max_eva_nf,k_x,k_y,qx,qy = lin_stab(U,V,beta,eta,Nx,Ny,rho,f0,Lx,Ly,params_nf.Qy) + +## plotting differences +z = -cumsum(H) + +fig = Figure(resolution=(1000, 600)) + +ax1_kwargs = (xlabel = "Qy", + ylabel = "z", + aspect = 1.) + +ax1 = Axis(fig[1, 1]; title = "Qy", ax1_kwargs...) + +ylims!(ax1,minimum(z), maximum(z)) + +lines!(ax1,params_nf.Qy[1,1,:],z,linewidth=3.,label="no fix") +lines!(ax1,params_sf.Qy[1,1,:],z,linewidth=3.,label="sign fix only") +lines!(ax1,params_f.Qy[1,1,:],z,linewidth=3.,label="sign fix and g' fix") + +axislegend(position=:lt) + +ax2_kwargs = (xlabel = "|ψ|", + ylabel = "z", + aspect = 1.) + +ax2 = Axis(fig[1, 2]; title = "|ψ|", ax2_kwargs...) + +ylims!(ax2,minimum(z), maximum(z)) + +lines!(ax2,abs.(max_eve_nf),z,linewidth=3.,label="no fix") +lines!(ax2,abs.(max_eve_sf),z,linewidth=3.,label="sign fix only") +lines!(ax2,abs.(max_eve_f),z,linewidth=3.,label="sign fix and g' fix") + +axislegend(position=:lt) + + + diff --git a/examples/matts_examples/mjcl_stab.jl b/examples/matts_examples/mjcl_stab.jl new file mode 100644 index 00000000..48e168d3 --- /dev/null +++ b/examples/matts_examples/mjcl_stab.jl @@ -0,0 +1,250 @@ +## define stability functions + +function lin_stab(U::Vector{Float64},V::Vector{Float64},beta,eta,Nx::Int64,Ny::Int64,rho::Vector{Float64},f0::Float64,Lx::Float64,Ly::Float64) + # U: (Nx x Nz) vector of zonal mean background velocity + # V: (Ny x Nz) vector of meridional mean background velocity + # beta: + # + + Nz = length(rho) + # define wavenumbers + k_x = reshape(fftfreq(Nx, 2π/Lx*Nx),(1,Nx)) + k_y = reshape(fftfreq(Ny, 2π/Ly*Ny),(1,Ny)) + + # k_x,k_y = wavenumber_grid(Nx,Ny,Lx,Ly) + + # k2 = k_x.^2 + k_y.^2 # this is for an isotropic wavenumber grid only (i.e. Nx=Ny) + + # define stretching matrix + S = calc_stretching_mat(Nz,rho,f0,H,rho[1]) + + # change dimensions of U and V to match domain size + U2 = zeros(1,Nz); U2[:] = U; U2 = repeat(U2,outer=(Ny,1)) + U = zeros(1,Ny,Nz); U[1,:,:] = U2 + + V2 = zeros(1,Nz); V2[:] = V; V2 = repeat(V2,outer=(Nx,1)) + V = zeros(1,Nx,Nz); V[1,:,:] = V2 + + # define background QG PV gradients + Qy = calc_PV_grad_y(U,beta,eta,Ny,Nz,k_y,S) + Qx = calc_PV_grad_x(V,eta,Nx,Nz,k_x,S) + + # perform linear stability analysis + evecs_all,evals_all = calc_lin_stab(Qy,Qx,U,V,S,k_x,k_y,Nz) + + # keep largest growth rates per wavenumber + evecs,evals,max_evec,max_eval = find_growth(evecs_all,evals_all,Nx,Ny,Nz) + + # def rad + r_d = sqrt(gp(rho[1:2],rho[1])*H[1])/f0 + + return fftshift(evecs),fftshift(evals),max_evec,max_eval,fftshift(k_x),fftshift(k_y),mean(Qx[1,:,:],dims=1),mean(Qy[1,:,:],dims=1) +end + +function lin_stab(U::Vector{Float64},V::Vector{Float64},beta,eta,Nx::Int64,Ny::Int64,rho::Vector{Float64},f0::Float64,Lx::Float64,Ly::Float64,Qy) + # Takes Qy as arg + # U: (Nx x Nz) vector of zonal mean background velocity + # V: (Ny x Nz) vector of meridional mean background velocity + # beta: + # + + Nz = length(rho) + # define wavenumbers + k_x = reshape(fftfreq(Nx, 2π/Lx*Nx),(1,Nx)) + k_y = reshape(fftfreq(Ny, 2π/Ly*Ny),(1,Ny)) + + # k_x,k_y = wavenumber_grid(Nx,Ny,Lx,Ly) + + # k2 = k_x.^2 + k_y.^2 # this is for an isotropic wavenumber grid only (i.e. Nx=Ny) + + # define stretching matrix + S = calc_stretching_mat(Nz,rho,f0,H,rho[1]) + + # change dimensions of U and V to match domain size + U2 = zeros(1,Nz); U2[:] = U; U2 = repeat(U2,outer=(Ny,1)) + U = zeros(1,Ny,Nz); U[1,:,:] = U2 + + V2 = zeros(1,Nz); V2[:] = V; V2 = repeat(V2,outer=(Nx,1)) + V = zeros(1,Nx,Nz); V[1,:,:] = V2 + + # define background QG PV gradients (TEMPORARY) + Qy = reshape(Qy[1,:,:],(1,Ny,Nz)) + Qx = zeros(size(Qy)) + + # perform linear stability analysis + evecs_all,evals_all = calc_lin_stab(Qy,Qx,U,V,S,k_x,k_y,Nz) + + # keep largest growth rates per wavenumber + evecs,evals,max_evec,max_eval = find_growth(evecs_all,evals_all,Nx,Ny,Nz) + + # def rad + r_d = sqrt(gp(rho[1:2],rho[1])*H[1])/f0 + + return fftshift(evecs),fftshift(evals),max_evec,max_eval,fftshift(k_x),fftshift(k_y),mean(Qx[1,:,:],dims=1),mean(Qy[1,:,:],dims=1) +end + +function wavenumber_grid(Nx,Ny,Lx,Ly) + # + # nk_x = div(Nx,2)+1; nk_y = div(Ny,2)+1 + + nk_x = Nx; nk_y = Ny + + k_x = reshape(LinRange(-2*pi/Lx*nk_x,2*pi/Lx*nk_x,nk_x),(1,Nx)) + k_y = reshape(LinRange(-2*pi/Ly*nk_y,2*pi/Ly*nk_y,nk_y),(1,Ny)) + + # k_x = LinRange(0.,2*pi/Lx*nk_x,nk_x) + # k_y = LinRange(0.,2*pi/Ly*nk_y,nk_y) + + return k_x,k_y +end + +function calc_PV_grad_y(U,beta,eta,Ny::Int64,Nz::Int64,k_y,S) + # calculates PV gradients in one meridional direction + # U is (Ny x Nz) + # k_y is (Ny x 1) + # + + Uyy = real.(ifft(-k_y.^2 .* fft(U))) + + # Uyy = repeat(Uyy, outer=(Nx, 1, 1)) + + # Q_y = zeros(Nx,Nz) + + F = zeros(size(U)) + for i=1:Ny + F[1,i,:] = S * U[1,i,:] + end + + Q_y = beta .- (Uyy .+ F) + + return Q_y +end + +function calc_PV_grad_x(V,eta,Nx::Int64,Nz::Int64,k_x,S) + # calculates PV gradients in one zonal direction + + Vxx = real.(ifft(k_x.^2 .* fft(V))) + + # Q_y = zeros(Nx,Nz) + + F = zeros(size(V)) + for i=1:Nx + F[1,i,:] = S * V[1,i,:] + end + + Q_x = Vxx .+ F + + return Q_x +end + +function gp(rho,rho0) + g = 9.81 + g_prime = g*(rho[2]-rho[1])/rho0 + + return g_prime +end + +function calc_stretching_mat(Nz,rho,f0,H,rho0) + # + S = zeros((Nz,Nz,)) + + alpha = 0 + + S[1,1] = -f0^2/H[1]/gp(rho[1:2],rho0) + alpha + S[1,2] = f0^2/H[1]/gp(rho[1:2],rho0) + for i = 2:Nz-1 + S[i,i-1] = f0^2/H[i]/gp(rho[i-1:i],rho0) + S[i,i] = -(f0^2/H[i]/gp(rho[i:i+1],rho0) + f0^2/H[i]/gp(rho[i-1:i],rho0)) + S[i,i+1] = f0^2/H[i]/gp(rho[i:i+1],rho0) + end + S[Nz,Nz-1] = f0^2/H[Nz]/gp(rho[Nz-1:Nz],rho0) + S[Nz,Nz] = -f0^2/H[Nz]/gp(rho[Nz-1:Nz],rho0) + + return S +end + +function calc_lin_stab(Qy,Qx,U,V,S,k_x,k_y,Nz) + # + # A = (k_x .* U .+ k_y .* V) + + # B = zeros((1,length(k2),Nz)) + # ell = zeros((Nz,Nz,length(k_x),length(k_y))) + # for i=eachindex(k_x) + # for j=eachindex(k_y) + # ell[:,:,i,j] = S .- (k_x[i]^2 + k_y[j]^2) * I + # B[1,i,j,:] = ell[:,:,i,j] * A[1,i,:] + # end + # end + + # C = k_x .* Qy .- k_y .* Qx + + # M = inv(ell) * (A .+ C) + + # evecs,evals = eig(M) + + + ################# + evecs = zeros(Nx,Ny,Nz,Nz) .+ 0im + evals = zeros(Nx,Ny,Nz) .+ 0im + k2 = zeros(Nx,Ny) + + for i=1:Nx + for j=1:Ny + if i==1 && j==1 + # do nothing + else + A = make_diag(k_x[i] .* U[:,j,:] .+ k_y[j] .* V[:,i,:]) + # ell_i = inv(S - (k_x[i]^2 + k_y[j]^2)*I) + # B = ell_i .* make_diag(k_x[i] * Qy[:,j,:] - k_y[j] * Qx[:,i,:]) .+ 0im + + # evecs[i,j,:,:] = eigvecs(A .+ B) + # evals[i,j,:] = eigvals(A .+ B) + + ell = (S - (k_x[i]^2 + k_y[j]^2) * I) + A2 = ell * A + ell_i = inv(ell) + Q2 = (k_x[i] * Qy[:,j,:] - k_y[j] * Qx[:,i,:]) .+ 0im + B2 = transpose(ell_i) * transpose(A2 .+ make_diag(Q2)) .+ 0im + + k2[i,j] = (k_x[i]^2 + k_y[j]^2) + + evecs[i,j,:,:] = eigvecs(B2) + evals[i,j,:] = eigvals(B2) + end + + end + end + + return evecs,evals +end + +function make_diag(array_in) + matrix_out = zeros(length(array_in),length(array_in)) + for i=eachindex(array_in) + matrix_out[i,i] = array_in[i] + end + return matrix_out +end + +function find_growth(evecs_all,evals_all,Nx,Ny,Nz) + # + evecs = zeros(Nx,Ny,Nz) .+ 0im; evals = zeros(Nx,Ny) .+ 0im + + for i=1:Nx + for j=1:Ny + indMax = argmax(imag(evals_all[i,j,:])) + evals[i,j] = evals_all[i,j,indMax] + evecs[i,j,:] = evecs_all[i,j,:,indMax] + end + end + + sigma = imag(evals) + + indMax = argmax(sigma) + + max_eval = sigma[indMax] + + max_evec = abs.(evecs[indMax,:]) + + return evecs,evals,max_evec,max_eval +end \ No newline at end of file diff --git a/examples/matts_examples/phillips_stab.jl b/examples/matts_examples/phillips_stab.jl new file mode 100644 index 00000000..b817deae --- /dev/null +++ b/examples/matts_examples/phillips_stab.jl @@ -0,0 +1,316 @@ +# a test space for building a linear stability analysis tool in julia +# for now it is separate from GeophysicalFlows.jl Model object, but might +# be merged in the future + +## load packages + +using GeophysicalFlows, CairoMakie, Printf, FFTW, LinearAlgebra, Statistics + +using Random: seed! + +include("./mjcl_stab.jl") + +## build basic model + +dev = CPU() # device (CPU) + +# numerical params +Ny = Nx = n = 128 # 2D resolution = n² +stepper = "FilteredRK4" # time stepping scheme +dt = 1e-2 # time step +nsteps = 60000 # total number of time-steps 20000 +nsubs = 50 # number of time-steps for plotting (nsteps must be multiple of nsubs) + +# physical params +L = 2π # domain size +μ = 5e-2 # bottom drag +beta = β = 0. #1.2130692965249345e-11 # the y-gradient of planetary PV + +nlayers = 8 # number of layers +f0 = f₀= 1. #0.0001236812857687059 # coriolis param [s^-1] +g = 9.81 +H = [0.1, 0.3, 1.] # the rest depths of each layer +rho = ρ = [4.0, 5.0, 5.1] # the density of each layer + +H=ones(8) +rho = ρ =collect(range(4.0,5.0,8)) +U = collect(range(1.0,0.0,8)) +V = zeros(8) + +# U = zeros(nlayers) # the imposed mean zonal flow in each layer +# U[1] = 1.0 +# U[2] = 0.05 +# U[3] = 0.005 + + +# V = zeros(nlayers) # the imposed mean zonal flow in each layer +# V[1] = 1.0 +# V[2] = 0.5 +# V[3] = 0.25 + +# setting up the ``problem'' +prob = MultiLayerQG.Problem(nlayers, dev; nx=n, Lx=L, f₀, g, H, ρ, U, μ, β, + dt, stepper, aliased_fraction=0) + +sol, clock, params, vars, grid = prob.sol, prob.clock, prob.params, prob.vars, prob.grid +x, y = grid.x, grid.y + +# initial conditions +seed!(1234) # reset of the random number generator for reproducibility +q₀ = 1e-2 * device_array(dev)(randn((grid.nx, grid.ny, nlayers))) +q₀h = prob.timestepper.filter .* rfft(q₀, (1, 2)) # apply rfft only in dims=1, 2 +q₀ = irfft(q₀h, grid.nx, (1, 2)) # apply irfft only in dims=1, 2 + +MultiLayerQG.set_q!(prob, q₀) + +# diagnostics +E = Diagnostic(MultiLayerQG.energies, prob; nsteps) +diags = [E] # A list of Diagnostics types passed to "stepforward!" will be updated every timestep. + +# output dirs +filepath = "." +plotpath = "./figs/plots_2layer" +plotname = "snapshots" +filename = joinpath(filepath, "2layer.jld2") + +# file management +if isfile(filename); rm(filename); end +if !isdir(plotpath); mkdir(plotpath); end + +# ``create output'' (?) +get_sol(prob) = prob.sol # extracts the Fourier-transformed solution + +function get_u(prob) + sol, params, vars, grid = prob.sol, prob.params, prob.vars, prob.grid + + @. vars.qh = sol + streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) + @. vars.uh = -im * grid.l * vars.ψh + invtransform!(vars.u, vars.uh, params) + + return vars.u +end + +out = Output(prob, filename, (:sol, get_sol), (:u, get_u)) + +# visualizing the simulation... + +Lx, Ly = grid.Lx, grid.Ly + +title_KE = Observable(@sprintf("μt = %.2f", μ * clock.t)) + +q1 = Observable(Array(vars.q[:, :, 1])) +q2 = Observable(Array(vars.q[:, :, 2])) +q3 = Observable(Array(vars.q[:, :, 3])) +q4 = Observable(Array(vars.q[:, :, 4])) +q5 = Observable(Array(vars.q[:, :, 5])) +q6 = Observable(Array(vars.q[:, :, 6])) +q7 = Observable(Array(vars.q[:, :, 7])) +q8 = Observable(Array(vars.q[:, :, 8])) + +# psi1 = Observable(Array(vars.ψ[:, :, 1])) +# psi2 = Observable(Array(vars.ψ[:, :, 2])) +# psi3 = Observable(Array(vars.ψ[:, :, 3])) +# psi4 = Observable(Array(vars.ψ[:, :, 4])) +# psi5 = Observable(Array(vars.ψ[:, :, 5])) +# psi6 = Observable(Array(vars.ψ[:, :, 6])) +# psi7 = Observable(Array(vars.ψ[:, :, 7])) +# psi8 = Observable(Array(vars.ψ[:, :, 8])) + +# function compute_levels(maxf, nlevels=8) +# # -max(|f|):...:max(|f|) +# levelsf = @lift collect(range(-$maxf, stop = $maxf, length=nlevels)) + +# # only positive +# levelsf⁺ = @lift collect(range($maxf/(nlevels-1), stop = $maxf, length=Int(nlevels/2))) + +# # only negative +# levelsf⁻ = @lift collect(range(-$maxf, stop = -$maxf/(nlevels-1), length=Int(nlevels/2))) + +# return levelsf, levelsf⁺, levelsf⁻ +# end + +# maxpsi1 = Observable(maximum(abs, vars.ψ[:, :, 1])) +# maxpsi2 = Observable(maximum(abs, vars.ψ[:, :, 2])) +# maxpsi3 = Observable(maximum(abs, vars.ψ[:, :, 3])) +# maxpsi4 = Observable(maximum(abs, vars.ψ[:, :, 4])) +# maxpsi5 = Observable(maximum(abs, vars.ψ[:, :, 5])) +# maxpsi6 = Observable(maximum(abs, vars.ψ[:, :, 6])) +# maxpsi7 = Observable(maximum(abs, vars.ψ[:, :, 7])) +# maxpsi8 = Observable(maximum(abs, vars.ψ[:, :, 8])) + + +# levelspsi1, levelspsi1⁺, levelspsi1⁻ = compute_levels(maxpsi1) +# levelspsi2, levelspsi2⁺, levelspsi2⁻ = compute_levels(maxpsi2) +# levelspsi3, levelspsi3⁺, levelspsi3⁻ = compute_levels(maxpsi3) +# levelspsi4, levelspsi4⁺, levelspsi4⁻ = compute_levels(maxpsi4) +# levelspsi5, levelspsi5⁺, levelspsi5⁻ = compute_levels(maxpsi5) +# levelspsi6, levelspsi6⁺, levelspsi6⁻ = compute_levels(maxpsi6) +# levelspsi7, levelspsi7⁺, levelspsi7⁻ = compute_levels(maxpsi7) +# levelspsi8, levelspsi8⁺, levelspsi8⁻ = compute_levels(maxpsi8) + +# KE₁ = Observable(Point2f[(μ * E.t[1], E.data[1][1][1])]) +# KE₂ = Observable(Point2f[(μ * E.t[1], E.data[1][1][end])]) +# PE = Observable(Point2f[(μ * E.t[1], E.data[1][2][end])]) + +fig = Figure(resolution=(1000, 600)) + +axis_kwargs = (xlabel = "x", + ylabel = "y", + aspect = 1, + limits = ((-Lx/2, Lx/2), (-Ly/2, Ly/2))) + +axq1 = Axis(fig[1, 1]; title = "q1", axis_kwargs...) + +# axpsi1 = Axis(fig[2, 1]; title = "psi1", axis_kwargs...) + +axq2 = Axis(fig[1, 2]; title = "q2", axis_kwargs...) + +# axpsi2 = Axis(fig[2, 2]; title = "psi2", axis_kwargs...) + +axq3 = Axis(fig[1, 3]; title = "q3", axis_kwargs...) +axq4 = Axis(fig[1, 4]; title = "q4", axis_kwargs...) +axq5 = Axis(fig[2, 1]; title = "q5", axis_kwargs...) +axq6 = Axis(fig[2, 2]; title = "q6", axis_kwargs...) +axq7 = Axis(fig[2, 3]; title = "q7", axis_kwargs...) +axq8 = Axis(fig[2, 4]; title = "q8", axis_kwargs...) + +# axKE = Axis(fig[1, 3], +# xlabel = "μ t", +# ylabel = "KE", +# title = title_KE, +# yscale = log10, +# limits = ((-0.1, 2.6), (1e-9, 5))) + +# axPE = Axis(fig[2, 3], +# xlabel = "μ t", +# ylabel = "PE", +# yscale = log10, +# limits = ((-0.1, 2.6), (1e-9, 5))) + +heatmap!(axq1, x, y, q1; colormap = :balance) +heatmap!(axq2, x, y, q2; colormap = :balance) +heatmap!(axq3, x, y, q3; colormap = :balance) +heatmap!(axq4, x, y, q4; colormap = :balance) +heatmap!(axq5, x, y, q5; colormap = :balance) +heatmap!(axq6, x, y, q6; colormap = :balance) +heatmap!(axq7, x, y, q7; colormap = :balance) +heatmap!(axq8, x, y, q8; colormap = :balance) + +# contourf!(axpsi1, x, y, psi1; +# levels = levelspsi1, colormap = :viridis, extendlow = :auto, extendhigh = :auto) +# contour!(axpsi1, x, y, psi1; +# levels = levelspsi1⁺, color=:black) +# contour!(axpsi1, x, y, psi1; +# levels = levelspsi1⁻, color=:black, linestyle = :dash) + +# contourf!(axpsi2, x, y, psi2; +# levels = levelspsi2, colormap = :viridis, extendlow = :auto, extendhigh = :auto) +# contour!(axpsi2, x, y, psi2; +# levels = levelspsi2⁺, color=:black) +# contour!(axpsi2, x, y, psi2; +# levels = levelspsi2⁻, color=:black, linestyle = :dash) + +# ke₁ = lines!(axKE, KE₁; linewidth = 3) +# ke₂ = lines!(axKE, KE₂; linewidth = 3) +# Legend(fig[1, 4], [ke₁, ke₂,], ["KE₁", "KE₂"]) + +# lines!(axPE, PE; linewidth = 3) + +fig + +# trying to run the model now +startwalltime = time() + +frames = 0:round(Int, nsteps / nsubs) + +record(fig, "multilayerqg_2layer.mp4", frames, framerate = 18) do j + if j % (1000 / nsubs) == 0 + cfl = clock.dt * maximum([maximum(vars.u) / grid.dx, maximum(vars.v) / grid.dy]) + + log = @sprintf("step: %04d, t: %.1f, cfl: %.2f, KE₁: %.3e, KE₂: %.3e, PE: %.3e, walltime: %.2f min", + clock.step, clock.t, cfl, E.data[E.i][1][1], E.data[E.i][1][2], E.data[E.i][2][1], (time()-startwalltime)/60) + + println(log) + end + + q1[] = vars.q[:, :, 1] + q2[] = vars.q[:, :, 2] + q3[] = vars.q[:, :, 3] + q4[] = vars.q[:, :, 4] + q5[] = vars.q[:, :, 5] + q6[] = vars.q[:, :, 6] + q7[] = vars.q[:, :, 7] + q8[] = vars.q[:, :, end] + + # psi1[] = vars.ψ[:, :, 1] + # psi2[] = vars.ψ[:, :, 2] + # psi3[] = vars.ψ[:, :, 3] + # psi4[] = vars.ψ[:, :, 4] + # psi5[] = vars.ψ[:, :, 5] + # psi6[] = vars.ψ[:, :, 6] + # psi7[] = vars.ψ[:, :, 7] + # psi8[] = vars.ψ[:, :, end] + + # maxpsi1[] = maximum(abs, vars.ψ[:, :, 1]) + # maxpsi2[] = maximum(abs, vars.ψ[:, :, 2]) + # maxpsi3[] = maximum(abs, vars.ψ[:, :, 3]) + # maxpsi4[] = maximum(abs, vars.ψ[:, :, 4]) + # maxpsi5[] = maximum(abs, vars.ψ[:, :, 5]) + # maxpsi6[] = maximum(abs, vars.ψ[:, :, 6]) + # maxpsi7[] = maximum(abs, vars.ψ[:, :, 7]) + # maxpsi8[] = maximum(abs, vars.ψ[:, :, end]) + + + # KE₁[] = push!(KE₁[], Point2f(μ * E.t[E.i], E.data[E.i][1][1])) + # KE₂[] = push!(KE₂[], Point2f(μ * E.t[E.i], E.data[E.i][1][end])) + # PE[] = push!(PE[] , Point2f(μ * E.t[E.i], E.data[E.i][2][end])) + + # title_KE[] = @sprintf("μ t = %.2f", μ * clock.t)] + + stepforward!(prob, diags, nsubs) + MultiLayerQG.updatevars!(prob) + + # savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), clock.step) + # savefig(savename) + +end + + +# savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), clock.step) +# savefig(savename) + + +## use stability functions +# Ny = div(Nx,2) + +eta = 0 + +eve,eva,k_x,k_y,qx,qy = lin_stab(U,V,beta,eta,Nx,Ny,rho,f0,Lx,Ly) + +k_xg = range(k_x[1],step=(k_x[2]-k_x[1]),length=length(k_x)) +k_yg = range(k_y[1],step=(k_y[2]-k_y[1]),length=length(k_y)) + + + +# Plots.contourf(transpose(imag(eva[60:70,60:70]))) + +# sigma = Observable(Array(imag(eva))) + +# fig = Figure(resolution=(1000, 600)) + +# axis_kwargs = (xlabel = "k_x", +# ylabel = "k_y", +# aspect = 1, +# limits = ((k_x[60], k_x[70]), (k_y[60],k_y[70]))) + +# ax_sig = Axis(fig[1, 1]; title = "σ", axis_kwargs...) + + +# contourf!(ax_sig, k_xg, k_yg, sigma; +# colormap = :viridis, extendlow = :auto, extendhigh = :auto) +# contour!(ax_sig, k_xg, k_yg, sigma; +# color=:black) + + + diff --git a/examples/matts_examples/prof_stab.jl b/examples/matts_examples/prof_stab.jl new file mode 100644 index 00000000..89aae07c --- /dev/null +++ b/examples/matts_examples/prof_stab.jl @@ -0,0 +1,89 @@ +# a test space for building a linear stability analysis tool in julia +# for now it is separate from GeophysicalFlows.jl Model object, but might +# be merged in the future + +## load packages + +using GeophysicalFlows, CairoMakie, Printf, FFTW, LinearAlgebra, Statistics + +using Random: seed! + +include("./mjcl_stab.jl") + +## build basic model + +dev = CPU() # device (CPU) + +# numerical params +Ny = Nx = n = 128 # 2D resolution = n² +stepper = "FilteredRK4" # time stepping scheme +dt = 1e-2 # time step +nsteps = 60000 # total number of time-steps 20000 +nsubs = 50 # number of time-steps for plotting (nsteps must be multiple of nsubs) + +# physical params +L = 2π # domain size +μ = 5e-2 # bottom drag +beta = β = 0. #1.2130692965249345e-11 # the y-gradient of planetary PV + +nlayers = 8 # number of layers +f0 = f₀= 1. #0.0001236812857687059 # coriolis param [s^-1] +g = 9.81 +H = [0.1, 0.3, 1.] # the rest depths of each layer +rho = ρ = [4.0, 5.0, 5.1] # the density of each layer + +H=ones(8) +rho = ρ =collect(range(4.0,5.0,8)) +U = collect(range(1.0,0.0,8)) +V = zeros(8) + +# U = zeros(nlayers) # the imposed mean zonal flow in each layer +# U[1] = 1.0 +# U[2] = 0.05 +# U[3] = 0.005 + + +# V = zeros(nlayers) # the imposed mean zonal flow in each layer +# V[1] = 1.0 +# V[2] = 0.5 +# V[3] = 0.25 + +# setting up the ``problem'' +prob = MultiLayerQG.Problem(nlayers, dev; nx=n, Lx=L, f₀, g, H, ρ, U, μ, β, + dt, stepper, aliased_fraction=0) + +sol, clock, params, vars, grid = prob.sol, prob.clock, prob.params, prob.vars, prob.grid +x, y = grid.x, grid.y + +Lx, Ly = grid.Lx, grid.Ly + +eta = 0 + +eve,eva,max_eve,max_eva,k_x,k_y,qx,qy = lin_stab(U,V,beta,eta,Nx,Ny,rho,f0,Lx,Ly) + +k_xg = range(k_x[1],step=(k_x[2]-k_x[1]),length=length(k_x)) +k_yg = range(k_y[1],step=(k_y[2]-k_y[1]),length=length(k_y)) + + + +# Plots.contourf(transpose(imag(eva[60:70,60:70]))) + +# sigma = Observable(Array(imag(eva))) + +# fig = Figure(resolution=(1000, 600)) + +# axis_kwargs = (xlabel = "k_x", +# ylabel = "k_y", +# aspect = 1, +# limits = ((k_x[60], k_x[70]), (k_y[60],k_y[70]))) + +# ax_sig = Axis(fig[1, 1]; title = "σ", axis_kwargs...) + + +# contourf!(ax_sig, k_xg, k_yg, sigma; +# colormap = :viridis, extendlow = :auto, extendhigh = :auto) +# contour!(ax_sig, k_xg, k_yg, sigma; +# color=:black) + + + diff --git a/examples/matts_examples/sphere_ex.jl b/examples/matts_examples/sphere_ex.jl new file mode 100644 index 00000000..e0e08a2e --- /dev/null +++ b/examples/matts_examples/sphere_ex.jl @@ -0,0 +1,33 @@ +# function to calculate the volume of a sphere +function sphere_vol(r) + # julia allows Unicode names (in UTF-8 encoding) + # so either "pi" or the symbol π can be used + return 4/3*pi*r^3 +end + +# functions can also be defined more succinctly +quadratic(a, sqr_term, b) = (-b + sqr_term) / 2a + +# calculates x for 0 = a*x^2+b*x+c, arguments types can be defined in function definitions +function quadratic2(a::Float64, b::Float64, c::Float64) + # unlike other languages 2a is equivalent to 2*a + # a^2 is used instead of a**2 or pow(a,2) + sqr_term = sqrt(b^2-4a*c) + r1 = quadratic(a, sqr_term, b) + r2 = quadratic(a, -sqr_term, b) + # multiple values can be returned from a function using tuples + # if the return keyword is omitted, the last term is returned + r1, r2 +end + +vol = sphere_vol(3) +# @printf allows number formatting but does not automatically append the \n to statements, see below +using Printf +@printf "volume = %0.3f\n" vol +#> volume = 113.097 + +quad1, quad2 = quadratic2(2.0, -2.0, -12.0) +println("result 1: ", quad1) +#> result 1: 3.0 +println("result 2: ", quad2) +#> result 2: -2.0 \ No newline at end of file diff --git a/src/multilayerqg_no_fixes.jl b/src/multilayerqg_no_fixes.jl new file mode 100644 index 00000000..687f5b29 --- /dev/null +++ b/src/multilayerqg_no_fixes.jl @@ -0,0 +1,1107 @@ +module MultiLayerQG_nf + +export + fwdtransform!, + invtransform!, + streamfunctionfrompv!, + pvfromstreamfunction!, + updatevars!, + + set_q!, + set_ψ!, + energies, + fluxes + +using + FFTW, + CUDA, + LinearAlgebra, + StaticArrays, + Reexport, + DocStringExtensions + +@reexport using FourierFlows + +using FourierFlows: parsevalsum, parsevalsum2, superzeros, plan_flows_rfft + +nothingfunction(args...) = nothing + +""" + Problem(nlayers :: Int, + dev = CPU(); + nx = 128, + ny = nx, + Lx = 2π, + Ly = Lx, + f₀ = 1.0, + β = 0.0, + g = 1.0, + U = zeros(nlayers), + H = 1/nlayers * ones(nlayers), + ρ = Array{Float64}(1:nlayers), + eta = nothing, + topographic_pv_gradient = (0, 0), + μ = 0, + ν = 0, + nν = 1, + dt = 0.01, + stepper = "RK4", + calcFq = nothingfunction, + stochastic = false, + linear = false, + aliased_fraction = 1/3, + T = Float64) + +Construct a multi-layer quasi-geostrophic problem with `nlayers` fluid layers on device `dev`. + +Arguments +========= +- `nlayers`: (required) Number of fluid layers. +- `dev`: (required) `CPU()` (default) or `GPU()`; computer architecture used to time-step `problem`. + +Keyword arguments +================= + - `nx`: Number of grid points in ``x``-domain. + - `ny`: Number of grid points in ``y``-domain. + - `Lx`: Extent of the ``x``-domain. + - `Ly`: Extent of the ``y``-domain. + - `f₀`: Constant planetary vorticity. + - `β`: Planetary vorticity ``y``-gradient. + - `g`: Gravitational acceleration constant. + - `U`: The imposed constant zonal flow ``U(y)`` in each fluid layer. + - `H`: Rest height of each fluid layer. + - `ρ`: Density of each fluid layer. + - `eta`: Periodic component of the topographic potential vorticity. + - `topographic_pv_gradient`: The ``(x, y)`` components of the topographic PV large-scale gradient. + - `μ`: Linear bottom drag coefficient. + - `ν`: Small-scale (hyper)-viscosity coefficient. + - `nν`: (Hyper)-viscosity order, `nν```≥ 1``. + - `dt`: Time-step. + - `stepper`: Time-stepping method. + - `calcF`: Function that calculates the Fourier transform of the forcing, ``F̂``. + - `stochastic`: `true` or `false` (default); boolean denoting whether `calcF` is temporally stochastic. + - `linear`: `true` or `false` (default); boolean denoting whether the linearized equations of motions are used. + - `aliased_fraction`: the fraction of high-wavenumbers that are zero-ed out by `dealias!()`. + - `T`: `Float32` or `Float64` (default); floating point type used for `problem` data. +""" +function Problem_nf(nlayers::Int, # number of fluid layers + dev = CPU(); + # Numerical parameters + nx = 128, + ny = nx, + Lx = 2π, + Ly = Lx, + # Physical parameters + f₀ = 1.0, # Coriolis parameter + β = 0.0, # y-gradient of Coriolis parameter + g = 1.0, # gravitational constant + U = zeros(nlayers), # imposed zonal flow U(y) in each layer + H = 1/nlayers * ones(nlayers), # rest fluid height of each layer + ρ = Array{Float64}(1:nlayers), # density of each layer + eta = nothing, # periodic component of the topographic PV + topographic_pv_gradient = (0, 0), # tuple with the ``(x, y)`` components of topographic PV large-scale gradient + # Bottom Drag and/or (hyper)-viscosity + μ = 0, + ν = 0, + nν = 1, + # Timestepper and equation options + dt = 0.01, + stepper = "RK4", + calcFq = nothingfunction, + stochastic = false, + linear = false, + # Float type and dealiasing + aliased_fraction = 1/3, + T = Float64) + + if dev == GPU() && nlayers > 2 + @warn """MultiLayerQG module is not optimized on the GPU yet for configurations with + 3 fluid layers or more! + + See issues on Github at https://github.com/FourierFlows/GeophysicalFlows.jl/issues/112 + and https://github.com/FourierFlows/GeophysicalFlows.jl/issues/267. + + To use MultiLayerQG with 3 fluid layers or more we suggest, for now, to restrict running + on CPU.""" + end + + if nlayers == 1 + @warn """MultiLayerQG module does work for single-layer configuration but may not be as + optimized. We suggest using SingleLayerQG module for single-layer QG simulation unless + you have reasons to use MultiLayerQG in a single-layer configuration, e.g., you want to + compare solutions with varying number of fluid layers.""" + end + + # topographic PV + eta === nothing && (eta = zeros(dev, T, (nx, ny))) + + grid = TwoDGrid(dev; nx, Lx, ny, Ly, aliased_fraction, T) + + params = Params_nf(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ, ν, nν, grid; calcFq) + + vars = calcFq == nothingfunction ? DecayingVars(grid, params) : (stochastic ? StochasticForcedVars(grid, params) : ForcedVars(grid, params)) + + equation = linear ? LinearEquation(params, grid) : Equation(params, grid) + + FourierFlows.Problem(equation, stepper, dt, grid, vars, params) +end + +""" + struct Params{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams + +The parameters for the `MultiLayerQG` problem. + +$(TYPEDFIELDS) +""" +struct Params_nf{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams + # prescribed params + "number of fluid layers" + nlayers :: Int + "gravitational constant" + g :: T + "constant planetary vorticity" + f₀ :: T + "planetary vorticity ``y``-gradient" + β :: T + "array with density of each fluid layer" + ρ :: Aphys3D + "array with rest height of each fluid layer" + H :: Aphys3D + "array with imposed constant zonal flow ``U(y)`` in each fluid layer" + U :: Aphys3D + "array containing the topographic PV" + eta :: Aphys2D + "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" + topographic_pv_gradient :: Tuple{T, T} + "linear bottom drag coefficient" + μ :: T + "small-scale (hyper)-viscosity coefficient" + ν :: T + "(hyper)-viscosity order, `nν```≥ 1``" + nν :: Int + "function that calculates the Fourier transform of the forcing, ``F̂``" + calcFq! :: Function + + # derived params + "array with the reduced gravity constants for each fluid interface" + g′ :: Aphys1D + "array containing ``x``-gradient of PV due to topographic PV in each fluid layer" + Qx :: Aphys3D + "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV in each fluid layer" + Qy :: Aphys3D + "array containing coeffients for getting PV from streamfunction" + S :: Atrans4D + "array containing coeffients for inverting PV to streamfunction" + S⁻¹ :: Atrans4D + "rfft plan for FFTs" + rfftplan :: Trfft +end + +""" + struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams + +The parameters for the a single-layer `MultiLayerQG` problem. + +$(TYPEDFIELDS) +""" +struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams + # prescribed params + "planetary vorticity ``y``-gradient" + β :: T + "array with imposed constant zonal flow ``U(y)``" + U :: Aphys3D + "array containing the periodic component of the topographic PV" + eta :: Aphys2D + "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" + topographic_pv_gradient :: Tuple{T, T} + "linear drag coefficient" + μ :: T + "small-scale (hyper)-viscosity coefficient" + ν :: T + "(hyper)-viscosity order, `nν```≥ 1``" + nν :: Int + "function that calculates the Fourier transform of the forcing, ``F̂``" + calcFq! :: Function + + # derived params + "array containing ``x``-gradient of PV due to topographic PV" + Qx :: Aphys3D + "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV" + Qy :: Aphys3D + "rfft plan for FFTs" + rfftplan :: Trfft +end + +""" + struct TwoLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams + +The parameters for the a two-layer `MultiLayerQG` problem. + +$(TYPEDFIELDS) +""" +struct TwoLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams + # prescribed params + "gravitational constant" + g :: T + "constant planetary vorticity" + f₀ :: T + "planetary vorticity ``y``-gradient" + β :: T + "array with density of each fluid layer" + ρ :: Aphys3D + "tuple with rest height of each fluid layer" + H :: Tuple + "array with imposed constant zonal flow ``U(y)`` in each fluid layer" + U :: Aphys3D + "array containing periodic component of the topographic PV" + eta :: Aphys2D + "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" + topographic_pv_gradient :: Tuple{T, T} + "linear bottom drag coefficient" + μ :: T + "small-scale (hyper)-viscosity coefficient" + ν :: T + "(hyper)-viscosity order, `nν```≥ 1``" + nν :: Int + "function that calculates the Fourier transform of the forcing, ``F̂``" + calcFq! :: Function + + # derived params + "the reduced gravity constants for the fluid interface" + g′ :: T + "array containing ``x``-gradient of PV due to topographic PV in each fluid layer" + Qx :: Aphys3D + "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV in each fluid layer" + Qy :: Aphys3D + "rfft plan for FFTs" + rfftplan :: Trfft +end + +function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 1}) where TU + T = eltype(grid) + if length(U) == nlayers + U_2D = zeros(dev, T, (1, nlayers)) + U_2D[:] = U + U_2D = repeat(U_2D, outer=(grid.ny, 1)) + else + U_2D = zeros(dev, T, (grid.ny, 1)) + U_2D[:] = U + end + U_3D = zeros(dev, T, (1, grid.ny, nlayers)) + @views U_3D[1, :, :] = U_2D + return U_3D +end + +function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 2}) where TU + T = eltype(grid) + U_3D = zeros(dev, T, (1, grid.ny, nlayers)) + @views U_3D[1, :, :] = U + return U_3D +end + +function convert_U_to_U3D(dev, nlayers, grid, U::Number) + T = eltype(grid) + A = device_array(dev) + U_3D = reshape(repeat([T(U)], outer=(grid.ny, 1)), (1, grid.ny, nlayers)) + return A(U_3D) +end + +function Params_nf(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ, ν, nν, grid; calcFq=nothingfunction, effort=FFTW.MEASURE) + dev = grid.device + T = eltype(grid) + A = device_array(dev) + + ny, nx = grid.ny , grid.nx + nkr, nl = grid.nkr, grid.nl + kr, l = grid.kr , grid.l + + U = convert_U_to_U3D(dev, nlayers, grid, U) + + Uyy = real.(ifft(-l.^2 .* fft(U))) + Uyy = CUDA.@allowscalar repeat(Uyy, outer=(nx, 1, 1)) + + # Calculate periodic components of the topographic PV gradients. + etah = rfft(A(eta)) + etax = irfft(im * kr .* etah, nx) # ∂η/∂x + etay = irfft(im * l .* etah, nx) # ∂η/∂y + + # Add topographic PV large-scale gradient + topographic_pv_gradient = T.(topographic_pv_gradient) + @. etax += topographic_pv_gradient[1] + @. etay += topographic_pv_gradient[2] + + Qx = zeros(dev, T, (nx, ny, nlayers)) + @views @. Qx[:, :, nlayers] += etax + + Qy = zeros(dev, T, (nx, ny, nlayers)) + Qy = T(β) .- Uyy # T(β) is needed to ensure that Qy remains same type as U + @views @. Qy[:, :, nlayers] += etay + + rfftplanlayered = plan_flows_rfft(A{T, 3}(undef, grid.nx, grid.ny, nlayers), [1, 2]; flags=effort) + + if nlayers==1 + return SingleLayerParams(T(β), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, Qx, Qy, rfftplanlayered) + + else # if nlayers≥2 + + ρ = reshape(T.(ρ), (1, 1, nlayers)) + H = reshape(T.(H), (1, 1, nlayers)) + + g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface + + Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) + Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) + + typeofSkl = SArray{Tuple{nlayers, nlayers}, T, 2, nlayers^2} # StaticArrays of type T and dims = (nlayers x nlayers) + + S = Array{typeofSkl, 2}(undef, (nkr, nl)) + calcS!(S, Fp, Fm, nlayers, grid) + + S⁻¹ = Array{typeofSkl, 2}(undef, (nkr, nl)) + calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) + + S, S⁻¹, Fp, Fm = A(S), A(S⁻¹), A(Fp), A(Fm) # convert to appropriate ArrayType + + CUDA.@allowscalar @views Qy[:, :, 1] = @. Qy[:, :, 1] - Fp[1] * (U[:, :, 2] - U[:, :, 1]) + for j = 2:nlayers-1 + CUDA.@allowscalar @views Qy[:, :, j] = @. Qy[:, :, j] - Fp[j] * (U[:, :, j+1] - U[:, :, j]) + Fm[j-1] * (U[:, :, j-1] - U[:, :, j]) + end + CUDA.@allowscalar @views Qy[:, :, nlayers] = @. Qy[:, :, nlayers] - Fm[nlayers-1] * (U[:, :, nlayers-1] - U[:, :, nlayers]) + + if nlayers==2 + return TwoLayerParams(T(g), T(f₀), T(β), A(ρ), (T(H[1]), T(H[2])), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, T(g′[1]), Qx, Qy, rfftplanlayered) + else # if nlayers>2 + return Params_nf(nlayers, T(g), T(f₀), T(β), A(ρ), A(H), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, A(g′), Qx, Qy, S, S⁻¹, rfftplanlayered) + end + end +end + +numberoflayers(params) = params.nlayers +numberoflayers(::SingleLayerParams) = 1 +numberoflayers(::TwoLayerParams) = 2 + +# --------- +# Equations +# --------- + +""" + hyperviscosity(params, grid) + +Return the linear operator `L` that corresponds to (hyper)-viscosity of order ``n_ν`` with +coefficient ``ν`` for ``n`` fluid layers. +```math +L_j = - ν |𝐤|^{2 n_ν}, \\ j = 1, ...,n . +``` +""" +function hyperviscosity(params, grid) + dev = grid.device + T = eltype(grid) + + L = device_array(dev){T}(undef, (grid.nkr, grid.nl, numberoflayers(params))) + @. L = - params.ν * grid.Krsq^params.nν + @views @. L[1, 1, :] = 0 + + return L +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 +`hyperviscosity(params, grid)`. + +The nonlinear term is computed via function `calcNlinear!`. +""" +function LinearEquation(params, grid) + L = hyperviscosity(params, grid) + + return FourierFlows.Equation(L, calcNlinear!, grid) +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 +`hyperviscosity(params, grid)`. + +The nonlinear term is computed via function `calcN!`. +""" +function Equation(params, grid) + L = hyperviscosity(params, grid) + + return FourierFlows.Equation(L, calcN!, grid) +end + + +# ---- +# Vars +# ---- + +""" + struct Vars{Aphys, Atrans, F, P} <: AbstractVars + +The variables for multi-layer QG problem. + +$(FIELDS) +""" +struct Vars{Aphys, Atrans, F, P} <: AbstractVars + "relative vorticity + vortex stretching" + q :: Aphys + "streamfunction" + ψ :: Aphys + "x-component of velocity" + u :: Aphys + "y-component of velocity" + v :: Aphys + "Fourier transform of relative vorticity + vortex stretching" + qh :: Atrans + "Fourier transform of streamfunction" + ψh :: Atrans + "Fourier transform of ``x``-component of velocity" + uh :: Atrans + "Fourier transform of ``y``-component of velocity" + vh :: Atrans + "Fourier transform of forcing" + Fqh :: F + "`sol` at previous time-step" + prevsol :: P +end + +const DecayingVars = Vars{<:AbstractArray, <:AbstractArray, Nothing, Nothing} +const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing} +const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray} + +""" + DecayingVars(grid, params) + +Return the variables for an unforced multi-layer QG problem with `grid` and `params`. +""" +function DecayingVars(grid, params) + Dev = typeof(grid.device) + T = eltype(grid) + nlayers = numberoflayers(params) + + @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v + @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh + + return Vars(q, ψ, u, v, qh, ψh, uh, vh, nothing, nothing) +end + +""" + ForcedVars(grid, params) + +Return the variables for a forced multi-layer QG problem with `grid` and `params`. +""" +function ForcedVars(grid, params) + Dev = typeof(grid.device) + T = eltype(grid) + nlayers = numberoflayers(params) + + @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v + @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh + + return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, nothing) +end + +""" + StochasticForcedVars(grid, params) + +Return the variables for a forced multi-layer QG problem with `grid` and `params`. +""" +function StochasticForcedVars(grid, params) + Dev = typeof(grid.device) + T = eltype(grid) + nlayers = numberoflayers(params) + + @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v + @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh prevsol + + return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, prevsol) +end + +""" + fwdtransform!(varh, var, params) + +Compute the Fourier transform of `var` and store it in `varh`. +""" +fwdtransform!(varh, var, params::AbstractParams) = mul!(varh, params.rfftplan, var) + +""" + invtransform!(var, varh, params) + +Compute the inverse Fourier transform of `varh` and store it in `var`. +""" +invtransform!(var, varh, params::AbstractParams) = ldiv!(var, params.rfftplan, varh) + +""" + pvfromstreamfunction!(qh, ψh, params, grid) + +Obtain the Fourier transform of the PV from the streamfunction `ψh` in each layer using +`qh = params.S * ψh`. +""" +function pvfromstreamfunction!(qh, ψh, params, grid) + for j=1:grid.nl, i=1:grid.nkr + CUDA.@allowscalar @views qh[i, j, :] .= params.S[i, j] * ψh[i, j, :] + end + + return nothing +end + +""" + pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) + +Obtain the Fourier transform of the PV from the streamfunction `ψh` for the special +case of a single fluid layer configuration. In this case, ``q̂ = - k² ψ̂``. +""" +function pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) + @. qh = -grid.Krsq * ψh + + return nothing +end + +""" + pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) + +Obtain the Fourier transform of the PV from the streamfunction `ψh` for the special +case of a two fluid layer configuration. In this case we have, + +```math +q̂₁ = - k² ψ̂₁ + f₀² / (g′ H₁) * (ψ̂₂ - ψ̂₁) , +``` + +```math +q̂₂ = - k² ψ̂₂ + f₀² / (g′ H₂) * (ψ̂₁ - ψ̂₂) . +``` + +(Here, the PV-streamfunction relationship is hard-coded to avoid scalar operations +on the GPU.) +""" +function pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) + f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] + + ψ1h, ψ2h = view(ψh, :, :, 1), view(ψh, :, :, 2) + + @views @. qh[:, :, 1] = - grid.Krsq * ψ1h + f₀^2 / (g′ * H₁) * (ψ2h - ψ1h) + @views @. qh[:, :, 2] = - grid.Krsq * ψ2h + f₀^2 / (g′ * H₂) * (ψ1h - ψ2h) + + return nothing +end + +""" + streamfunctionfrompv!(ψh, qh, params, grid) + +Invert the PV to obtain the Fourier transform of the streamfunction `ψh` in each layer from +`qh` using `ψh = params.S⁻¹ qh`. +""" +function streamfunctionfrompv!(ψh, qh, params, grid) + for j=1:grid.nl, i=1:grid.nkr + CUDA.@allowscalar @views ψh[i, j, :] .= params.S⁻¹[i, j] * qh[i, j, :] + end + + return nothing +end + +""" + streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) + +Invert the PV to obtain the Fourier transform of the streamfunction `ψh` for the special +case of a single fluid layer configuration. In this case, ``ψ̂ = - k⁻² q̂``. +""" +function streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) + @. ψh = -grid.invKrsq * qh + + return nothing +end + +""" + streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) + +Invert the PV to obtain the Fourier transform of the streamfunction `ψh` for the special +case of a two fluid layer configuration. In this case we have, + +```math +ψ̂₁ = - [k² q̂₁ + (f₀² / g′) (q̂₁ / H₂ + q̂₂ / H₁)] / Δ , +``` + +```math +ψ̂₂ = - [k² q̂₂ + (f₀² / g′) (q̂₁ / H₂ + q̂₂ / H₁)] / Δ , +``` + +where ``Δ = k² [k² + f₀² (H₁ + H₂) / (g′ H₁ H₂)]``. + +(Here, the PV-streamfunction relationship is hard-coded to avoid scalar operations +on the GPU.) +""" +function streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) + f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] + + q1h, q2h = view(qh, :, :, 1), view(qh, :, :, 2) + + @views @. ψh[:, :, 1] = - grid.Krsq * q1h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) + @views @. ψh[:, :, 2] = - grid.Krsq * q2h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) + + for j in 1:2 + @views @. ψh[:, :, j] *= grid.invKrsq / (grid.Krsq + f₀^2 / g′ * (H₁ + H₂) / (H₁ * H₂)) + end + + return nothing +end + +""" + calcS!(S, Fp, Fm, nlayers, grid) + +Construct the array ``𝕊``, which consists of `nlayer` x `nlayer` static arrays ``𝕊_𝐤`` that +relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``q̂_𝐤 = 𝕊_𝐤 ψ̂_𝐤``. +""" +function calcS!(S, Fp, Fm, nlayers, grid) + F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) + + for n=1:grid.nl, m=1:grid.nkr + k² = CUDA.@allowscalar grid.Krsq[m, n] + Skl = SMatrix{nlayers, nlayers}(- k² * I + F) + S[m, n] = Skl + end + + return nothing +end + +""" + calcS⁻¹!(S, Fp, Fm, nlayers, grid) + +Construct the array ``𝕊⁻¹``, which consists of `nlayer` x `nlayer` static arrays ``(𝕊_𝐤)⁻¹`` +that relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``ψ̂_𝐤 = (𝕊_𝐤)⁻¹ q̂_𝐤``. +""" +function calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) + F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) + + for n=1:grid.nl, m=1:grid.nkr + k² = CUDA.@allowscalar grid.Krsq[m, n] == 0 ? 1 : grid.Krsq[m, n] + Skl = - k² * I + F + S⁻¹[m, n] = SMatrix{nlayers, nlayers}(I / Skl) + end + + T = eltype(grid) + S⁻¹[1, 1] = SMatrix{nlayers, nlayers}(zeros(T, (nlayers, nlayers))) + + return nothing +end + + +# ------- +# Solvers +# ------- + +""" + calcN!(N, sol, t, clock, vars, params, grid) + +Compute the nonlinear term, that is the advection term, the bottom drag, and the forcing: + +```math +N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} + + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . +``` +""" +function calcN!(N, sol, t, clock, vars, params, grid) + nlayers = numberoflayers(params) + + dealias!(sol, grid) + + calcN_advection!(N, sol, vars, params, grid) + + @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag + + addforcing!(N, sol, t, clock, vars, params, grid) + + return nothing +end + +""" + calcNlinear!(N, sol, t, clock, vars, params, grid) + +Compute the nonlinear term of the linearized equations: + +```math +N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} + \\widehat{(∂_y ψ_j)(∂_x Q_j)} +- \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . +``` +""" +function calcNlinear!(N, sol, t, clock, vars, params, grid) + nlayers = numberoflayers(params) + + calcN_linearadvection!(N, sol, vars, params, grid) + @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag + addforcing!(N, sol, t, clock, vars, params, grid) + + return nothing +end + +""" + calcN_advection!(N, sol, vars, params, grid) + +Compute the advection term and stores it in `N`: + +```math +N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} + + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} . +``` +""" +function calcN_advection!(N, sol, vars, params, grid) + @. vars.qh = sol + + streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) + + @. vars.uh = -im * grid.l * vars.ψh + @. vars.vh = im * grid.kr * vars.ψh + + invtransform!(vars.u, vars.uh, params) + @. vars.u += params.U # add the imposed zonal flow U + + uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables + @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x + fwdtransform!(uQxh, uQx, params) + @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} + + invtransform!(vars.v, vars.vh, params) + + vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables + @. vQy = vars.v * params.Qy # v*∂Q/∂y + fwdtransform!(vQyh, vQy, params) + @. N -= vQyh # -\hat{v*∂Q/∂y} + + invtransform!(vars.q, vars.qh, params) + + uq , vq = vars.u , vars.v # use vars.u and vars.v as scratch variables + uqh, vqh = vars.uh, vars.vh # use vars.uh and vars.vh as scratch variables + @. uq *= vars.q # (U+u)*q + @. vq *= vars.q # v*q + + fwdtransform!(uqh, uq, params) + fwdtransform!(vqh, vq, params) + + @. N -= im * grid.kr * uqh + im * grid.l * vqh # -\hat{∂[(U+u)q]/∂x} - \hat{∂[vq]/∂y} + + return nothing +end + + +""" + calcN_linearadvection!(N, sol, vars, params, grid) + +Compute the advection term of the linearized equations and stores it in `N`: + +```math +N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} + + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} . +``` +""" +function calcN_linearadvection!(N, sol, vars, params, grid) + @. vars.qh = sol + + streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) + + @. vars.uh = -im * grid.l * vars.ψh + @. vars.vh = im * grid.kr * vars.ψh + + invtransform!(vars.u, vars.uh, params) + @. vars.u += params.U # add the imposed zonal flow U + uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables + @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x + fwdtransform!(uQxh, uQx, params) + @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} + + invtransform!(vars.v, vars.vh, params) + + vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables + + @. vQy = vars.v * params.Qy # v*∂Q/∂y + fwdtransform!(vQyh, vQy, params) + @. N -= vQyh # -\hat{v*∂Q/∂y} + + invtransform!(vars.q, vars.qh, params) + + @. vars.u = params.U + Uq , Uqh = vars.u , vars.uh # use vars.u and vars.uh as scratch variables + @. Uq *= vars.q # U*q + + fwdtransform!(Uqh, Uq, params) + + @. N -= im * grid.kr * Uqh # -\hat{∂[U*q]/∂x} + + return nothing +end + + +""" + addforcing!(N, sol, t, clock, vars, params, grid) + +When the problem includes forcing, calculate the forcing term ``F̂`` for each layer and add +it to the nonlinear term ``N``. +""" +addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing + +function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) + params.calcFq!(vars.Fqh, sol, t, clock, vars, params, grid) + @. N += vars.Fqh + + return nothing +end + + +# ---------------- +# Helper functions +# ---------------- + +""" + updatevars!(vars, params, grid, sol) + updatevars!(prob) + +Update all problem variables using `sol`. +""" +function updatevars!(vars, params, grid, sol) + dealias!(sol, grid) + + @. vars.qh = sol + streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) + @. vars.uh = -im * grid.l * vars.ψh + @. vars.vh = im * grid.kr * vars.ψh + + invtransform!(vars.q, deepcopy(vars.qh), params) + invtransform!(vars.ψ, deepcopy(vars.ψh), params) + invtransform!(vars.u, deepcopy(vars.uh), params) + invtransform!(vars.v, deepcopy(vars.vh), params) + + return nothing +end + +updatevars!(prob) = updatevars!(prob.vars, prob.params, prob.grid, prob.sol) + + +""" + set_q!(sol, params, vars, grid, q) + set_q!(prob, q) + +Set the solution `prob.sol` as the transform of `q` and update variables. +""" +function set_q!(sol, params, vars, grid, q) + A = typeof(vars.q) + fwdtransform!(vars.qh, A(q), params) + @. vars.qh[1, 1, :] = 0 + @. sol = vars.qh + updatevars!(vars, params, grid, sol) + + return nothing +end + +function set_q!(sol, params::SingleLayerParams, vars, grid, q::AbstractArray{T, 2}) where T + A = typeof(vars.q[:, :, 1]) + q_3D = vars.q + @views q_3D[:, :, 1] = A(q) + set_q!(sol, params, vars, grid, q_3D) + + return nothing +end + +set_q!(prob, q) = set_q!(prob.sol, prob.params, prob.vars, prob.grid, q) + + +""" + set_ψ!(params, vars, grid, sol, ψ) + set_ψ!(prob, ψ) + +Set the solution `prob.sol` to the transform `qh` that corresponds to streamfunction `ψ` +and update variables. +""" +function set_ψ!(sol, params, vars, grid, ψ) + A = typeof(vars.q) + fwdtransform!(vars.ψh, A(ψ), params) + pvfromstreamfunction!(vars.qh, vars.ψh, params, grid) + invtransform!(vars.q, vars.qh, params) + + set_q!(sol, params, vars, grid, vars.q) + + return nothing +end + +function set_ψ!(sol, params::SingleLayerParams, vars, grid, ψ::AbstractArray{T, 2}) where T + A = typeof(vars.ψ[:, :, 1]) + ψ_3D = vars.ψ + @views ψ_3D[:, :, 1] = A(ψ) + + set_ψ!(sol, params, vars, grid, ψ_3D) + + return nothing +end + +set_ψ!(prob, ψ) = set_ψ!(prob.sol, prob.params, prob.vars, prob.grid, ψ) + + +""" + energies(vars, params, grid, sol) + energies(prob) + +Return the kinetic energy of each fluid layer KE``_1, ...,`` KE``_{n}``, and the +potential energy of each fluid interface PE``_{3/2}, ...,`` PE``_{n-1/2}``, where ``n`` +is the number of layers in the fluid. (When ``n=1``, only the kinetic energy is returned.) + +The kinetic energy at the ``j``-th fluid layer is + +```math +𝖪𝖤_j = \\frac{H_j}{H} \\int \\frac1{2} |{\\bf ∇} ψ_j|^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{H_j}{H} \\sum_{𝐤} |𝐤|² |ψ̂_j|², \\ j = 1, ..., n , +``` + +while the potential energy that corresponds to the interface ``j+1/2`` (i.e., the interface +between the ``j``-th and ``(j+1)``-th fluid layer) is + +```math +𝖯𝖤_{j+1/2} = \\int \\frac1{2} \\frac{f₀^2}{g'_{j+1/2} H} (ψ_j - ψ_{j+1})^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{f₀^2}{g'_{j+1/2} H} \\sum_{𝐤} |ψ̂_j - ψ̂_{j+1}|², \\ j = 1, ..., n-1 . +``` +""" +function energies(vars, params, grid, sol) + nlayers = numberoflayers(params) + KE, PE = zeros(nlayers), zeros(nlayers-1) + + @. vars.qh = sol + streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) + + abs²∇𝐮h = vars.uh # use vars.uh as scratch variable + @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) + + for j = 1:nlayers + CUDA.@allowscalar KE[j] = 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) + end + + for j = 1:nlayers-1 + CUDA.@allowscalar PE[j] = 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′[j] * parsevalsum(abs2.(vars.ψh[:, :, j+1] .- vars.ψh[:, :, j]), grid) + end + + return KE, PE +end + +function energies(vars, params::TwoLayerParams, grid, sol) + nlayers = numberoflayers(params) + KE, PE = zeros(nlayers), zeros(nlayers-1) + + @. vars.qh = sol + streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) + + abs²∇𝐮h = vars.uh # use vars.uh as scratch variable + @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) + + ψ1h, ψ2h = view(vars.ψh, :, :, 1), view(vars.ψh, :, :, 2) + + for j = 1:nlayers + CUDA.@allowscalar KE[j] = @views 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) + end + + PE = @views 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′ * parsevalsum(abs2.(ψ2h .- ψ1h), grid) + + return KE, PE +end + +function energies(vars, params::SingleLayerParams, grid, sol) + @. vars.qh = sol + streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) + + abs²∇𝐮h = vars.uh # use vars.uh as scratch variable + @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) + + return 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h, grid) +end + +energies(prob) = energies(prob.vars, prob.params, prob.grid, prob.sol) + +""" + fluxes(vars, params, grid, sol) + fluxes(prob) + +Return the lateral eddy fluxes within each fluid layer, lateralfluxes``_1,...,``lateralfluxes``_n`` +and also the vertical eddy fluxes at each fluid interface, +verticalfluxes``_{3/2},...,``verticalfluxes``_{n-1/2}``, where ``n`` is the total number of layers in the fluid. +(When ``n=1``, only the lateral fluxes are returned.) + +The lateral eddy fluxes within the ``j``-th fluid layer are + +```math +\\textrm{lateralfluxes}_j = \\frac{H_j}{H} \\int U_j v_j ∂_y u_j +\\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n , +``` + +while the vertical eddy fluxes at the ``j+1/2``-th fluid interface (i.e., interface between +the ``j``-th and ``(j+1)``-th fluid layer) are + +```math +\\textrm{verticalfluxes}_{j+1/2} = \\int \\frac{f₀²}{g'_{j+1/2} H} (U_j - U_{j+1}) \\, +v_{j+1} ψ_{j} \\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n-1. +``` +""" +function fluxes(vars, params, grid, sol) + nlayers = numberoflayers(params) + + lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) + + updatevars!(vars, params, grid, sol) + + ∂u∂yh = vars.uh # use vars.uh as scratch variable + ∂u∂y = vars.u # use vars.u as scratch variable + + @. ∂u∂yh = im * grid.l * vars.uh + invtransform!(∂u∂y, ∂u∂yh, params) + + lateralfluxes = (sum(@. params.H * params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] + lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) + + for j = 1:nlayers-1 + CUDA.@allowscalar verticalfluxes[j] = sum(@views @. params.f₀^2 / params.g′[j] * (params.U[: ,:, j] - params.U[:, :, j+1]) * vars.v[:, :, j+1] * vars.ψ[:, :, j]; dims=(1, 2))[1] + CUDA.@allowscalar verticalfluxes[j] *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) + end + + return lateralfluxes, verticalfluxes +end + +function fluxes(vars, params::TwoLayerParams, grid, sol) + nlayers = numberoflayers(params) + + lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) + + updatevars!(vars, params, grid, sol) + + ∂u∂yh = vars.uh # use vars.uh as scratch variable + ∂u∂y = vars.u # use vars.u as scratch variable + + @. ∂u∂yh = im * grid.l * vars.uh + invtransform!(∂u∂y, ∂u∂yh, params) + + lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] + @. lateralfluxes *= params.H + lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) + + U₁, U₂ = view(params.U, :, :, 1), view(params.U, :, :, 2) + ψ₁ = view(vars.ψ, :, :, 1) + v₂ = view(vars.v, :, :, 2) + + verticalfluxes = sum(@views @. params.f₀^2 / params.g′ * (U₁ - U₂) * v₂ * ψ₁; dims=(1, 2)) + verticalfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) + + return lateralfluxes, verticalfluxes +end + +function fluxes(vars, params::SingleLayerParams, grid, sol) + updatevars!(vars, params, grid, sol) + + ∂u∂yh = vars.uh # use vars.uh as scratch variable + ∂u∂y = vars.u # use vars.u as scratch variable + + @. ∂u∂yh = im * grid.l * vars.uh + invtransform!(∂u∂y, ∂u∂yh, params) + + lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] + lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly) + + return lateralfluxes +end + +fluxes(prob) = fluxes(prob.vars, prob.params, prob.grid, prob.sol) + +end # module diff --git a/src/multilayerqg_sign_fix_only.jl b/src/multilayerqg_sign_fix_only.jl new file mode 100644 index 00000000..6d3a27d3 --- /dev/null +++ b/src/multilayerqg_sign_fix_only.jl @@ -0,0 +1,1107 @@ +module MultiLayerQG_sf + +export + fwdtransform!, + invtransform!, + streamfunctionfrompv!, + pvfromstreamfunction!, + updatevars!, + + set_q!, + set_ψ!, + energies, + fluxes + +using + FFTW, + CUDA, + LinearAlgebra, + StaticArrays, + Reexport, + DocStringExtensions + +@reexport using FourierFlows + +using FourierFlows: parsevalsum, parsevalsum2, superzeros, plan_flows_rfft + +nothingfunction(args...) = nothing + +""" + Problem(nlayers :: Int, + dev = CPU(); + nx = 128, + ny = nx, + Lx = 2π, + Ly = Lx, + f₀ = 1.0, + β = 0.0, + g = 1.0, + U = zeros(nlayers), + H = 1/nlayers * ones(nlayers), + ρ = Array{Float64}(1:nlayers), + eta = nothing, + topographic_pv_gradient = (0, 0), + μ = 0, + ν = 0, + nν = 1, + dt = 0.01, + stepper = "RK4", + calcFq = nothingfunction, + stochastic = false, + linear = false, + aliased_fraction = 1/3, + T = Float64) + +Construct a multi-layer quasi-geostrophic problem with `nlayers` fluid layers on device `dev`. + +Arguments +========= +- `nlayers`: (required) Number of fluid layers. +- `dev`: (required) `CPU()` (default) or `GPU()`; computer architecture used to time-step `problem`. + +Keyword arguments +================= + - `nx`: Number of grid points in ``x``-domain. + - `ny`: Number of grid points in ``y``-domain. + - `Lx`: Extent of the ``x``-domain. + - `Ly`: Extent of the ``y``-domain. + - `f₀`: Constant planetary vorticity. + - `β`: Planetary vorticity ``y``-gradient. + - `g`: Gravitational acceleration constant. + - `U`: The imposed constant zonal flow ``U(y)`` in each fluid layer. + - `H`: Rest height of each fluid layer. + - `ρ`: Density of each fluid layer. + - `eta`: Periodic component of the topographic potential vorticity. + - `topographic_pv_gradient`: The ``(x, y)`` components of the topographic PV large-scale gradient. + - `μ`: Linear bottom drag coefficient. + - `ν`: Small-scale (hyper)-viscosity coefficient. + - `nν`: (Hyper)-viscosity order, `nν```≥ 1``. + - `dt`: Time-step. + - `stepper`: Time-stepping method. + - `calcF`: Function that calculates the Fourier transform of the forcing, ``F̂``. + - `stochastic`: `true` or `false` (default); boolean denoting whether `calcF` is temporally stochastic. + - `linear`: `true` or `false` (default); boolean denoting whether the linearized equations of motions are used. + - `aliased_fraction`: the fraction of high-wavenumbers that are zero-ed out by `dealias!()`. + - `T`: `Float32` or `Float64` (default); floating point type used for `problem` data. +""" +function Problem_sf(nlayers::Int, # number of fluid layers + dev = CPU(); + # Numerical parameters + nx = 128, + ny = nx, + Lx = 2π, + Ly = Lx, + # Physical parameters + f₀ = 1.0, # Coriolis parameter + β = 0.0, # y-gradient of Coriolis parameter + g = 1.0, # gravitational constant + U = zeros(nlayers), # imposed zonal flow U(y) in each layer + H = 1/nlayers * ones(nlayers), # rest fluid height of each layer + ρ = Array{Float64}(1:nlayers), # density of each layer + eta = nothing, # periodic component of the topographic PV + topographic_pv_gradient = (0, 0), # tuple with the ``(x, y)`` components of topographic PV large-scale gradient + # Bottom Drag and/or (hyper)-viscosity + μ = 0, + ν = 0, + nν = 1, + # Timestepper and equation options + dt = 0.01, + stepper = "RK4", + calcFq = nothingfunction, + stochastic = false, + linear = false, + # Float type and dealiasing + aliased_fraction = 1/3, + T = Float64) + + if dev == GPU() && nlayers > 2 + @warn """MultiLayerQG module is not optimized on the GPU yet for configurations with + 3 fluid layers or more! + + See issues on Github at https://github.com/FourierFlows/GeophysicalFlows.jl/issues/112 + and https://github.com/FourierFlows/GeophysicalFlows.jl/issues/267. + + To use MultiLayerQG with 3 fluid layers or more we suggest, for now, to restrict running + on CPU.""" + end + + if nlayers == 1 + @warn """MultiLayerQG module does work for single-layer configuration but may not be as + optimized. We suggest using SingleLayerQG module for single-layer QG simulation unless + you have reasons to use MultiLayerQG in a single-layer configuration, e.g., you want to + compare solutions with varying number of fluid layers.""" + end + + # topographic PV + eta === nothing && (eta = zeros(dev, T, (nx, ny))) + + grid = TwoDGrid(dev; nx, Lx, ny, Ly, aliased_fraction, T) + + params = Params_sf(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ, ν, nν, grid; calcFq) + + vars = calcFq == nothingfunction ? DecayingVars(grid, params) : (stochastic ? StochasticForcedVars(grid, params) : ForcedVars(grid, params)) + + equation = linear ? LinearEquation(params, grid) : Equation(params, grid) + + FourierFlows.Problem(equation, stepper, dt, grid, vars, params) +end + +""" + struct Params{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams + +The parameters for the `MultiLayerQG` problem. + +$(TYPEDFIELDS) +""" +struct Params_sf{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams + # prescribed params + "number of fluid layers" + nlayers :: Int + "gravitational constant" + g :: T + "constant planetary vorticity" + f₀ :: T + "planetary vorticity ``y``-gradient" + β :: T + "array with density of each fluid layer" + ρ :: Aphys3D + "array with rest height of each fluid layer" + H :: Aphys3D + "array with imposed constant zonal flow ``U(y)`` in each fluid layer" + U :: Aphys3D + "array containing the topographic PV" + eta :: Aphys2D + "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" + topographic_pv_gradient :: Tuple{T, T} + "linear bottom drag coefficient" + μ :: T + "small-scale (hyper)-viscosity coefficient" + ν :: T + "(hyper)-viscosity order, `nν```≥ 1``" + nν :: Int + "function that calculates the Fourier transform of the forcing, ``F̂``" + calcFq! :: Function + + # derived params + "array with the reduced gravity constants for each fluid interface" + g′ :: Aphys1D + "array containing ``x``-gradient of PV due to topographic PV in each fluid layer" + Qx :: Aphys3D + "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV in each fluid layer" + Qy :: Aphys3D + "array containing coeffients for getting PV from streamfunction" + S :: Atrans4D + "array containing coeffients for inverting PV to streamfunction" + S⁻¹ :: Atrans4D + "rfft plan for FFTs" + rfftplan :: Trfft +end + +""" + struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams + +The parameters for the a single-layer `MultiLayerQG` problem. + +$(TYPEDFIELDS) +""" +struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams + # prescribed params + "planetary vorticity ``y``-gradient" + β :: T + "array with imposed constant zonal flow ``U(y)``" + U :: Aphys3D + "array containing the periodic component of the topographic PV" + eta :: Aphys2D + "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" + topographic_pv_gradient :: Tuple{T, T} + "linear drag coefficient" + μ :: T + "small-scale (hyper)-viscosity coefficient" + ν :: T + "(hyper)-viscosity order, `nν```≥ 1``" + nν :: Int + "function that calculates the Fourier transform of the forcing, ``F̂``" + calcFq! :: Function + + # derived params + "array containing ``x``-gradient of PV due to topographic PV" + Qx :: Aphys3D + "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV" + Qy :: Aphys3D + "rfft plan for FFTs" + rfftplan :: Trfft +end + +""" + struct TwoLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams + +The parameters for the a two-layer `MultiLayerQG` problem. + +$(TYPEDFIELDS) +""" +struct TwoLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams + # prescribed params + "gravitational constant" + g :: T + "constant planetary vorticity" + f₀ :: T + "planetary vorticity ``y``-gradient" + β :: T + "array with density of each fluid layer" + ρ :: Aphys3D + "tuple with rest height of each fluid layer" + H :: Tuple + "array with imposed constant zonal flow ``U(y)`` in each fluid layer" + U :: Aphys3D + "array containing periodic component of the topographic PV" + eta :: Aphys2D + "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" + topographic_pv_gradient :: Tuple{T, T} + "linear bottom drag coefficient" + μ :: T + "small-scale (hyper)-viscosity coefficient" + ν :: T + "(hyper)-viscosity order, `nν```≥ 1``" + nν :: Int + "function that calculates the Fourier transform of the forcing, ``F̂``" + calcFq! :: Function + + # derived params + "the reduced gravity constants for the fluid interface" + g′ :: T + "array containing ``x``-gradient of PV due to topographic PV in each fluid layer" + Qx :: Aphys3D + "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV in each fluid layer" + Qy :: Aphys3D + "rfft plan for FFTs" + rfftplan :: Trfft +end + +function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 1}) where TU + T = eltype(grid) + if length(U) == nlayers + U_2D = zeros(dev, T, (1, nlayers)) + U_2D[:] = U + U_2D = repeat(U_2D, outer=(grid.ny, 1)) + else + U_2D = zeros(dev, T, (grid.ny, 1)) + U_2D[:] = U + end + U_3D = zeros(dev, T, (1, grid.ny, nlayers)) + @views U_3D[1, :, :] = U_2D + return U_3D +end + +function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 2}) where TU + T = eltype(grid) + U_3D = zeros(dev, T, (1, grid.ny, nlayers)) + @views U_3D[1, :, :] = U + return U_3D +end + +function convert_U_to_U3D(dev, nlayers, grid, U::Number) + T = eltype(grid) + A = device_array(dev) + U_3D = reshape(repeat([T(U)], outer=(grid.ny, 1)), (1, grid.ny, nlayers)) + return A(U_3D) +end + +function Params_sf(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ, ν, nν, grid; calcFq=nothingfunction, effort=FFTW.MEASURE) + dev = grid.device + T = eltype(grid) + A = device_array(dev) + + ny, nx = grid.ny , grid.nx + nkr, nl = grid.nkr, grid.nl + kr, l = grid.kr , grid.l + + U = convert_U_to_U3D(dev, nlayers, grid, U) + + Uyy = real.(ifft(-l.^2 .* fft(U))) + Uyy = CUDA.@allowscalar repeat(Uyy, outer=(nx, 1, 1)) + + # Calculate periodic components of the topographic PV gradients. + etah = rfft(A(eta)) + etax = irfft(im * kr .* etah, nx) # ∂η/∂x + etay = irfft(im * l .* etah, nx) # ∂η/∂y + + # Add topographic PV large-scale gradient + topographic_pv_gradient = T.(topographic_pv_gradient) + @. etax += topographic_pv_gradient[1] + @. etay += topographic_pv_gradient[2] + + Qx = zeros(dev, T, (nx, ny, nlayers)) + @views @. Qx[:, :, nlayers] += etax + + Qy = zeros(dev, T, (nx, ny, nlayers)) + Qy = T(β) .- Uyy # T(β) is needed to ensure that Qy remains same type as U + @views @. Qy[:, :, nlayers] += etay + + rfftplanlayered = plan_flows_rfft(A{T, 3}(undef, grid.nx, grid.ny, nlayers), [1, 2]; flags=effort) + + if nlayers==1 + return SingleLayerParams(T(β), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, Qx, Qy, rfftplanlayered) + + else # if nlayers≥2 + + ρ = reshape(T.(ρ), (1, 1, nlayers)) + H = reshape(T.(H), (1, 1, nlayers)) + + g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface + + Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) + Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) + + typeofSkl = SArray{Tuple{nlayers, nlayers}, T, 2, nlayers^2} # StaticArrays of type T and dims = (nlayers x nlayers) + + S = Array{typeofSkl, 2}(undef, (nkr, nl)) + calcS!(S, Fp, Fm, nlayers, grid) + + S⁻¹ = Array{typeofSkl, 2}(undef, (nkr, nl)) + calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) + + S, S⁻¹, Fp, Fm = A(S), A(S⁻¹), A(Fp), A(Fm) # convert to appropriate ArrayType + + CUDA.@allowscalar @views Qy[:, :, 1] = @. Qy[:, :, 1] - Fp[1] * (U[:, :, 2] - U[:, :, 1]) + for j = 2:nlayers-1 + CUDA.@allowscalar @views Qy[:, :, j] = @. Qy[:, :, j] + Fp[j] * (U[:, :, j+1] - U[:, :, j]) + Fm[j-1] * (U[:, :, j-1] - U[:, :, j]) + end + CUDA.@allowscalar @views Qy[:, :, nlayers] = @. Qy[:, :, nlayers] - Fm[nlayers-1] * (U[:, :, nlayers-1] - U[:, :, nlayers]) + + if nlayers==2 + return TwoLayerParams(T(g), T(f₀), T(β), A(ρ), (T(H[1]), T(H[2])), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, T(g′[1]), Qx, Qy, rfftplanlayered) + else # if nlayers>2 + return Params_sf(nlayers, T(g), T(f₀), T(β), A(ρ), A(H), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, A(g′), Qx, Qy, S, S⁻¹, rfftplanlayered) + end + end +end + +numberoflayers(params) = params.nlayers +numberoflayers(::SingleLayerParams) = 1 +numberoflayers(::TwoLayerParams) = 2 + +# --------- +# Equations +# --------- + +""" + hyperviscosity(params, grid) + +Return the linear operator `L` that corresponds to (hyper)-viscosity of order ``n_ν`` with +coefficient ``ν`` for ``n`` fluid layers. +```math +L_j = - ν |𝐤|^{2 n_ν}, \\ j = 1, ...,n . +``` +""" +function hyperviscosity(params, grid) + dev = grid.device + T = eltype(grid) + + L = device_array(dev){T}(undef, (grid.nkr, grid.nl, numberoflayers(params))) + @. L = - params.ν * grid.Krsq^params.nν + @views @. L[1, 1, :] = 0 + + return L +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 +`hyperviscosity(params, grid)`. + +The nonlinear term is computed via function `calcNlinear!`. +""" +function LinearEquation(params, grid) + L = hyperviscosity(params, grid) + + return FourierFlows.Equation(L, calcNlinear!, grid) +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 +`hyperviscosity(params, grid)`. + +The nonlinear term is computed via function `calcN!`. +""" +function Equation(params, grid) + L = hyperviscosity(params, grid) + + return FourierFlows.Equation(L, calcN!, grid) +end + + +# ---- +# Vars +# ---- + +""" + struct Vars{Aphys, Atrans, F, P} <: AbstractVars + +The variables for multi-layer QG problem. + +$(FIELDS) +""" +struct Vars{Aphys, Atrans, F, P} <: AbstractVars + "relative vorticity + vortex stretching" + q :: Aphys + "streamfunction" + ψ :: Aphys + "x-component of velocity" + u :: Aphys + "y-component of velocity" + v :: Aphys + "Fourier transform of relative vorticity + vortex stretching" + qh :: Atrans + "Fourier transform of streamfunction" + ψh :: Atrans + "Fourier transform of ``x``-component of velocity" + uh :: Atrans + "Fourier transform of ``y``-component of velocity" + vh :: Atrans + "Fourier transform of forcing" + Fqh :: F + "`sol` at previous time-step" + prevsol :: P +end + +const DecayingVars = Vars{<:AbstractArray, <:AbstractArray, Nothing, Nothing} +const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing} +const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray} + +""" + DecayingVars(grid, params) + +Return the variables for an unforced multi-layer QG problem with `grid` and `params`. +""" +function DecayingVars(grid, params) + Dev = typeof(grid.device) + T = eltype(grid) + nlayers = numberoflayers(params) + + @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v + @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh + + return Vars(q, ψ, u, v, qh, ψh, uh, vh, nothing, nothing) +end + +""" + ForcedVars(grid, params) + +Return the variables for a forced multi-layer QG problem with `grid` and `params`. +""" +function ForcedVars(grid, params) + Dev = typeof(grid.device) + T = eltype(grid) + nlayers = numberoflayers(params) + + @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v + @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh + + return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, nothing) +end + +""" + StochasticForcedVars(grid, params) + +Return the variables for a forced multi-layer QG problem with `grid` and `params`. +""" +function StochasticForcedVars(grid, params) + Dev = typeof(grid.device) + T = eltype(grid) + nlayers = numberoflayers(params) + + @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v + @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh prevsol + + return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, prevsol) +end + +""" + fwdtransform!(varh, var, params) + +Compute the Fourier transform of `var` and store it in `varh`. +""" +fwdtransform!(varh, var, params::AbstractParams) = mul!(varh, params.rfftplan, var) + +""" + invtransform!(var, varh, params) + +Compute the inverse Fourier transform of `varh` and store it in `var`. +""" +invtransform!(var, varh, params::AbstractParams) = ldiv!(var, params.rfftplan, varh) + +""" + pvfromstreamfunction!(qh, ψh, params, grid) + +Obtain the Fourier transform of the PV from the streamfunction `ψh` in each layer using +`qh = params.S * ψh`. +""" +function pvfromstreamfunction!(qh, ψh, params, grid) + for j=1:grid.nl, i=1:grid.nkr + CUDA.@allowscalar @views qh[i, j, :] .= params.S[i, j] * ψh[i, j, :] + end + + return nothing +end + +""" + pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) + +Obtain the Fourier transform of the PV from the streamfunction `ψh` for the special +case of a single fluid layer configuration. In this case, ``q̂ = - k² ψ̂``. +""" +function pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) + @. qh = -grid.Krsq * ψh + + return nothing +end + +""" + pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) + +Obtain the Fourier transform of the PV from the streamfunction `ψh` for the special +case of a two fluid layer configuration. In this case we have, + +```math +q̂₁ = - k² ψ̂₁ + f₀² / (g′ H₁) * (ψ̂₂ - ψ̂₁) , +``` + +```math +q̂₂ = - k² ψ̂₂ + f₀² / (g′ H₂) * (ψ̂₁ - ψ̂₂) . +``` + +(Here, the PV-streamfunction relationship is hard-coded to avoid scalar operations +on the GPU.) +""" +function pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) + f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] + + ψ1h, ψ2h = view(ψh, :, :, 1), view(ψh, :, :, 2) + + @views @. qh[:, :, 1] = - grid.Krsq * ψ1h + f₀^2 / (g′ * H₁) * (ψ2h - ψ1h) + @views @. qh[:, :, 2] = - grid.Krsq * ψ2h + f₀^2 / (g′ * H₂) * (ψ1h - ψ2h) + + return nothing +end + +""" + streamfunctionfrompv!(ψh, qh, params, grid) + +Invert the PV to obtain the Fourier transform of the streamfunction `ψh` in each layer from +`qh` using `ψh = params.S⁻¹ qh`. +""" +function streamfunctionfrompv!(ψh, qh, params, grid) + for j=1:grid.nl, i=1:grid.nkr + CUDA.@allowscalar @views ψh[i, j, :] .= params.S⁻¹[i, j] * qh[i, j, :] + end + + return nothing +end + +""" + streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) + +Invert the PV to obtain the Fourier transform of the streamfunction `ψh` for the special +case of a single fluid layer configuration. In this case, ``ψ̂ = - k⁻² q̂``. +""" +function streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) + @. ψh = -grid.invKrsq * qh + + return nothing +end + +""" + streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) + +Invert the PV to obtain the Fourier transform of the streamfunction `ψh` for the special +case of a two fluid layer configuration. In this case we have, + +```math +ψ̂₁ = - [k² q̂₁ + (f₀² / g′) (q̂₁ / H₂ + q̂₂ / H₁)] / Δ , +``` + +```math +ψ̂₂ = - [k² q̂₂ + (f₀² / g′) (q̂₁ / H₂ + q̂₂ / H₁)] / Δ , +``` + +where ``Δ = k² [k² + f₀² (H₁ + H₂) / (g′ H₁ H₂)]``. + +(Here, the PV-streamfunction relationship is hard-coded to avoid scalar operations +on the GPU.) +""" +function streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) + f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] + + q1h, q2h = view(qh, :, :, 1), view(qh, :, :, 2) + + @views @. ψh[:, :, 1] = - grid.Krsq * q1h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) + @views @. ψh[:, :, 2] = - grid.Krsq * q2h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) + + for j in 1:2 + @views @. ψh[:, :, j] *= grid.invKrsq / (grid.Krsq + f₀^2 / g′ * (H₁ + H₂) / (H₁ * H₂)) + end + + return nothing +end + +""" + calcS!(S, Fp, Fm, nlayers, grid) + +Construct the array ``𝕊``, which consists of `nlayer` x `nlayer` static arrays ``𝕊_𝐤`` that +relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``q̂_𝐤 = 𝕊_𝐤 ψ̂_𝐤``. +""" +function calcS!(S, Fp, Fm, nlayers, grid) + F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) + + for n=1:grid.nl, m=1:grid.nkr + k² = CUDA.@allowscalar grid.Krsq[m, n] + Skl = SMatrix{nlayers, nlayers}(- k² * I + F) + S[m, n] = Skl + end + + return nothing +end + +""" + calcS⁻¹!(S, Fp, Fm, nlayers, grid) + +Construct the array ``𝕊⁻¹``, which consists of `nlayer` x `nlayer` static arrays ``(𝕊_𝐤)⁻¹`` +that relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``ψ̂_𝐤 = (𝕊_𝐤)⁻¹ q̂_𝐤``. +""" +function calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) + F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) + + for n=1:grid.nl, m=1:grid.nkr + k² = CUDA.@allowscalar grid.Krsq[m, n] == 0 ? 1 : grid.Krsq[m, n] + Skl = - k² * I + F + S⁻¹[m, n] = SMatrix{nlayers, nlayers}(I / Skl) + end + + T = eltype(grid) + S⁻¹[1, 1] = SMatrix{nlayers, nlayers}(zeros(T, (nlayers, nlayers))) + + return nothing +end + + +# ------- +# Solvers +# ------- + +""" + calcN!(N, sol, t, clock, vars, params, grid) + +Compute the nonlinear term, that is the advection term, the bottom drag, and the forcing: + +```math +N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} + + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . +``` +""" +function calcN!(N, sol, t, clock, vars, params, grid) + nlayers = numberoflayers(params) + + dealias!(sol, grid) + + calcN_advection!(N, sol, vars, params, grid) + + @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag + + addforcing!(N, sol, t, clock, vars, params, grid) + + return nothing +end + +""" + calcNlinear!(N, sol, t, clock, vars, params, grid) + +Compute the nonlinear term of the linearized equations: + +```math +N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} + \\widehat{(∂_y ψ_j)(∂_x Q_j)} +- \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . +``` +""" +function calcNlinear!(N, sol, t, clock, vars, params, grid) + nlayers = numberoflayers(params) + + calcN_linearadvection!(N, sol, vars, params, grid) + @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag + addforcing!(N, sol, t, clock, vars, params, grid) + + return nothing +end + +""" + calcN_advection!(N, sol, vars, params, grid) + +Compute the advection term and stores it in `N`: + +```math +N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} + + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} . +``` +""" +function calcN_advection!(N, sol, vars, params, grid) + @. vars.qh = sol + + streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) + + @. vars.uh = -im * grid.l * vars.ψh + @. vars.vh = im * grid.kr * vars.ψh + + invtransform!(vars.u, vars.uh, params) + @. vars.u += params.U # add the imposed zonal flow U + + uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables + @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x + fwdtransform!(uQxh, uQx, params) + @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} + + invtransform!(vars.v, vars.vh, params) + + vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables + @. vQy = vars.v * params.Qy # v*∂Q/∂y + fwdtransform!(vQyh, vQy, params) + @. N -= vQyh # -\hat{v*∂Q/∂y} + + invtransform!(vars.q, vars.qh, params) + + uq , vq = vars.u , vars.v # use vars.u and vars.v as scratch variables + uqh, vqh = vars.uh, vars.vh # use vars.uh and vars.vh as scratch variables + @. uq *= vars.q # (U+u)*q + @. vq *= vars.q # v*q + + fwdtransform!(uqh, uq, params) + fwdtransform!(vqh, vq, params) + + @. N -= im * grid.kr * uqh + im * grid.l * vqh # -\hat{∂[(U+u)q]/∂x} - \hat{∂[vq]/∂y} + + return nothing +end + + +""" + calcN_linearadvection!(N, sol, vars, params, grid) + +Compute the advection term of the linearized equations and stores it in `N`: + +```math +N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} + + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} . +``` +""" +function calcN_linearadvection!(N, sol, vars, params, grid) + @. vars.qh = sol + + streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) + + @. vars.uh = -im * grid.l * vars.ψh + @. vars.vh = im * grid.kr * vars.ψh + + invtransform!(vars.u, vars.uh, params) + @. vars.u += params.U # add the imposed zonal flow U + uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables + @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x + fwdtransform!(uQxh, uQx, params) + @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} + + invtransform!(vars.v, vars.vh, params) + + vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables + + @. vQy = vars.v * params.Qy # v*∂Q/∂y + fwdtransform!(vQyh, vQy, params) + @. N -= vQyh # -\hat{v*∂Q/∂y} + + invtransform!(vars.q, vars.qh, params) + + @. vars.u = params.U + Uq , Uqh = vars.u , vars.uh # use vars.u and vars.uh as scratch variables + @. Uq *= vars.q # U*q + + fwdtransform!(Uqh, Uq, params) + + @. N -= im * grid.kr * Uqh # -\hat{∂[U*q]/∂x} + + return nothing +end + + +""" + addforcing!(N, sol, t, clock, vars, params, grid) + +When the problem includes forcing, calculate the forcing term ``F̂`` for each layer and add +it to the nonlinear term ``N``. +""" +addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing + +function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) + params.calcFq!(vars.Fqh, sol, t, clock, vars, params, grid) + @. N += vars.Fqh + + return nothing +end + + +# ---------------- +# Helper functions +# ---------------- + +""" + updatevars!(vars, params, grid, sol) + updatevars!(prob) + +Update all problem variables using `sol`. +""" +function updatevars!(vars, params, grid, sol) + dealias!(sol, grid) + + @. vars.qh = sol + streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) + @. vars.uh = -im * grid.l * vars.ψh + @. vars.vh = im * grid.kr * vars.ψh + + invtransform!(vars.q, deepcopy(vars.qh), params) + invtransform!(vars.ψ, deepcopy(vars.ψh), params) + invtransform!(vars.u, deepcopy(vars.uh), params) + invtransform!(vars.v, deepcopy(vars.vh), params) + + return nothing +end + +updatevars!(prob) = updatevars!(prob.vars, prob.params, prob.grid, prob.sol) + + +""" + set_q!(sol, params, vars, grid, q) + set_q!(prob, q) + +Set the solution `prob.sol` as the transform of `q` and update variables. +""" +function set_q!(sol, params, vars, grid, q) + A = typeof(vars.q) + fwdtransform!(vars.qh, A(q), params) + @. vars.qh[1, 1, :] = 0 + @. sol = vars.qh + updatevars!(vars, params, grid, sol) + + return nothing +end + +function set_q!(sol, params::SingleLayerParams, vars, grid, q::AbstractArray{T, 2}) where T + A = typeof(vars.q[:, :, 1]) + q_3D = vars.q + @views q_3D[:, :, 1] = A(q) + set_q!(sol, params, vars, grid, q_3D) + + return nothing +end + +set_q!(prob, q) = set_q!(prob.sol, prob.params, prob.vars, prob.grid, q) + + +""" + set_ψ!(params, vars, grid, sol, ψ) + set_ψ!(prob, ψ) + +Set the solution `prob.sol` to the transform `qh` that corresponds to streamfunction `ψ` +and update variables. +""" +function set_ψ!(sol, params, vars, grid, ψ) + A = typeof(vars.q) + fwdtransform!(vars.ψh, A(ψ), params) + pvfromstreamfunction!(vars.qh, vars.ψh, params, grid) + invtransform!(vars.q, vars.qh, params) + + set_q!(sol, params, vars, grid, vars.q) + + return nothing +end + +function set_ψ!(sol, params::SingleLayerParams, vars, grid, ψ::AbstractArray{T, 2}) where T + A = typeof(vars.ψ[:, :, 1]) + ψ_3D = vars.ψ + @views ψ_3D[:, :, 1] = A(ψ) + + set_ψ!(sol, params, vars, grid, ψ_3D) + + return nothing +end + +set_ψ!(prob, ψ) = set_ψ!(prob.sol, prob.params, prob.vars, prob.grid, ψ) + + +""" + energies(vars, params, grid, sol) + energies(prob) + +Return the kinetic energy of each fluid layer KE``_1, ...,`` KE``_{n}``, and the +potential energy of each fluid interface PE``_{3/2}, ...,`` PE``_{n-1/2}``, where ``n`` +is the number of layers in the fluid. (When ``n=1``, only the kinetic energy is returned.) + +The kinetic energy at the ``j``-th fluid layer is + +```math +𝖪𝖤_j = \\frac{H_j}{H} \\int \\frac1{2} |{\\bf ∇} ψ_j|^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{H_j}{H} \\sum_{𝐤} |𝐤|² |ψ̂_j|², \\ j = 1, ..., n , +``` + +while the potential energy that corresponds to the interface ``j+1/2`` (i.e., the interface +between the ``j``-th and ``(j+1)``-th fluid layer) is + +```math +𝖯𝖤_{j+1/2} = \\int \\frac1{2} \\frac{f₀^2}{g'_{j+1/2} H} (ψ_j - ψ_{j+1})^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{f₀^2}{g'_{j+1/2} H} \\sum_{𝐤} |ψ̂_j - ψ̂_{j+1}|², \\ j = 1, ..., n-1 . +``` +""" +function energies(vars, params, grid, sol) + nlayers = numberoflayers(params) + KE, PE = zeros(nlayers), zeros(nlayers-1) + + @. vars.qh = sol + streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) + + abs²∇𝐮h = vars.uh # use vars.uh as scratch variable + @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) + + for j = 1:nlayers + CUDA.@allowscalar KE[j] = 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) + end + + for j = 1:nlayers-1 + CUDA.@allowscalar PE[j] = 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′[j] * parsevalsum(abs2.(vars.ψh[:, :, j+1] .- vars.ψh[:, :, j]), grid) + end + + return KE, PE +end + +function energies(vars, params::TwoLayerParams, grid, sol) + nlayers = numberoflayers(params) + KE, PE = zeros(nlayers), zeros(nlayers-1) + + @. vars.qh = sol + streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) + + abs²∇𝐮h = vars.uh # use vars.uh as scratch variable + @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) + + ψ1h, ψ2h = view(vars.ψh, :, :, 1), view(vars.ψh, :, :, 2) + + for j = 1:nlayers + CUDA.@allowscalar KE[j] = @views 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) + end + + PE = @views 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′ * parsevalsum(abs2.(ψ2h .- ψ1h), grid) + + return KE, PE +end + +function energies(vars, params::SingleLayerParams, grid, sol) + @. vars.qh = sol + streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) + + abs²∇𝐮h = vars.uh # use vars.uh as scratch variable + @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) + + return 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h, grid) +end + +energies(prob) = energies(prob.vars, prob.params, prob.grid, prob.sol) + +""" + fluxes(vars, params, grid, sol) + fluxes(prob) + +Return the lateral eddy fluxes within each fluid layer, lateralfluxes``_1,...,``lateralfluxes``_n`` +and also the vertical eddy fluxes at each fluid interface, +verticalfluxes``_{3/2},...,``verticalfluxes``_{n-1/2}``, where ``n`` is the total number of layers in the fluid. +(When ``n=1``, only the lateral fluxes are returned.) + +The lateral eddy fluxes within the ``j``-th fluid layer are + +```math +\\textrm{lateralfluxes}_j = \\frac{H_j}{H} \\int U_j v_j ∂_y u_j +\\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n , +``` + +while the vertical eddy fluxes at the ``j+1/2``-th fluid interface (i.e., interface between +the ``j``-th and ``(j+1)``-th fluid layer) are + +```math +\\textrm{verticalfluxes}_{j+1/2} = \\int \\frac{f₀²}{g'_{j+1/2} H} (U_j - U_{j+1}) \\, +v_{j+1} ψ_{j} \\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n-1. +``` +""" +function fluxes(vars, params, grid, sol) + nlayers = numberoflayers(params) + + lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) + + updatevars!(vars, params, grid, sol) + + ∂u∂yh = vars.uh # use vars.uh as scratch variable + ∂u∂y = vars.u # use vars.u as scratch variable + + @. ∂u∂yh = im * grid.l * vars.uh + invtransform!(∂u∂y, ∂u∂yh, params) + + lateralfluxes = (sum(@. params.H * params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] + lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) + + for j = 1:nlayers-1 + CUDA.@allowscalar verticalfluxes[j] = sum(@views @. params.f₀^2 / params.g′[j] * (params.U[: ,:, j] - params.U[:, :, j+1]) * vars.v[:, :, j+1] * vars.ψ[:, :, j]; dims=(1, 2))[1] + CUDA.@allowscalar verticalfluxes[j] *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) + end + + return lateralfluxes, verticalfluxes +end + +function fluxes(vars, params::TwoLayerParams, grid, sol) + nlayers = numberoflayers(params) + + lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) + + updatevars!(vars, params, grid, sol) + + ∂u∂yh = vars.uh # use vars.uh as scratch variable + ∂u∂y = vars.u # use vars.u as scratch variable + + @. ∂u∂yh = im * grid.l * vars.uh + invtransform!(∂u∂y, ∂u∂yh, params) + + lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] + @. lateralfluxes *= params.H + lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) + + U₁, U₂ = view(params.U, :, :, 1), view(params.U, :, :, 2) + ψ₁ = view(vars.ψ, :, :, 1) + v₂ = view(vars.v, :, :, 2) + + verticalfluxes = sum(@views @. params.f₀^2 / params.g′ * (U₁ - U₂) * v₂ * ψ₁; dims=(1, 2)) + verticalfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) + + return lateralfluxes, verticalfluxes +end + +function fluxes(vars, params::SingleLayerParams, grid, sol) + updatevars!(vars, params, grid, sol) + + ∂u∂yh = vars.uh # use vars.uh as scratch variable + ∂u∂y = vars.u # use vars.u as scratch variable + + @. ∂u∂yh = im * grid.l * vars.uh + invtransform!(∂u∂y, ∂u∂yh, params) + + lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] + lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly) + + return lateralfluxes +end + +fluxes(prob) = fluxes(prob.vars, prob.params, prob.grid, prob.sol) + +end # module From 98572577f28897a428f9ff85428a9200a2f52846 Mon Sep 17 00:00:00 2001 From: Matt Lobo Date: Fri, 9 Jun 2023 18:29:09 -0400 Subject: [PATCH 12/16] fixed sign error in multilayerqg.jl ONLY --- examples/matts_examples/GFjl_bug_plots.jl | 102 ------- examples/matts_examples/mjcl_stab.jl | 250 ----------------- examples/matts_examples/phillips_stab.jl | 316 ---------------------- examples/matts_examples/prof_stab.jl | 89 ------ examples/matts_examples/sphere_ex.jl | 33 --- src/multilayerqg.jl | 4 +- 6 files changed, 2 insertions(+), 792 deletions(-) delete mode 100644 examples/matts_examples/GFjl_bug_plots.jl delete mode 100644 examples/matts_examples/mjcl_stab.jl delete mode 100644 examples/matts_examples/phillips_stab.jl delete mode 100644 examples/matts_examples/prof_stab.jl delete mode 100644 examples/matts_examples/sphere_ex.jl diff --git a/examples/matts_examples/GFjl_bug_plots.jl b/examples/matts_examples/GFjl_bug_plots.jl deleted file mode 100644 index 1d84f861..00000000 --- a/examples/matts_examples/GFjl_bug_plots.jl +++ /dev/null @@ -1,102 +0,0 @@ -# showing how two different bug fixes in GeophysicalFlows.jl -# modify vertical profile of interior PV - -## load packages - -using GeophysicalFlows, CairoMakie, Printf, FFTW, LinearAlgebra, Statistics - -here = "/home/matt/Desktop/research/QG/julia_QG/" -include(here*"mjcl_stab.jl") - -## build basic model - -dev = CPU() # device (CPU) - -# numerical params -Ny = Nx = n = 128 # 2D resolution = n² -stepper = "FilteredRK4" # time stepping scheme -dt = 1e-2 # time step -nsteps = 60000 # total number of time-steps 20000 -nsubs = 50 # number of time-steps for plotting (nsteps must be multiple of nsubs) - -# physical params -L = 2π # domain size -μ = 5e-2 # bottom drag -beta = β = 0. # no gradient of planetary PV - -nlayers = 6 # number of layers -f0 = f₀ = 1. # 0.0001236812857687059 # coriolis param [s^-1] -g = 9.81 # gravity - -H = ones(nlayers) # even depths -rho = ρ = collect(range(4.0,5.0,nlayers)) # constant N^2 -U = collect(range(1.0,0.0,nlayers)) # constant shear -V = zeros(nlayers) # zonal flow only - -## setting up problems -# setting up the problen with bug fixes (denoted ``f'') -prob_f = MultiLayerQG.Problem(nlayers, dev; nx=n, Lx=L, f₀, g, H, ρ, U, μ, β, - dt, stepper, aliased_fraction=0) - -sol_f, clock_f, params_f, vars_f, grid_f = prob_f.sol, prob_f.clock, prob_f.params, prob_f.vars, prob_f.grid - -# setting up the problen with bug fixes (denoted ``sf'') -include(here*"../gfjl/GeophysicalFlows.jl/src/multilayerqg_sign_fix_only.jl") - -prob_sf = MultiLayerQG_sf.Problem_sf(nlayers, dev; nx=n, Lx=L, f₀, g, H, ρ, U, μ, β, - dt, stepper, aliased_fraction=0) - -sol_sf, clock_sf, params_sf, vars_sf, grid_sf = prob_sf.sol, prob_sf.clock, prob_sf.params, prob_sf.vars, prob_sf.grid - -# setting up the problen with no bug fixes (denoted ``nf'') -include(here*"../gfjl/GeophysicalFlows.jl/src/multilayerqg_no_fixes.jl") - -prob_nf = MultiLayerQG_nf.Problem_nf(nlayers, dev; nx=n, Lx=L, f₀, g, H, ρ, U, μ, β, - dt, stepper, aliased_fraction=0) - -sol_nf, clock_nf, params_nf, vars_nf, grid_nf = prob_nf.sol, prob_nf.clock, prob_nf.params, prob_nf.vars, prob_nf.grid - -# linear stability analysis -Lx, Ly = grid_f.Lx, grid_f.Ly - -eta = 0 - -eve,eva,max_eve_f,max_eva_f,k_x,k_y,qx,qy = lin_stab(U,V,beta,eta,Nx,Ny,rho,f0,Lx,Ly,params_f.Qy) -eve,eva,max_eve_sf,max_eva_sf,k_x,k_y,qx,qy = lin_stab(U,V,beta,eta,Nx,Ny,rho,f0,Lx,Ly,params_sf.Qy) -eve,eva,max_eve_nf,max_eva_nf,k_x,k_y,qx,qy = lin_stab(U,V,beta,eta,Nx,Ny,rho,f0,Lx,Ly,params_nf.Qy) - -## plotting differences -z = -cumsum(H) - -fig = Figure(resolution=(1000, 600)) - -ax1_kwargs = (xlabel = "Qy", - ylabel = "z", - aspect = 1.) - -ax1 = Axis(fig[1, 1]; title = "Qy", ax1_kwargs...) - -ylims!(ax1,minimum(z), maximum(z)) - -lines!(ax1,params_nf.Qy[1,1,:],z,linewidth=3.,label="no fix") -lines!(ax1,params_sf.Qy[1,1,:],z,linewidth=3.,label="sign fix only") -lines!(ax1,params_f.Qy[1,1,:],z,linewidth=3.,label="sign fix and g' fix") - -axislegend(position=:lt) - -ax2_kwargs = (xlabel = "|ψ|", - ylabel = "z", - aspect = 1.) - -ax2 = Axis(fig[1, 2]; title = "|ψ|", ax2_kwargs...) - -ylims!(ax2,minimum(z), maximum(z)) - -lines!(ax2,abs.(max_eve_nf),z,linewidth=3.,label="no fix") -lines!(ax2,abs.(max_eve_sf),z,linewidth=3.,label="sign fix only") -lines!(ax2,abs.(max_eve_f),z,linewidth=3.,label="sign fix and g' fix") - -axislegend(position=:lt) - - - diff --git a/examples/matts_examples/mjcl_stab.jl b/examples/matts_examples/mjcl_stab.jl deleted file mode 100644 index 48e168d3..00000000 --- a/examples/matts_examples/mjcl_stab.jl +++ /dev/null @@ -1,250 +0,0 @@ -## define stability functions - -function lin_stab(U::Vector{Float64},V::Vector{Float64},beta,eta,Nx::Int64,Ny::Int64,rho::Vector{Float64},f0::Float64,Lx::Float64,Ly::Float64) - # U: (Nx x Nz) vector of zonal mean background velocity - # V: (Ny x Nz) vector of meridional mean background velocity - # beta: - # - - Nz = length(rho) - # define wavenumbers - k_x = reshape(fftfreq(Nx, 2π/Lx*Nx),(1,Nx)) - k_y = reshape(fftfreq(Ny, 2π/Ly*Ny),(1,Ny)) - - # k_x,k_y = wavenumber_grid(Nx,Ny,Lx,Ly) - - # k2 = k_x.^2 + k_y.^2 # this is for an isotropic wavenumber grid only (i.e. Nx=Ny) - - # define stretching matrix - S = calc_stretching_mat(Nz,rho,f0,H,rho[1]) - - # change dimensions of U and V to match domain size - U2 = zeros(1,Nz); U2[:] = U; U2 = repeat(U2,outer=(Ny,1)) - U = zeros(1,Ny,Nz); U[1,:,:] = U2 - - V2 = zeros(1,Nz); V2[:] = V; V2 = repeat(V2,outer=(Nx,1)) - V = zeros(1,Nx,Nz); V[1,:,:] = V2 - - # define background QG PV gradients - Qy = calc_PV_grad_y(U,beta,eta,Ny,Nz,k_y,S) - Qx = calc_PV_grad_x(V,eta,Nx,Nz,k_x,S) - - # perform linear stability analysis - evecs_all,evals_all = calc_lin_stab(Qy,Qx,U,V,S,k_x,k_y,Nz) - - # keep largest growth rates per wavenumber - evecs,evals,max_evec,max_eval = find_growth(evecs_all,evals_all,Nx,Ny,Nz) - - # def rad - r_d = sqrt(gp(rho[1:2],rho[1])*H[1])/f0 - - return fftshift(evecs),fftshift(evals),max_evec,max_eval,fftshift(k_x),fftshift(k_y),mean(Qx[1,:,:],dims=1),mean(Qy[1,:,:],dims=1) -end - -function lin_stab(U::Vector{Float64},V::Vector{Float64},beta,eta,Nx::Int64,Ny::Int64,rho::Vector{Float64},f0::Float64,Lx::Float64,Ly::Float64,Qy) - # Takes Qy as arg - # U: (Nx x Nz) vector of zonal mean background velocity - # V: (Ny x Nz) vector of meridional mean background velocity - # beta: - # - - Nz = length(rho) - # define wavenumbers - k_x = reshape(fftfreq(Nx, 2π/Lx*Nx),(1,Nx)) - k_y = reshape(fftfreq(Ny, 2π/Ly*Ny),(1,Ny)) - - # k_x,k_y = wavenumber_grid(Nx,Ny,Lx,Ly) - - # k2 = k_x.^2 + k_y.^2 # this is for an isotropic wavenumber grid only (i.e. Nx=Ny) - - # define stretching matrix - S = calc_stretching_mat(Nz,rho,f0,H,rho[1]) - - # change dimensions of U and V to match domain size - U2 = zeros(1,Nz); U2[:] = U; U2 = repeat(U2,outer=(Ny,1)) - U = zeros(1,Ny,Nz); U[1,:,:] = U2 - - V2 = zeros(1,Nz); V2[:] = V; V2 = repeat(V2,outer=(Nx,1)) - V = zeros(1,Nx,Nz); V[1,:,:] = V2 - - # define background QG PV gradients (TEMPORARY) - Qy = reshape(Qy[1,:,:],(1,Ny,Nz)) - Qx = zeros(size(Qy)) - - # perform linear stability analysis - evecs_all,evals_all = calc_lin_stab(Qy,Qx,U,V,S,k_x,k_y,Nz) - - # keep largest growth rates per wavenumber - evecs,evals,max_evec,max_eval = find_growth(evecs_all,evals_all,Nx,Ny,Nz) - - # def rad - r_d = sqrt(gp(rho[1:2],rho[1])*H[1])/f0 - - return fftshift(evecs),fftshift(evals),max_evec,max_eval,fftshift(k_x),fftshift(k_y),mean(Qx[1,:,:],dims=1),mean(Qy[1,:,:],dims=1) -end - -function wavenumber_grid(Nx,Ny,Lx,Ly) - # - # nk_x = div(Nx,2)+1; nk_y = div(Ny,2)+1 - - nk_x = Nx; nk_y = Ny - - k_x = reshape(LinRange(-2*pi/Lx*nk_x,2*pi/Lx*nk_x,nk_x),(1,Nx)) - k_y = reshape(LinRange(-2*pi/Ly*nk_y,2*pi/Ly*nk_y,nk_y),(1,Ny)) - - # k_x = LinRange(0.,2*pi/Lx*nk_x,nk_x) - # k_y = LinRange(0.,2*pi/Ly*nk_y,nk_y) - - return k_x,k_y -end - -function calc_PV_grad_y(U,beta,eta,Ny::Int64,Nz::Int64,k_y,S) - # calculates PV gradients in one meridional direction - # U is (Ny x Nz) - # k_y is (Ny x 1) - # - - Uyy = real.(ifft(-k_y.^2 .* fft(U))) - - # Uyy = repeat(Uyy, outer=(Nx, 1, 1)) - - # Q_y = zeros(Nx,Nz) - - F = zeros(size(U)) - for i=1:Ny - F[1,i,:] = S * U[1,i,:] - end - - Q_y = beta .- (Uyy .+ F) - - return Q_y -end - -function calc_PV_grad_x(V,eta,Nx::Int64,Nz::Int64,k_x,S) - # calculates PV gradients in one zonal direction - - Vxx = real.(ifft(k_x.^2 .* fft(V))) - - # Q_y = zeros(Nx,Nz) - - F = zeros(size(V)) - for i=1:Nx - F[1,i,:] = S * V[1,i,:] - end - - Q_x = Vxx .+ F - - return Q_x -end - -function gp(rho,rho0) - g = 9.81 - g_prime = g*(rho[2]-rho[1])/rho0 - - return g_prime -end - -function calc_stretching_mat(Nz,rho,f0,H,rho0) - # - S = zeros((Nz,Nz,)) - - alpha = 0 - - S[1,1] = -f0^2/H[1]/gp(rho[1:2],rho0) + alpha - S[1,2] = f0^2/H[1]/gp(rho[1:2],rho0) - for i = 2:Nz-1 - S[i,i-1] = f0^2/H[i]/gp(rho[i-1:i],rho0) - S[i,i] = -(f0^2/H[i]/gp(rho[i:i+1],rho0) + f0^2/H[i]/gp(rho[i-1:i],rho0)) - S[i,i+1] = f0^2/H[i]/gp(rho[i:i+1],rho0) - end - S[Nz,Nz-1] = f0^2/H[Nz]/gp(rho[Nz-1:Nz],rho0) - S[Nz,Nz] = -f0^2/H[Nz]/gp(rho[Nz-1:Nz],rho0) - - return S -end - -function calc_lin_stab(Qy,Qx,U,V,S,k_x,k_y,Nz) - # - # A = (k_x .* U .+ k_y .* V) - - # B = zeros((1,length(k2),Nz)) - # ell = zeros((Nz,Nz,length(k_x),length(k_y))) - # for i=eachindex(k_x) - # for j=eachindex(k_y) - # ell[:,:,i,j] = S .- (k_x[i]^2 + k_y[j]^2) * I - # B[1,i,j,:] = ell[:,:,i,j] * A[1,i,:] - # end - # end - - # C = k_x .* Qy .- k_y .* Qx - - # M = inv(ell) * (A .+ C) - - # evecs,evals = eig(M) - - - ################# - evecs = zeros(Nx,Ny,Nz,Nz) .+ 0im - evals = zeros(Nx,Ny,Nz) .+ 0im - k2 = zeros(Nx,Ny) - - for i=1:Nx - for j=1:Ny - if i==1 && j==1 - # do nothing - else - A = make_diag(k_x[i] .* U[:,j,:] .+ k_y[j] .* V[:,i,:]) - # ell_i = inv(S - (k_x[i]^2 + k_y[j]^2)*I) - # B = ell_i .* make_diag(k_x[i] * Qy[:,j,:] - k_y[j] * Qx[:,i,:]) .+ 0im - - # evecs[i,j,:,:] = eigvecs(A .+ B) - # evals[i,j,:] = eigvals(A .+ B) - - ell = (S - (k_x[i]^2 + k_y[j]^2) * I) - A2 = ell * A - ell_i = inv(ell) - Q2 = (k_x[i] * Qy[:,j,:] - k_y[j] * Qx[:,i,:]) .+ 0im - B2 = transpose(ell_i) * transpose(A2 .+ make_diag(Q2)) .+ 0im - - k2[i,j] = (k_x[i]^2 + k_y[j]^2) - - evecs[i,j,:,:] = eigvecs(B2) - evals[i,j,:] = eigvals(B2) - end - - end - end - - return evecs,evals -end - -function make_diag(array_in) - matrix_out = zeros(length(array_in),length(array_in)) - for i=eachindex(array_in) - matrix_out[i,i] = array_in[i] - end - return matrix_out -end - -function find_growth(evecs_all,evals_all,Nx,Ny,Nz) - # - evecs = zeros(Nx,Ny,Nz) .+ 0im; evals = zeros(Nx,Ny) .+ 0im - - for i=1:Nx - for j=1:Ny - indMax = argmax(imag(evals_all[i,j,:])) - evals[i,j] = evals_all[i,j,indMax] - evecs[i,j,:] = evecs_all[i,j,:,indMax] - end - end - - sigma = imag(evals) - - indMax = argmax(sigma) - - max_eval = sigma[indMax] - - max_evec = abs.(evecs[indMax,:]) - - return evecs,evals,max_evec,max_eval -end \ No newline at end of file diff --git a/examples/matts_examples/phillips_stab.jl b/examples/matts_examples/phillips_stab.jl deleted file mode 100644 index b817deae..00000000 --- a/examples/matts_examples/phillips_stab.jl +++ /dev/null @@ -1,316 +0,0 @@ -# a test space for building a linear stability analysis tool in julia -# for now it is separate from GeophysicalFlows.jl Model object, but might -# be merged in the future - -## load packages - -using GeophysicalFlows, CairoMakie, Printf, FFTW, LinearAlgebra, Statistics - -using Random: seed! - -include("./mjcl_stab.jl") - -## build basic model - -dev = CPU() # device (CPU) - -# numerical params -Ny = Nx = n = 128 # 2D resolution = n² -stepper = "FilteredRK4" # time stepping scheme -dt = 1e-2 # time step -nsteps = 60000 # total number of time-steps 20000 -nsubs = 50 # number of time-steps for plotting (nsteps must be multiple of nsubs) - -# physical params -L = 2π # domain size -μ = 5e-2 # bottom drag -beta = β = 0. #1.2130692965249345e-11 # the y-gradient of planetary PV - -nlayers = 8 # number of layers -f0 = f₀= 1. #0.0001236812857687059 # coriolis param [s^-1] -g = 9.81 -H = [0.1, 0.3, 1.] # the rest depths of each layer -rho = ρ = [4.0, 5.0, 5.1] # the density of each layer - -H=ones(8) -rho = ρ =collect(range(4.0,5.0,8)) -U = collect(range(1.0,0.0,8)) -V = zeros(8) - -# U = zeros(nlayers) # the imposed mean zonal flow in each layer -# U[1] = 1.0 -# U[2] = 0.05 -# U[3] = 0.005 - - -# V = zeros(nlayers) # the imposed mean zonal flow in each layer -# V[1] = 1.0 -# V[2] = 0.5 -# V[3] = 0.25 - -# setting up the ``problem'' -prob = MultiLayerQG.Problem(nlayers, dev; nx=n, Lx=L, f₀, g, H, ρ, U, μ, β, - dt, stepper, aliased_fraction=0) - -sol, clock, params, vars, grid = prob.sol, prob.clock, prob.params, prob.vars, prob.grid -x, y = grid.x, grid.y - -# initial conditions -seed!(1234) # reset of the random number generator for reproducibility -q₀ = 1e-2 * device_array(dev)(randn((grid.nx, grid.ny, nlayers))) -q₀h = prob.timestepper.filter .* rfft(q₀, (1, 2)) # apply rfft only in dims=1, 2 -q₀ = irfft(q₀h, grid.nx, (1, 2)) # apply irfft only in dims=1, 2 - -MultiLayerQG.set_q!(prob, q₀) - -# diagnostics -E = Diagnostic(MultiLayerQG.energies, prob; nsteps) -diags = [E] # A list of Diagnostics types passed to "stepforward!" will be updated every timestep. - -# output dirs -filepath = "." -plotpath = "./figs/plots_2layer" -plotname = "snapshots" -filename = joinpath(filepath, "2layer.jld2") - -# file management -if isfile(filename); rm(filename); end -if !isdir(plotpath); mkdir(plotpath); end - -# ``create output'' (?) -get_sol(prob) = prob.sol # extracts the Fourier-transformed solution - -function get_u(prob) - sol, params, vars, grid = prob.sol, prob.params, prob.vars, prob.grid - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - @. vars.uh = -im * grid.l * vars.ψh - invtransform!(vars.u, vars.uh, params) - - return vars.u -end - -out = Output(prob, filename, (:sol, get_sol), (:u, get_u)) - -# visualizing the simulation... - -Lx, Ly = grid.Lx, grid.Ly - -title_KE = Observable(@sprintf("μt = %.2f", μ * clock.t)) - -q1 = Observable(Array(vars.q[:, :, 1])) -q2 = Observable(Array(vars.q[:, :, 2])) -q3 = Observable(Array(vars.q[:, :, 3])) -q4 = Observable(Array(vars.q[:, :, 4])) -q5 = Observable(Array(vars.q[:, :, 5])) -q6 = Observable(Array(vars.q[:, :, 6])) -q7 = Observable(Array(vars.q[:, :, 7])) -q8 = Observable(Array(vars.q[:, :, 8])) - -# psi1 = Observable(Array(vars.ψ[:, :, 1])) -# psi2 = Observable(Array(vars.ψ[:, :, 2])) -# psi3 = Observable(Array(vars.ψ[:, :, 3])) -# psi4 = Observable(Array(vars.ψ[:, :, 4])) -# psi5 = Observable(Array(vars.ψ[:, :, 5])) -# psi6 = Observable(Array(vars.ψ[:, :, 6])) -# psi7 = Observable(Array(vars.ψ[:, :, 7])) -# psi8 = Observable(Array(vars.ψ[:, :, 8])) - -# function compute_levels(maxf, nlevels=8) -# # -max(|f|):...:max(|f|) -# levelsf = @lift collect(range(-$maxf, stop = $maxf, length=nlevels)) - -# # only positive -# levelsf⁺ = @lift collect(range($maxf/(nlevels-1), stop = $maxf, length=Int(nlevels/2))) - -# # only negative -# levelsf⁻ = @lift collect(range(-$maxf, stop = -$maxf/(nlevels-1), length=Int(nlevels/2))) - -# return levelsf, levelsf⁺, levelsf⁻ -# end - -# maxpsi1 = Observable(maximum(abs, vars.ψ[:, :, 1])) -# maxpsi2 = Observable(maximum(abs, vars.ψ[:, :, 2])) -# maxpsi3 = Observable(maximum(abs, vars.ψ[:, :, 3])) -# maxpsi4 = Observable(maximum(abs, vars.ψ[:, :, 4])) -# maxpsi5 = Observable(maximum(abs, vars.ψ[:, :, 5])) -# maxpsi6 = Observable(maximum(abs, vars.ψ[:, :, 6])) -# maxpsi7 = Observable(maximum(abs, vars.ψ[:, :, 7])) -# maxpsi8 = Observable(maximum(abs, vars.ψ[:, :, 8])) - - -# levelspsi1, levelspsi1⁺, levelspsi1⁻ = compute_levels(maxpsi1) -# levelspsi2, levelspsi2⁺, levelspsi2⁻ = compute_levels(maxpsi2) -# levelspsi3, levelspsi3⁺, levelspsi3⁻ = compute_levels(maxpsi3) -# levelspsi4, levelspsi4⁺, levelspsi4⁻ = compute_levels(maxpsi4) -# levelspsi5, levelspsi5⁺, levelspsi5⁻ = compute_levels(maxpsi5) -# levelspsi6, levelspsi6⁺, levelspsi6⁻ = compute_levels(maxpsi6) -# levelspsi7, levelspsi7⁺, levelspsi7⁻ = compute_levels(maxpsi7) -# levelspsi8, levelspsi8⁺, levelspsi8⁻ = compute_levels(maxpsi8) - -# KE₁ = Observable(Point2f[(μ * E.t[1], E.data[1][1][1])]) -# KE₂ = Observable(Point2f[(μ * E.t[1], E.data[1][1][end])]) -# PE = Observable(Point2f[(μ * E.t[1], E.data[1][2][end])]) - -fig = Figure(resolution=(1000, 600)) - -axis_kwargs = (xlabel = "x", - ylabel = "y", - aspect = 1, - limits = ((-Lx/2, Lx/2), (-Ly/2, Ly/2))) - -axq1 = Axis(fig[1, 1]; title = "q1", axis_kwargs...) - -# axpsi1 = Axis(fig[2, 1]; title = "psi1", axis_kwargs...) - -axq2 = Axis(fig[1, 2]; title = "q2", axis_kwargs...) - -# axpsi2 = Axis(fig[2, 2]; title = "psi2", axis_kwargs...) - -axq3 = Axis(fig[1, 3]; title = "q3", axis_kwargs...) -axq4 = Axis(fig[1, 4]; title = "q4", axis_kwargs...) -axq5 = Axis(fig[2, 1]; title = "q5", axis_kwargs...) -axq6 = Axis(fig[2, 2]; title = "q6", axis_kwargs...) -axq7 = Axis(fig[2, 3]; title = "q7", axis_kwargs...) -axq8 = Axis(fig[2, 4]; title = "q8", axis_kwargs...) - -# axKE = Axis(fig[1, 3], -# xlabel = "μ t", -# ylabel = "KE", -# title = title_KE, -# yscale = log10, -# limits = ((-0.1, 2.6), (1e-9, 5))) - -# axPE = Axis(fig[2, 3], -# xlabel = "μ t", -# ylabel = "PE", -# yscale = log10, -# limits = ((-0.1, 2.6), (1e-9, 5))) - -heatmap!(axq1, x, y, q1; colormap = :balance) -heatmap!(axq2, x, y, q2; colormap = :balance) -heatmap!(axq3, x, y, q3; colormap = :balance) -heatmap!(axq4, x, y, q4; colormap = :balance) -heatmap!(axq5, x, y, q5; colormap = :balance) -heatmap!(axq6, x, y, q6; colormap = :balance) -heatmap!(axq7, x, y, q7; colormap = :balance) -heatmap!(axq8, x, y, q8; colormap = :balance) - -# contourf!(axpsi1, x, y, psi1; -# levels = levelspsi1, colormap = :viridis, extendlow = :auto, extendhigh = :auto) -# contour!(axpsi1, x, y, psi1; -# levels = levelspsi1⁺, color=:black) -# contour!(axpsi1, x, y, psi1; -# levels = levelspsi1⁻, color=:black, linestyle = :dash) - -# contourf!(axpsi2, x, y, psi2; -# levels = levelspsi2, colormap = :viridis, extendlow = :auto, extendhigh = :auto) -# contour!(axpsi2, x, y, psi2; -# levels = levelspsi2⁺, color=:black) -# contour!(axpsi2, x, y, psi2; -# levels = levelspsi2⁻, color=:black, linestyle = :dash) - -# ke₁ = lines!(axKE, KE₁; linewidth = 3) -# ke₂ = lines!(axKE, KE₂; linewidth = 3) -# Legend(fig[1, 4], [ke₁, ke₂,], ["KE₁", "KE₂"]) - -# lines!(axPE, PE; linewidth = 3) - -fig - -# trying to run the model now -startwalltime = time() - -frames = 0:round(Int, nsteps / nsubs) - -record(fig, "multilayerqg_2layer.mp4", frames, framerate = 18) do j - if j % (1000 / nsubs) == 0 - cfl = clock.dt * maximum([maximum(vars.u) / grid.dx, maximum(vars.v) / grid.dy]) - - log = @sprintf("step: %04d, t: %.1f, cfl: %.2f, KE₁: %.3e, KE₂: %.3e, PE: %.3e, walltime: %.2f min", - clock.step, clock.t, cfl, E.data[E.i][1][1], E.data[E.i][1][2], E.data[E.i][2][1], (time()-startwalltime)/60) - - println(log) - end - - q1[] = vars.q[:, :, 1] - q2[] = vars.q[:, :, 2] - q3[] = vars.q[:, :, 3] - q4[] = vars.q[:, :, 4] - q5[] = vars.q[:, :, 5] - q6[] = vars.q[:, :, 6] - q7[] = vars.q[:, :, 7] - q8[] = vars.q[:, :, end] - - # psi1[] = vars.ψ[:, :, 1] - # psi2[] = vars.ψ[:, :, 2] - # psi3[] = vars.ψ[:, :, 3] - # psi4[] = vars.ψ[:, :, 4] - # psi5[] = vars.ψ[:, :, 5] - # psi6[] = vars.ψ[:, :, 6] - # psi7[] = vars.ψ[:, :, 7] - # psi8[] = vars.ψ[:, :, end] - - # maxpsi1[] = maximum(abs, vars.ψ[:, :, 1]) - # maxpsi2[] = maximum(abs, vars.ψ[:, :, 2]) - # maxpsi3[] = maximum(abs, vars.ψ[:, :, 3]) - # maxpsi4[] = maximum(abs, vars.ψ[:, :, 4]) - # maxpsi5[] = maximum(abs, vars.ψ[:, :, 5]) - # maxpsi6[] = maximum(abs, vars.ψ[:, :, 6]) - # maxpsi7[] = maximum(abs, vars.ψ[:, :, 7]) - # maxpsi8[] = maximum(abs, vars.ψ[:, :, end]) - - - # KE₁[] = push!(KE₁[], Point2f(μ * E.t[E.i], E.data[E.i][1][1])) - # KE₂[] = push!(KE₂[], Point2f(μ * E.t[E.i], E.data[E.i][1][end])) - # PE[] = push!(PE[] , Point2f(μ * E.t[E.i], E.data[E.i][2][end])) - - # title_KE[] = @sprintf("μ t = %.2f", μ * clock.t)] - - stepforward!(prob, diags, nsubs) - MultiLayerQG.updatevars!(prob) - - # savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), clock.step) - # savefig(savename) - -end - - -# savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), clock.step) -# savefig(savename) - - -## use stability functions -# Ny = div(Nx,2) - -eta = 0 - -eve,eva,k_x,k_y,qx,qy = lin_stab(U,V,beta,eta,Nx,Ny,rho,f0,Lx,Ly) - -k_xg = range(k_x[1],step=(k_x[2]-k_x[1]),length=length(k_x)) -k_yg = range(k_y[1],step=(k_y[2]-k_y[1]),length=length(k_y)) - - - -# Plots.contourf(transpose(imag(eva[60:70,60:70]))) - -# sigma = Observable(Array(imag(eva))) - -# fig = Figure(resolution=(1000, 600)) - -# axis_kwargs = (xlabel = "k_x", -# ylabel = "k_y", -# aspect = 1, -# limits = ((k_x[60], k_x[70]), (k_y[60],k_y[70]))) - -# ax_sig = Axis(fig[1, 1]; title = "σ", axis_kwargs...) - - -# contourf!(ax_sig, k_xg, k_yg, sigma; -# colormap = :viridis, extendlow = :auto, extendhigh = :auto) -# contour!(ax_sig, k_xg, k_yg, sigma; -# color=:black) - - - diff --git a/examples/matts_examples/prof_stab.jl b/examples/matts_examples/prof_stab.jl deleted file mode 100644 index 89aae07c..00000000 --- a/examples/matts_examples/prof_stab.jl +++ /dev/null @@ -1,89 +0,0 @@ -# a test space for building a linear stability analysis tool in julia -# for now it is separate from GeophysicalFlows.jl Model object, but might -# be merged in the future - -## load packages - -using GeophysicalFlows, CairoMakie, Printf, FFTW, LinearAlgebra, Statistics - -using Random: seed! - -include("./mjcl_stab.jl") - -## build basic model - -dev = CPU() # device (CPU) - -# numerical params -Ny = Nx = n = 128 # 2D resolution = n² -stepper = "FilteredRK4" # time stepping scheme -dt = 1e-2 # time step -nsteps = 60000 # total number of time-steps 20000 -nsubs = 50 # number of time-steps for plotting (nsteps must be multiple of nsubs) - -# physical params -L = 2π # domain size -μ = 5e-2 # bottom drag -beta = β = 0. #1.2130692965249345e-11 # the y-gradient of planetary PV - -nlayers = 8 # number of layers -f0 = f₀= 1. #0.0001236812857687059 # coriolis param [s^-1] -g = 9.81 -H = [0.1, 0.3, 1.] # the rest depths of each layer -rho = ρ = [4.0, 5.0, 5.1] # the density of each layer - -H=ones(8) -rho = ρ =collect(range(4.0,5.0,8)) -U = collect(range(1.0,0.0,8)) -V = zeros(8) - -# U = zeros(nlayers) # the imposed mean zonal flow in each layer -# U[1] = 1.0 -# U[2] = 0.05 -# U[3] = 0.005 - - -# V = zeros(nlayers) # the imposed mean zonal flow in each layer -# V[1] = 1.0 -# V[2] = 0.5 -# V[3] = 0.25 - -# setting up the ``problem'' -prob = MultiLayerQG.Problem(nlayers, dev; nx=n, Lx=L, f₀, g, H, ρ, U, μ, β, - dt, stepper, aliased_fraction=0) - -sol, clock, params, vars, grid = prob.sol, prob.clock, prob.params, prob.vars, prob.grid -x, y = grid.x, grid.y - -Lx, Ly = grid.Lx, grid.Ly - -eta = 0 - -eve,eva,max_eve,max_eva,k_x,k_y,qx,qy = lin_stab(U,V,beta,eta,Nx,Ny,rho,f0,Lx,Ly) - -k_xg = range(k_x[1],step=(k_x[2]-k_x[1]),length=length(k_x)) -k_yg = range(k_y[1],step=(k_y[2]-k_y[1]),length=length(k_y)) - - - -# Plots.contourf(transpose(imag(eva[60:70,60:70]))) - -# sigma = Observable(Array(imag(eva))) - -# fig = Figure(resolution=(1000, 600)) - -# axis_kwargs = (xlabel = "k_x", -# ylabel = "k_y", -# aspect = 1, -# limits = ((k_x[60], k_x[70]), (k_y[60],k_y[70]))) - -# ax_sig = Axis(fig[1, 1]; title = "σ", axis_kwargs...) - - -# contourf!(ax_sig, k_xg, k_yg, sigma; -# colormap = :viridis, extendlow = :auto, extendhigh = :auto) -# contour!(ax_sig, k_xg, k_yg, sigma; -# color=:black) - - - diff --git a/examples/matts_examples/sphere_ex.jl b/examples/matts_examples/sphere_ex.jl deleted file mode 100644 index e0e08a2e..00000000 --- a/examples/matts_examples/sphere_ex.jl +++ /dev/null @@ -1,33 +0,0 @@ -# function to calculate the volume of a sphere -function sphere_vol(r) - # julia allows Unicode names (in UTF-8 encoding) - # so either "pi" or the symbol π can be used - return 4/3*pi*r^3 -end - -# functions can also be defined more succinctly -quadratic(a, sqr_term, b) = (-b + sqr_term) / 2a - -# calculates x for 0 = a*x^2+b*x+c, arguments types can be defined in function definitions -function quadratic2(a::Float64, b::Float64, c::Float64) - # unlike other languages 2a is equivalent to 2*a - # a^2 is used instead of a**2 or pow(a,2) - sqr_term = sqrt(b^2-4a*c) - r1 = quadratic(a, sqr_term, b) - r2 = quadratic(a, -sqr_term, b) - # multiple values can be returned from a function using tuples - # if the return keyword is omitted, the last term is returned - r1, r2 -end - -vol = sphere_vol(3) -# @printf allows number formatting but does not automatically append the \n to statements, see below -using Printf -@printf "volume = %0.3f\n" vol -#> volume = 113.097 - -quad1, quad2 = quadratic2(2.0, -2.0, -12.0) -println("result 1: ", quad1) -#> result 1: 3.0 -println("result 2: ", quad2) -#> result 2: -2.0 \ No newline at end of file diff --git a/src/multilayerqg.jl b/src/multilayerqg.jl index e50206ad..fd6ca09d 100644 --- a/src/multilayerqg.jl +++ b/src/multilayerqg.jl @@ -347,7 +347,7 @@ function Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ ρ = reshape(T.(ρ), (1, 1, nlayers)) H = reshape(T.(H), (1, 1, nlayers)) - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[1] # reduced gravity at each interface + g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) @@ -364,7 +364,7 @@ function Params(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ CUDA.@allowscalar @views Qy[:, :, 1] = @. Qy[:, :, 1] - Fp[1] * (U[:, :, 2] - U[:, :, 1]) for j = 2:nlayers-1 - CUDA.@allowscalar @views Qy[:, :, j] = @. Qy[:, :, j] + Fp[j] * (U[:, :, j+1] - U[:, :, j]) + Fm[j-1] * (U[:, :, j-1] - U[:, :, j]) + CUDA.@allowscalar @views Qy[:, :, j] = @. Qy[:, :, j] - Fp[j] * (U[:, :, j+1] - U[:, :, j]) - Fm[j-1] * (U[:, :, j-1] - U[:, :, j]) end CUDA.@allowscalar @views Qy[:, :, nlayers] = @. Qy[:, :, nlayers] - Fm[nlayers-1] * (U[:, :, nlayers-1] - U[:, :, nlayers]) From ef43c785187dd7b379f9000b69a3e8ab3cbc8644 Mon Sep 17 00:00:00 2001 From: Matt Lobo Date: Fri, 9 Jun 2023 18:37:48 -0400 Subject: [PATCH 13/16] removed various multilayerqg files in src --- src/multilayerqg_no_fixes.jl | 1107 ----------------------------- src/multilayerqg_sign_fix_only.jl | 1107 ----------------------------- 2 files changed, 2214 deletions(-) delete mode 100644 src/multilayerqg_no_fixes.jl delete mode 100644 src/multilayerqg_sign_fix_only.jl diff --git a/src/multilayerqg_no_fixes.jl b/src/multilayerqg_no_fixes.jl deleted file mode 100644 index 687f5b29..00000000 --- a/src/multilayerqg_no_fixes.jl +++ /dev/null @@ -1,1107 +0,0 @@ -module MultiLayerQG_nf - -export - fwdtransform!, - invtransform!, - streamfunctionfrompv!, - pvfromstreamfunction!, - updatevars!, - - set_q!, - set_ψ!, - energies, - fluxes - -using - FFTW, - CUDA, - LinearAlgebra, - StaticArrays, - Reexport, - DocStringExtensions - -@reexport using FourierFlows - -using FourierFlows: parsevalsum, parsevalsum2, superzeros, plan_flows_rfft - -nothingfunction(args...) = nothing - -""" - Problem(nlayers :: Int, - dev = CPU(); - nx = 128, - ny = nx, - Lx = 2π, - Ly = Lx, - f₀ = 1.0, - β = 0.0, - g = 1.0, - U = zeros(nlayers), - H = 1/nlayers * ones(nlayers), - ρ = Array{Float64}(1:nlayers), - eta = nothing, - topographic_pv_gradient = (0, 0), - μ = 0, - ν = 0, - nν = 1, - dt = 0.01, - stepper = "RK4", - calcFq = nothingfunction, - stochastic = false, - linear = false, - aliased_fraction = 1/3, - T = Float64) - -Construct a multi-layer quasi-geostrophic problem with `nlayers` fluid layers on device `dev`. - -Arguments -========= -- `nlayers`: (required) Number of fluid layers. -- `dev`: (required) `CPU()` (default) or `GPU()`; computer architecture used to time-step `problem`. - -Keyword arguments -================= - - `nx`: Number of grid points in ``x``-domain. - - `ny`: Number of grid points in ``y``-domain. - - `Lx`: Extent of the ``x``-domain. - - `Ly`: Extent of the ``y``-domain. - - `f₀`: Constant planetary vorticity. - - `β`: Planetary vorticity ``y``-gradient. - - `g`: Gravitational acceleration constant. - - `U`: The imposed constant zonal flow ``U(y)`` in each fluid layer. - - `H`: Rest height of each fluid layer. - - `ρ`: Density of each fluid layer. - - `eta`: Periodic component of the topographic potential vorticity. - - `topographic_pv_gradient`: The ``(x, y)`` components of the topographic PV large-scale gradient. - - `μ`: Linear bottom drag coefficient. - - `ν`: Small-scale (hyper)-viscosity coefficient. - - `nν`: (Hyper)-viscosity order, `nν```≥ 1``. - - `dt`: Time-step. - - `stepper`: Time-stepping method. - - `calcF`: Function that calculates the Fourier transform of the forcing, ``F̂``. - - `stochastic`: `true` or `false` (default); boolean denoting whether `calcF` is temporally stochastic. - - `linear`: `true` or `false` (default); boolean denoting whether the linearized equations of motions are used. - - `aliased_fraction`: the fraction of high-wavenumbers that are zero-ed out by `dealias!()`. - - `T`: `Float32` or `Float64` (default); floating point type used for `problem` data. -""" -function Problem_nf(nlayers::Int, # number of fluid layers - dev = CPU(); - # Numerical parameters - nx = 128, - ny = nx, - Lx = 2π, - Ly = Lx, - # Physical parameters - f₀ = 1.0, # Coriolis parameter - β = 0.0, # y-gradient of Coriolis parameter - g = 1.0, # gravitational constant - U = zeros(nlayers), # imposed zonal flow U(y) in each layer - H = 1/nlayers * ones(nlayers), # rest fluid height of each layer - ρ = Array{Float64}(1:nlayers), # density of each layer - eta = nothing, # periodic component of the topographic PV - topographic_pv_gradient = (0, 0), # tuple with the ``(x, y)`` components of topographic PV large-scale gradient - # Bottom Drag and/or (hyper)-viscosity - μ = 0, - ν = 0, - nν = 1, - # Timestepper and equation options - dt = 0.01, - stepper = "RK4", - calcFq = nothingfunction, - stochastic = false, - linear = false, - # Float type and dealiasing - aliased_fraction = 1/3, - T = Float64) - - if dev == GPU() && nlayers > 2 - @warn """MultiLayerQG module is not optimized on the GPU yet for configurations with - 3 fluid layers or more! - - See issues on Github at https://github.com/FourierFlows/GeophysicalFlows.jl/issues/112 - and https://github.com/FourierFlows/GeophysicalFlows.jl/issues/267. - - To use MultiLayerQG with 3 fluid layers or more we suggest, for now, to restrict running - on CPU.""" - end - - if nlayers == 1 - @warn """MultiLayerQG module does work for single-layer configuration but may not be as - optimized. We suggest using SingleLayerQG module for single-layer QG simulation unless - you have reasons to use MultiLayerQG in a single-layer configuration, e.g., you want to - compare solutions with varying number of fluid layers.""" - end - - # topographic PV - eta === nothing && (eta = zeros(dev, T, (nx, ny))) - - grid = TwoDGrid(dev; nx, Lx, ny, Ly, aliased_fraction, T) - - params = Params_nf(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ, ν, nν, grid; calcFq) - - vars = calcFq == nothingfunction ? DecayingVars(grid, params) : (stochastic ? StochasticForcedVars(grid, params) : ForcedVars(grid, params)) - - equation = linear ? LinearEquation(params, grid) : Equation(params, grid) - - FourierFlows.Problem(equation, stepper, dt, grid, vars, params) -end - -""" - struct Params{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams - -The parameters for the `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct Params_nf{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams - # prescribed params - "number of fluid layers" - nlayers :: Int - "gravitational constant" - g :: T - "constant planetary vorticity" - f₀ :: T - "planetary vorticity ``y``-gradient" - β :: T - "array with density of each fluid layer" - ρ :: Aphys3D - "array with rest height of each fluid layer" - H :: Aphys3D - "array with imposed constant zonal flow ``U(y)`` in each fluid layer" - U :: Aphys3D - "array containing the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear bottom drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "array with the reduced gravity constants for each fluid interface" - g′ :: Aphys1D - "array containing ``x``-gradient of PV due to topographic PV in each fluid layer" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV in each fluid layer" - Qy :: Aphys3D - "array containing coeffients for getting PV from streamfunction" - S :: Atrans4D - "array containing coeffients for inverting PV to streamfunction" - S⁻¹ :: Atrans4D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -""" - struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - -The parameters for the a single-layer `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - # prescribed params - "planetary vorticity ``y``-gradient" - β :: T - "array with imposed constant zonal flow ``U(y)``" - U :: Aphys3D - "array containing the periodic component of the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "array containing ``x``-gradient of PV due to topographic PV" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV" - Qy :: Aphys3D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -""" - struct TwoLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - -The parameters for the a two-layer `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct TwoLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - # prescribed params - "gravitational constant" - g :: T - "constant planetary vorticity" - f₀ :: T - "planetary vorticity ``y``-gradient" - β :: T - "array with density of each fluid layer" - ρ :: Aphys3D - "tuple with rest height of each fluid layer" - H :: Tuple - "array with imposed constant zonal flow ``U(y)`` in each fluid layer" - U :: Aphys3D - "array containing periodic component of the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear bottom drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "the reduced gravity constants for the fluid interface" - g′ :: T - "array containing ``x``-gradient of PV due to topographic PV in each fluid layer" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV in each fluid layer" - Qy :: Aphys3D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 1}) where TU - T = eltype(grid) - if length(U) == nlayers - U_2D = zeros(dev, T, (1, nlayers)) - U_2D[:] = U - U_2D = repeat(U_2D, outer=(grid.ny, 1)) - else - U_2D = zeros(dev, T, (grid.ny, 1)) - U_2D[:] = U - end - U_3D = zeros(dev, T, (1, grid.ny, nlayers)) - @views U_3D[1, :, :] = U_2D - return U_3D -end - -function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 2}) where TU - T = eltype(grid) - U_3D = zeros(dev, T, (1, grid.ny, nlayers)) - @views U_3D[1, :, :] = U - return U_3D -end - -function convert_U_to_U3D(dev, nlayers, grid, U::Number) - T = eltype(grid) - A = device_array(dev) - U_3D = reshape(repeat([T(U)], outer=(grid.ny, 1)), (1, grid.ny, nlayers)) - return A(U_3D) -end - -function Params_nf(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ, ν, nν, grid; calcFq=nothingfunction, effort=FFTW.MEASURE) - dev = grid.device - T = eltype(grid) - A = device_array(dev) - - ny, nx = grid.ny , grid.nx - nkr, nl = grid.nkr, grid.nl - kr, l = grid.kr , grid.l - - U = convert_U_to_U3D(dev, nlayers, grid, U) - - Uyy = real.(ifft(-l.^2 .* fft(U))) - Uyy = CUDA.@allowscalar repeat(Uyy, outer=(nx, 1, 1)) - - # Calculate periodic components of the topographic PV gradients. - etah = rfft(A(eta)) - etax = irfft(im * kr .* etah, nx) # ∂η/∂x - etay = irfft(im * l .* etah, nx) # ∂η/∂y - - # Add topographic PV large-scale gradient - topographic_pv_gradient = T.(topographic_pv_gradient) - @. etax += topographic_pv_gradient[1] - @. etay += topographic_pv_gradient[2] - - Qx = zeros(dev, T, (nx, ny, nlayers)) - @views @. Qx[:, :, nlayers] += etax - - Qy = zeros(dev, T, (nx, ny, nlayers)) - Qy = T(β) .- Uyy # T(β) is needed to ensure that Qy remains same type as U - @views @. Qy[:, :, nlayers] += etay - - rfftplanlayered = plan_flows_rfft(A{T, 3}(undef, grid.nx, grid.ny, nlayers), [1, 2]; flags=effort) - - if nlayers==1 - return SingleLayerParams(T(β), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, Qx, Qy, rfftplanlayered) - - else # if nlayers≥2 - - ρ = reshape(T.(ρ), (1, 1, nlayers)) - H = reshape(T.(H), (1, 1, nlayers)) - - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface - - Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) - Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) - - typeofSkl = SArray{Tuple{nlayers, nlayers}, T, 2, nlayers^2} # StaticArrays of type T and dims = (nlayers x nlayers) - - S = Array{typeofSkl, 2}(undef, (nkr, nl)) - calcS!(S, Fp, Fm, nlayers, grid) - - S⁻¹ = Array{typeofSkl, 2}(undef, (nkr, nl)) - calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) - - S, S⁻¹, Fp, Fm = A(S), A(S⁻¹), A(Fp), A(Fm) # convert to appropriate ArrayType - - CUDA.@allowscalar @views Qy[:, :, 1] = @. Qy[:, :, 1] - Fp[1] * (U[:, :, 2] - U[:, :, 1]) - for j = 2:nlayers-1 - CUDA.@allowscalar @views Qy[:, :, j] = @. Qy[:, :, j] - Fp[j] * (U[:, :, j+1] - U[:, :, j]) + Fm[j-1] * (U[:, :, j-1] - U[:, :, j]) - end - CUDA.@allowscalar @views Qy[:, :, nlayers] = @. Qy[:, :, nlayers] - Fm[nlayers-1] * (U[:, :, nlayers-1] - U[:, :, nlayers]) - - if nlayers==2 - return TwoLayerParams(T(g), T(f₀), T(β), A(ρ), (T(H[1]), T(H[2])), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, T(g′[1]), Qx, Qy, rfftplanlayered) - else # if nlayers>2 - return Params_nf(nlayers, T(g), T(f₀), T(β), A(ρ), A(H), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, A(g′), Qx, Qy, S, S⁻¹, rfftplanlayered) - end - end -end - -numberoflayers(params) = params.nlayers -numberoflayers(::SingleLayerParams) = 1 -numberoflayers(::TwoLayerParams) = 2 - -# --------- -# Equations -# --------- - -""" - hyperviscosity(params, grid) - -Return the linear operator `L` that corresponds to (hyper)-viscosity of order ``n_ν`` with -coefficient ``ν`` for ``n`` fluid layers. -```math -L_j = - ν |𝐤|^{2 n_ν}, \\ j = 1, ...,n . -``` -""" -function hyperviscosity(params, grid) - dev = grid.device - T = eltype(grid) - - L = device_array(dev){T}(undef, (grid.nkr, grid.nl, numberoflayers(params))) - @. L = - params.ν * grid.Krsq^params.nν - @views @. L[1, 1, :] = 0 - - return L -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 -`hyperviscosity(params, grid)`. - -The nonlinear term is computed via function `calcNlinear!`. -""" -function LinearEquation(params, grid) - L = hyperviscosity(params, grid) - - return FourierFlows.Equation(L, calcNlinear!, grid) -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 -`hyperviscosity(params, grid)`. - -The nonlinear term is computed via function `calcN!`. -""" -function Equation(params, grid) - L = hyperviscosity(params, grid) - - return FourierFlows.Equation(L, calcN!, grid) -end - - -# ---- -# Vars -# ---- - -""" - struct Vars{Aphys, Atrans, F, P} <: AbstractVars - -The variables for multi-layer QG problem. - -$(FIELDS) -""" -struct Vars{Aphys, Atrans, F, P} <: AbstractVars - "relative vorticity + vortex stretching" - q :: Aphys - "streamfunction" - ψ :: Aphys - "x-component of velocity" - u :: Aphys - "y-component of velocity" - v :: Aphys - "Fourier transform of relative vorticity + vortex stretching" - qh :: Atrans - "Fourier transform of streamfunction" - ψh :: Atrans - "Fourier transform of ``x``-component of velocity" - uh :: Atrans - "Fourier transform of ``y``-component of velocity" - vh :: Atrans - "Fourier transform of forcing" - Fqh :: F - "`sol` at previous time-step" - prevsol :: P -end - -const DecayingVars = Vars{<:AbstractArray, <:AbstractArray, Nothing, Nothing} -const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing} -const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray} - -""" - DecayingVars(grid, params) - -Return the variables for an unforced multi-layer QG problem with `grid` and `params`. -""" -function DecayingVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, nothing, nothing) -end - -""" - ForcedVars(grid, params) - -Return the variables for a forced multi-layer QG problem with `grid` and `params`. -""" -function ForcedVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, nothing) -end - -""" - StochasticForcedVars(grid, params) - -Return the variables for a forced multi-layer QG problem with `grid` and `params`. -""" -function StochasticForcedVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh prevsol - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, prevsol) -end - -""" - fwdtransform!(varh, var, params) - -Compute the Fourier transform of `var` and store it in `varh`. -""" -fwdtransform!(varh, var, params::AbstractParams) = mul!(varh, params.rfftplan, var) - -""" - invtransform!(var, varh, params) - -Compute the inverse Fourier transform of `varh` and store it in `var`. -""" -invtransform!(var, varh, params::AbstractParams) = ldiv!(var, params.rfftplan, varh) - -""" - pvfromstreamfunction!(qh, ψh, params, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` in each layer using -`qh = params.S * ψh`. -""" -function pvfromstreamfunction!(qh, ψh, params, grid) - for j=1:grid.nl, i=1:grid.nkr - CUDA.@allowscalar @views qh[i, j, :] .= params.S[i, j] * ψh[i, j, :] - end - - return nothing -end - -""" - pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` for the special -case of a single fluid layer configuration. In this case, ``q̂ = - k² ψ̂``. -""" -function pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) - @. qh = -grid.Krsq * ψh - - return nothing -end - -""" - pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` for the special -case of a two fluid layer configuration. In this case we have, - -```math -q̂₁ = - k² ψ̂₁ + f₀² / (g′ H₁) * (ψ̂₂ - ψ̂₁) , -``` - -```math -q̂₂ = - k² ψ̂₂ + f₀² / (g′ H₂) * (ψ̂₁ - ψ̂₂) . -``` - -(Here, the PV-streamfunction relationship is hard-coded to avoid scalar operations -on the GPU.) -""" -function pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) - f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] - - ψ1h, ψ2h = view(ψh, :, :, 1), view(ψh, :, :, 2) - - @views @. qh[:, :, 1] = - grid.Krsq * ψ1h + f₀^2 / (g′ * H₁) * (ψ2h - ψ1h) - @views @. qh[:, :, 2] = - grid.Krsq * ψ2h + f₀^2 / (g′ * H₂) * (ψ1h - ψ2h) - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` in each layer from -`qh` using `ψh = params.S⁻¹ qh`. -""" -function streamfunctionfrompv!(ψh, qh, params, grid) - for j=1:grid.nl, i=1:grid.nkr - CUDA.@allowscalar @views ψh[i, j, :] .= params.S⁻¹[i, j] * qh[i, j, :] - end - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` for the special -case of a single fluid layer configuration. In this case, ``ψ̂ = - k⁻² q̂``. -""" -function streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) - @. ψh = -grid.invKrsq * qh - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` for the special -case of a two fluid layer configuration. In this case we have, - -```math -ψ̂₁ = - [k² q̂₁ + (f₀² / g′) (q̂₁ / H₂ + q̂₂ / H₁)] / Δ , -``` - -```math -ψ̂₂ = - [k² q̂₂ + (f₀² / g′) (q̂₁ / H₂ + q̂₂ / H₁)] / Δ , -``` - -where ``Δ = k² [k² + f₀² (H₁ + H₂) / (g′ H₁ H₂)]``. - -(Here, the PV-streamfunction relationship is hard-coded to avoid scalar operations -on the GPU.) -""" -function streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) - f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] - - q1h, q2h = view(qh, :, :, 1), view(qh, :, :, 2) - - @views @. ψh[:, :, 1] = - grid.Krsq * q1h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) - @views @. ψh[:, :, 2] = - grid.Krsq * q2h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) - - for j in 1:2 - @views @. ψh[:, :, j] *= grid.invKrsq / (grid.Krsq + f₀^2 / g′ * (H₁ + H₂) / (H₁ * H₂)) - end - - return nothing -end - -""" - calcS!(S, Fp, Fm, nlayers, grid) - -Construct the array ``𝕊``, which consists of `nlayer` x `nlayer` static arrays ``𝕊_𝐤`` that -relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``q̂_𝐤 = 𝕊_𝐤 ψ̂_𝐤``. -""" -function calcS!(S, Fp, Fm, nlayers, grid) - F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) - - for n=1:grid.nl, m=1:grid.nkr - k² = CUDA.@allowscalar grid.Krsq[m, n] - Skl = SMatrix{nlayers, nlayers}(- k² * I + F) - S[m, n] = Skl - end - - return nothing -end - -""" - calcS⁻¹!(S, Fp, Fm, nlayers, grid) - -Construct the array ``𝕊⁻¹``, which consists of `nlayer` x `nlayer` static arrays ``(𝕊_𝐤)⁻¹`` -that relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``ψ̂_𝐤 = (𝕊_𝐤)⁻¹ q̂_𝐤``. -""" -function calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) - F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) - - for n=1:grid.nl, m=1:grid.nkr - k² = CUDA.@allowscalar grid.Krsq[m, n] == 0 ? 1 : grid.Krsq[m, n] - Skl = - k² * I + F - S⁻¹[m, n] = SMatrix{nlayers, nlayers}(I / Skl) - end - - T = eltype(grid) - S⁻¹[1, 1] = SMatrix{nlayers, nlayers}(zeros(T, (nlayers, nlayers))) - - return nothing -end - - -# ------- -# Solvers -# ------- - -""" - calcN!(N, sol, t, clock, vars, params, grid) - -Compute the nonlinear term, that is the advection term, the bottom drag, and the forcing: - -```math -N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . -``` -""" -function calcN!(N, sol, t, clock, vars, params, grid) - nlayers = numberoflayers(params) - - dealias!(sol, grid) - - calcN_advection!(N, sol, vars, params, grid) - - @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag - - addforcing!(N, sol, t, clock, vars, params, grid) - - return nothing -end - -""" - calcNlinear!(N, sol, t, clock, vars, params, grid) - -Compute the nonlinear term of the linearized equations: - -```math -N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} + \\widehat{(∂_y ψ_j)(∂_x Q_j)} -- \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . -``` -""" -function calcNlinear!(N, sol, t, clock, vars, params, grid) - nlayers = numberoflayers(params) - - calcN_linearadvection!(N, sol, vars, params, grid) - @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag - addforcing!(N, sol, t, clock, vars, params, grid) - - return nothing -end - -""" - calcN_advection!(N, sol, vars, params, grid) - -Compute the advection term and stores it in `N`: - -```math -N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} . -``` -""" -function calcN_advection!(N, sol, vars, params, grid) - @. vars.qh = sol - - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.u, vars.uh, params) - @. vars.u += params.U # add the imposed zonal flow U - - uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables - @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x - fwdtransform!(uQxh, uQx, params) - @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} - - invtransform!(vars.v, vars.vh, params) - - vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables - @. vQy = vars.v * params.Qy # v*∂Q/∂y - fwdtransform!(vQyh, vQy, params) - @. N -= vQyh # -\hat{v*∂Q/∂y} - - invtransform!(vars.q, vars.qh, params) - - uq , vq = vars.u , vars.v # use vars.u and vars.v as scratch variables - uqh, vqh = vars.uh, vars.vh # use vars.uh and vars.vh as scratch variables - @. uq *= vars.q # (U+u)*q - @. vq *= vars.q # v*q - - fwdtransform!(uqh, uq, params) - fwdtransform!(vqh, vq, params) - - @. N -= im * grid.kr * uqh + im * grid.l * vqh # -\hat{∂[(U+u)q]/∂x} - \hat{∂[vq]/∂y} - - return nothing -end - - -""" - calcN_linearadvection!(N, sol, vars, params, grid) - -Compute the advection term of the linearized equations and stores it in `N`: - -```math -N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} . -``` -""" -function calcN_linearadvection!(N, sol, vars, params, grid) - @. vars.qh = sol - - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.u, vars.uh, params) - @. vars.u += params.U # add the imposed zonal flow U - uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables - @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x - fwdtransform!(uQxh, uQx, params) - @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} - - invtransform!(vars.v, vars.vh, params) - - vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables - - @. vQy = vars.v * params.Qy # v*∂Q/∂y - fwdtransform!(vQyh, vQy, params) - @. N -= vQyh # -\hat{v*∂Q/∂y} - - invtransform!(vars.q, vars.qh, params) - - @. vars.u = params.U - Uq , Uqh = vars.u , vars.uh # use vars.u and vars.uh as scratch variables - @. Uq *= vars.q # U*q - - fwdtransform!(Uqh, Uq, params) - - @. N -= im * grid.kr * Uqh # -\hat{∂[U*q]/∂x} - - return nothing -end - - -""" - addforcing!(N, sol, t, clock, vars, params, grid) - -When the problem includes forcing, calculate the forcing term ``F̂`` for each layer and add -it to the nonlinear term ``N``. -""" -addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing - -function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) - params.calcFq!(vars.Fqh, sol, t, clock, vars, params, grid) - @. N += vars.Fqh - - return nothing -end - - -# ---------------- -# Helper functions -# ---------------- - -""" - updatevars!(vars, params, grid, sol) - updatevars!(prob) - -Update all problem variables using `sol`. -""" -function updatevars!(vars, params, grid, sol) - dealias!(sol, grid) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.q, deepcopy(vars.qh), params) - invtransform!(vars.ψ, deepcopy(vars.ψh), params) - invtransform!(vars.u, deepcopy(vars.uh), params) - invtransform!(vars.v, deepcopy(vars.vh), params) - - return nothing -end - -updatevars!(prob) = updatevars!(prob.vars, prob.params, prob.grid, prob.sol) - - -""" - set_q!(sol, params, vars, grid, q) - set_q!(prob, q) - -Set the solution `prob.sol` as the transform of `q` and update variables. -""" -function set_q!(sol, params, vars, grid, q) - A = typeof(vars.q) - fwdtransform!(vars.qh, A(q), params) - @. vars.qh[1, 1, :] = 0 - @. sol = vars.qh - updatevars!(vars, params, grid, sol) - - return nothing -end - -function set_q!(sol, params::SingleLayerParams, vars, grid, q::AbstractArray{T, 2}) where T - A = typeof(vars.q[:, :, 1]) - q_3D = vars.q - @views q_3D[:, :, 1] = A(q) - set_q!(sol, params, vars, grid, q_3D) - - return nothing -end - -set_q!(prob, q) = set_q!(prob.sol, prob.params, prob.vars, prob.grid, q) - - -""" - set_ψ!(params, vars, grid, sol, ψ) - set_ψ!(prob, ψ) - -Set the solution `prob.sol` to the transform `qh` that corresponds to streamfunction `ψ` -and update variables. -""" -function set_ψ!(sol, params, vars, grid, ψ) - A = typeof(vars.q) - fwdtransform!(vars.ψh, A(ψ), params) - pvfromstreamfunction!(vars.qh, vars.ψh, params, grid) - invtransform!(vars.q, vars.qh, params) - - set_q!(sol, params, vars, grid, vars.q) - - return nothing -end - -function set_ψ!(sol, params::SingleLayerParams, vars, grid, ψ::AbstractArray{T, 2}) where T - A = typeof(vars.ψ[:, :, 1]) - ψ_3D = vars.ψ - @views ψ_3D[:, :, 1] = A(ψ) - - set_ψ!(sol, params, vars, grid, ψ_3D) - - return nothing -end - -set_ψ!(prob, ψ) = set_ψ!(prob.sol, prob.params, prob.vars, prob.grid, ψ) - - -""" - energies(vars, params, grid, sol) - energies(prob) - -Return the kinetic energy of each fluid layer KE``_1, ...,`` KE``_{n}``, and the -potential energy of each fluid interface PE``_{3/2}, ...,`` PE``_{n-1/2}``, where ``n`` -is the number of layers in the fluid. (When ``n=1``, only the kinetic energy is returned.) - -The kinetic energy at the ``j``-th fluid layer is - -```math -𝖪𝖤_j = \\frac{H_j}{H} \\int \\frac1{2} |{\\bf ∇} ψ_j|^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{H_j}{H} \\sum_{𝐤} |𝐤|² |ψ̂_j|², \\ j = 1, ..., n , -``` - -while the potential energy that corresponds to the interface ``j+1/2`` (i.e., the interface -between the ``j``-th and ``(j+1)``-th fluid layer) is - -```math -𝖯𝖤_{j+1/2} = \\int \\frac1{2} \\frac{f₀^2}{g'_{j+1/2} H} (ψ_j - ψ_{j+1})^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{f₀^2}{g'_{j+1/2} H} \\sum_{𝐤} |ψ̂_j - ψ̂_{j+1}|², \\ j = 1, ..., n-1 . -``` -""" -function energies(vars, params, grid, sol) - nlayers = numberoflayers(params) - KE, PE = zeros(nlayers), zeros(nlayers-1) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - for j = 1:nlayers - CUDA.@allowscalar KE[j] = 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) - end - - for j = 1:nlayers-1 - CUDA.@allowscalar PE[j] = 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′[j] * parsevalsum(abs2.(vars.ψh[:, :, j+1] .- vars.ψh[:, :, j]), grid) - end - - return KE, PE -end - -function energies(vars, params::TwoLayerParams, grid, sol) - nlayers = numberoflayers(params) - KE, PE = zeros(nlayers), zeros(nlayers-1) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - ψ1h, ψ2h = view(vars.ψh, :, :, 1), view(vars.ψh, :, :, 2) - - for j = 1:nlayers - CUDA.@allowscalar KE[j] = @views 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) - end - - PE = @views 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′ * parsevalsum(abs2.(ψ2h .- ψ1h), grid) - - return KE, PE -end - -function energies(vars, params::SingleLayerParams, grid, sol) - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - return 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h, grid) -end - -energies(prob) = energies(prob.vars, prob.params, prob.grid, prob.sol) - -""" - fluxes(vars, params, grid, sol) - fluxes(prob) - -Return the lateral eddy fluxes within each fluid layer, lateralfluxes``_1,...,``lateralfluxes``_n`` -and also the vertical eddy fluxes at each fluid interface, -verticalfluxes``_{3/2},...,``verticalfluxes``_{n-1/2}``, where ``n`` is the total number of layers in the fluid. -(When ``n=1``, only the lateral fluxes are returned.) - -The lateral eddy fluxes within the ``j``-th fluid layer are - -```math -\\textrm{lateralfluxes}_j = \\frac{H_j}{H} \\int U_j v_j ∂_y u_j -\\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n , -``` - -while the vertical eddy fluxes at the ``j+1/2``-th fluid interface (i.e., interface between -the ``j``-th and ``(j+1)``-th fluid layer) are - -```math -\\textrm{verticalfluxes}_{j+1/2} = \\int \\frac{f₀²}{g'_{j+1/2} H} (U_j - U_{j+1}) \\, -v_{j+1} ψ_{j} \\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n-1. -``` -""" -function fluxes(vars, params, grid, sol) - nlayers = numberoflayers(params) - - lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) - - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.H * params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - for j = 1:nlayers-1 - CUDA.@allowscalar verticalfluxes[j] = sum(@views @. params.f₀^2 / params.g′[j] * (params.U[: ,:, j] - params.U[:, :, j+1]) * vars.v[:, :, j+1] * vars.ψ[:, :, j]; dims=(1, 2))[1] - CUDA.@allowscalar verticalfluxes[j] *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - end - - return lateralfluxes, verticalfluxes -end - -function fluxes(vars, params::TwoLayerParams, grid, sol) - nlayers = numberoflayers(params) - - lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) - - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - @. lateralfluxes *= params.H - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - U₁, U₂ = view(params.U, :, :, 1), view(params.U, :, :, 2) - ψ₁ = view(vars.ψ, :, :, 1) - v₂ = view(vars.v, :, :, 2) - - verticalfluxes = sum(@views @. params.f₀^2 / params.g′ * (U₁ - U₂) * v₂ * ψ₁; dims=(1, 2)) - verticalfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - return lateralfluxes, verticalfluxes -end - -function fluxes(vars, params::SingleLayerParams, grid, sol) - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly) - - return lateralfluxes -end - -fluxes(prob) = fluxes(prob.vars, prob.params, prob.grid, prob.sol) - -end # module diff --git a/src/multilayerqg_sign_fix_only.jl b/src/multilayerqg_sign_fix_only.jl deleted file mode 100644 index 6d3a27d3..00000000 --- a/src/multilayerqg_sign_fix_only.jl +++ /dev/null @@ -1,1107 +0,0 @@ -module MultiLayerQG_sf - -export - fwdtransform!, - invtransform!, - streamfunctionfrompv!, - pvfromstreamfunction!, - updatevars!, - - set_q!, - set_ψ!, - energies, - fluxes - -using - FFTW, - CUDA, - LinearAlgebra, - StaticArrays, - Reexport, - DocStringExtensions - -@reexport using FourierFlows - -using FourierFlows: parsevalsum, parsevalsum2, superzeros, plan_flows_rfft - -nothingfunction(args...) = nothing - -""" - Problem(nlayers :: Int, - dev = CPU(); - nx = 128, - ny = nx, - Lx = 2π, - Ly = Lx, - f₀ = 1.0, - β = 0.0, - g = 1.0, - U = zeros(nlayers), - H = 1/nlayers * ones(nlayers), - ρ = Array{Float64}(1:nlayers), - eta = nothing, - topographic_pv_gradient = (0, 0), - μ = 0, - ν = 0, - nν = 1, - dt = 0.01, - stepper = "RK4", - calcFq = nothingfunction, - stochastic = false, - linear = false, - aliased_fraction = 1/3, - T = Float64) - -Construct a multi-layer quasi-geostrophic problem with `nlayers` fluid layers on device `dev`. - -Arguments -========= -- `nlayers`: (required) Number of fluid layers. -- `dev`: (required) `CPU()` (default) or `GPU()`; computer architecture used to time-step `problem`. - -Keyword arguments -================= - - `nx`: Number of grid points in ``x``-domain. - - `ny`: Number of grid points in ``y``-domain. - - `Lx`: Extent of the ``x``-domain. - - `Ly`: Extent of the ``y``-domain. - - `f₀`: Constant planetary vorticity. - - `β`: Planetary vorticity ``y``-gradient. - - `g`: Gravitational acceleration constant. - - `U`: The imposed constant zonal flow ``U(y)`` in each fluid layer. - - `H`: Rest height of each fluid layer. - - `ρ`: Density of each fluid layer. - - `eta`: Periodic component of the topographic potential vorticity. - - `topographic_pv_gradient`: The ``(x, y)`` components of the topographic PV large-scale gradient. - - `μ`: Linear bottom drag coefficient. - - `ν`: Small-scale (hyper)-viscosity coefficient. - - `nν`: (Hyper)-viscosity order, `nν```≥ 1``. - - `dt`: Time-step. - - `stepper`: Time-stepping method. - - `calcF`: Function that calculates the Fourier transform of the forcing, ``F̂``. - - `stochastic`: `true` or `false` (default); boolean denoting whether `calcF` is temporally stochastic. - - `linear`: `true` or `false` (default); boolean denoting whether the linearized equations of motions are used. - - `aliased_fraction`: the fraction of high-wavenumbers that are zero-ed out by `dealias!()`. - - `T`: `Float32` or `Float64` (default); floating point type used for `problem` data. -""" -function Problem_sf(nlayers::Int, # number of fluid layers - dev = CPU(); - # Numerical parameters - nx = 128, - ny = nx, - Lx = 2π, - Ly = Lx, - # Physical parameters - f₀ = 1.0, # Coriolis parameter - β = 0.0, # y-gradient of Coriolis parameter - g = 1.0, # gravitational constant - U = zeros(nlayers), # imposed zonal flow U(y) in each layer - H = 1/nlayers * ones(nlayers), # rest fluid height of each layer - ρ = Array{Float64}(1:nlayers), # density of each layer - eta = nothing, # periodic component of the topographic PV - topographic_pv_gradient = (0, 0), # tuple with the ``(x, y)`` components of topographic PV large-scale gradient - # Bottom Drag and/or (hyper)-viscosity - μ = 0, - ν = 0, - nν = 1, - # Timestepper and equation options - dt = 0.01, - stepper = "RK4", - calcFq = nothingfunction, - stochastic = false, - linear = false, - # Float type and dealiasing - aliased_fraction = 1/3, - T = Float64) - - if dev == GPU() && nlayers > 2 - @warn """MultiLayerQG module is not optimized on the GPU yet for configurations with - 3 fluid layers or more! - - See issues on Github at https://github.com/FourierFlows/GeophysicalFlows.jl/issues/112 - and https://github.com/FourierFlows/GeophysicalFlows.jl/issues/267. - - To use MultiLayerQG with 3 fluid layers or more we suggest, for now, to restrict running - on CPU.""" - end - - if nlayers == 1 - @warn """MultiLayerQG module does work for single-layer configuration but may not be as - optimized. We suggest using SingleLayerQG module for single-layer QG simulation unless - you have reasons to use MultiLayerQG in a single-layer configuration, e.g., you want to - compare solutions with varying number of fluid layers.""" - end - - # topographic PV - eta === nothing && (eta = zeros(dev, T, (nx, ny))) - - grid = TwoDGrid(dev; nx, Lx, ny, Ly, aliased_fraction, T) - - params = Params_sf(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ, ν, nν, grid; calcFq) - - vars = calcFq == nothingfunction ? DecayingVars(grid, params) : (stochastic ? StochasticForcedVars(grid, params) : ForcedVars(grid, params)) - - equation = linear ? LinearEquation(params, grid) : Equation(params, grid) - - FourierFlows.Problem(equation, stepper, dt, grid, vars, params) -end - -""" - struct Params{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams - -The parameters for the `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct Params_sf{T, Aphys3D, Aphys2D, Aphys1D, Atrans4D, Trfft} <: AbstractParams - # prescribed params - "number of fluid layers" - nlayers :: Int - "gravitational constant" - g :: T - "constant planetary vorticity" - f₀ :: T - "planetary vorticity ``y``-gradient" - β :: T - "array with density of each fluid layer" - ρ :: Aphys3D - "array with rest height of each fluid layer" - H :: Aphys3D - "array with imposed constant zonal flow ``U(y)`` in each fluid layer" - U :: Aphys3D - "array containing the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear bottom drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "array with the reduced gravity constants for each fluid interface" - g′ :: Aphys1D - "array containing ``x``-gradient of PV due to topographic PV in each fluid layer" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV in each fluid layer" - Qy :: Aphys3D - "array containing coeffients for getting PV from streamfunction" - S :: Atrans4D - "array containing coeffients for inverting PV to streamfunction" - S⁻¹ :: Atrans4D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -""" - struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - -The parameters for the a single-layer `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct SingleLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - # prescribed params - "planetary vorticity ``y``-gradient" - β :: T - "array with imposed constant zonal flow ``U(y)``" - U :: Aphys3D - "array containing the periodic component of the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "array containing ``x``-gradient of PV due to topographic PV" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV" - Qy :: Aphys3D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -""" - struct TwoLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - -The parameters for the a two-layer `MultiLayerQG` problem. - -$(TYPEDFIELDS) -""" -struct TwoLayerParams{T, Aphys3D, Aphys2D, Trfft} <: AbstractParams - # prescribed params - "gravitational constant" - g :: T - "constant planetary vorticity" - f₀ :: T - "planetary vorticity ``y``-gradient" - β :: T - "array with density of each fluid layer" - ρ :: Aphys3D - "tuple with rest height of each fluid layer" - H :: Tuple - "array with imposed constant zonal flow ``U(y)`` in each fluid layer" - U :: Aphys3D - "array containing periodic component of the topographic PV" - eta :: Aphys2D - "tuple containing the ``(x, y)`` components of topographic PV large-scale gradient" - topographic_pv_gradient :: Tuple{T, T} - "linear bottom drag coefficient" - μ :: T - "small-scale (hyper)-viscosity coefficient" - ν :: T - "(hyper)-viscosity order, `nν```≥ 1``" - nν :: Int - "function that calculates the Fourier transform of the forcing, ``F̂``" - calcFq! :: Function - - # derived params - "the reduced gravity constants for the fluid interface" - g′ :: T - "array containing ``x``-gradient of PV due to topographic PV in each fluid layer" - Qx :: Aphys3D - "array containing ``y``-gradient of PV due to ``β``, ``U``, and topographic PV in each fluid layer" - Qy :: Aphys3D - "rfft plan for FFTs" - rfftplan :: Trfft -end - -function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 1}) where TU - T = eltype(grid) - if length(U) == nlayers - U_2D = zeros(dev, T, (1, nlayers)) - U_2D[:] = U - U_2D = repeat(U_2D, outer=(grid.ny, 1)) - else - U_2D = zeros(dev, T, (grid.ny, 1)) - U_2D[:] = U - end - U_3D = zeros(dev, T, (1, grid.ny, nlayers)) - @views U_3D[1, :, :] = U_2D - return U_3D -end - -function convert_U_to_U3D(dev, nlayers, grid, U::AbstractArray{TU, 2}) where TU - T = eltype(grid) - U_3D = zeros(dev, T, (1, grid.ny, nlayers)) - @views U_3D[1, :, :] = U - return U_3D -end - -function convert_U_to_U3D(dev, nlayers, grid, U::Number) - T = eltype(grid) - A = device_array(dev) - U_3D = reshape(repeat([T(U)], outer=(grid.ny, 1)), (1, grid.ny, nlayers)) - return A(U_3D) -end - -function Params_sf(nlayers, g, f₀, β, ρ, H, U, eta, topographic_pv_gradient, μ, ν, nν, grid; calcFq=nothingfunction, effort=FFTW.MEASURE) - dev = grid.device - T = eltype(grid) - A = device_array(dev) - - ny, nx = grid.ny , grid.nx - nkr, nl = grid.nkr, grid.nl - kr, l = grid.kr , grid.l - - U = convert_U_to_U3D(dev, nlayers, grid, U) - - Uyy = real.(ifft(-l.^2 .* fft(U))) - Uyy = CUDA.@allowscalar repeat(Uyy, outer=(nx, 1, 1)) - - # Calculate periodic components of the topographic PV gradients. - etah = rfft(A(eta)) - etax = irfft(im * kr .* etah, nx) # ∂η/∂x - etay = irfft(im * l .* etah, nx) # ∂η/∂y - - # Add topographic PV large-scale gradient - topographic_pv_gradient = T.(topographic_pv_gradient) - @. etax += topographic_pv_gradient[1] - @. etay += topographic_pv_gradient[2] - - Qx = zeros(dev, T, (nx, ny, nlayers)) - @views @. Qx[:, :, nlayers] += etax - - Qy = zeros(dev, T, (nx, ny, nlayers)) - Qy = T(β) .- Uyy # T(β) is needed to ensure that Qy remains same type as U - @views @. Qy[:, :, nlayers] += etay - - rfftplanlayered = plan_flows_rfft(A{T, 3}(undef, grid.nx, grid.ny, nlayers), [1, 2]; flags=effort) - - if nlayers==1 - return SingleLayerParams(T(β), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, Qx, Qy, rfftplanlayered) - - else # if nlayers≥2 - - ρ = reshape(T.(ρ), (1, 1, nlayers)) - H = reshape(T.(H), (1, 1, nlayers)) - - g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface - - Fm = @. T(f₀^2 / (g′ * H[2:nlayers])) - Fp = @. T(f₀^2 / (g′ * H[1:nlayers-1])) - - typeofSkl = SArray{Tuple{nlayers, nlayers}, T, 2, nlayers^2} # StaticArrays of type T and dims = (nlayers x nlayers) - - S = Array{typeofSkl, 2}(undef, (nkr, nl)) - calcS!(S, Fp, Fm, nlayers, grid) - - S⁻¹ = Array{typeofSkl, 2}(undef, (nkr, nl)) - calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) - - S, S⁻¹, Fp, Fm = A(S), A(S⁻¹), A(Fp), A(Fm) # convert to appropriate ArrayType - - CUDA.@allowscalar @views Qy[:, :, 1] = @. Qy[:, :, 1] - Fp[1] * (U[:, :, 2] - U[:, :, 1]) - for j = 2:nlayers-1 - CUDA.@allowscalar @views Qy[:, :, j] = @. Qy[:, :, j] + Fp[j] * (U[:, :, j+1] - U[:, :, j]) + Fm[j-1] * (U[:, :, j-1] - U[:, :, j]) - end - CUDA.@allowscalar @views Qy[:, :, nlayers] = @. Qy[:, :, nlayers] - Fm[nlayers-1] * (U[:, :, nlayers-1] - U[:, :, nlayers]) - - if nlayers==2 - return TwoLayerParams(T(g), T(f₀), T(β), A(ρ), (T(H[1]), T(H[2])), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, T(g′[1]), Qx, Qy, rfftplanlayered) - else # if nlayers>2 - return Params_sf(nlayers, T(g), T(f₀), T(β), A(ρ), A(H), U, eta, topographic_pv_gradient, T(μ), T(ν), nν, calcFq, A(g′), Qx, Qy, S, S⁻¹, rfftplanlayered) - end - end -end - -numberoflayers(params) = params.nlayers -numberoflayers(::SingleLayerParams) = 1 -numberoflayers(::TwoLayerParams) = 2 - -# --------- -# Equations -# --------- - -""" - hyperviscosity(params, grid) - -Return the linear operator `L` that corresponds to (hyper)-viscosity of order ``n_ν`` with -coefficient ``ν`` for ``n`` fluid layers. -```math -L_j = - ν |𝐤|^{2 n_ν}, \\ j = 1, ...,n . -``` -""" -function hyperviscosity(params, grid) - dev = grid.device - T = eltype(grid) - - L = device_array(dev){T}(undef, (grid.nkr, grid.nl, numberoflayers(params))) - @. L = - params.ν * grid.Krsq^params.nν - @views @. L[1, 1, :] = 0 - - return L -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 -`hyperviscosity(params, grid)`. - -The nonlinear term is computed via function `calcNlinear!`. -""" -function LinearEquation(params, grid) - L = hyperviscosity(params, grid) - - return FourierFlows.Equation(L, calcNlinear!, grid) -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 -`hyperviscosity(params, grid)`. - -The nonlinear term is computed via function `calcN!`. -""" -function Equation(params, grid) - L = hyperviscosity(params, grid) - - return FourierFlows.Equation(L, calcN!, grid) -end - - -# ---- -# Vars -# ---- - -""" - struct Vars{Aphys, Atrans, F, P} <: AbstractVars - -The variables for multi-layer QG problem. - -$(FIELDS) -""" -struct Vars{Aphys, Atrans, F, P} <: AbstractVars - "relative vorticity + vortex stretching" - q :: Aphys - "streamfunction" - ψ :: Aphys - "x-component of velocity" - u :: Aphys - "y-component of velocity" - v :: Aphys - "Fourier transform of relative vorticity + vortex stretching" - qh :: Atrans - "Fourier transform of streamfunction" - ψh :: Atrans - "Fourier transform of ``x``-component of velocity" - uh :: Atrans - "Fourier transform of ``y``-component of velocity" - vh :: Atrans - "Fourier transform of forcing" - Fqh :: F - "`sol` at previous time-step" - prevsol :: P -end - -const DecayingVars = Vars{<:AbstractArray, <:AbstractArray, Nothing, Nothing} -const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing} -const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray} - -""" - DecayingVars(grid, params) - -Return the variables for an unforced multi-layer QG problem with `grid` and `params`. -""" -function DecayingVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, nothing, nothing) -end - -""" - ForcedVars(grid, params) - -Return the variables for a forced multi-layer QG problem with `grid` and `params`. -""" -function ForcedVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, nothing) -end - -""" - StochasticForcedVars(grid, params) - -Return the variables for a forced multi-layer QG problem with `grid` and `params`. -""" -function StochasticForcedVars(grid, params) - Dev = typeof(grid.device) - T = eltype(grid) - nlayers = numberoflayers(params) - - @devzeros Dev T (grid.nx, grid.ny, nlayers) q ψ u v - @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) qh ψh uh vh Fqh prevsol - - return Vars(q, ψ, u, v, qh, ψh, uh, vh, Fqh, prevsol) -end - -""" - fwdtransform!(varh, var, params) - -Compute the Fourier transform of `var` and store it in `varh`. -""" -fwdtransform!(varh, var, params::AbstractParams) = mul!(varh, params.rfftplan, var) - -""" - invtransform!(var, varh, params) - -Compute the inverse Fourier transform of `varh` and store it in `var`. -""" -invtransform!(var, varh, params::AbstractParams) = ldiv!(var, params.rfftplan, varh) - -""" - pvfromstreamfunction!(qh, ψh, params, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` in each layer using -`qh = params.S * ψh`. -""" -function pvfromstreamfunction!(qh, ψh, params, grid) - for j=1:grid.nl, i=1:grid.nkr - CUDA.@allowscalar @views qh[i, j, :] .= params.S[i, j] * ψh[i, j, :] - end - - return nothing -end - -""" - pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` for the special -case of a single fluid layer configuration. In this case, ``q̂ = - k² ψ̂``. -""" -function pvfromstreamfunction!(qh, ψh, params::SingleLayerParams, grid) - @. qh = -grid.Krsq * ψh - - return nothing -end - -""" - pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) - -Obtain the Fourier transform of the PV from the streamfunction `ψh` for the special -case of a two fluid layer configuration. In this case we have, - -```math -q̂₁ = - k² ψ̂₁ + f₀² / (g′ H₁) * (ψ̂₂ - ψ̂₁) , -``` - -```math -q̂₂ = - k² ψ̂₂ + f₀² / (g′ H₂) * (ψ̂₁ - ψ̂₂) . -``` - -(Here, the PV-streamfunction relationship is hard-coded to avoid scalar operations -on the GPU.) -""" -function pvfromstreamfunction!(qh, ψh, params::TwoLayerParams, grid) - f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] - - ψ1h, ψ2h = view(ψh, :, :, 1), view(ψh, :, :, 2) - - @views @. qh[:, :, 1] = - grid.Krsq * ψ1h + f₀^2 / (g′ * H₁) * (ψ2h - ψ1h) - @views @. qh[:, :, 2] = - grid.Krsq * ψ2h + f₀^2 / (g′ * H₂) * (ψ1h - ψ2h) - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` in each layer from -`qh` using `ψh = params.S⁻¹ qh`. -""" -function streamfunctionfrompv!(ψh, qh, params, grid) - for j=1:grid.nl, i=1:grid.nkr - CUDA.@allowscalar @views ψh[i, j, :] .= params.S⁻¹[i, j] * qh[i, j, :] - end - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` for the special -case of a single fluid layer configuration. In this case, ``ψ̂ = - k⁻² q̂``. -""" -function streamfunctionfrompv!(ψh, qh, params::SingleLayerParams, grid) - @. ψh = -grid.invKrsq * qh - - return nothing -end - -""" - streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) - -Invert the PV to obtain the Fourier transform of the streamfunction `ψh` for the special -case of a two fluid layer configuration. In this case we have, - -```math -ψ̂₁ = - [k² q̂₁ + (f₀² / g′) (q̂₁ / H₂ + q̂₂ / H₁)] / Δ , -``` - -```math -ψ̂₂ = - [k² q̂₂ + (f₀² / g′) (q̂₁ / H₂ + q̂₂ / H₁)] / Δ , -``` - -where ``Δ = k² [k² + f₀² (H₁ + H₂) / (g′ H₁ H₂)]``. - -(Here, the PV-streamfunction relationship is hard-coded to avoid scalar operations -on the GPU.) -""" -function streamfunctionfrompv!(ψh, qh, params::TwoLayerParams, grid) - f₀, g′, H₁, H₂ = params.f₀, params.g′, params.H[1], params.H[2] - - q1h, q2h = view(qh, :, :, 1), view(qh, :, :, 2) - - @views @. ψh[:, :, 1] = - grid.Krsq * q1h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) - @views @. ψh[:, :, 2] = - grid.Krsq * q2h - f₀^2 / g′ * (q1h / H₂ + q2h / H₁) - - for j in 1:2 - @views @. ψh[:, :, j] *= grid.invKrsq / (grid.Krsq + f₀^2 / g′ * (H₁ + H₂) / (H₁ * H₂)) - end - - return nothing -end - -""" - calcS!(S, Fp, Fm, nlayers, grid) - -Construct the array ``𝕊``, which consists of `nlayer` x `nlayer` static arrays ``𝕊_𝐤`` that -relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``q̂_𝐤 = 𝕊_𝐤 ψ̂_𝐤``. -""" -function calcS!(S, Fp, Fm, nlayers, grid) - F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) - - for n=1:grid.nl, m=1:grid.nkr - k² = CUDA.@allowscalar grid.Krsq[m, n] - Skl = SMatrix{nlayers, nlayers}(- k² * I + F) - S[m, n] = Skl - end - - return nothing -end - -""" - calcS⁻¹!(S, Fp, Fm, nlayers, grid) - -Construct the array ``𝕊⁻¹``, which consists of `nlayer` x `nlayer` static arrays ``(𝕊_𝐤)⁻¹`` -that relate the ``q̂_j``'s and ``ψ̂_j``'s for every wavenumber: ``ψ̂_𝐤 = (𝕊_𝐤)⁻¹ q̂_𝐤``. -""" -function calcS⁻¹!(S⁻¹, Fp, Fm, nlayers, grid) - F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp)) - - for n=1:grid.nl, m=1:grid.nkr - k² = CUDA.@allowscalar grid.Krsq[m, n] == 0 ? 1 : grid.Krsq[m, n] - Skl = - k² * I + F - S⁻¹[m, n] = SMatrix{nlayers, nlayers}(I / Skl) - end - - T = eltype(grid) - S⁻¹[1, 1] = SMatrix{nlayers, nlayers}(zeros(T, (nlayers, nlayers))) - - return nothing -end - - -# ------- -# Solvers -# ------- - -""" - calcN!(N, sol, t, clock, vars, params, grid) - -Compute the nonlinear term, that is the advection term, the bottom drag, and the forcing: - -```math -N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . -``` -""" -function calcN!(N, sol, t, clock, vars, params, grid) - nlayers = numberoflayers(params) - - dealias!(sol, grid) - - calcN_advection!(N, sol, vars, params, grid) - - @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag - - addforcing!(N, sol, t, clock, vars, params, grid) - - return nothing -end - -""" - calcNlinear!(N, sol, t, clock, vars, params, grid) - -Compute the nonlinear term of the linearized equations: - -```math -N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} + \\widehat{(∂_y ψ_j)(∂_x Q_j)} -- \\widehat{(∂_x ψ_j)(∂_y Q_j)} + δ_{j, n} μ |𝐤|^2 ψ̂_n + F̂_j . -``` -""" -function calcNlinear!(N, sol, t, clock, vars, params, grid) - nlayers = numberoflayers(params) - - calcN_linearadvection!(N, sol, vars, params, grid) - @views @. N[:, :, nlayers] += params.μ * grid.Krsq * vars.ψh[:, :, nlayers] # bottom linear drag - addforcing!(N, sol, t, clock, vars, params, grid) - - return nothing -end - -""" - calcN_advection!(N, sol, vars, params, grid) - -Compute the advection term and stores it in `N`: - -```math -N_j = - \\widehat{𝖩(ψ_j, q_j)} - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} . -``` -""" -function calcN_advection!(N, sol, vars, params, grid) - @. vars.qh = sol - - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.u, vars.uh, params) - @. vars.u += params.U # add the imposed zonal flow U - - uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables - @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x - fwdtransform!(uQxh, uQx, params) - @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} - - invtransform!(vars.v, vars.vh, params) - - vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables - @. vQy = vars.v * params.Qy # v*∂Q/∂y - fwdtransform!(vQyh, vQy, params) - @. N -= vQyh # -\hat{v*∂Q/∂y} - - invtransform!(vars.q, vars.qh, params) - - uq , vq = vars.u , vars.v # use vars.u and vars.v as scratch variables - uqh, vqh = vars.uh, vars.vh # use vars.uh and vars.vh as scratch variables - @. uq *= vars.q # (U+u)*q - @. vq *= vars.q # v*q - - fwdtransform!(uqh, uq, params) - fwdtransform!(vqh, vq, params) - - @. N -= im * grid.kr * uqh + im * grid.l * vqh # -\hat{∂[(U+u)q]/∂x} - \hat{∂[vq]/∂y} - - return nothing -end - - -""" - calcN_linearadvection!(N, sol, vars, params, grid) - -Compute the advection term of the linearized equations and stores it in `N`: - -```math -N_j = - \\widehat{U_j ∂_x Q_j} - \\widehat{U_j ∂_x q_j} - + \\widehat{(∂_y ψ_j)(∂_x Q_j)} - \\widehat{(∂_x ψ_j)(∂_y Q_j)} . -``` -""" -function calcN_linearadvection!(N, sol, vars, params, grid) - @. vars.qh = sol - - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.u, vars.uh, params) - @. vars.u += params.U # add the imposed zonal flow U - uQx, uQxh = vars.q, vars.uh # use vars.q and vars.uh as scratch variables - @. uQx = vars.u * params.Qx # (U+u)*∂Q/∂x - fwdtransform!(uQxh, uQx, params) - @. N = - uQxh # -\hat{(U+u)*∂Q/∂x} - - invtransform!(vars.v, vars.vh, params) - - vQy, vQyh = vars.q, vars.vh # use vars.q and vars.vh as scratch variables - - @. vQy = vars.v * params.Qy # v*∂Q/∂y - fwdtransform!(vQyh, vQy, params) - @. N -= vQyh # -\hat{v*∂Q/∂y} - - invtransform!(vars.q, vars.qh, params) - - @. vars.u = params.U - Uq , Uqh = vars.u , vars.uh # use vars.u and vars.uh as scratch variables - @. Uq *= vars.q # U*q - - fwdtransform!(Uqh, Uq, params) - - @. N -= im * grid.kr * Uqh # -\hat{∂[U*q]/∂x} - - return nothing -end - - -""" - addforcing!(N, sol, t, clock, vars, params, grid) - -When the problem includes forcing, calculate the forcing term ``F̂`` for each layer and add -it to the nonlinear term ``N``. -""" -addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing - -function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) - params.calcFq!(vars.Fqh, sol, t, clock, vars, params, grid) - @. N += vars.Fqh - - return nothing -end - - -# ---------------- -# Helper functions -# ---------------- - -""" - updatevars!(vars, params, grid, sol) - updatevars!(prob) - -Update all problem variables using `sol`. -""" -function updatevars!(vars, params, grid, sol) - dealias!(sol, grid) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - @. vars.uh = -im * grid.l * vars.ψh - @. vars.vh = im * grid.kr * vars.ψh - - invtransform!(vars.q, deepcopy(vars.qh), params) - invtransform!(vars.ψ, deepcopy(vars.ψh), params) - invtransform!(vars.u, deepcopy(vars.uh), params) - invtransform!(vars.v, deepcopy(vars.vh), params) - - return nothing -end - -updatevars!(prob) = updatevars!(prob.vars, prob.params, prob.grid, prob.sol) - - -""" - set_q!(sol, params, vars, grid, q) - set_q!(prob, q) - -Set the solution `prob.sol` as the transform of `q` and update variables. -""" -function set_q!(sol, params, vars, grid, q) - A = typeof(vars.q) - fwdtransform!(vars.qh, A(q), params) - @. vars.qh[1, 1, :] = 0 - @. sol = vars.qh - updatevars!(vars, params, grid, sol) - - return nothing -end - -function set_q!(sol, params::SingleLayerParams, vars, grid, q::AbstractArray{T, 2}) where T - A = typeof(vars.q[:, :, 1]) - q_3D = vars.q - @views q_3D[:, :, 1] = A(q) - set_q!(sol, params, vars, grid, q_3D) - - return nothing -end - -set_q!(prob, q) = set_q!(prob.sol, prob.params, prob.vars, prob.grid, q) - - -""" - set_ψ!(params, vars, grid, sol, ψ) - set_ψ!(prob, ψ) - -Set the solution `prob.sol` to the transform `qh` that corresponds to streamfunction `ψ` -and update variables. -""" -function set_ψ!(sol, params, vars, grid, ψ) - A = typeof(vars.q) - fwdtransform!(vars.ψh, A(ψ), params) - pvfromstreamfunction!(vars.qh, vars.ψh, params, grid) - invtransform!(vars.q, vars.qh, params) - - set_q!(sol, params, vars, grid, vars.q) - - return nothing -end - -function set_ψ!(sol, params::SingleLayerParams, vars, grid, ψ::AbstractArray{T, 2}) where T - A = typeof(vars.ψ[:, :, 1]) - ψ_3D = vars.ψ - @views ψ_3D[:, :, 1] = A(ψ) - - set_ψ!(sol, params, vars, grid, ψ_3D) - - return nothing -end - -set_ψ!(prob, ψ) = set_ψ!(prob.sol, prob.params, prob.vars, prob.grid, ψ) - - -""" - energies(vars, params, grid, sol) - energies(prob) - -Return the kinetic energy of each fluid layer KE``_1, ...,`` KE``_{n}``, and the -potential energy of each fluid interface PE``_{3/2}, ...,`` PE``_{n-1/2}``, where ``n`` -is the number of layers in the fluid. (When ``n=1``, only the kinetic energy is returned.) - -The kinetic energy at the ``j``-th fluid layer is - -```math -𝖪𝖤_j = \\frac{H_j}{H} \\int \\frac1{2} |{\\bf ∇} ψ_j|^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{H_j}{H} \\sum_{𝐤} |𝐤|² |ψ̂_j|², \\ j = 1, ..., n , -``` - -while the potential energy that corresponds to the interface ``j+1/2`` (i.e., the interface -between the ``j``-th and ``(j+1)``-th fluid layer) is - -```math -𝖯𝖤_{j+1/2} = \\int \\frac1{2} \\frac{f₀^2}{g'_{j+1/2} H} (ψ_j - ψ_{j+1})^2 \\frac{𝖽x 𝖽y}{L_x L_y} = \\frac1{2} \\frac{f₀^2}{g'_{j+1/2} H} \\sum_{𝐤} |ψ̂_j - ψ̂_{j+1}|², \\ j = 1, ..., n-1 . -``` -""" -function energies(vars, params, grid, sol) - nlayers = numberoflayers(params) - KE, PE = zeros(nlayers), zeros(nlayers-1) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - for j = 1:nlayers - CUDA.@allowscalar KE[j] = 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) - end - - for j = 1:nlayers-1 - CUDA.@allowscalar PE[j] = 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′[j] * parsevalsum(abs2.(vars.ψh[:, :, j+1] .- vars.ψh[:, :, j]), grid) - end - - return KE, PE -end - -function energies(vars, params::TwoLayerParams, grid, sol) - nlayers = numberoflayers(params) - KE, PE = zeros(nlayers), zeros(nlayers-1) - - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - ψ1h, ψ2h = view(vars.ψh, :, :, 1), view(vars.ψh, :, :, 2) - - for j = 1:nlayers - CUDA.@allowscalar KE[j] = @views 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h[:, :, j], grid) * params.H[j] / sum(params.H) - end - - PE = @views 1 / (2 * grid.Lx * grid.Ly * sum(params.H)) * params.f₀^2 / params.g′ * parsevalsum(abs2.(ψ2h .- ψ1h), grid) - - return KE, PE -end - -function energies(vars, params::SingleLayerParams, grid, sol) - @. vars.qh = sol - streamfunctionfrompv!(vars.ψh, vars.qh, params, grid) - - abs²∇𝐮h = vars.uh # use vars.uh as scratch variable - @. abs²∇𝐮h = grid.Krsq * abs2(vars.ψh) - - return 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs²∇𝐮h, grid) -end - -energies(prob) = energies(prob.vars, prob.params, prob.grid, prob.sol) - -""" - fluxes(vars, params, grid, sol) - fluxes(prob) - -Return the lateral eddy fluxes within each fluid layer, lateralfluxes``_1,...,``lateralfluxes``_n`` -and also the vertical eddy fluxes at each fluid interface, -verticalfluxes``_{3/2},...,``verticalfluxes``_{n-1/2}``, where ``n`` is the total number of layers in the fluid. -(When ``n=1``, only the lateral fluxes are returned.) - -The lateral eddy fluxes within the ``j``-th fluid layer are - -```math -\\textrm{lateralfluxes}_j = \\frac{H_j}{H} \\int U_j v_j ∂_y u_j -\\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n , -``` - -while the vertical eddy fluxes at the ``j+1/2``-th fluid interface (i.e., interface between -the ``j``-th and ``(j+1)``-th fluid layer) are - -```math -\\textrm{verticalfluxes}_{j+1/2} = \\int \\frac{f₀²}{g'_{j+1/2} H} (U_j - U_{j+1}) \\, -v_{j+1} ψ_{j} \\frac{𝖽x 𝖽y}{L_x L_y} , \\ j = 1, ..., n-1. -``` -""" -function fluxes(vars, params, grid, sol) - nlayers = numberoflayers(params) - - lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) - - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.H * params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - for j = 1:nlayers-1 - CUDA.@allowscalar verticalfluxes[j] = sum(@views @. params.f₀^2 / params.g′[j] * (params.U[: ,:, j] - params.U[:, :, j+1]) * vars.v[:, :, j+1] * vars.ψ[:, :, j]; dims=(1, 2))[1] - CUDA.@allowscalar verticalfluxes[j] *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - end - - return lateralfluxes, verticalfluxes -end - -function fluxes(vars, params::TwoLayerParams, grid, sol) - nlayers = numberoflayers(params) - - lateralfluxes, verticalfluxes = zeros(nlayers), zeros(nlayers-1) - - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - @. lateralfluxes *= params.H - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - U₁, U₂ = view(params.U, :, :, 1), view(params.U, :, :, 2) - ψ₁ = view(vars.ψ, :, :, 1) - v₂ = view(vars.v, :, :, 2) - - verticalfluxes = sum(@views @. params.f₀^2 / params.g′ * (U₁ - U₂) * v₂ * ψ₁; dims=(1, 2)) - verticalfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly * sum(params.H)) - - return lateralfluxes, verticalfluxes -end - -function fluxes(vars, params::SingleLayerParams, grid, sol) - updatevars!(vars, params, grid, sol) - - ∂u∂yh = vars.uh # use vars.uh as scratch variable - ∂u∂y = vars.u # use vars.u as scratch variable - - @. ∂u∂yh = im * grid.l * vars.uh - invtransform!(∂u∂y, ∂u∂yh, params) - - lateralfluxes = (sum(@. params.U * vars.v * ∂u∂y; dims=(1, 2)))[1, 1, :] - lateralfluxes *= grid.dx * grid.dy / (grid.Lx * grid.Ly) - - return lateralfluxes -end - -fluxes(prob) = fluxes(prob.vars, prob.params, prob.grid, prob.sol) - -end # module From 957805e3351a77fe37141f259d27bb69d31f20bd Mon Sep 17 00:00:00 2001 From: Matt Lobo <104666665+mjclobo@users.noreply.github.com> Date: Fri, 9 Jun 2023 19:50:29 -0400 Subject: [PATCH 14/16] Update Project.toml Updating version number for patch release --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 9e2daa16..8f14cd3a 100644 --- a/Project.toml +++ b/Project.toml @@ -2,7 +2,7 @@ name = "GeophysicalFlows" uuid = "44ee3b1c-bc02-53fa-8355-8e347616e15e" license = "MIT" authors = ["Navid C. Constantinou ", "Gregory L. Wagner ", "and co-contributors"] -version = "0.15.2" +version = "0.15.3" [deps] CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" From 80364d0d3ba4ef65115d2a03e38ee2874e539469 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Sat, 10 Jun 2023 18:22:02 +1000 Subject: [PATCH 15/16] add nonlinear advection test for 3 layers --- test/runtests.jl | 5 +- test/test_multilayerqg.jl | 246 +++++++++++++++++++++++++++++++------- 2 files changed, 203 insertions(+), 48 deletions(-) diff --git a/test/runtests.jl b/test/runtests.jl index 6ab2dbce..6d542b55 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -121,7 +121,8 @@ for dev in devices @test test_pvtofromstreamfunction_2layer(dev) @test test_pvtofromstreamfunction_3layer(dev) @test test_mqg_rossbywave("RK4", 1e-2, 20, dev) - @test test_mqg_nonlinearadvection(0.005, "ForwardEuler", dev) + @test test_mqg_nonlinearadvection_2layers(0.005, "ForwardEuler", dev) + @test test_mqg_nonlinearadvection_3layers(0.005, "ForwardEuler", dev) @test test_mqg_linearadvection(0.005, "ForwardEuler", dev) @test test_mqg_energies(dev) @test test_mqg_energysinglelayer(dev) @@ -132,7 +133,7 @@ for dev in devices @test test_mqg_paramsconstructor(dev) @test test_mqg_stochasticforcedproblemconstructor(dev) @test test_mqg_problemtype(dev, Float32) - @test MultiLayerQG.nothingfunction() == nothing + @test MultiLayerQG.nothingfunction() === nothing end end end # time diff --git a/test/test_multilayerqg.jl b/test/test_multilayerqg.jl index 9e8f717c..3085aa6e 100644 --- a/test/test_multilayerqg.jl +++ b/test/test_multilayerqg.jl @@ -10,55 +10,109 @@ function constructtestfields_2layer(gr) k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers # a set of streafunctions ψ1 and ψ2, ... - ψ1 = @. 1e-3 * ( 1/4*cos(2k₀*x)*cos(5l₀*y) + 1/3*cos(3k₀*x)*cos(3l₀*y) ) - ψ2 = @. 1e-3 * ( cos(3k₀*x)*cos(4l₀*y) + 1/2*cos(4k₀*x)*cos(2l₀*y) ) + ψ1 = @. 1e-3 * ( 1/4 * cos(2k₀*x) * cos(5l₀*y) + + 1/3 * cos(3k₀*x) * cos(3l₀*y) ) + ψ2 = @. 1e-3 * ( cos(3k₀*x) * cos(4l₀*y) + + 1/2 * cos(4k₀*x) * cos(2l₀*y) ) + + Δψ1 = @. 1e-3 * ( - 1/4 * ((2k₀)^2 + (5l₀)^2) * cos(2k₀*x) * cos(5l₀*y) + + - 1/3 * ((3k₀)^2 + (3l₀)^2) * cos(3k₀*x) * cos(3l₀*y) ) + Δψ2 = @. 1e-3 * ( - ((3k₀)^2 + (4l₀)^2) * cos(3k₀*x) * cos(4l₀*y) + + - 1/2 * ((4k₀)^2 + (2l₀)^2) * cos(4k₀*x) * cos(2l₀*y) ) + + Δ²ψ1 = @. 1e-3 * ( + 1/4 * ((2k₀)^2 + (5l₀)^2)^2 * cos(2k₀*x) * cos(5l₀*y) + + + 1/3 * ((3k₀)^2 + (3l₀)^2)^2 * cos(3k₀*x) * cos(3l₀*y) ) + Δ²ψ2 = @. 1e-3 * ( + ((3k₀)^2 + (4l₀)^2)^2 * cos(3k₀*x) * cos(4l₀*y) + + + 1/2 * ((4k₀)^2 + (2l₀)^2)^2 * cos(4k₀*x) * cos(2l₀*y) ) # ... their corresponding PVs q1, q2, - q1 = @. 1e-3 * ( 1/6 *( 75cos(4k₀*x)*cos(2l₀*y) + 2cos(3k₀*x)*( -43cos(3l₀*y) + 75cos(4l₀*y) ) - 81cos(2k₀*x)*cos(5l₀*y) ) ) - q2 = @. 1e-3 * ( 1/48*( -630cos(4k₀*x)*cos(2l₀*y) + 100cos(3k₀*x)*( cos(3l₀*y) - 15cos(4l₀*y) ) + 75cos(2k₀*x)*cos(5l₀*y) ) ) + q1 = @. Δψ1 + 25 * ψ2 - 25 * ψ1 + q2 = @. Δψ2 + 25/4 * ψ1 - 25/4 * ψ2 - # ... and various derived fields, e.g., ∂ψ1/∂x, - ψ1x = @. 1e-3 * ( -k₀/2*sin(2k₀*x)*cos(5l₀*y) - k₀*sin(3k₀*x)*cos(3l₀*y) ) - ψ2x = @. 1e-3 * ( -3k₀*sin(3k₀*x)*cos(4l₀*y) - 2k₀*sin(4k₀*x)*cos(2l₀*y) ) - Δψ2 = @. 1e-3 * ( -25*k₀*l₀*cos(3k₀*x)*cos(4l₀*y) - 10*k₀*l₀*cos(4k₀*x)*cos(2l₀*y) ) + # ... and various derived fields, e.g., ∂ψ1/∂x,s + ψ1x = @. 1e-3 * ( - 1/4 * 2k₀ * sin(2k₀*x) * cos(5l₀*y) + + - 1/3 * 3k₀ * sin(3k₀*x) * cos(3l₀*y) ) + ψ2x = @. 1e-3 * ( - 3k₀ * sin(3k₀*x) * cos(4l₀*y) + + - 1/2 * 4k₀ * sin(4k₀*x) * cos(2l₀*y) ) - q1x = @. 1e-3 * ( 1/6 *( -4k₀*75sin(4k₀*x)*cos(2l₀*y) - 3k₀*2sin(3k₀*x)*( -43cos(3l₀*y) + 75cos(4l₀*y) ) + 2k₀*81sin(2k₀*x)*cos(5l₀*y) ) ) - q2x = @. 1e-3 * ( 1/48*( 4k₀*630sin(4k₀*x)*cos(2l₀*y) - 3k₀*100sin(3k₀*x)*( cos(3l₀*y) - 15cos(4l₀*y) ) - 2k₀*75sin(2k₀*x)*cos(5l₀*y) ) ) + Δψ1x = @. 1e-3 * ( + 1/4 * 2k₀ * ((2k₀)^2 + (5l₀)^2) * sin(2k₀*x) * cos(5l₀*y) + + + 1/3 * 3k₀ * ((3k₀)^2 + (3l₀)^2) * sin(3k₀*x) * cos(3l₀*y) ) + Δψ2x = @. 1e-3 * ( + 3k₀ * ((3k₀)^2 + (4l₀)^2) * sin(3k₀*x) * cos(4l₀*y) + + + 1/2 * 4k₀ * ((4k₀)^2 + (2l₀)^2) * sin(4k₀*x) * cos(2l₀*y) ) - Δq1 = @. 1e-3 * (k₀*l₀)*( 1/6 *( -20* 75cos(4k₀*x)*cos(2l₀*y) + 2cos(3k₀*x)*( +18*43cos(3l₀*y) - 25*75cos(4l₀*y) ) +29*81cos(2k₀*x)*cos(5l₀*y) ) ) - Δq2 = @. 1e-3 * (k₀*l₀)*( 1/48*( +20*630cos(4k₀*x)*cos(2l₀*y) + 100cos(3k₀*x)*( -18cos(3l₀*y) + 25*15cos(4l₀*y) ) -29*75cos(2k₀*x)*cos(5l₀*y) ) ) + # ... their corresponding PVs q1, q2, q3, + Δq1 = @. Δ²ψ1 + 25 * Δψ2 - 25 * Δψ1 + Δq2 = @. Δ²ψ2 + 25/4 * Δψ1 - 25/4 * Δψ2 + + q1x = @. Δψ1x + 25 * ψ2x - 25 * ψ1x + q2x = @. Δψ2x + 25/4 * ψ1x - 25/4 * ψ2x return ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2 end - """ constructtestfields_3layer(gr) Constructs flow fields for a 3-layer problem with parameters such that q1 = Δψ1 + 20ψ2 - 20ψ1, q2 = Δψ2 + 20ψ1 - 44ψ2 + 24ψ3, - q2 = Δψ2 + 12ψ2 - 12ψ3. + q3 = Δψ3 + 12ψ2 - 12ψ3. """ function constructtestfields_3layer(gr) x, y = gridpoints(gr) k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers # a set of streafunctions ψ1, ψ2, ψ3, ... - ψ1 = @. 1e-3 * ( 1/4*cos(2k₀*x)*cos(5l₀*y) + 1/3*cos(3k₀*x)*cos(3l₀*y) ) - ψ2 = @. 1e-3 * ( cos(3k₀*x)*cos(4l₀*y) + 1/2*cos(4k₀*x)*cos(2l₀*y) ) - ψ3 = @. 1e-3 * ( cos(1k₀*x)*cos(3l₀*y) + 1/2*cos(2k₀*x)*cos(2l₀*y) ) - - Δψ1 = @. -1e-3 * ( 1/4*((2k₀)^2+(5l₀)^2)*cos(2k₀*x)*cos(5l₀*y) + 1/3*((3k₀)^2+(3l₀)^2)*cos(3k₀*x)*cos(3l₀*y) ) - Δψ2 = @. -1e-3 * ( ((3k₀)^2+(4l₀)^2)*cos(3k₀*x)*cos(4l₀*y) + 1/2*((4k₀)^2+(2l₀)^2)*cos(4k₀*x)*cos(2l₀*y) ) - Δψ3 = @. -1e-3 * ( ((1k₀)^2+(3l₀)^2)*cos(1k₀*x)*cos(3l₀*y) + 1/2*((2k₀)^2+(2l₀)^2)*cos(2k₀*x)*cos(2l₀*y) ) + ψ1 = @. 1e-3 * ( 1/4 * cos(2k₀*x) * cos(5l₀*y) + + 1/3 * cos(3k₀*x) * cos(3l₀*y) ) + ψ2 = @. 1e-3 * ( cos(3k₀*x) * cos(4l₀*y) + + 1/2 * cos(4k₀*x) * cos(2l₀*y) ) + ψ3 = @. 1e-3 * ( cos(1k₀*x) * cos(3l₀*y) + + 1/2 * cos(2k₀*x) * cos(2l₀*y) ) + + Δψ1 = @. 1e-3 * ( - 1/4 * ((2k₀)^2 + (5l₀)^2) * cos(2k₀*x) * cos(5l₀*y) + + - 1/3 * ((3k₀)^2 + (3l₀)^2) * cos(3k₀*x) * cos(3l₀*y) ) + Δψ2 = @. 1e-3 * ( - ((3k₀)^2 + (4l₀)^2) * cos(3k₀*x) * cos(4l₀*y) + + - 1/2 * ((4k₀)^2 + (2l₀)^2) * cos(4k₀*x) * cos(2l₀*y) ) + Δψ3 = @. 1e-3 * ( - ((1k₀)^2 + (3l₀)^2) * cos(1k₀*x) * cos(3l₀*y) + + - 1/2 * ((2k₀)^2 + (2l₀)^2) * cos(2k₀*x) * cos(2l₀*y) ) + + Δ²ψ1 = @. 1e-3 * ( + 1/4 * ((2k₀)^2 + (5l₀)^2)^2 * cos(2k₀*x) * cos(5l₀*y) + + + 1/3 * ((3k₀)^2 + (3l₀)^2)^2 * cos(3k₀*x) * cos(3l₀*y) ) + Δ²ψ2 = @. 1e-3 * ( + ((3k₀)^2 + (4l₀)^2)^2 * cos(3k₀*x) * cos(4l₀*y) + + + 1/2 * ((4k₀)^2 + (2l₀)^2)^2 * cos(4k₀*x) * cos(2l₀*y) ) + Δ²ψ3 = @. 1e-3 * ( + ((1k₀)^2 + (3l₀)^2)^2 * cos(1k₀*x) * cos(3l₀*y) + + + 1/2 * ((2k₀)^2 + (2l₀)^2)^2 * cos(2k₀*x) * cos(2l₀*y) ) # ... their corresponding PVs q1, q2, q3, - q1 = @. Δψ1 + 20ψ2 - 20ψ1 - q2 = @. Δψ2 + 20ψ1 - 44ψ2 + 24ψ3 - q3 = @. Δψ3 + 12ψ2 - 12ψ3 + q1 = @. Δψ1 + 20ψ2 - 20ψ1 + q2 = @. Δψ2 + 20ψ1 - 44ψ2 + 24ψ3 + q3 = @. Δψ3 + 12ψ2 - 12ψ3 - return ψ1, ψ2, ψ3, q1, q2, q3 + # ... and various derived fields, e.g., ∂ψ1/∂x, + ψ1x = @. 1e-3 * ( - 1/4 * 2k₀ * sin(2k₀*x) * cos(5l₀*y) + + - 1/3 * 3k₀ * sin(3k₀*x) * cos(3l₀*y) ) + ψ2x = @. 1e-3 * ( - 3k₀ * sin(3k₀*x) * cos(4l₀*y) + + - 1/2 * 4k₀ * sin(4k₀*x) * cos(2l₀*y) ) + ψ3x = @. 1e-3 * ( - 1k₀ * sin(1k₀*x) * cos(3l₀*y) + + - 1/2 * 2k₀ * sin(2k₀*x) * cos(2l₀*y) ) + + Δψ1x = @. 1e-3 * ( + 1/4 * 2k₀ * ((2k₀)^2 + (5l₀)^2) * sin(2k₀*x) * cos(5l₀*y) + + + 1/3 * 3k₀ * ((3k₀)^2 + (3l₀)^2) * sin(3k₀*x) * cos(3l₀*y) ) + Δψ2x = @. 1e-3 * ( + 3k₀ * ((3k₀)^2 + (4l₀)^2) * sin(3k₀*x) * cos(4l₀*y) + + + 1/2 * 4k₀ * ((4k₀)^2 + (2l₀)^2) * sin(4k₀*x) * cos(2l₀*y) ) + Δψ3x = @. 1e-3 * ( + 1k₀ * ((1k₀)^2 + (3l₀)^2) * sin(1k₀*x) * cos(3l₀*y) + + + 1/2 * 2k₀ * ((2k₀)^2 + (2l₀)^2) * sin(2k₀*x) * cos(2l₀*y) ) + + q1x = @. Δψ1x + 20ψ2x - 20ψ1x + q2x = @. Δψ2x + 20ψ1x - 44ψ2x + 24ψ3x + q3x = @. Δψ3x + 12ψ2x - 12ψ3x + + Δq1 = @. Δ²ψ1 + 20Δψ2 - 20Δψ1 + Δq2 = @. Δ²ψ2 + 20Δψ1 - 44Δψ2 + 24Δψ3 + Δq3 = @. Δ²ψ3 + 12Δψ2 - 12Δψ3 + + return ψ1, ψ2, ψ3, q1, q2, q3, ψ1x, ψ2x, ψ3x, q1x, q2x, q3x, Δq1, Δq2, Δq3, Δψ3 end @@ -127,7 +181,7 @@ function test_pvtofromstreamfunction_3layer(dev::Device=CPU()) prob = MultiLayerQG.Problem(nlayers, dev; nx=n, Lx=L, f₀, g, H, ρ) sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid - ψ1, ψ2, ψ3, q1, q2, q3 = constructtestfields_3layer(gr) + ψ1, ψ2, ψ3, q1, q2, q3, ψ1x, ψ2x, ψ3x, q1x, q2x, q3x, Δq1, Δq2, Δq3, Δψ3 = constructtestfields_3layer(gr) vs.ψh[:, :, 1] .= rfft(ψ1) vs.ψh[:, :, 2] .= rfft(ψ2) @@ -153,7 +207,7 @@ end """ - test_mqg_nonlinearadvection(dt, stepper, dev; kwargs...) + test_mqg_nonlinearadvection_2layers(dt, stepper, dev::Device=CPU()) Tests the advection term by timestepping a test problem with timestep dt and timestepper identified by the string stepper. The test 2-layer problem is @@ -162,15 +216,14 @@ the advection terms J(ψn, qn) are non-zero. Next, a forcing Ff is derived such that a solution to the problem forced by this Ff is then qf. (This solution may not be realized, at least at long times, if it is unstable.) """ -function test_mqg_nonlinearadvection(dt, stepper, dev::Device=CPU(); - n=128, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1) +function test_mqg_nonlinearadvection_2layers(dt, stepper, dev::Device=CPU()) A = device_array(dev) - tf = 0.5 + tf = 200dt nt = round(Int, tf/dt) - nx, ny = 64, 66 + nx, ny = 128, 126 Lx, Ly = 2π, 2π gr = TwoDGrid(dev; nx, Lx, ny, Ly) @@ -184,9 +237,11 @@ function test_mqg_nonlinearadvection(dt, stepper, dev::Device=CPU(); β = 0.35 - U1, U2 = 0.1, 0.05 - u1 = @. 0.5sech(gr.y/(Ly/15))^2; u1 = A(reshape(u1, (1, gr.ny))) - u2 = @. 0.02cos(3l₀*gr.y); u2 = A(reshape(u2, (1, gr.ny))) + U1, U2 = 0.02, 0.025 + + u1 = @. 0.05sech(gr.y / (Ly/15))^2; u1 = A(reshape(u1, (1, gr.ny))) + u2 = @. 0.02cos(3l₀*gr.y); u2 = A(reshape(u2, (1, gr.ny))) + uyy1 = real.(ifft( -gr.l.^2 .* fft(u1) )) uyy2 = real.(ifft( -gr.l.^2 .* fft(u2) )) @@ -197,13 +252,14 @@ function test_mqg_nonlinearadvection(dt, stepper, dev::Device=CPU(); μ, ν, nν = 0.1, 0.05, 1 η0, σx, σy = 1.0, Lx/25, Ly/20 - η = @. η0*exp( -(x+Lx/8)^2/(2σx^2) - (y-Ly/8)^2/(2σy^2) ) - ηx = @. -(x + Lx/8)/(σx^2) * η + η = @. η0 * exp( -(x + Lx/8)^2 / 2σx^2 - (y - Ly/8)^2 / 2σy^2) + ηx = @. - (x + Lx/8) / σx^2 * η ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2 = constructtestfields_2layer(gr) - Ff1 = FourierFlows.jacobian(ψ1, q1, gr) + @. (β - uyy1 - 25*(U2+u2-U1-u1) )*ψ1x + (U1+u1)*q1x - ν*Δq1 - Ff2 = FourierFlows.jacobian(ψ2, q2 + η, gr) + @. (β - uyy2 - 25/4*(U1+u1-U2-u2) )*ψ2x + (U2+u2)*(q2x + ηx) + μ*Δψ2 - ν*Δq2 + Ff1 = FourierFlows.jacobian(ψ1, q1, gr) + @. (β - uyy1 - 25 * ((U2+u2) - (U1+u1)) ) * ψ1x + (U1+u1) * q1x - ν * Δq1 + Ff2 = FourierFlows.jacobian(ψ2, q2, gr) + @. (β - uyy2 - 25/4 * ((U1+u1) - (U2+u2)) ) * ψ2x + (U2+u2) * q2x - ν * Δq2 + Ff2 .+= FourierFlows.jacobian(ψ2, η, gr) + @. (U2+u2) * ηx + μ * Δψ2 T = eltype(gr) @@ -217,7 +273,7 @@ function test_mqg_nonlinearadvection(dt, stepper, dev::Device=CPU(); function calcFq!(Fqh, sol, t, cl, v, p, g) Fqh .= Ffh - nothing + return nothing end prob = MultiLayerQG.Problem(nlayers, dev; nx, ny, Lx, Ly, f₀, g, H, ρ, U, @@ -242,6 +298,101 @@ function test_mqg_nonlinearadvection(dt, stepper, dev::Device=CPU(); isapprox(vs.ψ, ψf, rtol=rtol_multilayerqg) end +""" + test_mqg_nonlinearadvection_3layers(dt, stepper, dev::Device=CPU() + +Same as `test_mqg_nonlinearadvection_2layers` but for 3 layers. +""" +function test_mqg_nonlinearadvection_3layers(dt, stepper, dev::Device=CPU()) + + A = device_array(dev) + + tf = 200*dt + nt = round(Int, tf/dt) + + nx, ny = 128, 126 + Lx, Ly = 2π, 5π + gr = TwoDGrid(dev; nx, Lx, ny, Ly) + + x, y = gridpoints(gr) + k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers + + nlayers = 3 # these choice of parameters give the + f₀, g = 1, 1 # desired PV-streamfunction relations + H = [0.25, 0.25, 0.5] # q1 = Δψ1 + 20ψ2 - 20ψ1, + ρ = [4.0, 5.0, 6.0] # q2 = Δψ2 + 20ψ1 - 44ψ2 + 24ψ3, + # q3 = Δψ3 + 12ψ2 - 12ψ3. + + β = 0.35 + + U1, U2, U3 = 0.02, 0.025, 0.01 + u1 = @. 0.05sech(gr.y / (Ly/15))^2; u1 = A(reshape(u1, (1, gr.ny))) + u2 = @. 0.02cos(3l₀ * gr.y); u2 = A(reshape(u2, (1, gr.ny))) + u3 = @. 0.01cos(3l₀ * gr.y); u3 = A(reshape(u3, (1, gr.ny))) + + uyy1 = real.(ifft( -gr.l.^2 .* fft(u1) )) + uyy2 = real.(ifft( -gr.l.^2 .* fft(u2) )) + uyy3 = real.(ifft( -gr.l.^2 .* fft(u3) )) + + U = zeros(ny, nlayers) + CUDA.@allowscalar U[:, 1] = u1 .+ U1 + CUDA.@allowscalar U[:, 2] = u2 .+ U2 + CUDA.@allowscalar U[:, 3] = u3 .+ U3 + + μ, ν, nν = 0.1, 0.05, 1 + + η0, σx, σy = 1.0, Lx/25, Ly/20 + η = @. η0 * exp( -(x + Lx/8)^2 / 2σx^2 - (y - Ly/8)^2 / 2σy^2) + ηx = @. - (x + Lx/8) / σx^2 * η + + ψ1, ψ2, ψ3, q1, q2, q3, ψ1x, ψ2x, ψ3x, q1x, q2x, q3x, Δq1, Δq2, Δq3, Δψ3 = constructtestfields_3layer(gr) + + Ff1 = FourierFlows.jacobian(ψ1, q1, gr) + @. (β - uyy1 - 20 * ((U2+u2) - (U1+u1))) * ψ1x + (U1+u1) * q1x - ν * Δq1 + Ff2 = FourierFlows.jacobian(ψ2, q2, gr) + @. (β - uyy2 - 20 * ((U1+u1) - (U2+u2)) - 24 * ((U3+u3) - (U2+u2))) * ψ2x + (U2+u2) * q2x - ν * Δq2 + Ff3 = FourierFlows.jacobian(ψ3, q3, gr) + @. (β - uyy3 - 12 * ((U2+u2) - (U3+u3)) ) * ψ3x + (U3+u3) * q3x - ν * Δq3 + Ff3 .+= FourierFlows.jacobian(ψ3, η, gr) + @. (U3+u3) * ηx + μ * Δψ3 + + T = eltype(gr) + + Ff = zeros(dev, T, (gr.nx, gr.ny, nlayers)) + @views Ff[:, :, 1] = Ff1 + @views Ff[:, :, 2] = Ff2 + @views Ff[:, :, 3] = Ff3 + + Ffh = zeros(dev, Complex{T}, (gr.nkr, gr.nl, nlayers)) + @views Ffh[:, :, 1] = rfft(Ff1) + @views Ffh[:, :, 2] = rfft(Ff2) + @views Ffh[:, :, 3] = rfft(Ff3) + + function calcFq!(Fqh, sol, t, cl, v, p, g) + Fqh .= Ffh + return nothing + end + + prob = MultiLayerQG.Problem(nlayers, dev; nx, ny, Lx, Ly, f₀, g, H, ρ, U, + eta=η, β, μ, ν, nν, calcFq=calcFq!, stepper, dt) + + sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid + + qf = zeros(dev, T, (gr.nx, gr.ny, nlayers)) + @views qf[:, :, 1] = q1 + @views qf[:, :, 2] = q2 + @views qf[:, :, 3] = q3 + + ψf = zeros(dev, T, (gr.nx, gr.ny, nlayers)) + @views ψf[:, :, 1] = ψ1 + @views ψf[:, :, 2] = ψ2 + @views ψf[:, :, 3] = ψ3 + + MultiLayerQG.set_q!(prob, qf) + + stepforward!(prob, nt) + MultiLayerQG.updatevars!(prob) + + return isapprox(vs.q, qf, rtol=rtol_multilayerqg) && + isapprox(vs.ψ, ψf, rtol=rtol_multilayerqg) +end + """ test_mqg_linearadvection(dt, stepper, dev; kwargs...) @@ -276,8 +427,10 @@ function test_mqg_linearadvection(dt, stepper, dev::Device=CPU(); β = 0.35 U1, U2 = 0.1, 0.05 - u1 = @. 0.5sech(gr.y/(Ly/15))^2; u1 = A(reshape(u1, (1, gr.ny))) - u2 = @. 0.02cos(3l₀*gr.y); u2 = A(reshape(u2, (1, gr.ny))) + + u1 = @. 0.5sech(gr.y / (Ly/15))^2; u1 = A(reshape(u1, (1, gr.ny))) + u2 = @. 0.02cos(3l₀ * gr.y); u2 = A(reshape(u2, (1, gr.ny))) + uyy1 = real.(ifft( -gr.l.^2 .* fft(u1) )) uyy2 = real.(ifft( -gr.l.^2 .* fft(u2) )) @@ -288,13 +441,14 @@ function test_mqg_linearadvection(dt, stepper, dev::Device=CPU(); μ, ν, nν = 0.1, 0.05, 1 η0, σx, σy = 1.0, Lx/25, Ly/20 - η = @. η0*exp( -(x + Lx/8)^2/(2σx^2) -(y-Ly/8)^2/(2σy^2) ) - ηx = @. -(x + Lx/8)/(σx^2) * η + η = @. η0 * exp( - (x + Lx/8)^2 / 2σx^2 - (y - Ly/8)^2 / 2σy^2) + ηx = @. -(x + Lx/8) / σx^2 * η ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2 = constructtestfields_2layer(gr) - Ff1 = (β .- uyy1 .- 25*(U2.+u2.-U1.-u1) ).*ψ1x + (U1.+u1).*q1x - ν*Δq1 - Ff2 = FourierFlows.jacobian(ψ2, η, gr) + (β .- uyy2 .- 25/4*(U1.+u1.-U2.-u2) ).*ψ2x + (U2.+u2).*(q2x + ηx) + μ*Δψ2 - ν*Δq2 + Ff1 = @. (β - uyy1 - 25 * ((U2+u2) - (U1+u1)) ) * ψ1x + (U1+u1) * q1x - ν * Δq1 + Ff2 = @. (β - uyy2 - 25/4 * ((U1+u1) - (U2+u2)) ) * ψ2x + (U2+u2) * q2x - ν * Δq2 + Ff2 .+= FourierFlows.jacobian(ψ2, η, gr) + @. (U2+u2) * ηx + μ * Δψ2 T = eltype(gr) @@ -308,7 +462,7 @@ function test_mqg_linearadvection(dt, stepper, dev::Device=CPU(); function calcFq!(Fqh, sol, t, cl, v, p, g) Fqh .= Ffh - nothing + return nothing end prob = MultiLayerQG.Problem(nlayers, dev; nx, ny, Lx, Ly, f₀, g, H, ρ, U, From 38e331767078797434ea6b105c798ca658db7bed Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Sun, 11 Jun 2023 07:40:10 +1000 Subject: [PATCH 16/16] simplify tests --- test/test_multilayerqg.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/test_multilayerqg.jl b/test/test_multilayerqg.jl index 3085aa6e..c698d075 100644 --- a/test/test_multilayerqg.jl +++ b/test/test_multilayerqg.jl @@ -223,7 +223,7 @@ function test_mqg_nonlinearadvection_2layers(dt, stepper, dev::Device=CPU()) tf = 200dt nt = round(Int, tf/dt) - nx, ny = 128, 126 + nx, ny = 64, 66 Lx, Ly = 2π, 2π gr = TwoDGrid(dev; nx, Lx, ny, Ly) @@ -257,8 +257,8 @@ function test_mqg_nonlinearadvection_2layers(dt, stepper, dev::Device=CPU()) ψ1, ψ2, q1, q2, ψ1x, ψ2x, q1x, q2x, Δψ2, Δq1, Δq2 = constructtestfields_2layer(gr) - Ff1 = FourierFlows.jacobian(ψ1, q1, gr) + @. (β - uyy1 - 25 * ((U2+u2) - (U1+u1)) ) * ψ1x + (U1+u1) * q1x - ν * Δq1 - Ff2 = FourierFlows.jacobian(ψ2, q2, gr) + @. (β - uyy2 - 25/4 * ((U1+u1) - (U2+u2)) ) * ψ2x + (U2+u2) * q2x - ν * Δq2 + Ff1 = FourierFlows.jacobian(ψ1, q1, gr) + @. (β - uyy1 - 25 * ((U2+u2) - (U1+u1))) * ψ1x + (U1+u1) * q1x - ν * Δq1 + Ff2 = FourierFlows.jacobian(ψ2, q2, gr) + @. (β - uyy2 - 25/4 * ((U1+u1) - (U2+u2))) * ψ2x + (U2+u2) * q2x - ν * Δq2 Ff2 .+= FourierFlows.jacobian(ψ2, η, gr) + @. (U2+u2) * ηx + μ * Δψ2 T = eltype(gr) @@ -310,7 +310,7 @@ function test_mqg_nonlinearadvection_3layers(dt, stepper, dev::Device=CPU()) tf = 200*dt nt = round(Int, tf/dt) - nx, ny = 128, 126 + nx, ny = 64, 66 Lx, Ly = 2π, 5π gr = TwoDGrid(dev; nx, Lx, ny, Ly)