Merge pull request #228 from ChevronETC/wask/sparse-coordinates-for-complex

Coordinates for complex SparseFunction

Author: jkwashbourne-oss

Project.toml

@@ -1,7 +1,7 @@
 name = "Devito"
 uuid = "06ed14b1-0e40-4084-abdf-764a285f8c42"
 authors = ["Sam Kaplan <Sam.Kaplan@chevron.com>"]
-version = "1.4.0"
+version = "1.4.1"
 
 [deps]
 PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"

src/Devito.jl

@@ -1086,10 +1086,12 @@ data(x::SubFunction{T,N,M}) where {T,N,M} = data_allocated(x)
     coordinates(x::SparseDiscreteFunction)
 
 Returns a Devito function associated with the coordinates of a sparse time function.
-Note that contrary to typical Julia convention, coordinate order is from slow-to-fast (Python ordering).
+Note 1: contrary to typical Julia convention, coordinate order is from slow-to-fast (Python ordering).
 Thus, for a 3D grid, the sparse time function coordinates would be ordered x,y,z.
+Note 2: we need to handle complex data types, because coordinats are purely real
 """
 coordinates(x::SparseDiscreteFunction{T,N,M}) where {T,N,M} = SubFunction{T,2,M}(x.o.coordinates)
+coordinates(x::SparseDiscreteFunction{Complex{T},N,M}) where {T,N,M} = SubFunction{T,2,M}(x.o.coordinates)
 
 """
     coordinates_data(x::SparseDiscreteFunction)

test/serialtests.jl

@@ -203,7 +203,7 @@ end
     @test size_with_halo(sf) == (npoint,)
 end
 
-@testset "Sparse function coordinates, n=$n" for n in ( (10,11), (10,11,12) )
+@testset "Sparse function coordinates, Real, n=$n" for n in ( (10,11), (10,11,12) )
     grid = Grid(shape=n, dtype=Float32)
     sf = SparseFunction(name="sf", npoint=10, grid=grid)
     @test typeof(coordinates(sf)) <: SubFunction{Float32,2}
@@ -216,6 +216,19 @@ end
     @test _sf_coords ≈ x
 end
 
+@testset "Sparse function coordinates, Complex, n=$n" for n in ( (10,11), (10,11,12) )
+    grid = Grid(shape=n, dtype=Float32)
+    sf = SparseFunction(name="sf", npoint=10, grid=grid, dtype=Complex{Float32})
+    @test typeof(coordinates(sf)) <: SubFunction{Float32,2}
+    sf_coords = coordinates_data(sf)
+    @test isa(sf_coords, Devito.DevitoArray)
+    @test size(sf_coords) == (length(n),10)
+    x = rand(length(n),10)
+    sf_coords .= x
+    _sf_coords = coordinates_data(sf)
+    @test _sf_coords ≈ x
+end
+
 @testset "SparseFunction from PyObject, T=$T, n=$n, npoint=$npoint" for T in (Float32, Float64), n in ((3,4),(3,4,5)), npoint in (1,5,10)
     g = Grid(shape=n, dtype=T)
     sf = SparseFunction(name="sf", grid=g, npoint=npoint)
@@ -520,11 +533,11 @@ end
     apply(op)
     for j in 1:div(size,factr)+1
         @test data(f)[j] == data(g)[(j-1)*factr+1]
-    end
-    if ENV["DEVITO_BRANCH"] in ("main", "devitopro")
-         @test data(f)[end] == data(g)[end]
-    else
-         @test_broken data(f)[end] == data(g)[end]
+        if get(ENV,"DEVITO_BRANCH","main") in ("main", "devitopro")
+            @test data(f)[end] == data(g)[end]
+        else
+            @test_broken data(f)[end] == data(g)[end]
+        end
     end
 end