Add Adapt reconstruction for SamplingOp

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
+  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

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
+
 gradient_test_data(::Type{T}, n) where {T <: Integer} = rand(T(-5):T(5), n)
 gradient_test_data(::Type{T}, n) where {T} = rand(T, n)
 
@@ -306,6 +358,87 @@ storage_vector_type(array_type, T) = typeof(array_type(zeros(T, 0)))
         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,
@@ -735,12 +868,16 @@ 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])
   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, sampling_index)
+    @test_throws ArgumentError Adapt.adapt(target, sampling_mask)
     @test_throws ArgumentError Adapt.adapt(target, radon_op)
   end
 end

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)