Merge branch 'master' into GPUfyTwoDTurb

README.md

@@ -10,8 +10,8 @@
     <a href="https://fourierflows.github.io/GeophysicalFlows.jl/stable/">
         <img src="https://img.shields.io/badge/docs-stable-blue.svg">
     </a>
-    <a href="https://fourierflows.github.io/GeophysicalFlows.jl/latest/">
-        <img src="https://img.shields.io/badge/docs-latest-blue.svg">
+    <a href="https://fourierflows.github.io/GeophysicalFlows.jl/dev/">
+        <img src="https://img.shields.io/badge/docs-dev-blue.svg">
     </a>
     <a href='https://coveralls.io/github/FourierFlows/GeophysicalFlows.jl?branch=master'><img src='https://coveralls.io/repos/github/FourierFlows/GeophysicalFlows.jl/badge.svg?branch=master' alt='Coverage Status' />
     </a>

docs/Project.toml

@@ -1,6 +1,3 @@
 [deps]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 GeophysicalFlows = "44ee3b1c-bc02-53fa-8355-8e347616e15e"
-
-[compat]
-Documenter = "~0.22"

src/barotropicqg.jl

@@ -61,7 +61,7 @@ function Problem(;
              T = Float64)
 
   # the grid
-  gr  = TwoDGrid(nx, Lx, ny, Ly)
+  gr  = TwoDGrid(nx, Lx, ny, Ly; T=T)
   x, y = gridpoints(gr)
 
   # topographic PV
