generalizes testdx/dy/dz and testgridpoints

Author: navidcy

test/runtests.jl

@@ -49,7 +49,7 @@ for dev in devices
     # Test 1D grid
     g₁ = OneDGrid(dev, nx, Lx)
     @test testnx(g₁, nx)
-    @test testdx(g₁)
+    @test testdx(dev, g₁)
     @test testdk(g₁)
     @test testx(g₁)
     @test testk(g₁)
@@ -61,16 +61,16 @@ for dev in devices
     g₂ = TwoDGrid(dev, nx, Lx, ny, Ly)
     @test testnx(g₂, nx)
     @test testny(g₂, ny)
-    @test testdx(g₂)
-    @test testdy(g₂)
+    @test testdx(dev, g₂)
+    @test testdy(dev, g₂)
     @test testdk(g₂)
     @test testdl(g₂)
     @test testx(g₂)
     @test testy(g₂)
     @test testk(g₂)
     @test testkr(g₂)
     @test testl(g₂)
-    @test testgridpoints(g₂)
+    @test testgridpoints(dev, g₂)
     @test testdealias(g₂)
     @test testmakefilter(dev, g₂)
 
@@ -79,9 +79,9 @@ for dev in devices
     @test testnx(g₃, nx)
     @test testny(g₃, ny)
     @test testnz(g₃, nz)
-    @test testdx(g₃)
-    @test testdy(g₃)
-    @test testdz(g₃)
+    @test testdx(dev, g₃)
+    @test testdy(dev, g₃)
+    @test testdz(dev, g₃)
     @test testdk(g₃)
     @test testdl(g₃)
     @test testdm(g₃)
@@ -92,7 +92,7 @@ for dev in devices
     @test testkr(g₃)
     @test testl(g₃)
     @test testm(g₃)
-    @test testgridpoints(g₃)
+    @test testgridpoints(dev, g₃)
     @test testdealias(g₃)
     @test testmakefilter(dev, g₃)
 

test/test_grid.jl

@@ -3,39 +3,21 @@ testny(g, ny) = isapprox(g.ny, ny)
 testnz(g, nz) = isapprox(g.nz, nz)
 
 # Physical grid tests
-function testdx(g::Union{OneDGrid{T,<:Array}, TwoDGrid{T,<:Array}, ThreeDGrid{T,<:Array}}) where T
+function testdx(dev, g::Union{OneDGrid{T}, TwoDGrid{T}, ThreeDGrid{T}}) where T
   dxgrid = @. g.x[2:end] - g.x[1:end-1]
-  dxones = g.dx*ones(T, size(dxgrid))
+  dxones = ArrayType(dev)(g.dx*ones(T, size(dxgrid)))
   isapprox(dxgrid, dxones)
 end
 
-function testdy(g::Union{TwoDGrid{T,<:Array}, ThreeDGrid{T,<:Array}}) where T
+function testdy(dev, g::Union{TwoDGrid{T}, ThreeDGrid{T}}) where T
   dygrid = @. g.y[2:end] - g.y[1:end-1]
-  dyones = g.dy*ones(T, size(dygrid))
+  dyones = ArrayType(dev)(g.dy*ones(T, size(dygrid)))
   isapprox(dygrid, dyones)
 end
 
-function testdz(g::ThreeDGrid{T,<:Array}) where T
+function testdz(dev, g::ThreeDGrid{T}) where T
   dzgrid = @. g.z[2:end] - g.z[1:end-1]
-  dzones = g.dz*ones(T, size(dzgrid))
-  isapprox(dzgrid, dzones)
-end
-
-function testdx(g::Union{OneDGrid{T,<:CuArray}, TwoDGrid{T,<:CuArray}, ThreeDGrid{T,<:CuArray}}) where T
-  dxgrid = @. g.x[2:end] - g.x[1:end-1]
-  dxones = CuArray(g.dx*ones(T, size(dxgrid)))
-  isapprox(dxgrid, dxones)
-end
-
-function testdy(g::Union{TwoDGrid{T,<:CuArray}, ThreeDGrid{T,<:CuArray}}) where T
-  dygrid = @. g.y[2:end] - g.y[1:end-1]
-  dyones = CuArray(g.dy*ones(T, size(dygrid)))
-  isapprox(dygrid, dyones)
-end
-
-function testdz(g::ThreeDGrid{T,<:CuArray}) where T
-  dzgrid = @. g.z[2:end] - g.z[1:end-1]
-  dzones = CuArray(g.dz*ones(T, size(dzgrid)))
+  dzones = ArrayType(dev)(g.dz*ones(T, size(dzgrid)))
   isapprox(dzgrid, dzones)
 end
 
@@ -65,23 +47,23 @@ testl(g::Union{TwoDGrid, ThreeDGrid}) = testwavenumberalignment(g.l, g.ny)
 testm(g::ThreeDGrid) = testwavenumberalignment(g.m, g.nz)
 testkr(g) = isapprox(cat(g.k[1:g.nkr-1], abs(g.k[g.nkr]), dims=1), g.kr)
 
-function testgridpoints(g::TwoDGrid{T}) where T
+function testgridpoints(dev::Device, g::TwoDGrid{T}) where T
   X, Y = gridpoints(g)
   dXgrid = @. X[2:end, :] - X[1:end-1, :]
   dYgrid = @. Y[:, 2:end] - Y[:, 1:end-1]
-  dXones = g.dx*ones(T, size(dXgrid))
-  dYones = g.dy*ones(T, size(dYgrid))
+  dXones = ArrayType(dev)(g.dx*ones(T, size(dXgrid)))
+  dYones = ArrayType(dev)(g.dy*ones(T, size(dYgrid)))
   isapprox(dXgrid, dXones) && isapprox(dYgrid, dYones)
 end
 
-function testgridpoints(g::ThreeDGrid{T}) where T
+function testgridpoints(dev::Device, g::ThreeDGrid{T}) where T
   X, Y, Z = gridpoints(g)
   dXgrid = @. X[2:end, :, :] - X[1:end-1, :, :]
   dYgrid = @. Y[:, 2:end, :] - Y[:, 1:end-1, :]
   dZgrid = @. Z[:, :, 2:end] - Z[:, :, 1:end-1]
-  dXones = g.dx*ones(T, size(dXgrid))
-  dYones = g.dy*ones(T, size(dYgrid))
-  dZones = g.dz*ones(T, size(dZgrid))
+  dXones = ArrayType(dev)(g.dx*ones(T, size(dXgrid)))
+  dYones = ArrayType(dev)(g.dy*ones(T, size(dYgrid)))
+  dZones = ArrayType(dev)(g.dz*ones(T, size(dZgrid)))
   isapprox(dXgrid, dXones) && isapprox(dYgrid, dYones) && isapprox(dZgrid, dZones)
 end