/
_fourier.py
253 lines (209 loc) · 9.4 KB
/
_fourier.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
import numpy
import cupy
from cupy import _core
from cupy._core import internal
from cupyx.scipy.ndimage import _util
from cupyx.scipy import special
def _get_output_fourier(output, input, complex_only=False):
types = [cupy.complex64, cupy.complex128]
if not complex_only:
types += [cupy.float32, cupy.float64]
if output is None:
if input.dtype in types:
output = cupy.zeros(input.shape, dtype=input.dtype)
else:
output = cupy.zeros(input.shape, dtype=types[-1])
elif type(output) is type:
if output not in types:
raise RuntimeError('output type not supported')
output = cupy.zeros(input.shape, dtype=output)
elif output.shape != input.shape:
raise RuntimeError('output shape not correct')
return output
def _reshape_nd(arr, ndim, axis):
"""Promote a 1d array to ndim with non-singleton size along axis."""
nd_shape = (1,) * axis + (arr.size,) + (1,) * (ndim - axis - 1)
return arr.reshape(nd_shape)
def fourier_gaussian(input, sigma, n=-1, axis=-1, output=None):
"""Multidimensional Gaussian shift filter.
The array is multiplied with the Fourier transform of a (separable)
Gaussian kernel.
Args:
input (cupy.ndarray): The input array.
sigma (float or sequence of float): The sigma of the Gaussian kernel.
If a float, `sigma` is the same for all axes. If a sequence,
`sigma` has to contain one value for each axis.
n (int, optional): If `n` is negative (default), then the input is
assumed to be the result of a complex fft. If `n` is larger than or
equal to zero, the input is assumed to be the result of a real fft,
and `n` gives the length of the array before transformation along
the real transform direction.
axis (int, optional): The axis of the real transform (only used when
``n > -1``).
output (cupy.ndarray, optional):
If given, the result of shifting the input is placed in this array.
Returns:
output (cupy.ndarray): The filtered output.
"""
ndim = input.ndim
output = _get_output_fourier(output, input)
axis = internal._normalize_axis_index(axis, ndim)
sigmas = _util._fix_sequence_arg(sigma, ndim, 'sigma')
_core.elementwise_copy(input, output)
for ax, (sigmak, ax_size) in enumerate(zip(sigmas, output.shape)):
# compute the frequency grid in Hz
if ax == axis and n > 0:
arr = cupy.arange(ax_size, dtype=output.real.dtype)
arr /= n
else:
arr = cupy.fft.fftfreq(ax_size)
arr = arr.astype(output.real.dtype, copy=False)
# compute the Gaussian weights
arr *= arr
scale = sigmak * sigmak / -2
arr *= (4 * numpy.pi * numpy.pi) * scale
cupy.exp(arr, out=arr)
# reshape for broadcasting
arr = _reshape_nd(arr, ndim=ndim, axis=ax)
output *= arr
return output
def fourier_uniform(input, size, n=-1, axis=-1, output=None):
"""Multidimensional uniform shift filter.
The array is multiplied with the Fourier transform of a box of given size.
Args:
input (cupy.ndarray): The input array.
size (float or sequence of float): The sigma of the box used for
filtering. If a float, `size` is the same for all axes. If a
sequence, `size` has to contain one value for each axis.
n (int, optional): If `n` is negative (default), then the input is
assumed to be the result of a complex fft. If `n` is larger than or
equal to zero, the input is assumed to be the result of a real fft,
and `n` gives the length of the array before transformation along
the real transform direction.
axis (int, optional): The axis of the real transform (only used when
``n > -1``).
output (cupy.ndarray, optional):
If given, the result of shifting the input is placed in this array.
Returns:
output (cupy.ndarray): The filtered output.
"""
ndim = input.ndim
output = _get_output_fourier(output, input)
axis = internal._normalize_axis_index(axis, ndim)
sizes = _util._fix_sequence_arg(size, ndim, 'size')
_core.elementwise_copy(input, output)
for ax, (size, ax_size) in enumerate(zip(sizes, output.shape)):
# compute the frequency grid in Hz
if ax == axis and n > 0:
arr = cupy.arange(ax_size, dtype=output.real.dtype)
arr /= n
else:
arr = cupy.fft.fftfreq(ax_size)
arr = arr.astype(output.real.dtype, copy=False)
# compute the uniform filter weights
arr *= size
cupy.sinc(arr, out=arr)
# reshape for broadcasting
arr = _reshape_nd(arr, ndim=ndim, axis=ax)
output *= arr
return output
def fourier_shift(input, shift, n=-1, axis=-1, output=None):
"""Multidimensional Fourier shift filter.
The array is multiplied with the Fourier transform of a shift operation.
Args:
input (cupy.ndarray): The input array. This should be in the Fourier
domain.
shift (float or sequence of float): The size of shift. If a float,
`shift` is the same for all axes. If a sequence, `shift` has to
contain one value for each axis.
n (int, optional): If `n` is negative (default), then the input is
assumed to be the result of a complex fft. If `n` is larger than or
equal to zero, the input is assumed to be the result of a real fft,
and `n` gives the length of the array before transformation along
the real transform direction.
axis (int, optional): The axis of the real transform (only used when
``n > -1``).
output (cupy.ndarray, optional):
If given, the result of shifting the input is placed in this array.
Returns:
output (cupy.ndarray): The shifted output (in the Fourier domain).
"""
ndim = input.ndim
output = _get_output_fourier(output, input, complex_only=True)
axis = internal._normalize_axis_index(axis, ndim)
shifts = _util._fix_sequence_arg(shift, ndim, 'shift')
_core.elementwise_copy(input, output)
for ax, (shiftk, ax_size) in enumerate(zip(shifts, output.shape)):
if shiftk == 0:
continue
if ax == axis and n > 0:
# cp.fft.rfftfreq(ax_size) * (-2j * numpy.pi * shiftk * ax_size/n)
arr = cupy.arange(ax_size, dtype=output.dtype)
arr *= -2j * numpy.pi * shiftk / n
else:
arr = cupy.fft.fftfreq(ax_size)
arr = arr * (-2j * numpy.pi * shiftk)
cupy.exp(arr, out=arr)
# reshape for broadcasting
arr = _reshape_nd(arr, ndim=ndim, axis=ax)
output *= arr
return output
def fourier_ellipsoid(input, size, n=-1, axis=-1, output=None):
"""Multidimensional ellipsoid Fourier filter.
The array is multiplied with the fourier transform of a ellipsoid of
given sizes.
Args:
input (cupy.ndarray): The input array.
size (float or sequence of float): The size of the box used for
filtering. If a float, `size` is the same for all axes. If a
sequence, `size` has to contain one value for each axis.
n (int, optional): If `n` is negative (default), then the input is
assumed to be the result of a complex fft. If `n` is larger than or
equal to zero, the input is assumed to be the result of a real fft,
and `n` gives the length of the array before transformation along
the real transform direction.
axis (int, optional): The axis of the real transform (only used when
``n > -1``).
output (cupy.ndarray, optional):
If given, the result of shifting the input is placed in this array.
Returns:
output (cupy.ndarray): The filtered output.
"""
ndim = input.ndim
if ndim == 1:
return fourier_uniform(input, size, n, axis, output)
if ndim > 3:
# Note: SciPy currently does not do any filtering on >=4d inputs, but
# does not warn about this!
raise NotImplementedError('Only 1d, 2d and 3d inputs are supported')
output = _get_output_fourier(output, input)
axis = internal._normalize_axis_index(axis, ndim)
sizes = _util._fix_sequence_arg(size, ndim, 'size')
_core.elementwise_copy(input, output)
# compute the distance from the origin for all samples in Fourier space
distance = 0
for ax, (size, ax_size) in enumerate(zip(sizes, output.shape)):
# compute the frequency grid in Hz
if ax == axis and n > 0:
arr = cupy.arange(ax_size, dtype=output.real.dtype)
arr *= numpy.pi * size / n
else:
arr = cupy.fft.fftfreq(ax_size)
arr *= numpy.pi * size
arr = arr.astype(output.real.dtype, copy=False)
arr *= arr
arr = _reshape_nd(arr, ndim=ndim, axis=ax)
distance = distance + arr
cupy.sqrt(distance, out=distance)
if ndim == 2:
special.j1(distance, out=output)
output *= 2
output /= distance
elif ndim == 3:
cupy.sin(distance, out=output)
output -= distance * cupy.cos(distance)
output *= 3
output /= distance ** 3
output[(0,) * ndim] = 1.0 # avoid NaN in corner at frequency=0 location
output *= input
return output