Merge pull request #330 from FourierFlows/ncc/grid-architecture

Grid refactor

.buildkite/pipeline.yml

@@ -14,10 +14,10 @@ steps:
     if: build.message !~ /\[skip tests\]/
     timeout_in_minutes: 60
 
-  - label: "🐫 Julia 1.7"
+  - label: "🐫 Julia 1.8"
     plugins:
       - JuliaCI/julia#v1:
-          version: 1.7
+          version: 1.8
       - JuliaCI/julia-test#v1: ~
       - JuliaCI/julia-coverage#v1:
           codecov: true

.github/workflows/main.yml

@@ -10,7 +10,7 @@ jobs:
     steps:
       - uses: julia-actions/setup-julia@latest
         with:
-          version: 1.7
+          version: 1.6
       - name: Pkg.add("CompatHelper")
         run: julia -e 'using Pkg; Pkg.add("CompatHelper")'
       - name: CompatHelper.main()

Project.toml

@@ -6,7 +6,7 @@ authors = ["Gregory L. Wagner <[email protected]>", "Navid C. Constantinou <
 description = "Tools for building fast, hackable, pseudospectral partial differential equation solvers on periodic domains."
 documentation = "https://fourierflows.github.io/FourierFlowsDocumentation/stable/"
 repository = "https://github.com/FourierFlows/FourierFlows.jl"
-version = "0.9.4"
+version = "0.10.0"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

README.md

@@ -102,7 +102,7 @@ and [Navid C. Constantinou][] (@navidcy).
 
 The code is citable via [zenodo](https://zenodo.org). Please cite as:
 
-> Gregory L. Wagner, Navid C. Constantinou, and contributors. (2022). FourierFlows/FourierFlows.jl: FourierFlows v0.9.4 (Version v0.9.4). Zenodo. [http://doi.org/10.5281/zenodo.1161724](http://doi.org/10.5281/zenodo.1161724)
+> Gregory L. Wagner, Navid C. Constantinou, and contributors. (2022). FourierFlows/FourierFlows.jl: FourierFlows v0.10.0 (Version v0.10.0). Zenodo.  [http://doi.org/10.5281/zenodo.1161724](http://doi.org/10.5281/zenodo.1161724)
 
 
 ## Contributing

appveyor.yml

@@ -1,7 +1,7 @@
 environment:
   matrix:
   - julia_version: 1.6
-  - julia_version: 1.7
+  - julia_version: 1.8
   - julia_version: nightly
 
 platform:

docs/make.jl

@@ -5,12 +5,7 @@ using
   Glob,
   FourierFlows
 
-# Gotta set this environment variable when using the GR run-time on CI machines.
-# This happens as examples will use Plots.jl to make plots and movies.
-# See: https://github.com/jheinen/GR.jl/issues/278
-ENV["GKSwstype"] = "100"
-
- #####
+#####
 ##### Generate examples
 #####
 
@@ -66,29 +61,29 @@ pages = [
         ],
 ]
 
-makedocs(
+makedocs(;
    sitename = "FourierFlows.jl",
-    authors = "Gregory L. Wagner and Navid C. Constantinou",
+    authors = "Gregory L. Wagner and Navid C. Constantinou and contributors",
     modules = [FourierFlows],
-     format = format,
-      pages = pages,
+     format,
+      pages,
     doctest = true,
      strict = :doctest,
       clean = true,
   checkdocs = :exports
 )
 
+@info "Cleaning up temporary .jld2 and .nc files created by doctests..."
+for file in vcat(glob("docs/*.jld2"))
+  rm(file)
+end
+
 withenv("GITHUB_REPOSITORY" => "FourierFlows/FourierFlowsDocumentation") do
   deploydocs(
             repo = "github.com/FourierFlows/FourierFlowsDocumentation.git",
         versions = ["stable" => "v^", "v#.#.#", "dev" => "dev"],
+       forcepush = true,
     push_preview = false,
        devbranch = "main"
   )
 end
-
-
-@info "Cleaning up temporary .jld2 and .nc files created by doctests..."
-for file in vcat(glob("docs/*.jld2"))
-  rm(file)
-end

docs/src/aliasing.md

@@ -25,7 +25,7 @@ on a grid ``x \in [-π, π)`` with ``n_x = 10`` grid points.
 using FourierFlows, CairoMakie
 
 nx, Lx = 10, 2π
-grid = OneDGrid(nx, Lx)
+grid = OneDGrid(; nx, Lx)
 
 f1(x) = cos(4x)
 f2(x) = cos(6x)
@@ -94,7 +94,7 @@ FourierFlows.dealias!
 If we construct a grid with `aliased_fraction=0`, e.g.,
 
 ```@example 1
-grid_nodealias = OneDGrid(nx, Lx; aliased_fraction=0)
+grid_nodealias = OneDGrid(; nx, Lx, aliased_fraction=0)
 ```
 
 then `dealias!(fh, grid_nodealias)` will have _no effect_ whatsoever on `fh`.

docs/src/diagnostics.md

@@ -8,7 +8,7 @@ DocTestSetup = quote
     using FourierFlows
     using LinearAlgebra: mul!, ldiv!
     nx, Lx = 32, 2.0
-    grid = OneDGrid(nx, Lx)
+    grid = OneDGrid(; nx, Lx)
     struct Params <: AbstractParams
     α :: Float64
     end
@@ -43,7 +43,7 @@ set_theme!(Theme(linewidth = 3, fontsize = 20))
 using LinearAlgebra: mul!
 
 nx, Lx = 32, 2.0
-grid = OneDGrid(nx, Lx)
+grid = OneDGrid(; nx, Lx)
 
 struct Params <: AbstractParams
   α :: Float64

docs/src/gpu.md

@@ -16,9 +16,8 @@ overload all methods to work with `GPU()` device as their argument (instead of t
 It's easy to construct a grid that lives on the GPU. Calling:
 
 ```julia
-dev = GPU()
-n, L = 16, 2.0
-grid = OneDGrid(dev, n, L)
+nx, Lx = 16, 2.0
+grid = OneDGrid(GPU(); nx, Lx)
 
 OneDimensionalGrid
   ├─────────── Device: GPU
@@ -30,22 +29,21 @@ OneDimensionalGrid
   └─ aliased fraction: 0.3333333333333333
 ```
 
-gives out a grid whose arrays are `CuArrays`. (Calling `OneDGrid(n, L)` defaults to CPU, i.e., 
-`OneDGrid(CPU(), n, L)`.)
+returns a grid whose arrays are `CuArrays`. (Calling `OneDGrid(; nx, Lx)` defaults to using
+`CPU()`, i.e., `OneDGrid(CPU(); nx, Lx)`.)
 
 When we construct the `Params`, `Vars`, and `Equation` for our problem we need to make sure
 that we create arrays on the appropriate device, i.e., `Arrays` for `CPU` or `CuArrays` for
-the `GPU`. Function `ArrayType` is useful in constructing appropriately chosen arrays.
+the `GPU`. Function `device_array` is useful in constructing appropriately chosen arrays.
 
 ```@docs
-ArrayType
+device_array
 ```
 
-The `FourierFlows.Problem` constructor then takes an optional positional argument 
-`dev::Device`. If not provided anything, the default values for `dev=CPU()`.
+The `FourierFlows.Problem` constructor infers the device from the `grid` that is provided.
 
 ```julia
-problem = Problem(equation, stepper, dt, grid, vars, params, GPU())
+problem = Problem(equation, stepper, dt, grid, vars, params)
 ```
 
 The `FourierFlows.Diffusion` module is written in a way such that switching from CPU to GPU 

docs/src/grids.md

@@ -19,7 +19,7 @@ using FourierFlows
 
 nx, Lx = 64, 2π
 
-grid = OneDGrid(nx, Lx)
+grid = OneDGrid(; nx, Lx)
 
 # output
 
@@ -48,7 +48,7 @@ using LinearAlgebra: mul!, ldiv!
 using CairoMakie
 set_theme!(Theme(linewidth = 3, fontsize = 20))
 nx, Lx = 64, 2π
-grid = OneDGrid(nx, Lx)
+grid = OneDGrid(; nx, Lx)
 ```
 
 ```@example 1
@@ -161,7 +161,7 @@ nothing # hide
 ```
 
 The grid contains the wavenumbers (both for real-value functions `grid.kr` and 
-for complex-valued functions `grid.k`). We populate array `∂ₓuh` is with ``i k \hat{u}``:
+for complex-valued functions `grid.k`). We populate array `∂ₓuh` is with ``\mathrm{i} k \hat{u}``:
 
 ```@example 1
 @. ∂ₓuh = im * grid.kr * uh

docs/src/output.md

@@ -9,7 +9,7 @@ DocTestSetup = quote
     using FourierFlows
     using LinearAlgebra: mul!, ldiv!
     nx, Lx = 32, 2.0
-    grid = OneDGrid(nx, Lx)
+    grid = OneDGrid(; nx, Lx)
     struct Params <: AbstractParams
     α :: Float64
     end
@@ -47,7 +47,7 @@ set_theme!(Theme(linewidth = 3, fontsize = 20))
 using LinearAlgebra: mul!, ldiv!
 
 nx, Lx = 32, 2.0
-grid = OneDGrid(nx, Lx)
+grid = OneDGrid(; nx, Lx)
 
 struct Params <: AbstractParams
   α :: Float64

docs/src/problem.md

@@ -27,7 +27,7 @@ using FourierFlows
 
 nx, Lx = 32, 2.0
 
-grid = OneDGrid(nx, Lx)
+grid = OneDGrid(; nx, Lx)
 ```
 
 Our problem has a parameter ``\alpha``. Thus, we create a `Params` as:

examples/OneDShallowWaterGeostrophicAdjustment.jl

@@ -77,10 +77,12 @@ nothing #hide
 # A constructor populates empty arrays based on the dimension of the `grid`
 # and then creates `Vars` struct.
 """
-    Vars(dev, grid)
+    Vars(grid)
 Constructs Vars for 1D shallow water based on the dimensions of arrays of the `grid`.
 """
-function Vars(::Dev, grid) where Dev
+function Vars(grid)
+  Dev = typeof(grid.device)
+
   T = eltype(grid)
   @devzeros Dev T grid.nx u v η
   @devzeros Dev Complex{T} grid.nkr uh vh ηh
@@ -125,19 +127,21 @@ nothing #hide
 # Next we construct the `Equation` struct:
 
 """
-    Equation(dev, params, grid)
+    Equation(params, grid)
 Construct the equation: the linear part, in this case the hyperviscous dissipation,
 and the nonlinear part, which is computed by `calcN!` function.
 """
-function Equation(dev, params, grid)
+function Equation(params, grid)
   T = eltype(grid)
+  dev = grid.device
+
   L = zeros(dev, T, (grid.nkr, 3))
   D = @. - params.ν * grid.kr^(2*params.nν)
-  
+
   L[:, 1] .= D # for u equation
   L[:, 2] .= D # for v equation
   L[:, 3] .= D # for η equation
-  
+
   return FourierFlows.Equation(L, calcN!, grid)
 end
 nothing #hide
@@ -223,16 +227,17 @@ nothing # hide
 
 
 # ## Construct the `struct`s and you are ready to go!
+
 # Create a `grid` and also `params`, `vars`, and the `equation` structs. Then 
 # give them all as input to the `FourierFlows.Problem()` constructor to get a
 # problem struct, `prob`, that contains all of the above.
 
-    grid = OneDGrid(dev, nx, Lx)
+    grid = OneDGrid(dev; nx, Lx)
   params = Params(ν, nν, g, H, f)
-    vars = Vars(dev, grid)
-equation = Equation(dev, params, grid)
+    vars = Vars(grid)
+equation = Equation(params, grid)
 
-    prob = FourierFlows.Problem(equation, stepper, dt, grid, vars, params, dev)
+    prob = FourierFlows.Problem(equation, stepper, dt, grid, vars, params)
 nothing #hide
 
 # ## Setting initial conditions

src/FourierFlows.jl

@@ -11,7 +11,7 @@ export
   Device,
   CPU,
   GPU,
-  ArrayType,
+  device_array,
 
   cxtype,
   fltype,
@@ -21,7 +21,6 @@ export
   AbstractParams,
 
   AbstractGrid,
-  ZeroDGrid,
   OneDGrid,
   TwoDGrid,
   ThreeDGrid,
@@ -43,7 +42,6 @@ export
   @devzeros,
   @createarrays,
   @superzeros,
-  devzeros,
   superzeros,
   supersize,
 
@@ -77,7 +75,7 @@ using Base: fieldnames
 using FFTW: fftfreq, rfftfreq
 
 "Abstract supertype for grids."
-abstract type AbstractGrid{T, A, Alias} end
+abstract type AbstractGrid{T, A, Talias, D} end
 
 "Abstract supertype for timesteppers."
 abstract type AbstractTimeStepper{T} end
@@ -91,7 +89,7 @@ abstract type AbstractVars end
 "Abstract supertype for diagnostics."
 abstract type AbstractDiagnostic end
 
-"Abstract supertype for device."
+"Abstract supertype for supported devices."
 abstract type Device end
 
 "CPU device."
@@ -111,9 +109,6 @@ include("timesteppers.jl")
 # Physics
 include("diffusion.jl")
 
-# CUDA functionality
-include("CuFourierFlows.jl")
-
 
 function __init__()
   threads = Threads.nthreads()

src/diagnostics.jl

@@ -117,10 +117,10 @@ getindex(d::Diagnostic, idx...) = getindex(d.data, idx...)
 
 show(io::IO, d::Diagnostic{T, N}) where {T, N} =
      print(io, "Diagnostic\n",
-               "  ├─── calc: ", d.calc, '\n',
-               "  ├─── prob: ", summary(d.prob), '\n', 
-               "  ├─── data: ", summary(d.data), '\n', 
-               "  ├────── t: ", summary(d.t), '\n', 
-               "  ├── steps: ", summary(d.steps), '\n', 
-               "  ├─── freq: ", d.freq, '\n', 
+               "  ├─── calc: ", d.calc, "\n",
+               "  ├─── prob: ", summary(d.prob), "\n",
+               "  ├─── data: ", summary(d.data), "\n",
+               "  ├────── t: ", summary(d.t), "\n",
+               "  ├── steps: ", summary(d.steps), "\n",
+               "  ├─── freq: ", d.freq, "\n",
                "  └────── i: ", d.i)