Merge pull request #56 from EHTJulia/rohandahale-multidomainnuft

NFFT for space, time and frequency domain

src/fourierdomain/nuft/dft_alg.jl

@@ -10,11 +10,13 @@ since it's easy to define derivatives.
 struct DFTAlg <: NUFT end
 
 # internal function that creates an DFT matrix/plan to use used for the img.
-function plan_nuft(::DFTAlg, imagegrid::AbstractRectiGrid, visdomain::UnstructuredDomain)
+function plan_nuft_spatial(::DFTAlg, imagegrid::AbstractRectiGrid,
+                           visdomain::UnstructuredDomain)
     visp = domainpoints(visdomain)
     (; X, Y) = imagegrid
     uv = domainpoints(visdomain)
-    dft = similar(Array{Complex{eltype(imagegrid)}}, length(visdomain), size(imagegrid)...)
+    dft = similar(Array{Complex{eltype(imagegrid)}}, length(visdomain),
+                  size(imagegrid)[1:2]...)
     @fastmath for i in eachindex(Y), j in eachindex(X), k in eachindex(visp)
         u = uv.U[k]
         v = uv.V[k]

src/fourierdomain/nuft/nfft_alg.jl

@@ -37,18 +37,41 @@ function applyft(p::AbstractNUFTPlan, img::AbstractArray)
     return vis
 end
 
-function plan_nuft(alg::NFFTAlg, imagegrid::AbstractRectiGrid,
-                   visdomain::UnstructuredDomain)
+# This a new function is overloaded to handle when NUFTPlan has plans
+# as dictionaries in the case of Ti or Fr case
+
+@inline function applyft(p::NUFTPlan{<:FourierTransform,<:AbstractDict},
+                         img::AbstractArray)
+    vis_list = similar(baseimage(img), Complex{eltype(img)}, p.totalvis)
+    plans = p.plan
+    iminds, visinds = p.indices
+    for i in eachindex(iminds, visinds)
+        imind = iminds[i]
+        visind = visinds[i]
+        # TODO
+        # If visinds are consecutive then we can use the in-place _nuft!:
+        # _nuft!(visind, plans[imind], @view(img[:, :, imind...])  
+        vis_inner = nuft(plans[imind], @view(img[:, :, imind]))
+        # After the todo this wont be required
+        vis_list[visind] .= vis_inner
+    end
+
+    vis_list .= vis_list .* p.phases
+    return vis_list
+end
+
+function plan_nuft_spatial(alg::NFFTAlg, imagegrid::AbstractRectiGrid,
+                           visdomain::UnstructuredDomain)
     visp = domainpoints(visdomain)
     uv2 = similar(visp.U, (2, length(visdomain)))
     dpx = pixelsizes(imagegrid)
     dx = dpx.X
     dy = dpx.Y
     # Here we flip the sign because the NFFT uses the -2pi convention
-    uv2[1, :] .= -visp.U * dx
-    uv2[2, :] .= -visp.V * dy
+    uv2[1, :] .= -visp.U .* dx
+    uv2[2, :] .= -visp.V .* dy
     (; reltol, precompute, fftflags) = alg
-    plan = plan_nfft(uv2, size(imagegrid); reltol, precompute, fftflags)
+    plan = plan_nfft(uv2, size(imagegrid)[1:2]; reltol, precompute, fftflags)
     return plan
 end
 

src/fourierdomain/nuft/nuft.jl

