Gpu support for FilledExtrapolation (#541)

Author: drewrobson

src/extrapolation/filled.jl

@@ -1,4 +1,4 @@
-mutable struct FilledExtrapolation{T,N,ITP<:AbstractInterpolation,IT,FT} <: AbstractExtrapolation{T,N,ITP,IT}
+struct FilledExtrapolation{T,N,ITP<:AbstractInterpolation,IT,FT} <: AbstractExtrapolation{T,N,ITP,IT}
     itp::ITP
     fillvalue::FT
 end

src/gpu_support.jl

@@ -41,6 +41,12 @@ function adapt_structure(to, itp::Extrapolation{T,N}) where {T,N}
     Extrapolation{eltype(itp′),N,typeof(itp′),itptype(itp),typeof(et)}(itp′, et)
 end
 
+function adapt_structure(to, itp::FilledExtrapolation{T,N}) where {T,N}
+    fillvalue = itp.fillvalue
+    itp′ = adapt(to, itp.itp)
+    FilledExtrapolation{eltype(itp′),N,typeof(itp′),itptype(itp),typeof(fillvalue)}(itp′, fillvalue)
+end
+
 import Base.Broadcast: broadcasted, BroadcastStyle
 using Base.Broadcast: broadcastable, combine_styles, AbstractArrayStyle
 function broadcasted(itp::AbstractInterpolation, args...)
@@ -58,6 +64,7 @@ Some array wrappers, like `OffsetArray`, should be skipped.
 """
 root_storage_type(::Type{T}) where {T<:AbstractInterpolation} = Array{eltype(T),ndims(T)} # fallback to `Array` by default.
 root_storage_type(::Type{T}) where {T<:Extrapolation} = root_storage_type(fieldtype(T, 1))
+root_storage_type(::Type{T}) where {T<:FilledExtrapolation} = root_storage_type(fieldtype(T, 1))
 root_storage_type(::Type{T}) where {T<:ScaledInterpolation} = root_storage_type(fieldtype(T, 1))
 root_storage_type(::Type{T}) where {T<:BSplineInterpolation} = root_storage_type(fieldtype(T, 1))
 root_storage_type(::Type{T}) where {T<:LanczosInterpolation} = root_storage_type(fieldtype(T, 1))

test/gpu_support.jl

@@ -31,6 +31,15 @@ JLArrays.allowscalar(false)
     @test gradient.(Ref(esitp), idx) ==
         collect(gradient.(Ref(jlesitp), idx)) ==
         collect(gradient.(Ref(jlesitp), jlidx))
+
+    esitp = extrapolate(sitp, 0.0)
+    jlesitp = jl(esitp)
+    idx = -1.0:0.84:41.0
+    jlidx = jl(collect(idx))
+    @test esitp.(idx) == collect(jlesitp.(idx)) == collect(jlesitp.(jlidx))
+    @test gradient.(Ref(esitp), idx) ==
+        collect(gradient.(Ref(jlesitp), idx)) ==
+        collect(gradient.(Ref(jlesitp), jlidx))
 end
 
 @testset "2d GPU Interpolation" begin
@@ -69,6 +78,16 @@ end
     @test gradient.(Ref(esitp), idx, idx') ==
         collect(gradient.(Ref(jlesitp), idx, idx')) ==
         collect(gradient.(Ref(jlesitp), jlidx, jlidx'))
+
+    esitp = extrapolate(sitp, 0.0)
+    jlesitp = jl(esitp)
+    idx = -1.0:0.84:41.0
+    jlidx = jl(collect(idx))
+    @test esitp.(idx, idx') == collect(jlesitp.(idx, idx')) == collect(jlesitp.(jlidx, jlidx'))
+    # gradient for `extrapolation` is currently broken under CUDA
+    @test gradient.(Ref(esitp), idx, idx') ==
+        collect(gradient.(Ref(jlesitp), idx, idx')) ==
+        collect(gradient.(Ref(jlesitp), jlidx, jlidx'))
 end
 
 @testset "Lanczos on gpu" begin
@@ -99,6 +118,7 @@ end
     @test eltype(adapt(Array{Real}, itp)) === Float64
     @test eltype(adapt(Array{Float32}, scale(itp, A_x))) === Float32
     @test eltype(adapt(Array{Float32}, extrapolate(scale(itp, A_x), Flat()))) === Float32
+    @test eltype(adapt(Array{Float32}, extrapolate(scale(itp, A_x), 0.0))) === Float32
     itp = interpolate((-1:0.2:1, -1:0.2:1), randn(11, 11), Gridded(Linear()))
     @test eltype(adapt(Array{Float32}, itp)) === Float32
     itp = interpolate((1.0:0.0, 1.:0.), randn(0, 0), Gridded(Linear()))