Merge pull request #30 from FourierFlows/FixTypeInstabilities

Fixes type instabilities in various modules

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

@@ -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, beta, Array{T}(rho), 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)
 

src/twodturb.jl

@@ -53,7 +53,7 @@ function Problem(;
            T = Float64
 )
 
-  gr = TwoDGrid(nx, Lx, ny, Ly)
+  gr = TwoDGrid(nx, Lx, ny, Ly; T=T)
   pr = Params{T}(nu, nnu, mu, nmu, calcF)
   vs = calcF == nothingfunction ? Vars(gr) : (stochastic ? StochasticForcedVars(gr) : ForcedVars(gr))
   eq = Equation(pr, gr)

test/runtests.jl

@@ -43,6 +43,8 @@ end
   @test test_twodturb_stochasticforcingbudgets()
   @test test_twodturb_deterministicforcingbudgets()
   @test test_twodturb_energyenstrophy()
+  @test test_twodturb_problemtype(Float32)
+  @test TwoDTurb.nothingfunction() == nothing
 end
 
 @testset "BarotropicQG" begin
@@ -62,6 +64,7 @@ end
   @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
 
@@ -80,6 +83,7 @@ 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
 
@@ -91,8 +95,9 @@ end
   @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

test/test_multilayerqg.jl

@@ -313,7 +313,7 @@ end
 Tests the set_q!() and set_psi!() functions that initialize sol with a flow with
 given `q` or `psi` respectively.
 """
-function test_setqsetpsi(;dt=0.001, stepper="ForwardEuler", n=64, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)
+function test_mqg_setqsetpsi(;dt=0.001, stepper="ForwardEuler", n=64, L=2π, nlayers=2, μ=0.0, ν=0.0, nν=1)
   nx, ny = 32, 34
   L = 2π
   gr = TwoDGrid(nx, L, ny, L)
@@ -355,7 +355,7 @@ end
 Tests that `Params` constructor works with both mean flow `U` being a floats
 (i.e., constant `U` in each layer) or vectors (i.e., `U(y)` in each layer).
 """
-function test_paramsconstructor(;dt=0.001, stepper="ForwardEuler")
+function test_mqg_paramsconstructor(;dt=0.001, stepper="ForwardEuler")
   nx, ny = 32, 34
   L = 2π
   gr = TwoDGrid(nx, L, ny, L)
@@ -380,3 +380,9 @@ function test_paramsconstructor(;dt=0.001, stepper="ForwardEuler")
 
   isapprox(probUfloats.params.U, probUvectors.params.U, rtol=rtol_multilayerqg)
 end
+
+function test_mqg_problemtype(T=Float32)
+  prob = MultilayerQG.Problem(nlayers=2, T=T)
+
+  (typeof(prob.sol)==Array{Complex{T},3} && typeof(prob.grid.Lx)==T && typeof(prob.grid.x)==Array{T,2} && typeof(prob.vars.u)==Array{T,3})
+end

test/test_twodturb.jl

@@ -186,23 +186,34 @@ function test_twodturb_energyenstrophy()
   nx, Lx  = 128, 2π
   ny, Ly  = 128, 3π
   gr = TwoDGrid(nx, Lx, ny, Ly)
-  k0, l0 = gr.k[2], gr.l[2] # fundamental wavenumbers
   x, y = gridpoints(gr)
 
+  k0, l0 = gr.k[2], gr.l[2] # fundamental wavenumbers
     psi0 = @. sin(2k0*x)*cos(2l0*y) + 2sin(k0*x)*cos(3l0*y)
    zeta0 = @. -((2k0)^2+(2l0)^2)*sin(2k0*x)*cos(2l0*y) - (k0^2+(3l0)^2)*2sin(k0*x)*cos(3l0*y)
 
   energy_calc = 29/9
   enstrophy_calc = 2701/162
 
   prob = TwoDTurb.Problem(nx=nx, Lx=Lx, ny=ny, Ly=Ly, stepper="ForwardEuler")
+
+  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid;
 
   TwoDTurb.set_zeta!(prob, zeta0)
   TwoDTurb.updatevars!(prob)
 
   energyzeta0 = TwoDTurb.energy(prob)
   enstrophyzeta0 = TwoDTurb.enstrophy(prob)
+
+  params = TwoDTurb.Params(p.nu, p.nnu)
+
+  (isapprox(energyzeta0, energy_calc, rtol=rtol_twodturb) &&
+   isapprox(enstrophyzeta0, enstrophy_calc, rtol=rtol_twodturb) &&
+   TwoDTurb.addforcing!(prob.timestepper.N, sol, cl.t, cl, v, p, g)==nothing && p == params)
+end
+
+function test_twodturb_problemtype(T=Float32)
+  prob = TwoDTurb.Problem(T=T)
 
-  (isapprox(energyzeta0, 29.0/9, rtol=rtol_twodturb) &&
-   isapprox(enstrophyzeta0, 2701.0/162, rtol=rtol_twodturb))
+  (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