@@ -8,21 +8,93 @@ Internal type used to store the cache for a non-uniform Fourier transform (NUFT)
 
 The user should instead create this using the [`FourierDualDomain`](@ref) function.
 """
-struct NUFTPlan{A,P,M} <: AbstractNUFTPlan
+struct NUFTPlan{A,P,M,I,T} <: AbstractNUFTPlan
     alg::A # which algorithm to use
     plan::P #NUFT matrix or plan
     phases::M #FT phases needed to phase center things
+    indices::I # imgdomain Ti/Fr indices mapped to visdomain indices
+    totalvis::T # Total number of visibility points
+end
+
+# We do this for speed an readability since all seems to be very slow 
+_compare(nv::NamedTuple{N}, val) where {N} = mapreduce(n -> (nv[n] == val[n]), *, N)
+
+# creates the indexing plan for the multidomain nuft. Returns a tuple with 
+# iminds whose elements defines the index in the image domain  
+# visinds whose elements are the indices of the visibilities that correspond to that imind  
+# The order of data is currently set by imgdomain. 
+
+# The order of data is currently set by imgdomain. 
+function plan_indices(imgdomain::AbstractRectiGrid, visdomain::UnstructuredDomain)
+    # TODO: Change the ordering so that visdomain is accessed in a constant stride so  
+    # we can utilize in-place nuft and save a bunch of allocations
+    spatialdims = ComradeBase.dims(imgdomain)[3:end]
+    nms = map(name, spatialdims)
+
+    # DimPoints stack overflows for an empty tuple  
+    isempty(spatialdims) && return (0, 0)
+
+    itr = pairs(DimPoints(spatialdims))
+    T = typeof(first(first(itr)))
+    iminds = T[]
+    visinds = Vector{Int}[]
+    visp = domainpoints(visdomain)
+    for (i, vals) in itr
+        nv = NamedTuple{nms}(vals)
+        push!(iminds, i)
+        push!(visinds, findall(p -> _compare(nv, p), visp))
+    end
+    return iminds, visinds
+end
+
+function plan_nuft(alg::NUFT, imagegrid::AbstractRectiGrid,
+                   visdomain::UnstructuredDomain, indices)
+    # check_image_uv(imagegrid, visdomain) 
+    # Check if Ti or Fr in visdomain are subset of imgdomain Ti or Fr if present
+    points = domainpoints(visdomain)
+    iminds, visinds = indices
+
+    uv = UnstructuredDomain(points[visinds[1]], executor(visdomain), header(visdomain))
+    tplan = plan_nuft_spatial(alg, imagegrid, uv)
+    plans = Dict{typeof(iminds[1]),typeof(tplan)}()
+
+    for i in eachindex(iminds, visinds)
+        imind = iminds[i]
+        visind = visinds[i]
+        uv = UnstructuredDomain(points[visind], executor(visdomain), header(visdomain))
+        plans[imind] = plan_nuft_spatial(alg, imagegrid, uv)
+    end
+    return plans
 end
 
 function create_forward_plan(algorithm::NUFT, imgdomain::AbstractRectiGrid,
                              visdomain::UnstructuredDomain)
-    plan = plan_nuft(algorithm, imgdomain, visdomain)
     phases = make_phases(algorithm, imgdomain, visdomain)
-    return NUFTPlan(algorithm, plan, phases)
+    indices = plan_indices(imgdomain, visdomain)
+    if hasproperty(imgdomain, :Ti) || hasproperty(imgdomain, :Fr)
+        plan = plan_nuft(algorithm, imgdomain, visdomain, indices)
+    else
+        plan = plan_nuft_spatial(algorithm, imgdomain, visdomain)
+    end
+    return NUFTPlan(algorithm, plan, phases, indices, size(visdomain)[1])
 end
 
 function inverse_plan(plan::NUFTPlan)
-    return NUFTPlan(plan.alg, plan.plan', inv.(plan.phases))
+    return NUFTPlan(plan.alg, plan.plan', inv.(plan.phases), plan.indices, plan.totalvis)
+end
+
+function inverse_plan(plan::NUFTPlan{<:FourierTransform,<:AbstractDict})
+    iminds, visinds = plan.indices
+
+    inverse_plans_t = plan.plan[iminds[1]]'
+    inverse_plans = Dict{typeof(iminds[1]),typeof(inverse_plans_t)}()
+
+    for i in eachindex(iminds, visinds)
+        imind = iminds[i]
+        inverse_plans[imind] = plan.plan[imind]'
+    end
+
+    return NUFTPlan(plan.alg, inverse_plans, inv.(plan.phases), plan.indices, plan.totalvis)
 end
 
 @inline function nuft(A, b::IntensityMap)

test/Project.toml

@@ -1,4 +1,5 @@
 [deps]
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"