Skip to content
New issue

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

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

Already on GitHub? Sign in to your account

Fixes type instabilities in various modules #30

Merged
merged 4 commits into from
Sep 10, 2019
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 13 additions & 13 deletions src/barotropicqg.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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


Expand All @@ -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


Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
30 changes: 15 additions & 15 deletions src/barotropicqgql.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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


Expand All @@ -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


Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/multilayerqg.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Copy link
Member

Choose a reason for hiding this comment

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

missing space

Copy link
Member Author

Choose a reason for hiding this comment

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

damn! you caught me!

vars = calcFq == nothingfunction ? Vars(grid, params) : ForcedVars(grid, params)
eqn = linear ? LinearEquation(params, grid) : Equation(params, grid)

Expand Down
2 changes: 1 addition & 1 deletion src/twodturb.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
9 changes: 7 additions & 2 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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

Expand All @@ -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

Expand Down
6 changes: 6 additions & 0 deletions test/test_barotropicqg.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
6 changes: 6 additions & 0 deletions test/test_barotropicqgql.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
10 changes: 8 additions & 2 deletions test/test_multilayerqg.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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
17 changes: 14 additions & 3 deletions test/test_twodturb.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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