def test_scale_nd_axes(self):
X = np.empty((4, 2, 16, 8), dtype='d')
X.flat[:] = np.cbrt(np.arange(X.size, dtype=X.dtype))
f = mkl_fft.fftn(X, axes=(0, 1, 2, 3))
f_scale = mkl_fft.fftn(X, axes=(0, 1, 2, 3), forward_scale=0.2)
r_tol, a_tol = _get_rtol_atol(X)
assert_allclose(f, 5*f_scale, rtol=r_tol, atol=a_tol)
def test_scale_nd(self):
X = np.empty((2, 4, 8, 16), dtype='d')
X.flat[:] = np.cbrt(np.arange(0, X.size, dtype=X.dtype))
f = mkl_fft.fftn(X)
f_scale = mkl_fft.fftn(X, forward_scale=0.2)
r_tol, a_tol = _get_rtol_atol(X)
assert_allclose(f, 5*f_scale, rtol=r_tol, atol=a_tol)
def test_matrix4(self):
"""fftn of strided array is same as fftn of a contiguous copy"""
for ar in [self.md, self.mz, self.mf, self.mc]:
r_tol, a_tol = _get_rtol_atol(ar)
d_strided = ar[::2, ::2]
d_contig = d_strided.copy()
t_strided = mkl_fft.fftn(d_strided)
t_contig = mkl_fft.fftn(d_contig)
assert_allclose(t_strided, t_contig, rtol=r_tol, atol=a_tol)
def test_matrix5(self):
"""fftn of strided array is same as fftn of a contiguous copy"""
rs = rnd.RandomState(1234)
x = rs.randn(6, 11, 12, 13)
y = x[::-2, :, :, ::3]
r_tol, a_tol = _get_rtol_atol(y)
f = mkl_fft.fftn(y, axes=(1,2))
for i0 in range(y.shape[0]):
for i3 in range(y.shape[3]):
assert_allclose(
f[i0, :, :, i3],
mkl_fft.fftn(y[i0, :, : , i3]),
rtol=r_tol, atol=a_tol
)
def test_matrix3(self):
"""fftn(ifftn(x)) is x"""
for ar in [self.md, self.mz, self.mf, self.mc]:
d = ar.copy()
r_tol, a_tol = _get_rtol_atol(d)
t = mkl_fft.fftn(mkl_fft.ifftn(d))
assert_allclose(d, t, rtol=r_tol, atol=a_tol, err_msg = "failed test for dtype {}, max abs diff: {}".format(d.dtype, np.max(np.abs(d-t))))
def test_matrix1(self):
"""fftn equals repeated fft"""
for ar in [self.md, self.mz, self.mf, self.mc]:
r_tol, a_tol = _get_rtol_atol(ar)
d = ar.copy()
t1 = mkl_fft.fftn(d)
t2 = mkl_fft.fft(mkl_fft.fft(d, axis=0), axis=1)
t3 = mkl_fft.fft(mkl_fft.fft(d, axis=1), axis=0)
assert_allclose(t1, t2, rtol=r_tol, atol=a_tol, err_msg = "failed test for dtype {}, max abs diff: {}".format(d.dtype, np.max(np.abs(t1-t2))))
assert_allclose(t1, t3, rtol=r_tol, atol=a_tol, err_msg = "failed test for dtype {}, max abs diff: {}".format(d.dtype, np.max(np.abs(t1-t3))))
def measure_mkl_fft(a, nrepeat, nthr):
import os
os.environ['OMP_NUM_THREADS'] = str(nthr)
import mkl_fft
tmin = 1e38
for i in range(nrepeat):
t0 = time()
b = mkl_fft.fftn(a)
t1 = time()
tmin = min(tmin, t1 - t0)
return tmin, b
def propagate(self):
self.work.fill(0)
self.work_view1[:] = self.field1.field
for W in self.W1:
self.work_view1 *= W
self.work_F = mkl_fft.fftn(self.work, overwrite_x=True)
for WF in self.W12F_pad:
self.work_F *= WF
self.work[:] = mkl_fft.ifftn(self.work, overwrite_x=True)
for W in self.W2:
self.work_view2 *= W
self.field2.field = self.work_view2 * self.Sf
def propagate_backward(self):
self.work.fill(0)
self.work_view2[:] = self.field2.field
for W in self.W2:
self.work_view2 *= W.conj()
self.work_F = mkl_fft.fftn(self.work,
axes=range(-self.ndim, 0),
overwrite_x=True)
for WF in self.W12F_pad:
self.work_F *= WF.conj()
self.work[:] = mkl_fft.ifftn(self.work,
axes=range(-self.ndim, 0),
overwrite_x=True)
for W in self.W1:
self.work_view1 *= W.conj()
self.field1.field = self.work_view1 * self.Sb
def _operation(self, field, adjoint):
'''The internal filtering operation.
Parameters
----------
field : Field
The input field.
adjoint : boolean
Whether to perform a forward or adjoint filter.
Returns
-------
Field
The filtered field.
'''
self._compute_functions(field)
if self.cutout is None:
f = field.shaped
else:
f = self.internal_array
f[:] = 0
c = tuple([slice(None)] * field.tensor_order) + self.cutout
f[c] = field.shaped
# Don't overwrite f if it's the input array.
if _use_mkl:
kwargs = {
'overwrite_x': self.cutout is not None and field.grid.ndim > 1
}
else:
kwargs = {}
f = _fft_module.fftn(f,
axes=tuple(range(-self.input_grid.ndim, 0)),
**kwargs)
if (self._transfer_function.ndim - self.internal_grid.ndim) == 2:
# The transfer function is a matrix field.
s1 = f.shape[:-self.internal_grid.ndim] + (
self.internal_grid.size, )
f = Field(f.reshape(s1), self.internal_grid)
s2 = self._transfer_function.shape[:-self.internal_grid.ndim] + (
self.internal_grid.size, )
tf = Field(self._transfer_function.reshape(s2), self.internal_grid)
if adjoint:
tf = field_conjugate_transpose(tf)
f = field_dot(tf, f).shaped
else:
# The transfer function is a scalar field.
if adjoint:
tf = self._transfer_function.conj()
else:
tf = self._transfer_function
f *= tf
if _use_mkl:
kwargs = {'overwrite_x': True}
else:
kwargs = {}
f = _fft_module.ifftn(f,
axes=tuple(range(-self.input_grid.ndim, 0)),
**kwargs)
s = f.shape[:-self.internal_grid.ndim] + (-1, )
if self.cutout is None:
res = f.reshape(s)
else:
res = f[c].reshape(s)
return Field(res, self.input_grid)
