Skip to content

Commit cb3cac0

Browse files
authored
Merge pull request #160 from JuliaMath/jf/sdc-real
Streamline `sdc` for real sample density compensation factors (DCFs)
2 parents d04aea8 + 012688b commit cb3cac0

File tree

9 files changed

+271
-101
lines changed

9 files changed

+271
-101
lines changed

AbstractNFFTs/src/interface.jl

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -189,7 +189,7 @@ The transformation is applied along `D-R+1` dimensions specified in the plan `p`
189189
size_in(p)
190190
191191
Size of the input array for the plan p (NFFT/NFCT/NFST/NNFFT).
192-
The returned tuple has `R` entries.
192+
The returned tuple has `D` entries.
193193
Note that this will be the output size for the transposed / adjoint operator.
194194
"""
195195
@mustimplement size_in(p::AbstractFTPlan{T,D,R}) where {T,D,R}
@@ -222,10 +222,10 @@ Change nodes `k` in the plan `p` operation and return the plan.
222222
convolve!(p::AbstractNFFTPlan{T,D,R}, g::AbstractArray, fHat::AbstractArray)
223223
224224
Overwrite `R`-dimensional array `fHat`
225-
(often R=1, and `fHat` has length `p.J`)
225+
(often R=1, and `fHat` has length `only(p.size_in)`)
226226
with the result of "convolution" (i.e., interpolation)
227227
between `D`-dimensional equispaced input array `g`
228-
(often of size `p.Ñ`)
228+
(often of size `p.Ñ` which is not part of abstract interface)
229229
and the interpolation kernel in NFFT plan `p`.
230230
"""
231231
@mustimplement convolve!(p::AbstractNFFTPlan, g::AbstractArray, fHat::AbstractArray)

NFFTTools/Project.toml

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,18 @@
11
name = "NFFTTools"
22
uuid = "7424e34d-94f7-41d6-98a0-85abaf1b6c91"
3-
author = ["Tobias Knopp <tobias@knoppweb.de>"]
4-
version = "0.2.7"
3+
author = ["Tobias Knopp <tobias@knoppweb.de> and contributors"]
4+
version = "0.3"
55

66
[deps]
77
AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
88
AbstractNFFTs = "7f219486-4aa7-41d6-80a7-e08ef20ceed7"
99
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
1010
FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
11+
NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d"
1112

1213
[compat]
1314
julia = "1.6"
1415
AbstractFFTs = "1.0"
1516
AbstractNFFTs = "0.6, 0.7, 0.8, 0.9"
1617
FFTW = "1"
18+
NFFT = "0.14"

NFFTTools/src/NFFTTools.jl

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,12 @@ module NFFTTools
33
export sdc
44
export calculateToeplitzKernel, calculateToeplitzKernel!, convolveToeplitzKernel!
55

6-
using AbstractNFFTs: AbstractNFFTPlan, plan_nfft, nodes!
6+
using AbstractNFFTs: AbstractNFFTPlan, plan_nfft, nodes!, size_in, size_out
77
using AbstractNFFTs: convolve!, convolve_transpose!
8+
#using NFFT: NFFTPlan
89
using FFTW: fftshift, plan_fft, plan_ifft
910
using LinearAlgebra: adjoint, mul!
10-
import FFTW # ESTIMATE (which is currently non-public)
11+
import FFTW # ESTIMATE
1112

1213
include("samplingDensity.jl")
1314
include("Toeplitz.jl")

NFFTTools/src/samplingDensity.jl

Lines changed: 142 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,49 +1,155 @@
1-
21
#=
32
using AbstractNFFTs: AbstractNFFTPlan, convolve!, convolve_transpose!
3+
using AbstractNFFTs: size_in, size_out
44
using LinearAlgebra: mul!
55
=#
66

