Merge pull request #40 from MagneticResonanceImaging/CoilMapsCG

Option to run a coilwise CG recon in the coil maps calculation.
MagneticResonanceImaging · Oct 4, 2024 · c1ddb60 · c1ddb60
2 parents ae488c3 + 0a7f145
commit c1ddb60
Show file tree

Hide file tree

Showing 5 changed files with 58 additions and 30 deletions.
diff --git a/Project.toml b/Project.toml
@@ -6,6 +6,7 @@ version = "0.8.1-DEV"
 [deps]
 ExponentialUtilities = "d4d017d3-3776-5f7e-afef-a10c40355c18"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearOperators = "5c8ed15e-5a4c-59e4-a42b-c7e8811fb125"
 MRICoilSensitivities = "c57eb701-aafc-44a2-a53c-128049758959"

diff --git a/src/BackProjection.jl b/src/BackProjection.jl
@@ -21,68 +21,72 @@ Calculate (filtered) backprojection
 # Notes
 - The type of the elements of the trajectory define if a gridded backprojection (eltype(trj[1]) or eltype(trj) <: Int) or a non-uniform (else) is performed.
 """
-function calculateBackProjection(data::AbstractVector{<:AbstractArray{cT}}, trj::AbstractVector{<:AbstractMatrix{T}}, img_shape::NTuple{N,Int}; U=I(length(data)), density_compensation=:none, verbose=false) where {T <: Real, cT <: Complex{T},N}
-    Ncoef = size(U,2)
+function calculateBackProjection(data::AbstractVector{<:AbstractArray{cT}}, trj::AbstractVector{<:AbstractMatrix{T}}, img_shape::NTuple{N,Int}; U=I(length(data)), density_compensation=:none, verbose=false) where {T<:Real,cT<:Complex{T},N}
+    Ncoef = size(U, 2)
 
     trj_v = reduce(hcat, trj)
     p = plan_nfft(trj_v, img_shape; precompute=TENSOR, blocking=true, fftflags=FFTW.MEASURE)
 
     Ncoil = size(data[1], 2)
     xbp = Array{cT}(undef, img_shape..., Ncoef, Ncoil)
 
-    trj_l = [size(trj[it],2) for it in eachindex(trj)]
-    data_temp = Vector{cT}(undef,sum(trj_l))
+    trj_l = [size(trj[it], 2) for it in eachindex(trj)]
+    data_temp = Vector{cT}(undef, sum(trj_l))
 
     img_idx = CartesianIndices(img_shape)
-    verbose && println("calculating backprojection..."); flush(stdout)
+    verbose && println("calculating backprojection...")
+    flush(stdout)
     for icoef ∈ axes(U, 2)
         t = @elapsed for icoil ∈ axes(data[1], 2)
             @simd for it in eachindex(data)
                 idx1 = sum(trj_l[1:it-1]) + 1
                 idx2 = sum(trj_l[1:it])
-                @views data_temp[idx1:idx2] .= data[it][:,icoil] .* conj(U[it,icoef])
+                @views data_temp[idx1:idx2] .= data[it][:, icoil] .* conj(U[it, icoef])
             end
             applyDensityCompensation!(data_temp, trj_v; density_compensation)
 
             @views mul!(xbp[img_idx, icoef, icoil], adjoint(p), data_temp)
         end
-        verbose && println("coefficient = $icoef: t = $t s"); flush(stdout)
+        verbose && println("coefficient = $icoef: t = $t s")
+        flush(stdout)
     end
     return xbp
 end
 
-function calculateBackProjection(data::AbstractVector{<:AbstractMatrix{cT}}, trj::AbstractVector{<:AbstractMatrix{T}}, cmaps::AbstractVector{<:AbstractArray{cT,N}}; U=I(length(data)), density_compensation=:none, verbose=false) where {T <: Real, cT <: Complex{T}, N}
+function calculateBackProjection(data::AbstractVector{<:AbstractMatrix{cT}}, trj::AbstractVector{<:AbstractMatrix{T}}, cmaps::AbstractVector{<:AbstractArray{cT,N}}; U=I(length(data)), density_compensation=:none, verbose=false) where {T<:Real,cT<:Complex{T},N}
     test_dimension(data, trj, U, cmaps)
 
     trj_v = reduce(hcat, trj)
     Ncoef = size(U, 2)
     img_shape = size(cmaps[1])
 
-    p = plan_nfft(trj_v, img_shape; precompute=TENSOR, blocking = true, fftflags = FFTW.MEASURE)
+    p = plan_nfft(trj_v, img_shape; precompute=TENSOR, blocking=true, fftflags=FFTW.MEASURE)
     xbp = zeros(cT, img_shape..., Ncoef)
     xtmp = Array{cT}(undef, img_shape)
 
-    trj_l = [size(trj[it],2) for it in eachindex(trj)]
-    data_temp = Vector{cT}(undef,sum(trj_l))
+    trj_l = [size(trj[it], 2) for it in eachindex(trj)]
+    data_temp = Vector{cT}(undef, sum(trj_l))
     img_idx = CartesianIndices(img_shape)
-    verbose && println("calculating backprojection..."); flush(stdout)
+    verbose && println("calculating backprojection...")
+    flush(stdout)
     for icoef ∈ axes(U, 2)
         t = @elapsed for icoil ∈ eachindex(cmaps)
             @simd for it in eachindex(data)
                 idx1 = sum(trj_l[1:it-1]) + 1
                 idx2 = sum(trj_l[1:it])
-                @views data_temp[idx1:idx2] .= data[it][:,icoil] .* conj(U[it,icoef])
+                @views data_temp[idx1:idx2] .= data[it][:, icoil] .* conj(U[it, icoef])
             end
             applyDensityCompensation!(data_temp, trj_v; density_compensation)
             mul!(xtmp, adjoint(p), data_temp)
-            xbp[img_idx,icoef] .+= conj.(cmaps[icoil]) .* xtmp
+            xbp[img_idx, icoef] .+= conj.(cmaps[icoil]) .* xtmp
         end
-        verbose && println("coefficient = $icoef: t = $t s"); flush(stdout)
+        verbose && println("coefficient = $icoef: t = $t s")
+        flush(stdout)
     end
     return xbp
 end
 
-function calculateBackProjection(data::AbstractArray{cT}, trj::AbstractMatrix{T}, cmaps_img_shape; density_compensation=:none, verbose=false) where {T <: Real, cT <: Complex{T}}
+function calculateBackProjection(data::AbstractArray{cT}, trj::AbstractMatrix{T}, cmaps_img_shape; density_compensation=:none, verbose=false) where {T<:Real,cT<:Complex{T}}
     return calculateBackProjection([data], [trj], cmaps_img_shape; U=I(1), density_compensation, verbose)
 end
 
@@ -112,6 +116,22 @@ function calculateBackProjection(data::AbstractVector{<:AbstractArray}, trj::Abs
     return xbp
 end
 
+
+function calculateCoilwiseCG(data::AbstractVector{<:AbstractArray{cT}}, trj::AbstractVector{<:AbstractMatrix{T}}, img_shape::NTuple{N,Int}; U=I(length(data)), maxiter=100, verbose=false) where {T<:Real,cT<:Complex{T},N}
+    Ncoil = size(data[1], 2)
+
+    AᴴA = NFFTNormalOp(img_shape, trj, U[:, 1]; verbose)
+    xbp = calculateBackProjection(data, trj, img_shape; U=U[:, 1], verbose)
+    x = zeros(cT, img_shape..., Ncoil)
+
+    for icoil = 1:Ncoil
+        bi = vec(@view xbp[CartesianIndices(img_shape), 1, icoil])
+        xi = vec(@view x[CartesianIndices(img_shape), icoil])
+        cg!(xi, AᴴA, bi; maxiter, verbose, reltol=0)
+    end
+    return x
+end
+
 ## ##########################################################################
 # Internal use
 #############################################################################

diff --git a/src/CoilMaps.jl b/src/CoilMaps.jl
@@ -24,31 +24,36 @@ Estimate coil sensitivity maps using ESPIRiT [1].
 # References
 [1] Uecker, M., Lai, P., Murphy, M.J., Virtue, P., Elad, M., Pauly, J.M., Vasanawala, S.S. and Lustig, M. (2014), ESPIRiT—an eigenvalue approach to autocalibrating parallel MRI: Where SENSE meets GRAPPA. Magn. Reson. Med., 71: 990-1001. https://doi.org/10.1002/mrm.24751
 """
