Add Adapt reconstruction for RadonOp

ext/LinearOperatorRadonKAExt/LinearOperatorRadonKAExt.jl

@@ -1,6 +1,7 @@
 module LinearOperatorRadonKAExt
 
 using LinearOperatorCollection, RadonKA
+using Adapt
 
 include("RadonOp.jl")
 

ext/LinearOperatorRadonKAExt/RadonOp.jl

@@ -10,6 +10,9 @@ Generates a `RadonOp` which evaluates the Radon transform operator and its adjoi
 * `geometry`                - Radon geometry descriptor (default: parallel beam circle)
 * `μ`                       - optional attenuation map (for attenuated Radon transform)
 * `S`                       - storage type for internal vectors (default: `Vector{T}`)
+
+`Adapt.adapt(target, op)` rebuilds the operator for `target` with the same shape and
+geometry; array `angles` and `μ` are adapted to `target`, ranges and scalars are kept.
 """
 function LinearOperatorCollection.RadonOp(::Type{T}; shape::NTuple{N, Int}, angles,
    geometry = RadonParallelCircle(shape[1], -(shape[1]-1)÷2:(shape[1]-1)÷2), μ = nothing, S = Vector{T}) where {T, N}
@@ -32,6 +35,7 @@ mutable struct RadonOpImpl{T, vecT <: AbstractVector{T}, vecT2, G, A} <: RadonOp
   const angles :: vecT2
   const geometry :: G
   const μ :: A
+  const shape :: Tuple{Vararg{Int}}
 end
 
 LinearOperators.storage_type(::RadonOpImpl{T, vecT}) where {T, vecT} = vecT
@@ -46,7 +50,13 @@ function RadonOpImpl(T::Type; shape::NTuple{N, Int64}, angles, geometry, μ, S)
   (res, x) -> prod_radon!(res, x, shape, angles, geometry, μ),
   nothing, 
   (res, x) -> ctprod_radon!(res, x, (N_sinogram, N_angles, d), angles, geometry, μ),
-  0, 0, 0, S(undef, 0), S(undef, 0), angles, geometry, μ)
+  0, 0, 0, S(undef, 0), S(undef, 0), angles, geometry, μ, shape)
+end
+
+function Adapt.adapt_structure(to, op::RadonOpImpl)
+  target_storage = LinearOperatorCollection._adapt_target_storage(to, op.Mv, "RadonOp")
+  return RadonOp(eltype(op); shape=op.shape, angles=Adapt.adapt(to, op.angles),
+                 geometry=op.geometry, μ=Adapt.adapt(to, op.μ), S=typeof(target_storage))
 end
 
 prod_radon!(res::vecT, x::vecT, shape, angles::vecT2, geometry::G, μ::A) where {vecT, vecT2, G, A} = copyto!(res, radon(reshape(x, shape), angles; geometry, μ))

test/testAdapt.jl

@@ -334,8 +334,15 @@ storage_vector_type(array_type, T) = typeof(array_type(zeros(T, 0)))
             compound_adapted = Adapt.adapt(array_type, compound)
             @test all(typeof(child.Mv) <: array_type for child in compound_adapted.ops)
             @test all(child.plan !== source.plan for child in compound_adapted.ops)
+            first_child, second_child = compound_adapted.ops
+            @test first_child !== second_child
+            @test first_child.plan !== second_child.plan
+            @test first_child.Mv !== second_child.Mv
             z = randn(ComplexF32, 96)
-            @test Array(compound_adapted * array_type(z)) ≈ compound * z rtol=1f-4 atol=1f-5
+            # The repeated source child races when DiagOp runs both blocks in
+            # parallel, so the reference applies it to each block in turn.
+            reference = vcat(source * z[1:48], source * z[49:96])
+            @test Array(compound_adapted * array_type(z)) ≈ reference rtol=1f-4 atol=1f-5
 
             round_trip = Adapt.adapt(Array, adapted)
             @test round_trip !== adapted
@@ -497,6 +504,91 @@ storage_vector_type(array_type, T) = typeof(array_type(zeros(T, 0)))
         end
       end
 
+      @testset "RadonOp" begin
+        N = 16
+        angle_range = range(0f0, Float32(π), 12)
+        default_geometry = RadonParallelCircle(N, -(N - 1) ÷ 2:(N - 1) ÷ 2)
+        in_height = -(N ÷ 4):(N ÷ 4)
+        geometry = RadonParallelCircle(N, in_height,
+                                       Float32.(1 .+ 0.5 .* rand(length(in_height))))
+        cases = (
+          (name = "2-D non-default geometry", shape = (N, N),
+           angles = collect(angle_range), geometry = geometry, μ = nothing),
+          (name = "3-D with range angles", shape = (N, N, 3),
+           angles = angle_range, geometry = default_geometry, μ = nothing),
+          (name = "array μ", shape = (N, N), angles = collect(angle_range),
+           geometry = default_geometry, μ = 0.02f0 .* rand(Float32, N, N)),
+          (name = "scalar μ", shape = (N, N), angles = collect(angle_range),
+           geometry = default_geometry, μ = 0.01f0),
+        )
+
+        for case in cases
+          @testset "$(case.name)" begin
+            source = RadonOp(Float32; shape = case.shape, angles = case.angles,
+                             geometry = case.geometry, μ = case.μ, S = Vector{Float32})
+            source_angles = copy(case.angles)
+            source_μ = case.μ isa AbstractArray ? copy(case.μ) : case.μ
+            adapted = Adapt.adapt(array_type, source)
+
+            @test adapted !== source
+            @test adapted.Mv !== source.Mv
+            @test adapted.Mtu !== source.Mtu
+            check_adapted_storage(adapted, array_type)
+            @test typeof(adapted.Mtu) <: array_type
+            @test size(adapted) == size(source)
+            @test adapted.shape == source.shape == case.shape
+            @test adapted.geometry === source.geometry
+            if case.angles isa AbstractRange
+              @test adapted.angles === case.angles
+            else
+              @test typeof(adapted.angles) <: array_type
+            end
+            if case.μ isa AbstractArray
+              @test typeof(adapted.μ) <: array_type
+            else
+              @test adapted.μ === case.μ
+            end
+            @test source.angles == source_angles
+            @test source.μ == source_μ
+
+            x = rand(Float32, prod(case.shape))
+            y = rand(Float32, size(source, 1))
+            if case.geometry === geometry
+              # The narrowed detector changes the sinogram size, and the
+              # products must match RadonKA called with the same geometry.
+              @test size(source, 1) == length(in_height) * length(case.angles)
+              @test length(in_height) != length(default_geometry.in_height)
+              @test source * x ≈ vec(radon(reshape(x, case.shape), case.angles; geometry)) rtol = 1f-5
+              @test adjoint(source) * y ≈ vec(backproject(reshape(y, length(in_height), :),
+                                                          case.angles; geometry)) rtol = 1f-5
+            end
+
+            if array_type === JLArray
+              # RadonKA calls KernelAbstractions.synchronize, which the
+              # JLArrays backend does not implement.
+              @test_broken Array(adapted * array_type(x)) ≈ source * x rtol = 1f-5
+              @test_broken Array(adjoint(adapted) * array_type(y)) ≈ adjoint(source) * y rtol = 1f-5
+            else
+              @test Array(adapted * array_type(x)) ≈ source * x rtol = 1f-5
+              @test Array(adjoint(adapted) * array_type(y)) ≈ adjoint(source) * y rtol = 1f-5
+            end
+
+            # A different backend copies mutable configuration arrays, so
+            # later source mutation does not reach the adapted operator.
+            if array_type !== Array
+              if !(case.angles isa AbstractRange)
+                source.angles .+= 1
+                @test Array(adapted.angles) == source_angles
+              end
+              if case.μ isa AbstractArray
+                source.μ .+= 1
+                @test Array(adapted.μ) == source_μ
+              end
+            end
+          end
+        end
+      end
+
       @testset "GradientOp rebuilds closures and scratch" begin
         cases = (((7,), nothing, (1,)), ((7,), 1, 1),
                  ((4, 3), nothing, (1, 2)), ((4, 3), 2, 2), ((4, 3), 1:2, (1, 2)),
@@ -643,11 +735,13 @@ struct UnsupportedAdaptTarget end
   toeplitz = with(nfft_backend => NFFT.backend()) do
     normalOperator(nfft)
   end
+  radon_op = RadonOp(Float32; shape=(8, 8), angles=collect(range(0f0, Float32(π), 4)))
 
   for target in (UnsupportedAdaptTarget, AbstractVector)
     @test_throws ArgumentError Adapt.adapt(target, fft)
     @test_throws ArgumentError Adapt.adapt(target, nfft)
     @test_throws ArgumentError Adapt.adapt(target, toeplitz)
+    @test_throws ArgumentError Adapt.adapt(target, radon_op)
   end
 end