def fftfreq(resolution, mode='vector', dtype=None):
"""
Returns the discrete Fourier transform sample frequencies.
These are the frequencies corresponding to the components of the result of `math.fft` on a tensor of shape `resolution`.
:param resolution: grid resolution measured in cells
:param mode: one of (None, 'vector', 'absolute', 'square')
:param dtype: data type of the returned tensor
:return: tensor holding the frequencies of the corresponding values computed by math.fft
"""
assert mode in ('vector', 'absolute', 'square')
k = np.meshgrid(*[np.fft.fftfreq(int(n)) for n in resolution],
indexing='ij')
k = math.expand_dims(math.stack(k, -1), 0)
if dtype is not None:
k = k.astype(dtype)
else:
k = math.to_float(k)
if mode == 'vector':
return k
k = math.sum(k**2, axis=-1, keepdims=True)
if mode == 'square':
return k
else:
return math.sqrt(k)
def fftfreq(resolution, mode='vector', dtype=np.float32):
assert mode in ('vector', 'absolute', 'square')
k = np.meshgrid(*[np.fft.fftfreq(int(n)) for n in resolution], indexing='ij')
k = math.expand_dims(math.stack(k, -1), 0)
k = k.astype(dtype)
if mode == 'vector':
return k
k = math.sum(k**2, axis=-1, keepdims=True)
if mode == 'square':
return k
else:
return math.sqrt(k)
def upsample2x(tensor, interpolation='linear'):
if struct.isstruct(tensor):
return struct.map(lambda s: upsample2x(s, interpolation), tensor, recursive=False)
if interpolation.lower() != 'linear':
raise ValueError('Only linear interpolation supported')
dims = range(spatial_rank(tensor))
vlen = tensor.shape[-1]
spatial_dims = tensor.shape[1:-1]
rank = spatial_rank(tensor)
tensor = math.pad(tensor, _get_pad_width(rank), 'replicate')
for dim in dims:
lower, center, upper = _dim_shifted(tensor, dim, (-1, 0, 1))
combined = math.stack([0.25 * lower + 0.75 * center, 0.75 * center + 0.25 * upper], axis=2 + dim)
tensor = math.reshape(combined, [-1] + [spatial_dims[dim] * 2 if i == dim else tensor.shape[i + 1] for i in dims] + [vlen])
return tensor