Add GPU support for sparse single patch recos

src/Algorithms/SinglePatchAlgorithms/SinglePatchAlgorithm.jl

@@ -1,19 +1,21 @@
 Base.@kwdef struct SinglePatchReconstructionParameter{L<:AbstractSystemMatrixLoadingParameter, SL<:AbstractLinearSolver,
-   SP<:AbstractSolverParameters{SL}, R<:AbstractRegularization, W<:AbstractWeightingParameters} <: AbstractSinglePatchReconstructionParameters
+   matT <: AbstractArray, SP<:AbstractSolverParameters{SL}, R<:AbstractRegularization, W<:AbstractWeightingParameters} <: AbstractSinglePatchReconstructionParameters
   # File
   sf::MPIFile
   sfLoad::Union{L, ProcessResultCache{L}}
+  arrayType::Type{matT} = Array
   # Solver
   solverParams::SP
   reg::Vector{R} = AbstractRegularization[]
   weightingParams::W = NoWeightingParameters()
 end
 
-Base.@kwdef mutable struct SinglePatchReconstructionAlgorithm{P} <: AbstractSinglePatchReconstructionAlgorithm where {P<:AbstractSinglePatchAlgorithmParameters}
+Base.@kwdef mutable struct SinglePatchReconstructionAlgorithm{P, matT <: AbstractArray} <: AbstractSinglePatchReconstructionAlgorithm where {P<:AbstractSinglePatchAlgorithmParameters}
   params::P
   # Could also do reconstruction progress meter here
   sf::Union{MPIFile, Vector{MPIFile}}
   S::AbstractArray
+  arrayType::Type{matT}
   grid::RegularGridPositions
   freqs::Vector{CartesianIndex{2}}
   output::Channel{Any}
