diff --git a/ext/LinearOperatorRadonKAExt/LinearOperatorRadonKAExt.jl b/ext/LinearOperatorRadonKAExt/LinearOperatorRadonKAExt.jl index 29caf35..0662a76 100644 --- a/ext/LinearOperatorRadonKAExt/LinearOperatorRadonKAExt.jl +++ b/ext/LinearOperatorRadonKAExt/LinearOperatorRadonKAExt.jl @@ -1,6 +1,7 @@ module LinearOperatorRadonKAExt using LinearOperatorCollection, RadonKA +using Adapt include("RadonOp.jl") diff --git a/ext/LinearOperatorRadonKAExt/RadonOp.jl b/ext/LinearOperatorRadonKAExt/RadonOp.jl index 5bd69ff..399b74b 100644 --- a/ext/LinearOperatorRadonKAExt/RadonOp.jl +++ b/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, μ)) diff --git a/test/testAdapt.jl b/test/testAdapt.jl index 66d151f..03fefa2 100644 --- a/test/testAdapt.jl +++ b/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