Merge pull request #30 from MagneticResonanceImaging/grog_interfacere…

…write

WIP: re-write GROG/Cartesian interface

src/FFTNormalOpBasisFunc.jl

@@ -1,22 +1,38 @@
-function calculateKernelBasis(img_shape, D::AbstractArray{G}, U::Matrix{Complex{T}}; verbose = false) where {G,T}
+function calculateKernelBasis(img_shape, trj, U)
+    Ncoeff = size(U, 2)
+    Nt = length(trj) # number of time points
+    @assert Nt == size(U, 1) "Mismatch between trajectory and basis"
+
+    Λ = zeros(eltype(U), Ncoeff, Ncoeff, img_shape...)
+
+    for it ∈ eachindex(trj), ix ∈ axes(trj[it], 2)
+        k_idx = ntuple(j -> mod1(Int(trj[it][j, ix]) - img_shape[j] ÷ 2, img_shape[j]), length(img_shape)) # incorporates ifftshift
+        k_idx = CartesianIndex(k_idx)
 
-    Ncoeff = size(U,2)
-    Λ = Array{Complex{T}}(undef, Ncoeff, Ncoeff, img_shape...)
-    t = @elapsed begin
-        Threads.@threads for i ∈ CartesianIndices(img_shape) # takes 0.5s for 2D
-            Λ[:,:,i] .= U' * (D[i,:] .* U) #U' * diagm(D) * U
+        for ic ∈ CartesianIndices((Ncoeff, Ncoeff))
+            Λ[ic[1], ic[2], k_idx] += conj(U[it, ic[1]]) * U[it, ic[2]]
         end
-        Λ .= ifftshift(Λ, 3:(3+length(img_shape)-1)) #could fftshift D first
     end
-    verbose && println("Kernel calculation: t = $t s"); flush(stdout)
+    return Λ
+end
+
+function calculateKernelBasis(D, U)
+    Ncoeff = size(U, 2)
+    img_shape = size(D)[1:end-1]
+    Λ = Array{eltype(U)}(undef, Ncoeff, Ncoeff, img_shape...)
+
+    D .= ifftshift(D, 1:length(img_shape))
+    Threads.@threads for i ∈ CartesianIndices(img_shape)
+        Λ[:, :, i] .= U' * (D[i, :] .* U) #U' * diagm(D) * U
+    end
 
     return Λ
 end
 
 ## ##########################################################################
-# FFTNormalOpBasisFunc
+# FFTNormalOpBasis
 #############################################################################
-struct FFTNormalOpBasisFunc{S,T,N,E,F,G}
+struct _FFTNormalOpBasis{S,T,N,E,F,G}
     shape::S
     Ncoeff::Int
     fftplan::E
@@ -28,31 +44,45 @@ struct FFTNormalOpBasisFunc{S,T,N,E,F,G}
     cmaps::G
 end
 
-function FFTNormalOpBasisFunc(
-    img_shape,
-    U::Matrix{Complex{T}};
-    cmaps = (1,),
-    verbose = false,
-    D::AbstractArray{G} = ones(Int8, img_shape..., size(U,1)),
-    Λ = calculateKernelBasis(img_shape, D, U; verbose = verbose),
-    ) where {G,T}
+function FFTNormalOpBasis(img_shape, trj, U; cmaps=(1,))
+    Λ = calculateKernelBasis(img_shape, trj, U)
+    return FFTNormalOpBasis(Λ; cmaps)
+end
 
-    Ncoeff = size(U, 2)
-    kL1 = Array{Complex{T}}(undef, img_shape..., Ncoeff)
+function FFTNormalOpBasis(D, U; cmaps=(1,))
+    Λ = calculateKernelBasis(D, U)
+    return FFTNormalOpBasis(Λ; cmaps)
+end
+
+function FFTNormalOpBasis(Λ; cmaps=(1,))
+    Ncoeff = size(Λ, 1)
+    img_shape = size(Λ)[3:end]
+    kL1 = Array{eltype(Λ)}(undef, img_shape..., Ncoeff)
     kL2 = similar(kL1)
 
-    @views kmask = (Λ[1,1,CartesianIndices(img_shape)] .!= 0)
+    @views kmask = (Λ[1, 1, CartesianIndices(img_shape)] .!= 0)
     kmask_indcs = findall(vec(kmask))
     Λ = reshape(Λ, Ncoeff, Ncoeff, :)
