Merge pull request #270 from FourierFlows/ncc/fix-bug-on_grid

Fix bug in `on_grid()`

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.15"
+version = "0.6.16"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

README.md

@@ -91,7 +91,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.15 (Version v0.6.15). 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.16 (Version v0.6.16). Zenodo.  [http://doi.org/10.5281/zenodo.1161724](http://doi.org/10.5281/zenodo.1161724)
 
 
 [Julia]: https://julialang.org/

src/utils.jl

@@ -248,17 +248,11 @@ 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) = CUDA.@allowscalar func.(grid.x)
+on_grid(func, grid::OneDGrid{T, A}) where {T, A} = CUDA.@allowscalar A([func(grid.x[i]) for i=1:grid.nx])
 
-function on_grid(func, grid::TwoDGrid)
-  x, y = gridpoints(grid)
-  return CUDA.@allowscalar func.(x, y)
-end
-
-function on_grid(func, grid::ThreeDGrid)
-  x, y, z = gridpoints(grid)
-  return CUDA.@allowscalar func.(x, y, z)
-end
+on_grid(func, grid::TwoDGrid{T, A}) where {T, A} = CUDA.@allowscalar A([func(grid.x[i], grid.y[j]) for i=1:grid.nx, j=1:grid.ny])
+ 
+on_grid(func, grid::ThreeDGrid{T, A}) where {T, A} = CUDA.@allowscalar A([func(grid.x[i], grid.y[j], grid.z[k]) for i=1:grid.nx, j=1:grid.ny, k=1:grid.nz])
 
 """
     ArrayType(::Device)

test/test_utils.jl

@@ -157,5 +157,5 @@ function test_ongrid(dev::Device)
   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₁) && CUDA.@allowscalar FourierFlows.on_grid(f₂, g₂) == f₂.(X₂, Y₂) && CUDA.@allowscalar FourierFlows.on_grid(f₃, g₃) == f₃.(X₃, Y₃, Z₃))
+  return (FourierFlows.on_grid(f₁, g₁) ≈ f₁.(X₁) && FourierFlows.on_grid(f₂, g₂) ≈ f₂.(X₂, Y₂) && FourierFlows.on_grid(f₃, g₃) ≈ f₃.(X₃, Y₃, Z₃))
 end