From 4a8f7103154837e21586dce2a558be3ff202b3b7 Mon Sep 17 00:00:00 2001 From: shukebeta Date: Mon, 2 Nov 2026 23:31:42 +1300 Subject: [PATCH] feat: add Adapt reconstruction for SamplingOp MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SamplingOp now returns a SamplingOpImpl wrapper that keeps the wrapped operator, the effective pattern, and the shape. Adapt and copy rebuild the operator from the effective pattern: integer indices are copied and stay on the host, the Boolean mask is copied to the target backend. The 5-arg adjoint product of the integer mode is marked broken in the tests: LinearOperators' opRestriction transpose kernel ignores α and β. Co-Authored-By: Claude Opus 5.5 --- src/SamplingOp.jl | 66 +++++++++++++++++--- test/testAdapt.jl | 137 ++++++++++++++++++++++++++++++++++++++++++ test/testOperators.jl | 52 ++++++++++++++++ 3 files changed, 246 insertions(+), 9 deletions(-) diff --git a/src/SamplingOp.jl b/src/SamplingOp.jl index e57b0d7..d48419a 100644 --- a/src/SamplingOp.jl +++ b/src/SamplingOp.jl @@ -1,14 +1,26 @@ export vectorizePattern """ - SamplingOp(pattern::Array{Int}, shape::Tuple) + SamplingOp(::Type{T}; pattern, shape::Tuple=(), S = Vector{T}) -builds a `LinearOperator` which only returns the vector elements at positions -indicated by pattern. +builds a linear operator that samples a vector at the positions given by `pattern`. -# Arguents -* `pattern::Array{Int}` - indices to sample -* `shape::Tuple` - size of the array to sample +Two modes are supported: +* integer indices with `shape`: returns the entries of a vectorized array of size + `shape` at the indices in `pattern`. A multidimensional `pattern` holds indices + into the first dimension, one column per remaining index, and is converted with + `vectorizePattern`. A one-dimensional `pattern` is used as is and not copied. +* Boolean mask without `shape`: multiplies the input elementwise with a copy of + `pattern` stored on the backend of `S`. + +`Adapt.adapt(to, op)` rebuilds the operator for the storage type `to` from its +effective pattern. Integer indices are copied and stay host arrays; the Boolean +mask is copied to the target backend. + +# Arguments +* `pattern` - integer indices or Boolean mask +* `shape::Tuple` - size of the array to sample (integer indices only) +* `S` - storage type of the operator """ function LinearOperatorCollection.SamplingOp(::Type{T}; pattern::P, shape::Tuple=(), S = Vector{T}) where {P} where T <: Number @@ -28,9 +40,20 @@ function vectorizePattern(idx::T, shape::Tuple) where T<:AbstractArray{Int} return [ floor(Int,(i-1)/size(idx,1))*shape[1]+idx[i] for i = 1:length(idx) ] end +struct SamplingOpImpl{T, vecT <: AbstractVector, O <: AbstractLinearOperator{T}, P <: AbstractArray, N} <: SamplingOp{T} + op::O + # effective pattern used by `op`: host integer indices or the owned Boolean mask + pattern::P + shape::NTuple{N, Int} + function SamplingOpImpl(op::O, pattern::P, shape::NTuple{N, Int}) where {T, O <: AbstractLinearOperator{T}, P <: AbstractArray, N} + return new{T, storage_type(op), O, P, N}(op, pattern, shape) + end +end + function SamplingOpImpl(T::Type{<:Number}, pattern::AbstractArray{Int}, shape::Tuple; S = Vector{T}) ndims(pattern)>1 ? idx = vectorizePattern(pattern, shape) : idx = pattern - return opEye(T,length(idx), S = S)*opRestriction(idx, prod(shape); S = S) + op = opEye(T,length(idx), S = S)*opRestriction(idx, prod(shape); S = S) + return SamplingOpImpl(op, idx, map(Int, shape)) end function SamplingOpImpl(T::Type{<:Number}, pattern::AbstractArray{Bool}; S = Vector{T}) @@ -40,5 +63,30 @@ function SamplingOpImpl(T::Type{<:Number}, pattern::AbstractArray{Bool}; S = Vec res .= pattern.*x end - return LinearOperator(T, length(pattern), length(pattern), true, false, prod!; S = S) -end \ No newline at end of file + op = LinearOperator(T, length(pattern), length(pattern), true, false, prod!; S = S) + return SamplingOpImpl(op, pattern, ()) +end + +function Base.getproperty(op::SamplingOpImpl, field::Symbol) + if in(field, (:op, :pattern, :shape)) + return getfield(op, field) + else + return getproperty(getfield(op, :op), field) + end +end +Base.setproperty!(op::SamplingOpImpl, field::Symbol, value) = setproperty!(op.op, field, value) + +storage_type(::SamplingOpImpl{T, vecT}) where {T, vecT} = vecT + +# Rebuild from the effective pattern: integer indices are copied, the constructor copies the mask. +_rebuild_sampling(op::SamplingOpImpl{T, V, O, <:AbstractArray{Bool}}, S) where {T, V, O} = + SamplingOp(T; pattern = op.pattern, S = S) +_rebuild_sampling(op::SamplingOpImpl{T}, S) where T = + SamplingOp(T; pattern = copy(op.pattern), shape = op.shape, S = S) + +Base.copy(op::SamplingOpImpl) = _rebuild_sampling(op, storage_type(op)) + +function Adapt.adapt_structure(to, op::SamplingOpImpl) + target_storage = _adapt_target_storage(to, op.Mv, "SamplingOp") + return _rebuild_sampling(op, typeof(target_storage)) +end diff --git a/test/testAdapt.jl b/test/testAdapt.jl index 8528559..e21990c 100644 --- a/test/testAdapt.jl +++ b/test/testAdapt.jl @@ -54,6 +54,58 @@ function custom_point_transform(x::AbstractFloat) end custom_point_transform(x::Tuple) = map(custom_point_transform, x) +const SamplingOpImpl = LinearOperatorCollection.SamplingOpImpl + +function init_sampling_workspaces!(op, array_type) + T = eltype(op) + m, n = size(op) + mul!(array_type(rand(T, m)), op, array_type(rand(T, n)), T(2), T(3)) + mul!(array_type(rand(T, n)), adjoint(op), array_type(rand(T, m)), T(2), T(3)) + return op +end + +snapshot_sampling(op) = (pattern = Array(op.pattern), Mv = Array(op.Mv), Mtu = Array(op.Mtu)) + +function check_sampling_snapshot(op, snapshot) + @test Array(op.pattern) == snapshot.pattern + @test Array(op.Mv) == snapshot.Mv + @test Array(op.Mtu) == snapshot.Mtu +end + +function check_sampling_action(op, array_type, x, y, forward, adjoint_result) + @test Array(op * array_type(x)) ≈ forward rtol = 1f-4 + @test Array(adjoint(op) * array_type(y)) ≈ adjoint_result rtol = 1f-4 +end + +function check_sampling_rebuild(adapted, source, array_type) + @test adapted isa SamplingOpImpl + @test adapted !== source + @test adapted.op !== source.op + @test adapted.prod! !== source.prod! + check_adapted_storage(adapted, array_type) + @test typeof(adapted.Mv) <: array_type + @test typeof(adapted.Mtu) <: array_type + @test adapted.pattern !== source.pattern + @test Array(adapted.pattern) == Array(source.pattern) + @test adapted.shape == source.shape + if eltype(source.pattern) == Bool + @test typeof(adapted.pattern) <: array_type + @test eltype(adapted.pattern) == Bool + else + @test adapted.pattern isa Vector{Int} + end +end + +# GPU-safe: whole-array broadcast for masks, host indices are plain vectors. +function mutate_sampling_pattern!(pattern, n) + if eltype(pattern) == Bool + pattern .= .!pattern + else + pattern[1] = mod1(pattern[1] + 1, n) + end + return pattern +end + @testset "Adapt reconstruction" begin N = 16 @@ -231,6 +283,87 @@ custom_point_transform(x::Tuple) = map(custom_point_transform, x) check_adapted_mul(source_normal, adapted_normal, x, array_type) end + @testset "SamplingOp $(mode) $(T)" for mode in (:index, :mask), T in (Float32, ComplexF32) + pattern_matrix = [mod1(3i + j, N) for i in 1:N÷2, j in 1:N] + caller = mode == :index ? pattern_matrix : [isodd(i ÷ 3) for i in 1:N^2] + source = mode == :index ? SamplingOp(T; pattern = caller, shape = (N, N)) : + SamplingOp(T; pattern = caller) + x = rand(T, size(source, 2)) + y = rand(T, size(source, 1)) + forward = source * x + adjoint_result = adjoint(source) * y + init_sampling_workspaces!(source, Array) + source_snapshot = snapshot_sampling(source) + + # The caller's matrix or mask was copied at construction. + mutate_sampling_pattern!(caller, N^2) + check_sampling_snapshot(source, source_snapshot) + + target = Adapt.adapt(array_type, source) + check_sampling_rebuild(target, source, array_type) + check_sampling_action(target, array_type, x, y, forward, adjoint_result) + check_sampling_snapshot(source, source_snapshot) + + # Same-backend reconstruction on the target. + init_sampling_workspaces!(target, array_type) + target_snapshot = snapshot_sampling(target) + same = Adapt.adapt(array_type, target) + check_sampling_rebuild(same, target, array_type) + check_sampling_action(same, array_type, x, y, forward, adjoint_result) + check_sampling_snapshot(target, target_snapshot) + mutate_sampling_pattern!(target.pattern, N^2) + check_sampling_action(same, array_type, x, y, forward, adjoint_result) + + round_trip = Adapt.adapt(Array, same) + check_sampling_rebuild(round_trip, same, Array) + check_sampling_action(round_trip, Array, x, y, forward, adjoint_result) + + # copy has its own operator, pattern, and workspaces. + copied = copy(source) + check_sampling_rebuild(copied, source, Array) + init_sampling_workspaces!(copied, Array) + isempty(source.Mv) || @test copied.Mv !== source.Mv + isempty(source.Mtu) || @test copied.Mtu !== source.Mtu + check_sampling_snapshot(source, source_snapshot) + check_sampling_action(copied, Array, x, y, forward, adjoint_result) + + # Mutating the source pattern after adaptation does not reach the target. + adapted = Adapt.adapt(array_type, source) + mutate_sampling_pattern!(source.pattern, N^2) + check_sampling_action(adapted, array_type, x, y, forward, adjoint_result) + end + + @testset "SamplingOp live index vector" begin + indices = collect(1:3:N^2) + source = SamplingOp(Float32; pattern = indices, shape = (N, N)) + @test source.pattern === indices + indices[1] = 2 + adapted = Adapt.adapt(array_type, source) + x = rand(Float32, N^2) + @test source * x == x[indices] + @test Array(adapted * array_type(x)) ≈ x[indices] + @test adapted.pattern == indices + @test adapted.pattern !== indices + end + + @testset "SamplingOp ∘ FFTOp" begin + pattern = [mod1(3i + j, N) for i in 1:N÷2, j in 1:N] + source = SamplingOp(ComplexF32; pattern, shape = (N, N)) ∘ + FFTOp(ComplexF32; shape = (N, N), S = Vector{ComplexF32}) + @test size(source) == (N^2 ÷ 2, N^2) + adapted = Adapt.adapt(array_type, source) + @test adapted isa ProdOp + @test adapted.A isa SamplingOpImpl + check_adapted_storage(adapted, array_type) + check_adapted_storage(adapted.A, array_type) + check_adapted_storage(adapted.B, array_type) + @test typeof(adapted.tmp) <: array_type + @test adapted.A.pattern == source.A.pattern + x = randn(ComplexF32, N^2) + y = randn(ComplexF32, N^2 ÷ 2) + check_sampling_action(adapted, array_type, x, y, source * x, adjoint(source) * y) + end + @testset "FFTOp rebuilds plans and preserves action" begin for shift in (false, true), unitary in (false, true) source = FFTOp(ComplexF32; shape=(6, 8), shift, unitary, @@ -460,11 +593,15 @@ struct UnsupportedAdaptTarget end toeplitz = with(nfft_backend => NFFT.backend()) do normalOperator(nfft) end + sampling_index = SamplingOp(Float32; pattern = [1, 3], shape = (4,)) + sampling_mask = SamplingOp(Float32; pattern = [true, false, true]) 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, sampling_index) + @test_throws ArgumentError Adapt.adapt(target, sampling_mask) end end diff --git a/test/testOperators.jl b/test/testOperators.jl index bd15ca8..a882e0c 100644 --- a/test/testOperators.jl +++ b/test/testOperators.jl @@ -184,6 +184,57 @@ function testDirectionalGradOp(N=64;arrayType = Array) @test norm(x2r-vec(x2r_ref)) / norm(x2r_ref) ≈ 0 atol=0.001 end +function testSamplingModes(; arrayType = Array) + shape = (4, 3) + pattern = [4 2 3; 1 4 2] + idx = [4, 1, 6, 8, 11, 10] + @test vectorizePattern(pattern, shape) == idx + mask = zeros(Bool, 12); mask[[2, 5, 9]] .= true + + for T in (Float32, ComplexF32) + S = typeof(arrayType(zeros(T, 1))) + x = rand(T, 12) + y = rand(T, 6) + ops = ( + (SamplingOp(T; pattern, shape, S), idx), + (createLinearOperator(SamplingOp{T}; pattern, shape, S), idx), + (SamplingOp(T; pattern = mask, S), nothing), + (createLinearOperator(SamplingOp{T}; pattern = mask, S), nothing), + ) + for (op, ref_idx) in ops + @test op isa SamplingOp{T} + @test eltype(op) == T + check_storage(op, arrayType) + if isnothing(ref_idx) + @test size(op) == (12, 12) + forward_ref = mask .* x + adjoint_in = x + adjoint_ref = mask .* x + else + @test size(op) == (6, 12) + forward_ref = x[ref_idx] + adjoint_in = y + adjoint_ref = zeros(T, 12); adjoint_ref[ref_idx] .= y + end + @test Array(op * arrayType(x)) ≈ forward_ref + @test Array(adjoint(op) * arrayType(adjoint_in)) ≈ adjoint_ref + res = rand(T, size(op, 1)) + res_op = arrayType(copy(res)) + mul!(res_op, op, arrayType(x), T(2), T(3)) + @test Array(res_op) ≈ 2 .* forward_ref .+ 3 .* res + res = rand(T, size(op, 2)) + res_op = arrayType(copy(res)) + mul!(res_op, adjoint(op), arrayType(adjoint_in), T(2), T(3)) + if isnothing(ref_idx) + @test Array(res_op) ≈ 2 .* adjoint_ref .+ 3 .* res + else + # LinearOperators' opRestriction transpose kernel ignores α and β. + @test_broken Array(res_op) ≈ 2 .* adjoint_ref .+ 3 .* res + end + end + end +end + function testSampling(N=64;arrayType = Array) x = rand(ComplexF64,N,N) xop = arrayType(vec(x)) @@ -465,6 +516,7 @@ end @testset testDirectionalGradOp(64;arrayType) @info "test SamplingOp: $arrayType" @testset testSampling(64;arrayType) + @testset testSamplingModes(;arrayType) @info "test WaveletOp: $arrayType" @testset testWavelet(64,64;arrayType) @testset testWavelet(64,60;arrayType)