-    Λ = Λ[:,:,kmask_indcs]
+    Λ = Λ[:, :, kmask_indcs]
+
+    ktmp = @view kL1[CartesianIndices(img_shape), 1]
+    fftplan = plan_fft!(ktmp; flags=FFTW.MEASURE, num_threads=round(Int, Threads.nthreads() / Ncoeff))
+    ifftplan = plan_ifft!(ktmp; flags=FFTW.MEASURE, num_threads=round(Int, Threads.nthreads() / Ncoeff))
+    A = _FFTNormalOpBasis(img_shape, Ncoeff, fftplan, ifftplan, Λ, kmask_indcs, kL1, kL2, cmaps)
 
-    ktmp = @view kL1[CartesianIndices(img_shape),1]
-    fftplan = plan_fft!( ktmp; flags = FFTW.MEASURE, num_threads=round(Int, Threads.nthreads()/Ncoeff))
-    ifftplan = plan_ifft!(ktmp; flags = FFTW.MEASURE, num_threads=round(Int, Threads.nthreads()/Ncoeff))
-    return FFTNormalOpBasisFunc(img_shape, Ncoeff, fftplan, ifftplan, Λ, kmask_indcs, kL1, kL2, cmaps)
+    return LinearOperator(
+        eltype(Λ),
+        prod(A.shape) * A.Ncoeff,
+        prod(A.shape) * A.Ncoeff,
+        true,
+        true,
+        (res, x, α, β) -> mul!(res, A, x, α, β),
+        nothing,
+        (res, x, α, β) -> mul!(res, A, x, α, β),
+    )
 end
 
-function LinearAlgebra.mul!(x::Vector{T}, S::FFTNormalOpBasisFunc, b, α, β) where {T}
+function LinearAlgebra.mul!(x::Vector{T}, S::_FFTNormalOpBasis, b, α, β) where {T}
     idx = CartesianIndices(S.shape)
 
     b = reshape(b, S.shape..., S.Ncoeff)
@@ -83,46 +113,11 @@ function LinearAlgebra.mul!(x::Vector{T}, S::FFTNormalOpBasisFunc, b, α, β) wh
 
             Threads.@threads for i ∈ 1:S.Ncoeff # multiply by C' and F'
                 @views S.ifftplan * S.kL2[idx, i]
-                @views xr[idx,i] .+= α .* conj.(cmap) .* S.kL2[idx,i]
+                @views xr[idx, i] .+= α .* conj.(cmap) .* S.kL2[idx, i]
             end
         end
     finally
         BLAS.set_num_threads(bthreads)
     end
     return x
-end
-
-Base.:*(S::FFTNormalOpBasisFunc, b::AbstractVector) = mul!(similar(b), S, b)
-Base.size(S::FFTNormalOpBasisFunc) = S.shape
-Base.size(S::FFTNormalOpBasisFunc, dim) = S.shape[dim]
-Base.eltype(::Type{FFTNormalOpBasisFunc{S,T,N,E,F,G}}) where {S,T,N,E,F,G} = T
-
-
-## ##########################################################################
-# LinearOperator of FFTNormalOpBasisFunc
-#############################################################################
-function FFTNormalOpBasisFuncLO(A::FFTNormalOpBasisFunc{S,T,N,E,F,G}) where {S,T,N,E,F,G}
-    return LinearOperator(
-        Complex{T},
-        prod(A.shape) * A.Ncoeff,
-        prod(A.shape) * A.Ncoeff,
-        true,
-        true,
-        (res, x, α, β) -> mul!(res, A, x, α, β),
-        nothing,
-        (res, x, α, β) -> mul!(res, A, x, α, β),
-    )
-end
-
-function FFTNormalOpBasisFuncLO(
-    img_shape,
-    U::Matrix{Complex{T}};
-    cmaps = (1,),
-    verbose = false,
-    D::AbstractArray{G} = ones(Int8, img_shape..., size(U,1)),
-    Λ = calculateKernelBasis(img_shape, D, U; verbose = verbose),
-    ) where {G,T}
-
-    S = FFTNormalOpBasisFunc(img_shape, U; cmaps = cmaps, D=D, Λ = Λ)
-    return FFTNormalOpBasisFuncLO(S)
-end
+end

src/GROG.jl

@@ -1,18 +1,17 @@
-function scGROG(data::AbstractArray{Complex{T}}, trj) where {T}
-    # self-calibrating radial GROG
-    # doi: 10.1002/mrm.21565
+function grog_calculatekernel(data, trj, Nr)
+    # self-calibrating radial GROG (http://doi.org/10.1002/mrm.21565)
 
