Merge pull request #305 from FourierFlows/ncc/bump-FourierFlows-0.10

Update to latest grid refactor

Project.toml

@@ -2,7 +2,7 @@ name = "GeophysicalFlows"
 uuid = "44ee3b1c-bc02-53fa-8355-8e347616e15e"
 license = "MIT"
 authors = ["Navid C. Constantinou <[email protected]>", "Gregory L. Wagner <[email protected]>", "and co-contributors"]
-version = "0.14.2"
+version = "0.15.0"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
@@ -21,7 +21,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 CUDA = "^1, ^2.4.2, 3.0.0 - 3.6.4, ^3.7.1"
 DocStringExtensions = "^0.8, 0.9"
 FFTW = "^1"
-FourierFlows = "^0.9.4"
+FourierFlows = "^0.10.0"
 JLD2 = "^0.1, ^0.2, ^0.3, ^0.4"
 Reexport = "^0.2, ^1"
 SpecialFunctions = "^0.10, ^1, 2"

docs/src/aliasing.md

@@ -11,7 +11,7 @@ in Fourier space before transforming to physical space to compute nonlinear term
     Users can construct a `grid` with different `aliased_fraction` via
 
     ```julia
-    julia> grid = OneDGrid(64, 2π, aliased_fraction=1/2)
+    julia> grid = OneDGrid(; nx=64, Lx=2π, aliased_fraction=1/2)
 
     julia> OneDimensionalGrid
              ├─────────── Device: CPU

examples/barotropicqgql_betaforced.jl

@@ -62,7 +62,7 @@ forcing_wavenumber = 14.0 * 2π/L  # the forcing wavenumber, `k_f`, for a spectr
 forcing_bandwidth  = 1.5  * 2π/L  # the width of the forcing spectrum, `δ_f`
 ε = 0.001                         # energy input rate by the forcing
 
-grid = TwoDGrid(dev, n, L)
+grid = TwoDGrid(dev; nx=n, Lx=L)
 
 K = @. sqrt(grid.Krsq)            # a 2D array with the total wavenumber
 
@@ -133,7 +133,7 @@ fig
 # ## Setting initial conditions
 
 # Our initial condition is simply fluid at rest.
-BarotropicQGQL.set_zeta!(prob, ArrayType(dev)(zeros(grid.nx, grid.ny)))
+BarotropicQGQL.set_zeta!(prob, device_array(dev)(zeros(grid.nx, grid.ny)))
 nothing # hide
 
 # ## Diagnostics

examples/multilayerqg_2layer.jl

@@ -76,11 +76,11 @@ nothing # hide
 # Our initial condition is some small-amplitude random noise. We smooth our initial
 # condidtion using the `timestepper`'s high-wavenumber `filter`.
 #
-# `ArrayType()` function returns the array type appropriate for the device, i.e., `Array` for
+# `device_array()` function returns the array type appropriate for the device, i.e., `Array` for
 # `dev = CPU()` and `CuArray` for `dev = GPU()`.
 
 seed!(1234) # reset of the random number generator for reproducibility
-q₀  = 1e-2 * ArrayType(dev)(randn((grid.nx, grid.ny, nlayers)))
+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
 

examples/singlelayerqg_betadecay.jl

@@ -64,15 +64,15 @@ nothing # hide
 
 # Our initial condition consist of a flow that has power only at wavenumbers with
 # ``6 < \frac{L}{2\pi} \sqrt{k_x^2 + k_y^2} < 10`` and initial energy ``E_0``.
-# `ArrayType()` function returns the array type appropriate for the device, i.e., `Array` for
+# `device_array()` function returns the array type appropriate for the device, i.e., `Array` for
 # `dev = CPU()` and `CuArray` for `dev = GPU()`.
 
 E₀ = 0.08 # energy of initial condition
 
 K = @. sqrt(grid.Krsq)                          # a 2D array with the total wavenumber
 
 Random.seed!(1234)
-q₀h = ArrayType(dev)(randn(Complex{eltype(grid)}, size(sol)))
+q₀h = device_array(dev)(randn(Complex{eltype(grid)}, size(sol)))
 @. q₀h = ifelse(K < 6  * 2π/L, 0, q₀h)
 @. q₀h = ifelse(K > 10 * 2π/L, 0, q₀h)
 @. q₀h[1, :] = 0    # remove any power from zonal wavenumber k=0

examples/singlelayerqg_betaforced.jl

@@ -63,7 +63,7 @@ forcing_wavenumber = 14.0 * 2π/L  # the forcing wavenumber, `k_f`, for a spectr
 forcing_bandwidth  = 1.5  * 2π/L  # the width of the forcing spectrum, `δ_f`
 ε = 0.001                         # energy input rate by the forcing
 
-grid = TwoDGrid(dev, n, L)
+grid = TwoDGrid(dev; nx=n, Lx=L)
 
 K = @. sqrt(grid.Krsq)            # a 2D array with the total wavenumber
 
@@ -132,7 +132,7 @@ fig
 # ## Setting initial conditions
 
 # Our initial condition is simply fluid at rest.
-SingleLayerQG.set_q!(prob, ArrayType(dev)(zeros(grid.nx, grid.ny)))
+SingleLayerQG.set_q!(prob, device_array(dev)(zeros(grid.nx, grid.ny)))
 
 
 # ## Diagnostics

examples/singlelayerqg_decaying_topography.jl

@@ -88,15 +88,15 @@ fig
 
 # Our initial condition consist of a flow that has power only at wavenumbers with
 # ``6 < \frac{L}{2\pi} \sqrt{k_x^2 + k_y^2} < 12`` and initial energy ``E_0``.
-# `ArrayType()` function returns the array type appropriate for the device, i.e., `Array` for
+# `device_array()` function returns the array type appropriate for the device, i.e., `Array` for
 # `dev = CPU()` and `CuArray` for `dev = GPU()`.
 
 E₀ = 0.04 # energy of initial condition
 
 K = @. sqrt(grid.Krsq)                             # a 2D array with the total wavenumber
 
 Random.seed!(1234)
-qih = ArrayType(dev)(randn(Complex{eltype(grid)}, size(sol)))
+qih = device_array(dev)(randn(Complex{eltype(grid)}, size(sol)))
 @. qih = ifelse(K < 6  * 2π/L, 0, qih)
 @. qih = ifelse(K > 12 * 2π/L, 0, qih)
 qih *= sqrt(E₀ / SingleLayerQG.energy(qih, vars, params, grid))  # normalize qi to have energy E₀

examples/twodnavierstokes_stochasticforcing.jl

@@ -55,7 +55,7 @@ forcing_wavenumber = 14.0 * 2π/L  # the forcing wavenumber, `k_f`, for a spectr
 forcing_bandwidth  = 1.5  * 2π/L  # the width of the forcing spectrum, `δ_f`
 ε = 0.1                           # energy input rate by the forcing
 
-grid = TwoDGrid(dev, n, L)
+grid = TwoDGrid(dev; nx=n, Lx=L)
 
 K = @. sqrt(grid.Krsq)             # a 2D array with the total wavenumber
 
@@ -125,7 +125,7 @@ fig
 # ## Setting initial conditions
 
 # Our initial condition is a fluid at rest.
-TwoDNavierStokes.set_ζ!(prob, ArrayType(dev)(zeros(grid.nx, grid.ny)))
+TwoDNavierStokes.set_ζ!(prob, device_array(dev)(zeros(grid.nx, grid.ny)))
 
 
 # ## Diagnostics

examples/twodnavierstokes_stochasticforcing_budgets.jl

@@ -57,7 +57,7 @@ forcing_wavenumber = 14.0 * 2π/L  # the forcing wavenumber, `k_f`, for a spectr
 forcing_bandwidth  = 1.5  * 2π/L  # the width of the forcing spectrum, `δ_f`
 ε = 0.1                           # energy input rate by the forcing
 
-grid = TwoDGrid(dev, n, L)
+grid = TwoDGrid(dev; nx=n, Lx=L)
 
 K = @. sqrt(grid.Krsq)             # a 2D array with the total wavenumber
 
@@ -128,7 +128,7 @@ fig
 # ## Setting initial conditions
 
 # Our initial condition is a fluid at rest.
-TwoDNavierStokes.set_ζ!(prob, ArrayType(dev)(zeros(grid.nx, grid.ny)))
+TwoDNavierStokes.set_ζ!(prob, device_array(dev)(zeros(grid.nx, grid.ny)))
 
 
 # ## Diagnostics

src/barotropicqgql.jl

@@ -90,21 +90,21 @@ function Problem(dev::Device=CPU();
                  T = Float64)
 
   # the grid
-  grid = TwoDGrid(dev, nx, Lx, ny, Ly; aliased_fraction=aliased_fraction, T=T)
+  grid = TwoDGrid(dev; nx, Lx, ny, Ly, aliased_fraction=aliased_fraction, T)
   x, y = gridpoints(grid)
 
   # topographic PV
   eta === nothing && ( eta = zeros(dev, T, (nx, ny)) )
 
-  params = !(typeof(eta)<:ArrayType(dev)) ?
+  params = !(typeof(eta)<:device_array(dev)) ?
            Params(grid, β, eta, μ, ν, nν, calcF) :
            Params(β, eta, rfft(eta), μ, ν, nν, calcF)
 
-  vars = calcF == nothingfunction ? DecayingVars(dev, grid) : stochastic ? StochasticForcedVars(dev, grid) : ForcedVars(dev, grid)
+  vars = calcF == nothingfunction ? DecayingVars(grid) : stochastic ? StochasticForcedVars(grid) : ForcedVars(grid)
 
   equation = BarotropicQGQL.Equation(params, grid)
 
-  FourierFlows.Problem(equation, stepper, dt, grid, vars, params, dev)
+  FourierFlows.Problem(equation, stepper, dt, grid, vars, params)
 end
 
 
@@ -232,44 +232,48 @@ const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothi
 const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray}
 
 """
-    DecayingVars(dev, grid)
+    DecayingVars(grid)
 
-Return the vars for unforced two-dimensional quasi-linear barotropic QG problem on device `dev` 
-and with `grid`.
+Return the vars for unforced two-dimensional quasi-linear barotropic QG problem on `grid`.
 """
-function DecayingVars(dev::Dev, grid::AbstractGrid) where Dev
+function DecayingVars(grid::AbstractGrid)
+  Dev = typeof(grid.device)
   T = eltype(grid)
+
   @devzeros Dev T (grid.nx, grid.ny) u v U uzeta vzeta zeta Zeta psi Psi
   @devzeros Dev Complex{T} (grid.nkr, grid.nl) N NZ uh vh Uh zetah Zetah psih Psih
 
   return Vars(u, v, U, uzeta, vzeta, zeta, Zeta, psi, Psi, N, NZ, uh, vh, Uh, zetah, Zetah, psih, Psih, nothing, nothing)
 end
 
 """
-    ForcedVars(dev, grid)
+    ForcedVars(grid)
 
-Return the `vars` for forced two-dimensional quasi-linear barotropic QG problem on device 
-`dev` and with `grid`.
+Return the `vars` for forced two-dimensional quasi-linear barotropic QG problem on `grid`.
 """
-function ForcedVars(dev::Dev, grid::AbstractGrid) where Dev
+function ForcedVars(grid::AbstractGrid)
+  Dev = typeof(grid.device)
   T = eltype(grid)
+
   @devzeros Dev T (grid.nx, grid.ny) u v U uzeta vzeta zeta Zeta psi Psi
   @devzeros Dev Complex{T} (grid.nkr, grid.nl) N NZ uh vh Uh zetah Zetah psih Psih Fh
 
   return Vars(u, v, U, uzeta, vzeta, zeta, Zeta, psi, Psi, N, NZ, uh, vh, Uh, zetah, Zetah, psih, Psih, Fh, nothing)
 end
 
 """
-    StochasticForcedVars(dev, grid)
+    StochasticForcedVars(grid)
 
 Return the `vars` for stochastically forced two-dimensional quasi-linear barotropic QG problem 
-on device `dev` and with `grid`.
+on `grid`.
 """
-function StochasticForcedVars(dev::Dev, grid::AbstractGrid) where Dev
+function StochasticForcedVars(grid::AbstractGrid)
+  Dev = typeof(grid.device)
   T = eltype(grid)
+
   @devzeros Dev T (grid.nx, grid.ny) u v U uzeta vzeta zeta Zeta psi Psi
   @devzeros Dev Complex{T} (grid.nkr, grid.nl) N NZ uh vh Uh zetah Zetah psih Psih Fh prevsol
-  
+
   return Vars(u, v, U, uzeta, vzeta, zeta, Zeta, psi, Psi, N, NZ, uh, vh, Uh, zetah, Zetah, psih, Psih, Fh, prevsol)
 end
 
@@ -335,11 +339,11 @@ N = - \\widehat{𝖩(ψ, ζ + η)}^{\\mathrm{QL}} + F̂ .
 """
 function calcN!(N, sol, t, clock, vars, params, grid)
   dealias!(sol, grid)
-  
+
   calcN_advection!(N, sol, t, clock, vars, params, grid)
-  
+
   addforcing!(N, sol, t, clock, vars, params, grid)
-  
+
   return nothing
 end
 

src/multilayerqg.jl

@@ -121,15 +121,15 @@ function Problem(nlayers::Int,                        # number of fluid layers
   # topographic PV
   eta === nothing && (eta = zeros(dev, T, (nx, ny)))
 
-  grid = TwoDGrid(dev, nx, Lx, ny, Ly; aliased_fraction=aliased_fraction, T=T)
+  grid = TwoDGrid(dev; nx, Lx, ny, Ly, aliased_fraction, T)
 
-  params = Params(nlayers, g, f₀, β, ρ, H, U, eta, μ, ν, nν, grid, calcFq=calcFq, dev=dev)   
+  params = Params(nlayers, g, f₀, β, ρ, H, U, eta, μ, ν, nν, grid, calcFq=calcFq)   
 
-  vars = calcFq == nothingfunction ? DecayingVars(dev, grid, params) : (stochastic ? StochasticForcedVars(dev, grid, params) : ForcedVars(dev, grid, params))
+  vars = calcFq == nothingfunction ? DecayingVars(grid, params) : (stochastic ? StochasticForcedVars(grid, params) : ForcedVars(grid, params))
 
-  equation = linear ? LinearEquation(dev, params, grid) : Equation(dev, params, grid)
+  equation = linear ? LinearEquation(params, grid) : Equation(params, grid)
 
-  FourierFlows.Problem(equation, stepper, dt, grid, vars, params, dev)
+  FourierFlows.Problem(equation, stepper, dt, grid, vars, params)
 end
 
 """
@@ -281,14 +281,15 @@ end
 
 function convert_U_to_U3D(dev, nlayers, grid, U::Number)
   T = eltype(grid)
-  A = ArrayType(dev)
+  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, μ, ν, nν, grid; calcFq=nothingfunction, effort=FFTW.MEASURE, dev::Device=CPU()) where TU
+function Params(nlayers, g, f₀, β, ρ, H, U, eta, μ, ν, nν, grid; calcFq=nothingfunction, effort=FFTW.MEASURE) where TU
+  dev = grid.device
   T = eltype(grid)
-  A = ArrayType(dev)
+  A = device_array(dev)
 
    ny, nx = grid.ny , grid.nx
   nkr, nl = grid.nkr, grid.nl
@@ -358,49 +359,51 @@ numberoflayers(::TwoLayerParams) = 2
 # ---------
 
 """
-    hyperviscosity(dev, params, grid)
+    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(dev, params, grid)
+function hyperviscosity(params, grid)
+  dev = grid.device
   T = eltype(grid)
-  L = ArrayType(dev){T}(undef, (grid.nkr, grid.nl, numberoflayers(params)))
+
+  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(dev, params, grid)
+    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(dev, params, grid)`.
+`hyperviscosity(params, grid)`.
 
 The nonlinear term is computed via function `calcNlinear!`.
 """
-function LinearEquation(dev, params, grid)
-  L = hyperviscosity(dev, params, grid)
+function LinearEquation(params, grid)
+  L = hyperviscosity(params, grid)
 
   return FourierFlows.Equation(L, calcNlinear!, grid)
 end
 
 """