-function calcCoilMaps(data::AbstractVector{<:AbstractMatrix{Complex{T}}}, trj::AbstractVector{<:AbstractMatrix{T}}, img_shape::NTuple{N,Int}; U = ones(length(data)), density_compensation=:radial_3D, kernel_size=ntuple(_ -> 6, N), calib_size=ntuple(_ -> 24, N), eigThresh_1=0.01, eigThresh_2=0.9, nmaps=1, verbose=false) where {N,T}
+function calcCoilMaps(data::AbstractVector{<:AbstractMatrix{Complex{T}}}, trj::AbstractVector{<:AbstractMatrix{T}}, img_shape::NTuple{N,Int}; U=ones(length(data)), kernel_size=ntuple(_ -> 6, N), calib_size=img_shape.÷(img_shape[1]÷32), eigThresh_1=0.01, eigThresh_2=0.9, nmaps=1, verbose=false) where {N,T}
     Ncoil = size(data[1], 2)
     Ndims = length(img_shape)
     imdims = ntuple(i -> i, Ndims)
+    Nt = length(trj)
 
-    xbp = calculateBackProjection(data, trj, img_shape; U=U[:,1], density_compensation, verbose)
-    xbp = dropdims(xbp, dims=ndims(xbp)-1)
+    calib_scale = img_shape ./ calib_size
 