-    # data should be passed with dimensions Nr x Ns x Ncoil
-    Nr = size(data, 1) #number of readout points
-    Ns = size(data, 2) # number of spokes across whole trajectory
     Ncoil = size(data, 3)
+    data = reshape(data, Nr, :, Ncoil)
+    Ns = size(data, 2) # number of spokes across whole trajectory
     Nd = size(trj[1], 1) # number of dimensions
+
     @assert Nr > Ncoil "Ncoil < Nr, problem is ill posed"
     @assert Ns > Ncoil^2 "Number of spokes < Ncoil^2, problem is ill posed"
 
     # preallocations
-    lnG = Array{Complex{T}}(undef, Nd, Ncoil, Ncoil) #matrix of GROG operators
-    vθ = Array{Complex{T}}(undef, Ns, Ncoil, Ncoil)
+    lnG = Array{eltype(data)}(undef, Nd, Ncoil, Ncoil) #matrix of GROG operators
+    vθ = Array{eltype(data)}(undef, Ns, Ncoil, Ncoil)
 
     # 1) Precompute n, m for the trajectory
     trjr = reshape(combinedimsview(trj), Nd, Nr, :)
@@ -34,80 +33,62 @@ function scGROG(data::AbstractArray{Complex{T}}, trj) where {T}
     return lnG
 end
 
-function griddedBackProjection(data::AbstractArray{Complex{T}}, lnG, trj, U::Matrix{Complex{T}}, cmaps; density=false, verbose=false) where {T}
-    # performs GROG gridding, returns backprojection and kernels
-    # assumes data is passed with dimensions Nr x NCyc*Nt x Ncoil
+function grog_griddata!(data, trj, Nr, img_shape)
+    lnG = grog_calculatekernel(data, trj, Nr)
 
-    img_shape = size(cmaps[1])
-    Nr = size(data, 1) #number of readout points
-    Nt = length(trj) # number of time points
-    @assert Nt == size(U, 1) "Mismatch between trajectory and basis"
-    Ncoeff = size(U, 2)
+    Ncoil = size(data, 3)
 
-    idx = CartesianIndices(img_shape)
-    Ncoil = length(cmaps)
+    Nt = length(trj) # number of time points
     data = reshape(data, :, Nt, Ncoil) # make sure data has correct size before gridding
 
     exp_method = ExpMethodHigham2005()
     cache = [ExponentialUtilities.alloc_mem(lnG[1], exp_method) for _ ∈ 1:Threads.nthreads()]
     lGcache = [similar(lnG[1]) for _ ∈ 1:Threads.nthreads()]
 
-    # gridding
-    t = @elapsed Threads.@threads for i ∈ CartesianIndices(@view data[:, :, 1])
-        idt = Threads.threadid()
-        for j ∈ length(img_shape):-1:1
+    Threads.@threads for i ∈ CartesianIndices(@view data[:, :, 1])
+        idt = Threads.threadid() # TODO: fix data race bug
+        for j ∈ eachindex(img_shape)
             trj_i = trj[i[2]][j, i[1]] * img_shape[j] + 1 / 2
             k_idx = round(trj_i)
             shift = (k_idx - trj_i) * Nr / img_shape[j]
+
+            # overwrite trj with rounded grid point index
             trj[i[2]][j, i[1]] = k_idx + img_shape[j] ÷ 2
 
+            # overwrite data with gridded data
             lGcache[idt] .= shift .* lnG[j]
             @views data[i, :] = exponential!(lGcache[idt], exp_method, cache[idt]) * data[i, :]
         end
     end
-    verbose && println("Gridding: t = $t s"); flush(stdout)
+end
 
-    # backprojection & kernel calculation
-    dataU = zeros(Complex{T}, img_shape..., Ncoil, Ncoeff)
-    Λ = zeros(Complex{T}, Ncoeff, Ncoeff, img_shape...)
-    if density
-        D = zeros(Int16, img_shape..., Nt)
-    end
+function calculateBackProjection_gridded(data, trj, U, cmaps)
+    Ncoil = length(cmaps)
+    Ncoeff = size(U, 2)
+    img_shape = size(cmaps[1])
+    img_idx = CartesianIndices(img_shape)
 
