From dc831bb9c5be3942a00c52bd5bc24af3240d4c81 Mon Sep 17 00:00:00 2001 From: shukebeta Date: Thu, 29 Oct 2026 20:19:31 +1300 Subject: [PATCH 1/2] feat: add Adapt reconstruction for RadonOp Store the image shape on RadonOpImpl and rebuild the operator through the RadonOp constructor for the target storage. Array angles and attenuation maps are adapted to the target; ranges, scalar attenuation, and the geometry are kept. Also give the FFTOp DiagOp compound test distinct children, since DiagOp evaluates them in parallel and a shared FFTOp races on its plan buffer under multiple threads. Co-Authored-By: Claude Opus 5.5 --- .../LinearOperatorRadonKAExt.jl | 1 + ext/LinearOperatorRadonKAExt/RadonOp.jl | 12 ++- test/testAdapt.jl | 91 ++++++++++++++++++- 3 files changed, 102 insertions(+), 2 deletions(-) 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 8528559..1ffef83 100644 --- a/test/testAdapt.jl +++ b/test/testAdapt.jl @@ -255,7 +255,9 @@ custom_point_transform(x::Tuple) = map(custom_point_transform, x) @test adjoint_result ≈ adjoint(source) * y rtol=1f-4 atol=1f-5 if shift && unitary - compound = DiagOp((source, source)) + # Distinct children: DiagOp runs them in parallel, and one FFTOp + # shares its plan and buffer across calls. + compound = DiagOp((source, copy(source))) 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) @@ -421,6 +423,91 @@ custom_point_transform(x::Tuple) = map(custom_point_transform, x) @test adapted.fftplan !== source_normal.fftplan 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 end end @@ -460,11 +547,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 From 226697be602acd74f1eb05ecfb51ff80a2204df3 Mon Sep 17 00:00:00 2001 From: shukebeta Date: Thu, 29 Oct 2026 20:25:17 +1300 Subject: [PATCH 2/2] test: keep repeated-child FFTOp compound adapt case Adapt the DiagOp with the same FFTOp twice, assert the adapted children have independent plans and workspaces, and compare against applying the source to each block in turn, which avoids the parallel race on the shared source child. Co-Authored-By: Claude Opus 5.5 --- test/testAdapt.jl | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/test/testAdapt.jl b/test/testAdapt.jl index 1ffef83..5a8bf1d 100644 --- a/test/testAdapt.jl +++ b/test/testAdapt.jl @@ -255,14 +255,19 @@ custom_point_transform(x::Tuple) = map(custom_point_transform, x) @test adjoint_result ≈ adjoint(source) * y rtol=1f-4 atol=1f-5 if shift && unitary - # Distinct children: DiagOp runs them in parallel, and one FFTOp - # shares its plan and buffer across calls. - compound = DiagOp((source, copy(source))) + compound = DiagOp((source, source)) 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