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