@@ -23,16 +25,16 @@ function SinglePatchReconstruction(params::SinglePatchParameters{<:AbstractMPIPr
   return SinglePatchReconstructionAlgorithm(params)
 end
 function SinglePatchReconstructionAlgorithm(params::SinglePatchParameters{<:AbstractMPIPreProcessingParameters, R, PT}) where {R<:AbstractSinglePatchReconstructionParameters, PT <:AbstractMPIPostProcessingParameters}
-  freqs, S, grid = prepareSystemMatrix(params.reco)
-  return SinglePatchReconstructionAlgorithm(params, params.reco.sf, S, grid, freqs, Channel{Any}(Inf))
+  freqs, S, grid, arrayType = prepareSystemMatrix(params.reco)
+  return SinglePatchReconstructionAlgorithm(params, params.reco.sf, S, arrayType, grid, freqs, Channel{Any}(Inf))
 end
 recoAlgorithmTypes(::Type{SinglePatchReconstruction}) = SystemMatrixBasedAlgorithm()
 AbstractImageReconstruction.parameter(algo::SinglePatchReconstructionAlgorithm) = algo.params
 
 function prepareSystemMatrix(reco::SinglePatchReconstructionParameter{L,S}) where {L<:AbstractSystemMatrixLoadingParameter, S<:AbstractLinearSolver}
   freqs, sf, grid = process(AbstractMPIRecoAlgorithm, reco.sfLoad, reco.sf)
-  sf, grid = prepareSF(S, sf, grid) 
-  return freqs, sf, grid
+  sf, grid = process(AbstractMPIRecoAlgorithm, reco.sfLoad, S, sf, grid, reco.arrayType)
+  return freqs, sf, grid, reco.arrayType
 end
 
 
@@ -44,13 +46,13 @@ function process(algo::SinglePatchReconstructionAlgorithm, params::Union{A, Proc
     @warn "System matrix and measurement have different element data type. Mapping measurment data to system matrix element type."
     result = map(eltype(algo.S),result)
   end
-  result = copyto!(similar(algo.S, size(result)...), result)
+  result = adapt(algo.arrayType, result)
   return result
 end
 
 
 function process(algo::SinglePatchReconstructionAlgorithm, params::SinglePatchReconstructionParameter, u)
-  weights = process(algo, params.weightingParams, u)
+  weights = adapt(algo.arrayType, process(algo, params.weightingParams, u))
 
   B = getLinearOperator(algo, params)
 
@@ -68,5 +70,5 @@ function getLinearOperator(algo::SinglePatchReconstructionAlgorithm, params::Sin
 end
 
 function getLinearOperator(algo::SinglePatchReconstructionAlgorithm, params::SinglePatchReconstructionParameter{<:SparseSystemMatrixLoadingParameter, S}) where {S}
-  return process(algo, params.sfLoad, eltype(algo.S), tuple(shape(algo.grid)...))
+  return process(algo, params.sfLoad, eltype(algo.S), algo.arrayType, tuple(shape(algo.grid)...))
 end

src/Algorithms/SinglePatchAlgorithms/SinglePatchBGEstimationAlgorithm.jl

@@ -1,9 +1,10 @@
 export SinglePatchBGEstimationAlgorithm, SinglePatchBGEstimationReconstructionParameter
 Base.@kwdef struct SinglePatchBGEstimationReconstructionParameter{L<:DenseSystemMatixLoadingParameter,
-  SP<:AbstractSolverParameters} <: AbstractSinglePatchReconstructionParameters
+  SP<:AbstractSolverParameters, matT <: AbstractArray} <: AbstractSinglePatchReconstructionParameters
   # File
   sf::MPIFile
   sfLoad::Union{L, ProcessResultCache{L}}
+  arrayType::Type{matT} = Array
   # Solver
   solverParams::SP
   λ::Float32
@@ -12,11 +13,12 @@ Base.@kwdef struct SinglePatchBGEstimationReconstructionParameter{L<:DenseSystem
   bgDict::BGDictParameter
 end
 
-Base.@kwdef mutable struct SinglePatchBGEstimationAlgorithm{P} <: AbstractSinglePatchReconstructionAlgorithm where {P<:AbstractSinglePatchAlgorithmParameters}
+Base.@kwdef mutable struct SinglePatchBGEstimationAlgorithm{P, matT <: AbstractArray} <: AbstractSinglePatchReconstructionAlgorithm where {P<:AbstractSinglePatchAlgorithmParameters}
   params::P
   # Could also do reconstruction progress meter here
   sf::Union{MPIFile,Vector{MPIFile}}
   S::AbstractArray
+  arrayType::Type{matT}
   bgDict::AbstractArray
   grid::RegularGridPositions
   freqs::Vector{CartesianIndex{2}}
@@ -27,16 +29,16 @@ function SinglePatchReconstruction(params::SinglePatchParameters{<:AbstractMPIPr
   return SinglePatchBGEstimationAlgorithm(params)
 end
 function SinglePatchBGEstimationAlgorithm(params::SinglePatchParameters{<:AbstractMPIPreProcessingParameters,R,PT}) where {R<:SinglePatchBGEstimationReconstructionParameter,PT<:AbstractMPIPostProcessingParameters}
-  freqs, S, grid = prepareSystemMatrix(params.reco)
-  return SinglePatchBGEstimationAlgorithm(params, params.reco.sf, S, process(SinglePatchBGEstimationAlgorithm, freqs, params.reco.bgDict), grid, freqs, Channel{Any}(Inf))
+  freqs, S, grid, arrayType = prepareSystemMatrix(params.reco)
+  return SinglePatchBGEstimationAlgorithm(params, params.reco.sf, S, arrayType, process(SinglePatchBGEstimationAlgorithm, freqs, params.reco.bgDict), grid, freqs, Channel{Any}(Inf))
 end
 recoAlgorithmTypes(::Type{SinglePatchBGEstimationAlgorithm}) = SystemMatrixBasedAlgorithm()
 AbstractImageReconstruction.parameter(algo::SinglePatchBGEstimationAlgorithm) = algo.origParam
 
 function prepareSystemMatrix(reco::SinglePatchBGEstimationReconstructionParameter{L}) where {L<:AbstractSystemMatrixLoadingParameter}
   freqs, sf, grid = process(AbstractMPIRecoAlgorithm, reco.sfLoad, reco.sf)
-  sf, grid = prepareSF(Kaczmarz, sf, grid)
-  return freqs, sf, grid
+  sf, grid = process(AbstractMPIRecoAlgorithm, reco.sfLoad, Kaczmarz, sf, grid, reco.arrayType)
+  return freqs, sf, grid, reco.arrayType
 end
 
 AbstractImageReconstruction.take!(algo::SinglePatchBGEstimationAlgorithm) = Base.take!(algo.output)
@@ -47,6 +49,7 @@ function process(algo::SinglePatchBGEstimationAlgorithm, params::AbstractMPIPreP
     @warn "System matrix and measurement have different element data type. Mapping measurment data to system matrix element type."
     result = map(eltype(algo.S), result)
   end
+  result = adapt(algo.arrayType, result)
   return result
 end
 

src/Algorithms/SinglePatchAlgorithms/SinglePatchTemporalRegularizationAlgorithm.jl

@@ -1,9 +1,10 @@
 export SinglePatchTemporalRegularizationAlgorithm, SinglePatchTemporalRegularizationReconstructionParameter
 Base.@kwdef struct SinglePatchTemporalRegularizationReconstructionParameter{L<:DenseSystemMatixLoadingParameter,
-  SP<:AbstractSolverParameters} <: AbstractSinglePatchReconstructionParameters
+  SP<:AbstractSolverParameters, matT <: AbstractArray} <: AbstractSinglePatchReconstructionParameters
   # File
   sf::MPIFile
   sfLoad::Union{L, ProcessResultCache{L}}
+  arrayType::Type{matT} = Array
   # Solver
   solverParams::SP
   λ::Float32
@@ -14,10 +15,11 @@ Base.@kwdef struct SinglePatchTemporalRegularizationReconstructionParameter{L<:D
   bgDict::BGDictParameter
 end
 
-Base.@kwdef mutable struct SinglePatchTemporalRegularizationAlgorithm{P} <: AbstractSinglePatchReconstructionAlgorithm where {P<:AbstractSinglePatchAlgorithmParameters}
+Base.@kwdef mutable struct SinglePatchTemporalRegularizationAlgorithm{P, matT <: AbstractArray} <: AbstractSinglePatchReconstructionAlgorithm where {P<:AbstractSinglePatchAlgorithmParameters}
   params::P
   sf::Union{MPIFile,Vector{MPIFile}}
   S::AbstractArray
+  arrayType::Type{matT}
   bgDict::AbstractArray
   idxFG::Union{Nothing, UnitRange{Int64}, Vector{Int64}} = nothing
   idxFG::Union{Nothing, UnitRange{Int64}, Vector{Int64}} = nothing
@@ -30,17 +32,17 @@ function SinglePatchReconstruction(params::SinglePatchParameters{<:AbstractMPIPr
   return SinglePatchTemporalRegularizationAlgorithm(params)
 end
 function SinglePatchTemporalRegularizationAlgorithm(params::SinglePatchParameters{<:AbstractMPIPreProcessingParameters,R,PT}) where {R<:SinglePatchTemporalRegularizationReconstructionParameter,PT<:AbstractMPIPostProcessingParameters}
-  freqs, S, grid = prepareSystemMatrix(params.reco)
-  return SinglePatchTemporalRegularizationAlgorithm(params, params.reco.sf, S, process(SinglePatchTemporalRegularizationAlgorithm, params.reco.bgDict, freqs)
+  freqs, S, grid, arrayType = prepareSystemMatrix(params.reco)
+  return SinglePatchTemporalRegularizationAlgorithm(params, params.reco.sf, S, arrayType, process(SinglePatchTemporalRegularizationAlgorithm, params.reco.bgDict, freqs)
     ,params.reco.idxFG, params.reco.idxBG, grid, freqs, Channel{Any}(Inf))
 end
 recoAlgorithmTypes(::Type{SinglePatchTemporalRegularizationAlgorithm}) = SystemMatrixBasedAlgorithm()
 AbstractImageReconstruction.parameter(algo::SinglePatchTemporalRegularizationAlgorithm) = algo.origParam
 
 function prepareSystemMatrix(reco::SinglePatchTemporalRegularizationReconstructionParameter{L}) where {L<:AbstractSystemMatrixLoadingParameter}
   freqs, sf, grid = process(AbstractMPIRecoAlgorithm, reco.sfLoad, reco.sf)
-  sf, grid = prepareSF(Kaczmarz, sf, grid)
-  return freqs, sf, grid
+  sf, grid = process(AbstractMPIRecoAlgorithm, reco.sfLoad, Kaczmarz, sf, grid, reco.arrayType)
+  return freqs, sf, grid, reco.arrayType
 end
 
 AbstractImageReconstruction.take!(algo::SinglePatchTemporalRegularizationAlgorithm) = Base.take!(algo.output)
@@ -51,6 +53,7 @@ function process(algo::SinglePatchTemporalRegularizationAlgorithm, params::Abstr
     @warn "System matrix and measurement have different element data type. Mapping measurment data to system matrix element type."
     result = map(eltype(algo.S), result)
   end
+  result = adapt(algo.arrayType, result)
   return result
 end
 

src/SystemMatrix/SystemMatrix.jl

@@ -67,7 +67,6 @@ end
 function process(t::Type{<:AbstractMPIRecoAlgorithm}, params::DenseSystemMatixLoadingParameter, sf::MPIFile, frequencies::Vector{CartesianIndex{2}})
   S, grid = getSF(sf, frequencies, nothing; toKwargs(params)...)
   @info "Loading SM"
-  S = adapt(params.arrayType, S)
   return S, grid
 end
 
@@ -90,8 +89,8 @@ function process(t::Type{<:AbstractMPIRecoAlgorithm}, params::SparseSystemMatrix
   S, grid = getSF(sf, frequencies, params.sparseTrafo; toKwargs(params)...)
   return S, grid
 end
-function process(t::Type{<:AbstractMPIRecoAlgorithm}, params::SparseSystemMatrixLoadingParameter, elType::Type{<:Number}, shape::NTuple{N, Int64}) where N
-  return createLinearOperator(params.sparseTrafo, elType; shape)
+function process(t::Type{<:AbstractMPIRecoAlgorithm}, params::SparseSystemMatrixLoadingParameter, elType::Type{<:Number}, arrayType, shape::NTuple{N, Int64}) where N
+  return createLinearOperator(params.sparseTrafo, elType; shape, S = typeof(arrayType{elType}(undef, 0)))
 end
 
 function converttoreal(S::AbstractArray{Complex{T}},f) where T
@@ -123,15 +122,30 @@ function getSF(bSF, frequencies, sparseTrafo, solver::AbstractString; kargs...)
   end
 end
 getSF(bSF, frequencies, sparseTrafo, solver::AbstractLinearSolver; kargs...) = getSF(bSF, frequencies, sparseTrafo, typeof(solver); kargs...)
-function getSF(bSF, frequencies, sparseTrafo, solver::Type{<:AbstractLinearSolver}; kargs...)
+function getSF(bSF, frequencies, sparseTrafo, solver::Type{<:AbstractLinearSolver}; arrayType = Array, kargs...)
   SF, grid = getSF(bSF, frequencies, sparseTrafo; kargs...)
-  return prepareSF(solver, SF, grid)
+  return prepareSF(solver, SF, grid, arrayType)
 end
 
+
+function AbstractImageReconstruction.process(::Type{<:AbstractMPIRecoAlgorithm}, params::AbstractSystemMatrixLoadingParameter, solver::Type{<:AbstractLinearSolver}, sf, grid, arrayType)
+  return prepareSF(solver, sf, grid, arrayType)
+end
+function prepareSF(solver::Type{<:AbstractLinearSolver}, SF, grid, arrayType)
+  SF, grid = prepareSF(solver, SF, grid)
+  # adapt(Array, Sparse-CPU) results in dense array
+  return arrayType != Array ? adapt(arrayType, SF) : SF, grid
+end
 prepareSF(solver::Type{Kaczmarz}, SF, grid) = transpose(SF), grid
 prepareSF(solver::Type{PseudoInverse}, SF, grid) = SVD(svd(transpose(SF))...), grid
-prepareSF(solver::Union{Type{<:RegularizedLeastSquares.AbstractLinearSolver}}, SF, grid) = copy(transpose(SF)), grid
 prepareSF(solver::Type{DirectSolver}, SF, grid) = RegularizedLeastSquares.tikhonovLU(copy(transpose(SF))), grid
+prepareSF(solver::Type{Kaczmarz}, SF::AbstractSparseArray, grid) = transpose(SF), grid
+prepareSF(solver::Type{PseudoInverse}, SF::AbstractSparseArray, grid) = SVD(svd(transpose(SF))...), grid
+prepareSF(solver::Type{DirectSolver}, SF::AbstractSparseArray, grid) = RegularizedLeastSquares.tikhonovLU(copy(transpose(SF))), grid
+
+prepareSF(solver, SF , grid) = copy(transpose(SF)), grid
+prepareSF(solver, SF::AbstractSparseArray, grid) = sparse(transpose(SF)), grid
+
 
 prepareNormalSF(solver::AbstractLinearSolver, SF) = prepareNormalSF(typeof(solver), SF)
 prepareNormalSF(solver::Type{<:RegularizedLeastSquares.AbstractLinearSolver}, SF) = nothing

test/ReconstructionGPU.jl

@@ -1,5 +1,5 @@
 for arrayType in arrayTypes
-  @testset "Reconstruction $arrayType" begin
+  @testset "Single Patch Reconstruction: $arrayType" begin
     bSF = MPIFile(joinpath(datadir, "calibrations", "12.mdf"))
     b = MPIFile(joinpath(datadir, "measurements", "20211226_203916_MultiPatch", "1.mdf"))
     names = (:color, :x, :y, :z, :time)
@@ -22,12 +22,33 @@ for arrayType in arrayTypes
     params[:solver] = CGNR
 
     c1 = reconstruct("SinglePatch", b; params...)
-    @test axisnames(c1) == names
-    @test axisvalues(c1) == values
 
     c2 = reconstruct("SinglePatch", b; params..., arrayType = arrayType)
-    @test axisnames(c2) == names
-    @test axisvalues(c2) == values
+
+    @test isapprox(arraydata(c1), arraydata(c2))
+  end
+
+  @testset "Sparse Single Patch Reconstruction: $arrayType" begin
+    bSF = MPIFile(joinpath(datadir, "calibrations", "7.mdf"))
+    b = MPIFile(joinpath(datadir, "measurements","20211226_204612_Dice", "1.mdf"))
+    redFactor = 0.01
+
+    params = Dict{Symbol, Any}()
+    params[:SNRThresh] = 2
+    params[:frames] = 1:100
+    params[:minFreq] = 80e3
+    params[:recChannels] = 1:2
+    params[:iterations] = 3
+    params[:spectralLeakageCorrection] = false
+    params[:sf] = bSF
+    params[:reg] = [L2Regularization(0.1f0)]
+    params[:numAverages] = 100
+    params[:solver] = CGNR
+    params[:redFactor] = redFactor
+
+    c1 = reconstruct("SinglePatchSparse", b; params..., sparseTrafo = "FFT", useDFFoV = false)
+
+    c2 = reconstruct("SinglePatchSparse", b; params..., sparseTrafo = "FFT", useDFFoV = false, arrayType = arrayType)
 
     @test isapprox(arraydata(c1), arraydata(c2))
   end