7-
function sdc(p::AbstractNFFTPlan{T,D,1}; iters=20) where {T,D}
8-
# Weights for sample density compensation.
9-
# Uses method of Pipe & Menon, 1999. Mag Reson Med, 186, 179.
10-
weights = similar(p.tmpVec, Complex{T}, p.J)
11-
weights .= one(Complex{T})
12-
weights_tmp = similar(weights)
13-
scaling_factor = zero(T)
7+
"""
8+
out = _fill_similar(array, v::T, dims) where T
9+
Return an allocated array similar to `array` filled with value `v`.
10+
11+
This 2-step initialization helper function
12+
is needed to accommodate GPU array types.
13+
The more obvious statement `weights = fill(v, dims)`
14+
can lead to the wrong array type and cause GPU tests to fail.
15+
"""
16+
function _fill_similar(array::AbstractArray, v::T, dims::Union{Integer,Dims}) where T
17+
weights = similar(array, T, dims)
18+
fill!(weights, v)
19+
return weights
20+
end
21+
22+
#=
23+
It would be desirable to use the memory pointed to by p.tmpVec
24+
as working buffer for sdc iterations,
25+
but this attempt led to intermittent corrupted outputs and errors.
26+
So it is left commented out for now.
27+
=#
28+
#=
29+
function _reinterpret_real(g::StridedArray{Complex{T}}) where {T <: Real}
30+
r1 = reinterpret(T, g)
31+
r2 = @view r1[1:2:end,:]
32+
return r2
33+
end
34+
=#
35+
36+
"""
37+
weights = sdc(plan::{Abstract}NFFTPlan; iters=20, ...)
38+
39+
Compute weights for sample density compensation for given NFFT `plan`
40+
Uses method of Pipe & Menon, Mag Reson Med, 44(1):179-186, Jan. 1999.
41+
DOI: 10.1002/(SICI)1522-2594(199901)41:1<179::AID-MRM25>3.0.CO;2-V
42+
The function applies `iters` iterations of that method.
43+
44+
The scaling here such that if the plan were 1D with N nodes equispaced
45+
from -0.5 to 0.5-1/N, then the returned weights are ≈ 1/N.
46+
47+
The weights are scaled such that ``A' diag(w) A 1_N ≈ 1_N``.
48+
49+
The returned vector is real, positive values of length `plan.J`.
50+
51+
This method _almost_ conforms to the `AbstractNFFT` interface
52+
except that it uses `p.Ñ` and `p.tmpVec` that are not part of that interface.
53+
54+
There are several named keyword arguments that are work buffers
55+
that are all mutated: `weights weights_tmp workg workf workv`.
56+
If the caller provides all of those,
57+
then this function should make only small allocations.
58+
"""
59+
function sdc(
60+
p::AbstractNFFTPlan{T,D,1};
61+
iters::Int = 20,
62+
# the type of p.tmpVec (if present) is needed for GPU arrays (JLArray):
63+
_array::AbstractArray = hasfield(typeof(p), :tmpVec) ?
64+
p.tmpVec : Array{T,D}(undef, zeros(Int, D)...),
65+
# the following are working buffers that are all mutated:
66+
weights::AbstractVector{T} = _fill_similar(_array, one(T), only(size_out(p))),
67+
weights_tmp::AbstractVector = similar(weights),
68+
# workg::AbstractArray = _reinterpret_real(p.tmpVec), # todo
69+
workg::AbstractArray = similar(_array, T, p.Ñ),
70+
workf::AbstractVector = similar(_array, Complex{T}, only(size_out(p))),
71+
workv::AbstractArray = similar(_array, Complex{T}, size_in(p)),
72+
) where {T <: Real, D}
73+
74+
return sdc!(
75+
p,
76+
iters,
77+
weights,
78+
weights_tmp,
79+
workg,
80+
workf,
81+
workv,
82+
)
83+
end
84+
85+
86+
"""
87+
weights = sdc!( p::AbstractNFFTPlan, iters, weights, weights_tmp, workg)
88+
Compute sampling density compensation `weights`
89+
without performing any final scaling step.
90+
91+
Ideally this function should be (nearly) non-allocating.
92+
"""
93+
function sdc!(
94+
p::AbstractNFFTPlan{T,D,1},
95+
iters::Int,
96+
# the following are working buffers that are all mutated:
97+
weights::AbstractVector{T},
98+
weights_tmp::AbstractVector,
99+
workg::AbstractArray,
100+
) where {T <: Real, D}
101+
102+
scaling_factor = missing # will be set below
14103

15104
# Pre-weighting to correct non-uniform sample density
16105
for i in 1:iters
17-
convolve_transpose!(p, weights, p.tmpVec)
18-
if i==1
19-
scaling_factor = maximum(abs.(p.tmpVec))
106+
convolve_transpose!(p, weights, workg)
107+
if i == 1
108+
scaling_factor = maximum(workg)
20109
end
21110

