Merge pull request #255 from FourierFlows/ncc/on_grid-patch

Fix bug in on_grid() utility function

Author: navidcy

Project.toml

@@ -6,7 +6,7 @@ authors = ["Gregory L. Wagner <wagner.greg@gmail.com>", "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.6.11"
+version = "0.6.12"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

README.md

@@ -61,7 +61,7 @@ and [Navid C. Constantinou][] (@navidcy).
 
 The code is citable via [zenodo](https://zenodo.org). Please cite as:
 
-> Navid C. Constantinou & Gregory L. Wagner. (2021). FourierFlows/FourierFlows.jl: FourierFlows v0.6.11 (Version v0.6.11). Zenodo.  [http://doi.org/10.5281/zenodo.1161724](http://doi.org/10.5281/zenodo.1161724)
+> Navid C. Constantinou & Gregory L. Wagner. (2021). FourierFlows/FourierFlows.jl: FourierFlows v0.6.12 (Version v0.6.12). Zenodo.  [http://doi.org/10.5281/zenodo.1161724](http://doi.org/10.5281/zenodo.1161724)
 
 
 [Julia]: https://julialang.org/

src/utils.jl

@@ -248,16 +248,16 @@ end
 Returns an array, of the ArrayType of the device `grid` lives on, that contains the values of
 function `func` evaluated on the `grid`.
 """
-on_grid(func, grid::OneDGrid) = @. func(grid.x)
+on_grid(func, grid::OneDGrid) = CUDA.@allowscalar func.(grid.x)
 
 function on_grid(func, grid::TwoDGrid)
   x, y = gridpoints(grid)
-  return func.(x, y)
+  return CUDA.@allowscalar func.(x, y)
 end
 
 function on_grid(func, grid::ThreeDGrid)
   x, y, z = gridpoints(grid)
-  return func.(x, y, z)
+  return CUDA.@allowscalar func.(x, y, z)
 end
 
 """

test/test_utils.jl

@@ -145,17 +145,17 @@ function test_ongrid(dev::Device)
   nx, ny, nz = 6, 8, 10
   Lx, Ly, Lz = 2π, 2.0, 3.0
 
-  g₁ = OneDGrid(nx, Lx)
+  g₁ = OneDGrid(dev, nx, Lx)
   X₁ = ArrayType(dev)(g₁.x)
   f₁(x) = x^2
 
-  g₂ = TwoDGrid(nx, Lx, ny, Ly)
+  g₂ = TwoDGrid(dev, nx, Lx, ny, Ly)
   X₂, Y₂ = gridpoints(g₂)
   f₂(x, y) = x^2 - y^3
 
-  g₃ = ThreeDGrid(nx, Lx, ny, Ly, nz, Lz)
+  g₃ = ThreeDGrid(dev, nx, Lx, ny, Ly, nz, Lz)
   X₃, Y₃, Z₃ = gridpoints(g₃)
   f₃(x, y, z) = x^2 - y^3 + sin(z)
 
-  return CUDA.@allowscalar (FourierFlows.on_grid(f₁, g₁) == f₁.(X₁) && FourierFlows.on_grid(f₂, g₂) == f₂.(X₂, Y₂) && FourierFlows.on_grid(f₃, g₃) == f₃.(X₃, Y₃, Z₃))
+  return (CUDA.@allowscalar FourierFlows.on_grid(f₁, g₁) == f₁.(X₁) && CUDA.@allowscalar FourierFlows.on_grid(f₂, g₂) == f₂.(X₂, Y₂) && CUDA.@allowscalar FourierFlows.on_grid(f₃, g₃) == f₃.(X₃, Y₃, Z₃))
 end