-    Equation(dev, params, grid)
+    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(dev, params, grid)`.
+`hyperviscosity(params, grid)`.
 
 The nonlinear term is computed via function `calcN!`.
 """
-function Equation(dev, params, grid)
-  L = hyperviscosity(dev, params, grid)
+function Equation(params, grid)
+  L = hyperviscosity(params, grid)
 
   return FourierFlows.Equation(L, calcN!, grid)
 end
@@ -445,11 +448,12 @@ const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothi
 const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray}
 
 """
-    DecayingVars(dev, grid, params)
+    DecayingVars(grid, params)
 
 Return the variables for an unforced multi-layer QG problem with `grid` and `params`.
 """
-function DecayingVars(dev::Dev, grid, params) where Dev
+function DecayingVars(grid, params)
+  Dev = typeof(grid.device)
   T = eltype(grid)
   nlayers = numberoflayers(params)
 
@@ -460,11 +464,12 @@ function DecayingVars(dev::Dev, grid, params) where Dev
 end
 
 """
-    ForcedVars(dev, grid, params)
+    ForcedVars(grid, params)
 
 Return the variables for a forced multi-layer QG problem with `grid` and `params`.
 """
-function ForcedVars(dev::Dev, grid, params) where Dev
+function ForcedVars(grid, params)
+  Dev = typeof(grid.device)
   T = eltype(grid)
   nlayers = numberoflayers(params)
 
@@ -475,11 +480,12 @@ function ForcedVars(dev::Dev, grid, params) where Dev
 end
 
 """
-    StochasticForcedVars(dev, rid, params)
+    StochasticForcedVars(rid, params)
 
 Return the variables for a forced multi-layer QG problem with `grid` and `params`.
 """
-function StochasticForcedVars(dev::Dev, grid, params) where Dev
+function StochasticForcedVars(grid, params)
+  Dev = typeof(grid.device)
   T = eltype(grid)
   nlayers = numberoflayers(params)