-    img_idx = CartesianIndices(img_shape)
-    kbp = fftshift(xbp, imdims)
+    trj_calib = Vector{Matrix{T}}(undef, Nt)
+    data_calib = Vector{Matrix{Complex{T}}}(undef, Nt)
+    Threads.@threads for it ∈ eachindex(trj)
+        trj_idx = [all(abs.(trj[it][:, i] .* calib_scale) .< T(0.5)) for i ∈ axes(trj[it], 2)]
+        trj_calib[it] = trj[it][:, trj_idx] .* calib_scale
+        data_calib[it] = data[it][trj_idx, :]
+    end
+    x = calculateCoilwiseCG(data_calib, trj_calib, calib_size; U)
+
+    kbp = fftshift(x, imdims)
     fft!(kbp, imdims)
     kbp = fftshift(kbp, imdims)
 
-    m = CartesianIndices(calib_size) .+ CartesianIndex((img_shape .- calib_size) .÷ 2)
-    kbp = kbp[m, :]
-
     t = @elapsed begin
-        cmaps = espirit(kbp, img_shape, kernel_size, eigThresh_1=eigThresh_1, eigThresh_2=eigThresh_2, nmaps=nmaps)
+        cmaps = espirit(kbp, img_shape, kernel_size; eigThresh_1, eigThresh_2, nmaps)
     end
     verbose && println("espirit: $t s")
 
-    cmaps = [cmaps[img_idx, ic, 1] for ic = 1:Ncoil]
+    cmaps = [cmaps[CartesianIndices(img_shape), ic, in] for ic = 1:Ncoil, in = (nmaps == 1 ? 1 : 1:nmaps)]
     return cmaps
 end
 
 function calcCoilMaps(data::AbstractMatrix{Complex{T}}, trj::AbstractMatrix{T}, img_shape::NTuple{N,Int}; density_compensation=:radial_3D, kernel_size=ntuple(_ -> 6, N), calib_size=ntuple(_ -> 24, N), eigThresh_1=0.01, eigThresh_2=0.9, nmaps=1, verbose=false) where {N,T}
-    calcCoilMaps([data], [trj], img_shape; U = I(1), density_compensation, kernel_size, calib_size, eigThresh_1, eigThresh_2, nmaps, verbose)
+    calcCoilMaps([data], [trj], img_shape; U=I(1), density_compensation, kernel_size, calib_size, eigThresh_1, eigThresh_2, nmaps, verbose)
 end
diff --git a/src/MRFingerprintingRecon.jl b/src/MRFingerprintingRecon.jl
@@ -8,6 +8,8 @@ using NFFTTools
 using MRICoilSensitivities
 using LinearOperators
 using ExponentialUtilities
+using IterativeSolvers
+
 
 export NFFTNormalOp, calcCoilMaps, calculateBackProjection, kooshball, kooshballGA, calcFilteredBackProjection
 export FFTNormalOp, radial_grog!

diff --git a/src/NFFTNormalOpBasisFunc.jl b/src/NFFTNormalOpBasisFunc.jl
@@ -22,7 +22,7 @@ Differentiate between functions exploiting a pre-calculated Toeplitz kernel basi
 function NFFTNormalOp(
     img_shape,
     trj::AbstractVector{<:AbstractMatrix{T}},
-    U::AbstractMatrix{Tc};
+    U::AbstractArray{Tc};
     cmaps=[ones(T, img_shape)],
     verbose = false,
     num_fft_threads = round(Int, Threads.nthreads()/size(U, 2)),
@@ -93,7 +93,7 @@ function calculate_kmask_indcs(img_shape_os, trj::AbstractVector{<:AbstractMatri
     return kmask_indcs
 end
 
-function calculateToeplitzKernelBasis(img_shape_os, trj::AbstractVector{<:AbstractMatrix{T}}, U::AbstractMatrix{Tc}; verbose = false) where {T, Tc <: Union{T, Complex{T}}}
+function calculateToeplitzKernelBasis(img_shape_os, trj::AbstractVector{<:AbstractMatrix{T}}, U::AbstractArray{Tc}; verbose = false) where {T, Tc <: Union{T, Complex{T}}}
 
     kmask_indcs = calculate_kmask_indcs(img_shape_os, trj)
     @assert all(kmask_indcs .> 0) # ensure that kmask is not out of bound
@@ -106,7 +106,7 @@ function calculateToeplitzKernelBasis(img_shape_os, trj::AbstractVector{<:Abstra
     λ2 = similar(λ)
     λ3 = similar(λ)
     Λ  = Array{Complex{T}}(undef, Ncoeff, Ncoeff, length(kmask_indcs))
-    
+
     trj_l = [size(trj[it],2) for it in eachindex(trj)]
     S  = Vector{Complex{T}}(undef, sum(trj_l))