def dft_postprocess_data(arr, real_grid, recip_grid, shift, axes,
interp, sign='-', op='multiply', out=None):
"""Post-process the Fourier-space data after DFT.
This function multiplies the given data with the separable
function::
q(xi) = exp(+- 1j * dot(x[0], xi)) * s * phi_hat(xi_bar)
where ``x[0]`` and ``s`` are the minimum point and the stride of
the real-space grid, respectively, and ``phi_hat(xi_bar)`` is the FT
of the interpolation kernel. The sign of the exponent depends on the
choice of ``sign``. Note that for ``op='divide'`` the
multiplication with ``s * phi_hat(xi_bar)`` is replaced by a
division with the same array.
In discretized form on the reciprocal grid, the exponential part
of this function becomes an array::
q[k] = exp(+- 1j * dot(x[0], xi[k]))
and the arguments ``xi_bar`` to the interpolation kernel
are the normalized frequencies::
for 'shift=True' : xi_bar[k] = -pi + pi * (2*k) / N
for 'shift=False' : xi_bar[k] = -pi + pi * (2*k+1) / N
See [Pre+2007], Section 13.9 "Computing Fourier Integrals Using
the FFT" for a similar approach.
Parameters
----------
arr : `array-like`
Array to be pre-processed. An array with real data type is
converted to its complex counterpart.
real_grid : uniform `RectGrid`
Real space grid in the transform.
recip_grid : uniform `RectGrid`
Reciprocal grid in the transform
shift : bool or sequence of bools
If ``True``, the grid is shifted by half a stride in the negative
direction in the corresponding axes. The sequence must have the
same length as ``axes``.
axes : int or sequence of ints
Dimensions along which to take the transform. The sequence must
have the same length as ``shifts``.
interp : string or sequence of strings
Interpolation scheme used in the real-space.
sign : {'-', '+'}, optional
Sign of the complex exponent.
op : {'multiply', 'divide'}, optional
Operation to perform with the stride times the interpolation
kernel FT
out : `numpy.ndarray`, optional
Array in which the result is stored. If ``out is arr``, an
in-place modification is performed.
Returns
-------
out : `numpy.ndarray`
Result of the post-processing. If ``out`` was given, the returned
object is a reference to it.
References
----------
[Pre+2007] Press, W H, Teukolsky, S A, Vetterling, W T, and Flannery, B P.
*Numerical Recipes in C - The Art of Scientific Computing* (Volume 3).
Cambridge University Press, 2007.
"""
arr = np.asarray(arr)
if is_real_floating_dtype(arr.dtype):
arr = arr.astype(complex_dtype(arr.dtype))
elif not is_complex_floating_dtype(arr.dtype):
raise ValueError('array data type {} is not a complex floating point '
'data type'.format(dtype_repr(arr.dtype)))
if out is None:
out = arr.copy()
elif out is not arr:
out[:] = arr
if axes is None:
axes = list(range(arr.ndim))
else:
try:
axes = [int(axes)]
except TypeError:
axes = list(axes)
shift_list = normalized_scalar_param_list(shift, length=len(axes),
param_conv=bool)
if sign == '-':
imag = -1j
elif sign == '+':
imag = 1j
else:
raise ValueError("`sign` '{}' not understood".format(sign))
op, op_in = str(op).lower(), op
if op not in ('multiply', 'divide'):
raise ValueError("kernel `op` '{}' not understood".format(op_in))
# Make a list from interp if that's not the case already
try:
# Duck-typed string check
interp + ''
except TypeError:
pass
else:
interp = [str(interp).lower()] * arr.ndim
onedim_arrs = []
for ax, shift, intp in zip(axes, shift_list, interp):
x = real_grid.min_pt[ax]
xi = recip_grid.coord_vectors[ax]
# First part: exponential array
onedim_arr = np.exp(imag * x * xi)
# Second part: interpolation kernel
len_dft = recip_grid.shape[ax]
len_orig = real_grid.shape[ax]
halfcomplex = (len_dft < len_orig)
odd = len_orig % 2
fmin = -0.5 if shift else -0.5 + 1.0 / (2 * len_orig)
if halfcomplex:
# maximum lies around 0, possibly half a cell left or right of it
if shift and odd:
fmax = - 1.0 / (2 * len_orig)
elif not shift and not odd:
fmax = 1.0 / (2 * len_orig)
else:
fmax = 0.0
else: # not halfcomplex
# maximum lies close to 0.5, half or full cell left of it
if shift:
# -0.5 + (N-1)/N = 0.5 - 1/N
fmax = 0.5 - 1.0 / len_orig
else:
# -0.5 + 1/(2*N) + (N-1)/N = 0.5 - 1/(2*N)
fmax = 0.5 - 1.0 / (2 * len_orig)
freqs = np.linspace(fmin, fmax, num=len_dft)
stride = real_grid.stride[ax]
interp_kernel = _interp_kernel_ft(freqs, intp)
interp_kernel *= stride
if op == 'multiply':
onedim_arr *= interp_kernel
else:
onedim_arr /= interp_kernel
onedim_arrs.append(onedim_arr.astype(out.dtype, copy=False))
fast_1d_tensor_mult(out, onedim_arrs, axes=axes, out=out)
return out
def test_fast_1d_tensor_mult_error():
shape = (2, 3, 4)
test_arr = np.ones(shape)
x, y, z = (np.arange(size, dtype='float64') for size in shape)
# No ndarray to operate on
with pytest.raises(TypeError):
fast_1d_tensor_mult([[0, 0], [0, 0]], [x, x])
# No 1d arrays given
with pytest.raises(ValueError):
fast_1d_tensor_mult(test_arr, [])
# Length or dimension mismatch
with pytest.raises(ValueError):
fast_1d_tensor_mult(test_arr, [x, y])
with pytest.raises(ValueError):
fast_1d_tensor_mult(test_arr, [x, y], (1, 2, 0))
with pytest.raises(ValueError):
fast_1d_tensor_mult(test_arr, [x, x, y, z])
# Axes out of bounds
with pytest.raises(ValueError):
fast_1d_tensor_mult(test_arr, [x, y, z], (1, 2, 3))
with pytest.raises(ValueError):
fast_1d_tensor_mult(test_arr, [x, y, z], (-2, -3, -4))
# Other than 1d arrays
with pytest.raises(ValueError):
fast_1d_tensor_mult(test_arr, [x, y, np.ones((4, 2))], (-2, -3, -4))
def dft_preprocess_data(arr, shift=True, axes=None, sign='-', out=None):
"""Pre-process the real-space data before DFT.
This function multiplies the given data with the separable
function::
p(x) = exp(+- 1j * dot(x - x[0], xi[0]))
where ``x[0]`` and ``xi[0]`` are the minimum coodinates of
the real-space and reciprocal grids, respectively. The sign of
the exponent depends on the choice of ``sign``. In discretized
form, this function becomes an array::
p[k] = exp(+- 1j * k * s * xi[0])
If the reciprocal grid is not shifted, i.e. symmetric around 0,
it is ``xi[0] = pi/s * (-1 + 1/N)``, hence::
p[k] = exp(-+ 1j * pi * k * (1 - 1/N))
For a shifted grid, we have :math:``xi[0] = -pi/s``, thus the
array is given by::
p[k] = (-1)**k
Parameters
----------
arr : `array-like`
Array to be pre-processed. If its data type is a real
non-floating type, it is converted to 'float64'.
shift : bool or or sequence of bools, optional
If ``True``, the grid is shifted by half a stride in the negative
direction. With a sequence, this option is applied separately on
each axis.
axes : int or sequence of ints, optional
Dimensions in which to calculate the reciprocal. The sequence
must have the same length as ``shift`` if the latter is given
as a sequence.
Default: all axes.
sign : {'-', '+'}, optional
Sign of the complex exponent.
out : `numpy.ndarray`, optional
Array in which the result is stored. If ``out is arr``,
an in-place modification is performed. For real data type,
this is only possible for ``shift=True`` since the factors are
complex otherwise.
Returns
-------
out : `numpy.ndarray`
Result of the pre-processing. If ``out`` was given, the returned
object is a reference to it.
Notes
-----
If ``out`` is not specified, the data type of the returned array
is the same as that of ``arr`` except when ``arr`` has real data
type and ``shift`` is not ``True``. In this case, the return type
is the complex counterpart of ``arr.dtype``.
"""
arr = np.asarray(arr)
if not is_scalar_dtype(arr.dtype):
raise ValueError('array has non-scalar data type {}'
''.format(dtype_repr(arr.dtype)))
elif is_real_dtype(arr.dtype) and not is_real_floating_dtype(arr.dtype):
arr = arr.astype('float64')
if axes is None:
axes = list(range(arr.ndim))
else:
try:
axes = [int(axes)]
except TypeError:
axes = list(axes)
shape = arr.shape
shift_list = normalized_scalar_param_list(shift, length=len(axes),
param_conv=bool)
# Make a copy of arr with correct data type if necessary, or copy values.
if out is None:
if is_real_dtype(arr.dtype) and not all(shift_list):
out = np.array(arr, dtype=complex_dtype(arr.dtype), copy=True)
else:
out = arr.copy()
else:
out[:] = arr
if is_real_dtype(out.dtype) and not shift:
raise ValueError('cannot pre-process real input in-place without '
'shift')
if sign == '-':
imag = -1j
elif sign == '+':
imag = 1j
else:
raise ValueError("`sign` '{}' not understood".format(sign))
def _onedim_arr(length, shift):
if shift:
# (-1)^indices
factor = np.ones(length, dtype=out.dtype)
factor[1::2] = -1
else:
factor = np.arange(length, dtype=out.dtype)
factor *= -imag * np.pi * (1 - 1.0 / length)
np.exp(factor, out=factor)
return factor.astype(out.dtype, copy=False)
onedim_arrs = []
for axis, shift in zip(axes, shift_list):
length = shape[axis]
onedim_arrs.append(_onedim_arr(length, shift))
fast_1d_tensor_mult(out, onedim_arrs, axes=axes, out=out)
return out
def test_fast_1d_tensor_mult():
# Full multiplication
def simple_mult_3(x, y, z):
return x[:, None, None] * y[None, :, None] * z[None, None, :]
shape = (2, 3, 4)
x, y, z = (np.arange(size, dtype='float64') for size in shape)
true_result = simple_mult_3(x, y, z)
# Standard call
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [x, y, z])
assert all_equal(out, true_result)
assert all_equal(test_arr, np.ones(shape)) # no changes to input
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [x, y, z], axes=(0, 1, 2))
assert all_equal(out, true_result)
# In-place with both same and different array as input
test_arr = np.ones(shape)
fast_1d_tensor_mult(test_arr, [x, y, z], out=test_arr)
assert all_equal(test_arr, true_result)
test_arr = np.ones(shape)
out = np.empty(shape)
fast_1d_tensor_mult(test_arr, [x, y, z], out=out)
assert all_equal(out, true_result)
# Different orderings
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [y, x, z], axes=(1, 0, 2))
assert all_equal(out, true_result)
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [x, z, y], axes=(0, 2, 1))
assert all_equal(out, true_result)
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [z, x, y], axes=(2, 0, 1))
assert all_equal(out, true_result)
# More arrays than dimensions also ok with explicit axes
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [z, x, y, np.ones(3)],
axes=(2, 0, 1, 1))
assert all_equal(out, true_result)
# Squeezable or extendable arrays also possible
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [x, y, z[None, :]])
assert all_equal(out, true_result)
shape = (1, 3, 4)
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [2, y, z])
assert all_equal(out, simple_mult_3(np.ones(1) * 2, y, z))
# Reduced multiplication, axis 0 not contained
def simple_mult_2(y, z, nx):
return np.ones((nx, 1, 1)) * y[None, :, None] * z[None, None, :]
shape = (2, 3, 4)
x, y, z = (np.arange(size, dtype='float64') for size in shape)
true_result = simple_mult_2(y, z, 2)
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [y, z], axes=(1, 2))
assert all_equal(out, true_result)
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [z, y], axes=(2, 1))
assert all_equal(out, true_result)
def dft_postprocess_data(arr, real_grid, recip_grid, shift, axes,
interp, sign='-', op='multiply', out=None):
"""Post-process the Fourier-space data after DFT.
This function multiplies the given data with the separable
function::
q(xi) = exp(+- 1j * dot(x[0], xi)) * s * phi_hat(xi_bar)
where ``x[0]`` and ``s`` are the minimum point and the stride of
the real-space grid, respectively, and ``phi_hat(xi_bar)`` is the FT
of the interpolation kernel. The sign of the exponent depends on the
choice of ``sign``. Note that for ``op='divide'`` the
multiplication with ``s * phi_hat(xi_bar)`` is replaced by a
division with the same array.
In discretized form on the reciprocal grid, the exponential part
of this function becomes an array::
q[k] = exp(+- 1j * dot(x[0], xi[k]))
and the arguments ``xi_bar`` to the interpolation kernel
are the normalized frequencies::
for 'shift=True' : xi_bar[k] = -pi + pi * (2*k) / N
for 'shift=False' : xi_bar[k] = -pi + pi * (2*k+1) / N
See [Pre+2007]_, Section 13.9 "Computing Fourier Integrals Using
the FFT" for a similar approach.
Parameters
----------
arr : `array-like`
Array to be pre-processed. An array with real data type is
converted to its complex counterpart.
real_grid : `RegularGrid`
Real space grid in the transform
recip_grid : `RegularGrid`
Reciprocal grid in the transform
shift : bool or sequence of bools
If ``True``, the grid is shifted by half a stride in the negative
direction in the corresponding axes. The sequence must have the
same length as ``axes``.
axes : int or sequence of ints
Dimensions along which to take the transform. The sequence must
have the same length as ``shifts``.
interp : string or sequence of strings
Interpolation scheme used in the real-space.
sign : {'-', '+'}, optional
Sign of the complex exponent.
op : {'multiply', 'divide'}
Operation to perform with the stride times the interpolation
kernel FT
out : `numpy.ndarray`, optional
Array in which the result is stored. If ``out is arr``, an
in-place modification is performed.
Returns
-------
out : `numpy.ndarray`
Result of the post-processing. If ``out`` was given, the returned
object is a reference to it.
"""
arr = np.asarray(arr)
if is_real_floating_dtype(arr.dtype):
arr = arr.astype(complex_dtype(arr.dtype))
elif not is_complex_floating_dtype(arr.dtype):
raise ValueError('array data type {} is not a complex floating point '
'data type'.format(dtype_repr(arr.dtype)))
if out is None:
out = arr.copy()
elif out is not arr:
out[:] = arr
if axes is None:
axes = list(range(arr.ndim))
else:
try:
axes = [int(axes)]
except TypeError:
axes = list(axes)
shift_list = normalized_scalar_param_list(shift, length=len(axes),
param_conv=bool)
if sign == '-':
imag = -1j
elif sign == '+':
imag = 1j
else:
raise ValueError("`sign` '{}' not understood".format(sign))
op, op_in = str(op).lower(), op
if op not in ('multiply', 'divide'):
raise ValueError("kernel `op` '{}' not understood".format(op_in))
# Make a list from interp if that's not the case already
try:
# Duck-typed string check
interp + ''
except TypeError:
pass
else:
interp = [str(interp).lower()] * arr.ndim
onedim_arrs = []
for ax, shift, intp in zip(axes, shift_list, interp):
x = real_grid.min_pt[ax]
xi = recip_grid.coord_vectors[ax]
# First part: exponential array
onedim_arr = np.exp(imag * x * xi)
# Second part: interpolation kernel
len_dft = recip_grid.shape[ax]
len_orig = real_grid.shape[ax]
halfcomplex = (len_dft < len_orig)
odd = len_orig % 2
fmin = -0.5 if shift else -0.5 + 1.0 / (2 * len_orig)
if halfcomplex:
# maximum lies around 0, possibly half a cell left or right of it
if shift and odd:
fmax = - 1.0 / (2 * len_orig)
elif not shift and not odd:
fmax = 1.0 / (2 * len_orig)
else:
fmax = 0.0
else: # not halfcomplex
# maximum lies close to 0.5, half or full cell left of it
if shift:
# -0.5 + (N-1)/N = 0.5 - 1/N
fmax = 0.5 - 1.0 / len_orig
else:
# -0.5 + 1/(2*N) + (N-1)/N = 0.5 - 1/(2*N)
fmax = 0.5 - 1.0 / (2 * len_orig)
freqs = np.linspace(fmin, fmax, num=len_dft)
stride = real_grid.stride[ax]
if op == 'multiply':
onedim_arr *= stride * _interp_kernel_ft(freqs, intp)
else:
onedim_arr /= stride * _interp_kernel_ft(freqs, intp)
onedim_arrs.append(onedim_arr.astype(out.dtype, copy=False))
fast_1d_tensor_mult(out, onedim_arrs, axes=axes, out=out)
return out
def dft_preprocess_data(arr, shift=True, axes=None, sign='-', out=None):
"""Pre-process the real-space data before DFT.
This function multiplies the given data with the separable
function::
p(x) = exp(+- 1j * dot(x - x[0], xi[0]))
where ``x[0]`` and ``xi[0]`` are the minimum coodinates of
the real-space and reciprocal grids, respectively. The sign of
the exponent depends on the choice of ``sign``. In discretized
form, this function becomes an array::
p[k] = exp(+- 1j * k * s * xi[0])
If the reciprocal grid is not shifted, i.e. symmetric around 0,
it is ``xi[0] = pi/s * (-1 + 1/N)``, hence::
p[k] = exp(-+ 1j * pi * k * (1 - 1/N))
For a shifted grid, we have :math:``xi[0] = -pi/s``, thus the
array is given by::
p[k] = (-1)**k
Parameters
----------
arr : `array-like`
Array to be pre-processed. If its data type is a real
non-floating type, it is converted to 'float64'.
shift : bool or or sequence of bools, optional
If ``True``, the grid is shifted by half a stride in the negative
direction. With a sequence, this option is applied separately on
each axis.
axes : int or sequence of ints, optional
Dimensions in which to calculate the reciprocal. The sequence
must have the same length as ``shift`` if the latter is given
as a sequence.
Default: all axes.
sign : {'-', '+'}, optional
Sign of the complex exponent.
out : `numpy.ndarray`, optional
Array in which the result is stored. If ``out is arr``,
an in-place modification is performed. For real data type,
this is only possible for ``shift=True`` since the factors are
complex otherwise.
Returns
-------
out : `numpy.ndarray`
Result of the pre-processing. If ``out`` was given, the returned
object is a reference to it.
Notes
-----
If ``out`` is not specified, the data type of the returned array
is the same as that of ``arr`` except when ``arr`` has real data
type and ``shift`` is not ``True``. In this case, the return type
is the complex counterpart of ``arr.dtype``.
"""
arr = np.asarray(arr)
if not is_scalar_dtype(arr.dtype):
raise ValueError('array has non-scalar data type {}'
''.format(dtype_repr(arr.dtype)))
elif is_real_dtype(arr.dtype) and not is_real_floating_dtype(arr.dtype):
arr = arr.astype('float64')
if axes is None:
axes = list(range(arr.ndim))
else:
try:
axes = [int(axes)]
except TypeError:
axes = list(axes)
shape = arr.shape
shift_list = normalized_scalar_param_list(shift, length=len(axes),
param_conv=bool)
# Make a copy of arr with correct data type if necessary, or copy values.
if out is None:
if is_real_dtype(arr.dtype) and not all(shift_list):
out = np.array(arr, dtype=complex_dtype(arr.dtype), copy=True)
else:
out = arr.copy()
else:
out[:] = arr
if is_real_dtype(out.dtype) and not shift:
raise ValueError('cannot pre-process real input in-place without '
'shift')
if sign == '-':
imag = -1j
elif sign == '+':
imag = 1j
else:
raise ValueError("`sign` '{}' not understood".format(sign))
def _onedim_arr(length, shift):
if shift:
# (-1)^indices
factor = np.ones(length, dtype=out.dtype)
factor[1::2] = -1
else:
factor = np.arange(length, dtype=out.dtype)
factor *= -imag * np.pi * (1 - 1.0 / length)
np.exp(factor, out=factor)
return factor.astype(out.dtype, copy=False)
onedim_arrs = []
for axis, shift in zip(axes, shift_list):
length = shape[axis]
onedim_arrs.append(_onedim_arr(length, shift))
fast_1d_tensor_mult(out, onedim_arrs, axes=axes, out=out)
return out
def test_fast_1d_tensor_mult_error():
shape = (2, 3, 4)
test_arr = np.ones(shape)
x, y, z = (np.arange(size, dtype='float64') for size in shape)
# No ndarray to operate on
with pytest.raises(TypeError):
fast_1d_tensor_mult([[0, 0], [0, 0]], [x, x])
# No 1d arrays given
with pytest.raises(ValueError):
fast_1d_tensor_mult(test_arr, [])
# Length or dimension mismatch
with pytest.raises(ValueError):
fast_1d_tensor_mult(test_arr, [x, y])
with pytest.raises(ValueError):
fast_1d_tensor_mult(test_arr, [x, y], (1, 2, 0))
with pytest.raises(ValueError):
fast_1d_tensor_mult(test_arr, [x, x, y, z])
# Axes out of bounds
with pytest.raises(ValueError):
fast_1d_tensor_mult(test_arr, [x, y, z], (1, 2, 3))
with pytest.raises(ValueError):
fast_1d_tensor_mult(test_arr, [x, y, z], (-2, -3, -4))
# Other than 1d arrays
with pytest.raises(ValueError):
fast_1d_tensor_mult(test_arr, [x, y, np.ones((4, 2))], (-2, -3, -4))
def test_fast_1d_tensor_mult():
# Full multiplication
def simple_mult_3(x, y, z):
return x[:, None, None] * y[None, :, None] * z[None, None, :]
shape = (2, 3, 4)
x, y, z = (np.arange(size, dtype='float64') for size in shape)
true_result = simple_mult_3(x, y, z)
# Standard call
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [x, y, z])
assert all_equal(out, true_result)
assert all_equal(test_arr, np.ones(shape)) # no changes to input
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [x, y, z], axes=(0, 1, 2))
assert all_equal(out, true_result)
# In-place with both same and different array as input
test_arr = np.ones(shape)
fast_1d_tensor_mult(test_arr, [x, y, z], out=test_arr)
assert all_equal(test_arr, true_result)
test_arr = np.ones(shape)
out = np.empty(shape)
fast_1d_tensor_mult(test_arr, [x, y, z], out=out)
assert all_equal(out, true_result)
# Different orderings
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [y, x, z], axes=(1, 0, 2))
assert all_equal(out, true_result)
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [x, z, y], axes=(0, 2, 1))
assert all_equal(out, true_result)
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [z, x, y], axes=(2, 0, 1))
assert all_equal(out, true_result)
# More arrays than dimensions also ok with explicit axes
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [z, x, y, np.ones(3)],
axes=(2, 0, 1, 1))
assert all_equal(out, true_result)
# Squeezable or extendable arrays also possible
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [x, y, z[None, :]])
assert all_equal(out, true_result)
shape = (1, 3, 4)
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [2, y, z])
assert all_equal(out, simple_mult_3(np.ones(1) * 2, y, z))
# Reduced multiplication, axis 0 not contained
def simple_mult_2(y, z, nx):
return np.ones((nx, 1, 1)) * y[None, :, None] * z[None, None, :]
shape = (2, 3, 4)
x, y, z = (np.arange(size, dtype='float64') for size in shape)
true_result = simple_mult_2(y, z, 2)
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [y, z], axes=(1, 2))
assert all_equal(out, true_result)
test_arr = np.ones(shape)
out = fast_1d_tensor_mult(test_arr, [z, y], axes=(2, 1))
assert all_equal(out, true_result)