@@ -70,10 +70,10 @@ function Problem(;
     etah = rfft(eta)
   end
 
-  if typeof(eta)!=Array{Float64,2} #this is true if eta was passes in Problem as a function
+  if typeof(eta)!=Array{T, 2} #this is true if eta was passes in Problem as a function
     pr = Params(gr, f0, beta, eta, mu, nu, nnu, calcFU, calcFq)
   else
-    pr = Params(f0, beta, eta, rfft(eta), mu, nu, nnu, calcFU, calcFq)
+    pr = Params{T}(f0, beta, eta, rfft(eta), mu, nu, nnu, calcFU, calcFq)
   end
 
   if calcFq == nothingfunction && calcFU == nothingfunction
@@ -120,11 +120,11 @@ end
 
 Constructor for Params that accepts a generating function for the topographic PV.
 """
-function Params(g, f0, beta, eta::Function, mu, nu, nnu, calcFU, calcFq)
+function Params(g::AbstractGrid{T}, f0, beta, eta::Function, mu, nu, nnu, calcFU, calcFq)  where T
   x, y = gridpoints(g)
   etagrid = eta(x, y)
   etah = rfft(etagrid)
-  Params(f0, beta, etagrid, etah, mu, nu, nnu, calcFU, calcFq)
+  Params{T}(f0, beta, etagrid, etah, mu, nu, nnu, calcFU, calcFq)
 end
 
 
@@ -138,9 +138,9 @@ end
 Returns the equation for two-dimensional barotropic QG problem with params p and grid g.
 """
 function Equation(p::Params, g::AbstractGrid{T}) where T
-  LC = @. -p.mu - p.nu*g.Krsq^p.nnu + im*p.beta*g.kr*g.invKrsq
-  LC[1, 1] = 0
-  FourierFlows.Equation(LC, calcN!, g)
+  L = @. -p.mu - p.nu*g.Krsq^p.nnu + im*p.beta*g.kr*g.invKrsq
+  L[1, 1] = 0
+  FourierFlows.Equation(L, calcN!, g)
 end
 
 
@@ -167,7 +167,7 @@ eval(structvarsexpr(:StochasticForcedVars, stochforcedvarspecs; parent=:Barotrop
 
 Returns the vars for unforced two-dimensional barotropic QG problem with grid g.
 """
-function Vars(g; T=typeof(g.Lx))
+function Vars(g::AbstractGrid{T}) where T
   U = Array{T,0}(undef, ); U[] = 0
   @createarrays T (g.nx, g.ny) q u v psi zeta
   @createarrays Complex{T} (g.nkr, g.nl) qh uh vh psih zetah
@@ -179,8 +179,8 @@ end
 
 Returns the vars for forced two-dimensional barotropic QG problem with grid g.
 """
-function ForcedVars(g; T=typeof(g.Lx))
-  v = Vars(g; T=T)
+function ForcedVars(g::AbstractGrid{T}) where T
+  v = Vars(g)
   Fqh = zeros(Complex{T}, (g.nkr, g.nl))
   ForcedVars(getfield.(Ref(v), fieldnames(typeof(v)))..., Fqh)
 end
@@ -190,8 +190,8 @@ end
 
 Returns the vars for stochastically forced two-dimensional barotropic QG problem with grid g.
 """
-function StochasticForcedVars(g; T=typeof(g.Lx))
-  v = ForcedVars(g; T=T)
+function StochasticForcedVars(g::AbstractGrid{T}) where T
+  v = ForcedVars(g)
   prevsol = zeros(Complex{T}, (g.nkr, g.nl))
   StochasticForcedVars(getfield.(Ref(v), fieldnames(typeof(v)))..., prevsol)
 end

src/barotropicqgql.jl

@@ -58,7 +58,7 @@ function Problem(;
              T = Float64)
 
   # the grid
-  gr  = BarotropicQGQL.TwoDGrid(nx, Lx, ny, Ly)
+  gr  = TwoDGrid(nx, Lx, ny, Ly; T=T)
   x, y = gridpoints(gr)
 
   # topographic PV
@@ -67,10 +67,10 @@ function Problem(;
     etah = rfft(eta)
   end
 
-  if typeof(eta)!=Array{Float64,2} #this is true if eta was passes in Problem as a function
+  if typeof(eta)!=Array{T, 2} #this is true if eta was passes in Problem as a function
     pr = Params(gr, f0, beta, eta, mu, nu, nnu, calcF)
   else
-    pr = Params(f0, beta, eta, rfft(eta), mu, nu, nnu, calcF)
+    pr = Params{T}(f0, beta, eta, rfft(eta), mu, nu, nnu, calcF)
   end
   vs = calcF == nothingfunction ? Vars(gr) : (stochastic ? StochasticForcedVars(gr) : ForcedVars(gr))
   eq = BarotropicQGQL.Equation(pr, gr)
@@ -93,8 +93,8 @@ Returns the params for an unforced two-dimensional barotropic QG problem.
 struct Params{T} <: AbstractParams
   f0::T                      # Constant planetary vorticity
   beta::T                    # Planetary vorticity y-gradient
-  eta::Array{T,2}            # Topographic PV
-  etah::Array{Complex{T},2}  # FFT of Topographic PV
+  eta::Array{T, 2}           # Topographic PV
+  etah::Array{Complex{T}, 2} # FFT of Topographic PV
   mu::T                      # Linear drag
   nu::T                      # Viscosity coefficient
   nnu::Int                   # Hyperviscous order (nnu=1 is plain old viscosity)
@@ -106,11 +106,11 @@ end
 
 Constructor for Params that accepts a generating function for the topographic PV.
 """
-function Params(g::TwoDGrid, f0, beta, eta::Function, mu, nu, nnu, calcF)
+function Params(g::AbstractGrid{T}, f0, beta, eta::Function, mu, nu, nnu, calcF) where T
   x, y = gridpoints(g)
   etagrid = eta(x, y)
   etah = rfft(etagrid)
-  Params(f0, beta, etagrid, etah, mu, nu, nnu, calcF)
+  Params{T}(f0, beta, etagrid, etah, mu, nu, nnu, calcF)
 end
 
 
@@ -124,9 +124,9 @@ end
 Returns the equation for two-dimensional barotropic QG problem with params p and grid g.
 """
 function Equation(p::Params, g::AbstractGrid{T}) where T
-  LC = @. -p.mu - p.nu*g.Krsq^p.nnu + im*p.beta*g.kr*g.invKrsq
-  LC[1, 1] = 0
-  FourierFlows.Equation(LC, calcN!, g)
+  L = @. -p.mu - p.nu*g.Krsq^p.nnu + im*p.beta*g.kr*g.invKrsq
+  L[1, 1] = 0
+  FourierFlows.Equation(L, calcN!, g)
 end
 
 
@@ -160,7 +160,7 @@ mutable struct Vars{T} <: BarotropicQGQLVars
   Psih::Array{Complex{T},2}
 end
 
-function Vars(g; T=typeof(g.Lx))
+function Vars(g::AbstractGrid{T}) where T
   @createarrays T (g.nx, g.ny) u v U uzeta vzeta zeta Zeta psi Psi
   @createarrays Complex{T} (g.nkr, g.nl) N NZ uh vh Uh zetah Zetah psih Psih
   Vars(u, v, U, uzeta, vzeta, zeta, Zeta, psi, Psi, N, NZ, uh, vh, Uh, zetah, Zetah, psih, Psih)
@@ -193,8 +193,8 @@ mutable struct ForcedVars{T} <: BarotropicQGQLForcedVars
   Fh::Array{Complex{T},2}
 end
 
-function ForcedVars(g; T=typeof(g.Lx))
-  v = Vars(g; T=T)
+function ForcedVars(g::AbstractGrid{T}) where T
+  v = Vars(g)
   Fh = zeros(Complex{T}, (g.nkr, g.nl))
   ForcedVars(getfield.(Ref(v), fieldnames(typeof(v)))..., Fh)
 end
@@ -223,8 +223,8 @@ mutable struct StochasticForcedVars{T} <: BarotropicQGQLForcedVars
   prevsol::Array{Complex{T},2}
 end
 
-function StochasticForcedVars(g; T=typeof(g.Lx))
-  v = ForcedVars(g; T=T)
+function StochasticForcedVars(g::AbstractGrid{T}) where T
+  v = ForcedVars(g)
   prevsol = zeros(Complex{T}, (g.nkr, g.nl))
   StochasticForcedVars(getfield.(Ref(v), fieldnames(typeof(v)))..., prevsol)
 end

src/multilayerqg.jl

@@ -42,11 +42,11 @@ function Problem(;
     # Physical parameters
      nlayers = 2,                       # number of fluid layers
           f0 = 1.0,                     # Coriolis parameter
-        beta = 0.0,                     # y-gradient of Coriolis parameter
+           β = 0.0,                     # y-gradient of Coriolis parameter
            g = 1.0,                     # gravitational constant
            U = zeros(ny, nlayers),      # imposed zonal flow U(y) in each layer
            H = [0.2, 0.8],              # rest fluid height of each layer
-         rho = [4.0, 5.0],              # density of each layer
+           ρ = [4.0, 5.0],              # density of each layer
          eta = zeros(nx, ny),           # topographic PV
     # Bottom Drag and/or (hyper)-viscosity
            μ = 0.0,
@@ -59,7 +59,7 @@ function Problem(;
            T = Float64)
 
    grid = TwoDGrid(nx, Lx, ny, Ly; T=T)
-   params = Params(nlayers, T(g), T(f0), T(beta), Array{T}(rho), Array{T}(H), Array{T}(U), Array{T}(eta), T(μ), T(ν), nν, grid, calcFq=calcFq)
+   params = Params(nlayers, g, f0, β, Array{T}(ρ), Array{T}(H), Array{T}(U), Array{T}(eta), μ, ν, nν, grid, calcFq=calcFq)
    vars = calcFq == nothingfunction ? Vars(grid, params) : ForcedVars(grid, params)
    eqn = linear ? LinearEquation(params, grid) : Equation(params, grid)
 
@@ -72,9 +72,9 @@ struct Params{T} <: AbstractParams
   # prescribed params
    nlayers :: Int            # Number of fluid layers
          g :: T              # Gravitational constant
-         f0:: T              # Constant planetary vorticity
-      beta :: T              # Planetary vorticity y-gradient
-       rho :: Array{T,3}     # Array with density of each fluid layer
+        f0 :: T              # Constant planetary vorticity
+         β :: T              # Planetary vorticity y-gradient
+         ρ :: Array{T,3}     # Array with density of each fluid layer
          H :: Array{T,3}     # Array with rest height of each fluid layer
          U :: Array{T,3}     # Array with imposed constant zonal flow U(y) in each fluid layer
        eta :: Array{T,2}     # Array containing topographic PV
@@ -85,8 +85,8 @@ struct Params{T} <: AbstractParams
 
   # derived params
         g′ :: Array{T,1}     # Array with the reduced gravity constants for each fluid interface
-        Qx :: Array{T,3}     # Array containing x-gradient of PV due to U and eta in each fluid layer
-        Qy :: Array{T,3}     # Array containing y-gradient of PV due to beta, U, and eta in each fluid layer
+        Qx :: Array{T,3}     # Array containing x-gradient of PV due to eta in each fluid layer
+        Qy :: Array{T,3}     # Array containing y-gradient of PV due to β, U, and eta in each fluid layer
          S :: Array{T,4}     # Array containing coeffients for getting PV from  streamfunction
       invS :: Array{T,4}     # Array containing coeffients for inverting PV to streamfunction
   rfftplan :: FFTW.rFFTWPlan{T,-1,false,3}  # rfft plan for FFTs
@@ -96,8 +96,8 @@ struct SingleLayerParams{T} <: BarotropicParams
   # prescribed params
          g :: T              # Gravitational constant
         f0 :: T              # Constant planetary vorticity
-      beta :: T              # Planetary vorticity y-gradient
-       rho :: T              # Fluid layer density
+         β :: T              # Planetary vorticity y-gradient
+         ρ :: T              # Fluid layer density
          H :: T              # Fluid layer rest height
          U :: Array{T,1}     # Imposed constant zonal flow U(y)
        eta :: Array{T,2}     # Array containing topographic PV
@@ -107,27 +107,27 @@ struct SingleLayerParams{T} <: BarotropicParams
    calcFq! :: Function       # Function that calculates the forcing on QGPV q
 
   # derived params
-        Qx :: Array{T,2}     # Array containing zonal PV gradient due to beta, U, and eta in each fluid layer
-        Qy :: Array{T,2}     # Array containing meridional PV gradient due to beta, U, and eta in each fluid layer
+        Qx :: Array{T,2}     # Array containing x-gradient of PV due to eta in each fluid layer
+        Qy :: Array{T,2}     # Array containing meridional PV gradient due to β, U, and eta in each fluid layer
   rfftplan :: FFTW.rFFTWPlan{T,-1,false,2}  # rfft plan for FFTs
 end
 
-function Params(nlayers, g, f0, beta, rho, H, U::Array{T,2}, eta, μ, ν, nν, grid::AbstractGrid{T}; calcFq=nothingfunction, effort=FFTW.MEASURE) where T
+function Params(nlayers, g, f0, β, ρ, H, U::Array{T,2}, eta, μ, ν, nν, grid::AbstractGrid{T}; calcFq=nothingfunction, effort=FFTW.MEASURE) where T
 
    ny, nx = grid.ny , grid.nx
   nkr, nl = grid.nkr, grid.nl
    kr, l  = grid.kr , grid.l
 
   U = reshape(U, (1, ny, nlayers))
 
-  # g′ = g*(rho[2:nlayers]-rho[1:nlayers-1]) ./ rho[1:nlayers-1] # definition match PYQG
-  g′ = g*(rho[2:nlayers]-rho[1:nlayers-1]) ./ rho[2:nlayers] # correct definition
+  # g′ = g*(ρ[2:nlayers]-ρ[1:nlayers-1]) ./ ρ[1:nlayers-1] # definition match PYQG
+  g′ = g*(ρ[2:nlayers]-ρ[1:nlayers-1]) ./ ρ[2:nlayers] # correct definition
 
   Fm = @. f0^2 / ( g′*H[2:nlayers  ] )
   Fp = @. f0^2 / ( g′*H[1:nlayers-1] )
 
-  rho = reshape(rho, (1,  1, nlayers))
-    H = reshape(H  , (1,  1, nlayers))
+   ρ = reshape(ρ, (1,  1, nlayers))
+   H = reshape(H, (1,  1, nlayers))
 
   Uyy = repeat(irfft( -l.^2 .* rfft(U, [1, 2]),  1, [1, 2]), outer=(1, 1, 1))
   Uyy = repeat(Uyy, outer=(nx, 1, 1))
@@ -140,11 +140,11 @@ function Params(nlayers, g, f0, beta, rho, H, U::Array{T,2}, eta, μ, ν, nν, g
   @views @. Qx[:, :, nlayers] += etax
 
   Qy = zeros(nx, ny, nlayers)
-  Qy[:, :, 1] = @. beta - Uyy[:, :, 1] - Fp[1]*( U[:, :, 2] - U[:, :, 1] )
+  Qy[:, :, 1] = @. β - Uyy[:, :, 1] - Fp[1]*( U[:, :, 2] - U[:, :, 1] )
   for j = 2:nlayers-1
-    Qy[:, :, j] = @. beta - Uyy[:, :, j] - Fp[j]*( U[:, :, j+1] - U[:, :, j] ) - Fm[j-1]*( U[:, :, j-1] - U[:, :, j] )
+    Qy[:, :, j] = @. β - Uyy[:, :, j] - Fp[j]*( U[:, :, j+1] - U[:, :, j] ) - Fm[j-1]*( U[:, :, j-1] - U[:, :, j] )
   end
-  @views Qy[:, :, nlayers] = @. beta - Uyy[:, :, nlayers] - Fm[nlayers-1]*( U[:, :, nlayers-1] - U[:, :, nlayers] )
+  @views Qy[:, :, nlayers] = @. β - Uyy[:, :, nlayers] - Fm[nlayers-1]*( U[:, :, nlayers-1] - U[:, :, nlayers] )
   @views @. Qy[:, :, nlayers] += etay
 
   S = Array{T}(undef, (nkr, nl, nlayers, nlayers))
@@ -156,16 +156,16 @@ function Params(nlayers, g, f0, beta, rho, H, U::Array{T,2}, eta, μ, ν, nν, g
   rfftplanlayered = plan_rfft(Array{T,3}(undef, grid.nx, grid.ny, nlayers), [1, 2]; flags=effort)
 
   if nlayers == 1
-    SingleLayerParams{T}(g, f0, beta, rho, H, U, eta, μ, ν, nν, calcFq, Qx, Qy, grid.rfftplan)
+    SingleLayerParams{T}(g, f0, β, ρ, H, U, eta, μ, ν, nν, calcFq, Qx, Qy, grid.rfftplan)
   else
-    Params{T}(nlayers, g, f0, beta, rho, H, U, eta, μ, ν, nν, calcFq, g′, Qx, Qy, S, invS, rfftplanlayered)
+    Params{T}(nlayers, g, f0, β, ρ, H, U, eta, μ, ν, nν, calcFq, g′, Qx, Qy, S, invS, rfftplanlayered)
   end
 end
 
-function Params(nlayers, g, f0, beta, rho, H, U::Array{T,1}, eta, μ, ν, nν, grid::AbstractGrid{T}; calcFq=nothingfunction, effort=FFTW.MEASURE) where T
+function Params(nlayers, g, f0, β, ρ, H, U::Array{T,1}, eta, μ, ν, nν, grid::AbstractGrid{T}; calcFq=nothingfunction, effort=FFTW.MEASURE) where T
   U = reshape(U, (1, nlayers))
   U = repeat(U, outer=(grid.ny, 1))
-  Params(nlayers, g, f0, beta, rho, H, U, eta, μ, ν, nν, grid; calcFq=calcFq, effort=effort)
+  Params(nlayers, g, f0, β, ρ, H, U, eta, μ, ν, nν, grid; calcFq=calcFq, effort=effort)
 end
 
 
@@ -191,7 +191,7 @@ function Equation(p, gr::AbstractGrid{T}) where T
 end
 
 function Equation(p::SingleLayerParams, gr::AbstractGrid{T}) where T
-  L = @. -p.μ - p.ν*gr.Krsq^p.nν + im*p.beta*gr.kr*gr.invKrsq
+  L = @. -p.μ - p.ν*gr.Krsq^p.nν + im*p.β*gr.kr*gr.invKrsq
   L[1, 1] = 0
   FourierFlows.Equation(L, calcN!, gr)
 end

src/twodturb.jl

@@ -67,7 +67,7 @@ end
 Returns the params for two-dimensional turbulence.
 """
 struct Params{T} <: AbstractParams
-      ν  :: T         # Vorticity viscosity
+       ν :: T         # Vorticity viscosity
       nν :: Int       # Vorticity hyperviscous order
        μ :: T         # Bottom drag or hypoviscosity
       nμ :: Int       # Order of hypodrag

test/runtests.jl

@@ -50,6 +50,7 @@ for dev in devices
     @test test_twodturb_stochasticforcingbudgets(dev)
     @test test_twodturb_deterministicforcingbudgets(dev)
     @test test_twodturb_energyenstrophy(dev)
+    @test test_twodturb_problemtype(Float32)
     @test TwoDTurb.nothingfunction() == nothing
   end
 
@@ -75,6 +76,7 @@ println("rest of tests only on CPU")
   @test test_bqg_formstress(0.01, "ForwardEuler")
   @test test_bqg_energyenstrophy()
   @test test_bqg_meanenergyenstrophy()
+  @test test_bqg_problemtype(Float32)
   @test BarotropicQG.nothingfunction() == nothing
 end
 
@@ -93,19 +95,22 @@ end
   @test test_bqgql_stochasticforcingbudgets()
   @test test_bqgql_advection(0.0005, "ForwardEuler")
   @test test_bqgql_energyenstrophy()
+  @test test_bqgql_problemtype(Float32)
   @test BarotropicQGQL.nothingfunction() == nothing
 end
 
 @testset "MultilayerQG" begin
   include("test_multilayerqg.jl")
 
-  @test test_pvtofromstreamfunction()
+  @test test_pvtofromstreamfunction_2layer()
+  @test test_pvtofromstreamfunction_3layer()
   @test test_mqg_nonlinearadvection(0.001, "ForwardEuler")
   @test test_mqg_linearadvection(0.001, "ForwardEuler")
   @test test_mqg_energies()
   @test test_mqg_fluxes()
-  @test test_setqsetpsi()
-  @test test_paramsconstructor()
+  @test test_mqg_setqsetpsi()
+  @test test_mqg_paramsconstructor()
+  @test test_mqg_problemtype(Float32)
   @test MultilayerQG.nothingfunction() == nothing
 end
 

test/test_barotropicqg.jl

@@ -334,3 +334,9 @@ function test_bqg_meanenergyenstrophy()
     isapprox(enstrophyzeta0, enstrophy_calc, rtol=1e-13)
   )
 end
+
+function test_bqg_problemtype(T=Float32)
+  prob = BarotropicQG.Problem(T=T)
+
+  (typeof(prob.sol)==Array{Complex{T},2} && typeof(prob.grid.Lx)==T && typeof(prob.grid.x)==Array{T,2} && typeof(prob.vars.u)==Array{T,2})
+end

test/test_barotropicqgql.jl

@@ -257,3 +257,9 @@ function test_bqgql_energyenstrophy()
 
   isapprox(energyzeta0, energy_calc, rtol=1e-13) && isapprox(enstrophyzeta0, enstrophy_calc, rtol=1e-13) && BarotropicQGQL.addforcing!(prob.timestepper.N, sol, cl.t, cl, v, p, g)==nothing
 end
+
+function test_bqgql_problemtype(T=Float32)
+  prob = BarotropicQGQL.Problem(T=T)
+
+  (typeof(prob.sol)==Array{Complex{T},2} && typeof(prob.grid.Lx)==T && typeof(prob.grid.x)==Array{T,2} && typeof(prob.vars.u)==Array{T,2})
+end