-    t = @elapsed for i ∈ CartesianIndices(@view data[:, :, 1])
-        k_idx = ntuple(j -> mod1(Int(trj[i[2]][j, i[1]]) - img_shape[j]÷2, img_shape[j]), length(img_shape)) # incorporates ifftshift
-        k_idx = CartesianIndex(k_idx)
+    Nt = length(trj)
+    @assert Nt == size(U, 1) "Mismatch between trajectory and basis"
+    data = reshape(data, :, Nt, Ncoil)
 
-        # multiply by basis for backprojection
-        for icoef ∈ axes(U, 2), icoil ∈ axes(data, 3)
-            @views dataU[k_idx, icoil, icoef] += data[i[1], i[2], icoil] * conj(U[i[2], icoef])
-        end
-        # add to kernel
-        for ic ∈ CartesianIndices((Ncoeff, Ncoeff))
-            Λ[ic[1], ic[2], k_idx] += conj(U[i[2], ic[1]]) * U[i[2], ic[2]]
-        end
-        if density
-            k_idx_D = CartesianIndex(ntuple(j -> Int(trj[i[2]][j, i[1]]), length(img_shape)))
-            D[k_idx_D, i[2]] += 1
-        end
-    end
-    verbose && println("Kernel calculation & back-projection time: t = $t s"); flush(stdout)
+    dataU = similar(data, img_shape..., Ncoeff)
+    xbp = zeros(eltype(data), img_shape..., Ncoeff)
 
-    # compute backprojection
-    xbp = zeros(Complex{T}, img_shape..., Ncoeff)
-    xbpci = [Array{Complex{T}}(undef, img_shape) for _ = 1:Threads.nthreads()]
     Threads.@threads for icoef ∈ axes(U, 2)
-        idt = Threads.threadid()
-        for icoil ∈ eachindex(cmaps)
-            @views ifft!(dataU[idx, icoil, icoef])
-            @views fftshift!(xbpci[idt], dataU[idx, icoil, icoef])
-            xbp[idx, icoef] .+= conj.(cmaps[icoil]) .* xbpci[idt]
-        end
-    end
+        for icoil ∈ axes(data, 3)
+            dataU[img_idx, icoef] .= 0
 
-    if density
-        return xbp, Λ, D
-    else
-        return xbp, Λ
+            for i ∈ CartesianIndices(@view data[:, :, 1])
+                k_idx = ntuple(j -> mod1(Int(trj[i[2]][j, i[1]]) - img_shape[j] ÷ 2, img_shape[j]), length(img_shape)) # incorporates ifftshift
+                k_idx = CartesianIndex(k_idx)
+
+                @views dataU[k_idx, icoef] += data[i[1], i[2], icoil] * conj(U[i[2], icoef])
+            end
+
+            @views ifft!(dataU[img_idx, icoef])
+            @views xbp[img_idx, icoef] .+= conj.(cmaps[icoil]) .* fftshift(dataU[img_idx, icoef])
+        end
     end
+    return xbp
 end

src/MRFingerprintingRecon.jl

@@ -10,7 +10,7 @@ using LinearOperators
 using SplitApplyCombine
 using ExponentialUtilities
 
-export FFTNormalOpBasisFunc, FFTNormalOpBasisFuncLO, scGROG, griddedBackProjection
+export FFTNormalOpBasis, grog_griddata!, calculateBackProjection_gridded
 export NFFTNormalOpBasisFunc, NFFTNormalOpBasisFuncLO, calcCoilMaps, calculateBackProjection, kooshball, kooshballGA
 
 function __init__()

test/reconstruct_cart.jl

@@ -52,9 +52,7 @@ for icoil = 1:Ncoil
     end
     data .= ifftshift(data, (1,2))
     fft!(data, (1,2))
-    data .= fftshift(data, (1,2))
-    data .*= D
-    data .= ifftshift(data, (1,2))
+    data .*= ifftshift(D, (1,2))
     ifft!(data, (1,2))
     data .= fftshift(data, (1,2))
     Threads.@threads for i ∈ CartesianIndices(@view x[:,:,1])
@@ -63,7 +61,7 @@ for icoil = 1:Ncoil
 end
 
 ## construct forward operator
-A = FFTNormalOpBasisFuncLO((Nx,Nx), U; cmaps=cmaps, D=D)
+A = FFTNormalOpBasis(D, U; cmaps)
 
 ## test forward operator is symmetric
 Λ = zeros(Complex{T}, Nc, Nc, Nx^2)