22-
p.tmpVec ./= scaling_factor
23-
convolve!(p, p.tmpVec, weights_tmp)
111+
workg ./= scaling_factor
112+
convolve!(p, workg, weights_tmp)
24113
weights_tmp ./= scaling_factor
25-
weights ./= (abs.(weights_tmp) .+ eps(T))
114+
any((0), weights_tmp) && throw("non-positive weights")
115+
weights ./= weights_tmp
26116
end
27-
# Post weights to correct image scaling
28-
# This finds c, where ||u - c*v||_2^2 = 0 and then uses
29-
# c to scale all weights by a scalar factor.
30-
u = similar(weights, Complex{T}, p.N)
31-
u .= one(Complex{T})
32-
33-
# conversion to Array is a workaround for CuNFFT. Without it we get strange
34-
# results that indicate some synchronization issue
35-
#f = Array( p * u )
36-
#b = f .* Array(weights) # apply weights from above
37-
#v = Array( adjoint(p) * convert(typeof(weights), b) )
38-
#c = vec(v) \ vec(Array(u)) # least squares diff
39-
#return abs.(convert(typeof(weights), c * Array(weights)))
40-
41-
# non converting version
42-
f = similar(p.tmpVec, Complex{T}, p.J)
43-
mul!(f, p, u)
44-
f .*= weights # apply weights from above
45-
v = similar(p.tmpVec, Complex{T}, p.N)
46-
mul!(v, adjoint(p), f)
47-
c = vec(v) \ vec(u) # least squares diff
48-
return abs.(c * weights)
117+
return weights
118+
end
119+
120+
121+
"""
122+
sdc!( ... )
123+
124+
This version scales the weights such that
125+
``A' D(w) A 1_N ≈ 1_N``.
126+
Find ``c ∈ ℝ`` that minimizes ``‖u - c * v‖₂``
127+
where ``u ≜ 1_N`` (vector of ones)
128+
and ``v ≜ A' D(w) A 1_N``, and then scale `w` by that `c`.
129+
The analytical solution is
130+
``c = real(u'v) / ‖v‖² = real(sum(v)) / ‖v‖².``
131+
132+
Ideally this function should be (nearly) non-allocating.
133+
"""
134+
function sdc!(
135+
p::AbstractNFFTPlan{T,D,1},
136+
iters::Int,
137+
# the following are working buffers that are all mutated:
138+
weights::AbstractVector{T},
139+
weights_tmp::AbstractVector,
140+
workg::AbstractArray,
141+
workf::AbstractVector,
142+
workv::AbstractArray,
143+
) where {T <: Real, D}
144+
145+
sdc!(p, iters, weights, weights_tmp, workg)
146+
147+
u = workv # trick to save memory
148+
fill!(u, one(T))
149+
mul!(workf, p, u)
150+
workf .*= weights # apply weights from above
151+
mul!(workv, adjoint(p), workf)
152+
c = real(sum(workv)) / sum(abs2, workv)
153+
weights .*= c
154+
return weights
49155
end

Project.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
name = "NFFT"
22
uuid = "efe261a4-0d2b-5849-be55-fc731d526b0d"
33
authors = ["Tobias Knopp <tobias@knoppweb.de> and contributors"]
4-
version = "0.14.2"
4+
version = "0.14.3"
55

66
[deps]
77
AbstractNFFTs = "7f219486-4aa7-41d6-80a7-e08ef20ceed7"

ext/NFFTGPUArraysExt/implementation.jl

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
const RealOrComplex = Union{T, Complex{T}} where {T <: Real}
2+
13
mutable struct GPU_NFFTPlan{T,D, arrTc <: AbstractGPUArray{Complex{T}, D}, vecI <: AbstractGPUVector{Int32}, FP, BP, INV, SM} <: AbstractNFFTPlan{T,D,1}
24
N::NTuple{D,Int64}
35
NOut::NTuple{1,Int64}
@@ -60,14 +62,21 @@ AbstractNFFTs.size_in(p::GPU_NFFTPlan) = p.N
6062
AbstractNFFTs.size_out(p::GPU_NFFTPlan) = p.NOut
6163

6264

63-
function AbstractNFFTs.convolve!(p::GPU_NFFTPlan{T,D, arrTc}, g::arrTc, fHat::arrH) where {D,T,arr<: AbstractGPUArray, arrTc <: arr, arrH <: arr}
65+
function AbstractNFFTs.convolve!(
66+
p::GPU_NFFTPlan{<:Real, D, <:AbstractGPUArray},
67+
g::AbstractGPUArray{<:RealOrComplex, D},
68+
fHat::AbstractGPUVector{<:RealOrComplex},
69+
) where {D}
6470
mul!(fHat, transpose(p.B), vec(g))
65-
return
6671
end
6772

68-
function AbstractNFFTs.convolve_transpose!(p::GPU_NFFTPlan{T,D, arrTc}, fHat::arrH, g::arrTc) where {D,T,arr<: AbstractGPUArray, arrTc <: arr, arrH <: arr}
73+
function AbstractNFFTs.convolve_transpose!(
74+
p::GPU_NFFTPlan{<:Real, D, <:AbstractGPUArray},
75+
fHat::AbstractGPUVector{<:RealOrComplex},
76+
g::AbstractGPUArray{<:RealOrComplex, D},
77+
) where {D}
6978
mul!(vec(g), p.B, fHat)
70-
return
79+
return g
7180
end
7281

7382
function AbstractNFFTs.deconvolve!(p::GPU_NFFTPlan{T,D, arrTc}, f::arrF, g::arrTc) where {D,T,arr<: AbstractGPUArray, arrTc <: arr, arrF <: arr}
@@ -125,4 +134,3 @@ function LinearAlgebra.mul!(f::arrF, pl::Adjoint{Complex{T},<:GPU_NFFTPlan{T,D,
125134

126135
return f
127136
end
128-

0 commit comments

Comments
 (0)