def grid_sample(self, resolution, size, batch_size=1, channels=None):
channels = channels or self.channels or len(size)
shape = (batch_size, ) + tuple(resolution) + (channels, )
rndj = math.to_complex(
self.math.random_normal(shape)) + 1j * math.to_complex(
self.math.random_normal(shape)) # Note: there is no complex32
k = math.fftfreq(
resolution) * resolution / size * self.scale # in physical units
k = math.sum(k**2, axis=-1, keepdims=True)
lowest_frequency = 0.1
weight_mask = 1 / (1 + math.exp(
(lowest_frequency - k) * 1e3)) # High pass filter
# --- Compute 1/k ---
k[(0, ) * len(k.shape)] = np.inf
inv_k = 1 / k
inv_k[(0, ) * len(k.shape)] = 0
# --- Compute result ---
fft = rndj * inv_k**self.smoothness * weight_mask
array = math.real(math.ifft(fft))
array /= math.std(array,
axis=tuple(range(1, math.ndims(array))),
keepdims=True)
array -= math.mean(array,
axis=tuple(range(1, math.ndims(array))),
keepdims=True)
array = math.to_float(array)
return array
def grid_sample(self,
resolution: math.Shape,
size,
shape: math.Shape = None):
shape = (self._shape if shape is None else shape) & resolution
for dim in channel(self._shape):
if dim.item_names[0] is None:
warnings.warn(
f"Please provide item names for Noise dim {dim} using {dim}='x,y,z'",
FutureWarning)
shape &= channel(**{dim.name: resolution.names})
rndj = math.to_complex(random_normal(shape)) + 1j * math.to_complex(
random_normal(shape)) # Note: there is no complex32
with math.NUMPY:
k = math.fftfreq(resolution) * resolution / math.tensor(
size) * math.tensor(self.scale) # in physical units
k = math.vec_squared(k)
lowest_frequency = 0.1
weight_mask = math.to_float(k > lowest_frequency)
# --- Compute 1/k ---
k._native[(0, ) * len(k.shape)] = np.inf
inv_k = 1 / k
inv_k._native[(0, ) * len(k.shape)] = 0
# --- Compute result ---
fft = rndj * inv_k**self.smoothness * weight_mask
array = math.real(math.ifft(fft))
array /= math.std(array, dim=array.shape.non_batch)
array -= math.mean(array, dim=array.shape.non_batch)
array = math.to_float(array)
return array
def grid_sample(self, resolution, size, batch_size=1, dtype=np.float32):
shape = (batch_size,) + tuple(resolution) + (self.channels,)
rndj = math.randn(shape, dtype) + 1j * math.randn(shape, dtype)
k = math.fftfreq(resolution) * resolution / size * self.scale # in physical units
k = math.sum(k ** 2, axis=-1, keepdims=True)
lowest_frequency = 0.1
weight_mask = 1 / (1 + math.exp((lowest_frequency - k) * 1e3)) # High pass filter
# --- Compute 1/k ---
k[(0,) * len(k.shape)] = np.inf
inv_k = 1 / k
inv_k[(0,) * len(k.shape)] = 0
# --- Compute result ---
fft = rndj * inv_k ** self.smoothness * weight_mask
array = math.real(math.ifft(fft)).astype(dtype)
array /= math.std(array, axis=tuple(range(1, math.ndims(array))), keepdims=True)
array -= math.mean(array, axis=tuple(range(1, math.ndims(array))), keepdims=True)
return array
def grid_sample(self, resolution: math.Shape, size, shape: math.Shape = None):
shape = (self._shape if shape is None else shape).combined(resolution)
rndj = math.to_complex(random_normal(shape)) + 1j * math.to_complex(random_normal(shape)) # Note: there is no complex32
with math.NUMPY_BACKEND:
k = math.fftfreq(resolution) * resolution / size * self.scale # in physical units
k = math.vec_squared(k)
lowest_frequency = 0.1
weight_mask = 1 / (1 + math.exp((lowest_frequency - k) * 1e3)) # High pass filter
# --- Compute 1/k ---
k.native()[(0,) * len(k.shape)] = np.inf
inv_k = 1 / k
inv_k.native()[(0,) * len(k.shape)] = 0
# --- Compute result ---
fft = rndj * inv_k ** self.smoothness * weight_mask
array = math.real(math.ifft(fft))
array /= math.std(array, dim=array.shape.non_batch)
array -= math.mean(array, dim=array.shape.non_batch)
array = math.to_float(array)
return array
def test_std_collapsed(self):
ones = math.ones(x=40000, y=30000)
std = math.std(ones)
self.assertEqual(0, std)
def test_std_collapsed(self):
ones = math.ones(spatial(x=4, y=3)) # hi-res disabled because the current implementation caches the tensor, causes out-of-memory
std = math.std(ones)
math.assert_close(0, std)
