def fourier_poisson(tensor, times=1):
""" Inverse operation to `fourier_laplace`. """
frequencies = math.fft(math.to_complex(tensor))
k = fftfreq(math.staticshape(tensor)[1:-1], mode='square')
fft_laplace = -(2 * np.pi)**2 * k
fft_laplace[(0, ) * math.ndims(k)] = np.inf
inv_fft_laplace = 1 / fft_laplace
inv_fft_laplace[(0, ) * math.ndims(k)] = 0
return math.real(math.ifft(frequencies * inv_fft_laplace**times))
def normalize_to(target, source=1, epsilon=1e-5, batch_dims=1):
"""
Multiplies the target so that its total content matches the source.
:param target: a tensor
:param source: a tensor or number
:param epsilon: small number to prevent division by zero or None.
:return: normalized tensor of the same shape as target
"""
target_total = math.sum(target, axis=tuple(range(batch_dims, math.ndims(target))), keepdims=True)
denominator = math.maximum(target_total, epsilon) if epsilon is not None else target_total
source_total = math.sum(source, axis=tuple(range(batch_dims, math.ndims(source))), keepdims=True)
return target * (source_total / denominator)
def batch_align_scalar(tensor, innate_spatial_dims, target):
if rank(tensor) == 0:
assert innate_spatial_dims == 0
return math.expand_dims(tensor, 0, len(math.staticshape(target)))
if math.staticshape(tensor)[-1] != 1 or math.ndims(tensor) <= 1:
tensor = math.expand_dims(tensor, -1)
result = batch_align(tensor, innate_spatial_dims + 1, target)
return result
def spatial_rank(tensor):
""" The spatial rank of a tensor is ndims - 2. """
return math.ndims(tensor) - 2
def is_scalar(tensor):
return math.ndims(tensor) == 0