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, 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 diffuse(field, amount, substeps=1):
assert isinstance(field, CenteredGrid)
if field.extrapolation == 'periodic':
frequencies = math.fft(math.to_complex(field.data))
k = math.fftfreq(field.resolution) / field.dx
k = math.sum(k**2, axis=-1, keepdims=True)
fft_laplace = -(2 * pi)**2 * k
diffuse_kernel = math.to_complex(math.exp(fft_laplace * amount))
data = math.ifft(frequencies * diffuse_kernel)
data = math.real(data)
else:
data = field.data
for i in range(substeps):
data += amount / substeps * field.laplace()
return field.with_data(data)
def diffuse(field, amount, substeps=1):
u"""
Simulate a finite-time diffusion process of the form dF/dt = α · ΔF on a given `Field` F with diffusion coefficient α.
If `field` is periodic (set via `extrapolation='periodic'`), diffusion may be simulated in Fourier space.
Otherwise, finite differencing is used to approximate the
:param field: CenteredGrid, StaggeredGrid or ConstantField
:param amount: number of Field, typically α · dt
:param substeps: number of iterations to use
:return: Field of same type as `field`
:rtype: Field
"""
if isinstance(field, ConstantField):
return field
if isinstance(field, StaggeredGrid):
return struct.map(
lambda grid: diffuse(grid, amount, substeps=substeps),
field,
leaf_condition=lambda x: isinstance(x, CenteredGrid))
assert isinstance(
field, CenteredGrid), "Cannot diffuse field of type '%s'" % type(field)
if field.extrapolation == 'periodic' and not isinstance(amount, Field):
fft_laplace = -(2 * pi)**2 * field.squared_frequencies
diffuse_kernel = math.exp(fft_laplace * amount)
return math.real(
math.ifft(field.fft() * math.to_complex(diffuse_kernel)))
else:
data = field.data
if isinstance(amount, Field):
amount = amount.at(field).data
else:
amount = math.batch_align(amount, 0, data)
for i in range(substeps):
data += amount / substeps * field.laplace().data
return field.with_data(data)
def test_cast(self):
for backend in BACKENDS:
with backend:
x = math.random_uniform(dtype=DType(float, 64))
self.assertEqual(DType(float, 32), math.to_float(x).dtype, msg=backend.name)
self.assertEqual(DType(complex, 64), math.to_complex(x).dtype, msg=backend.name)
with math.precision(64):
self.assertEqual(DType(float, 64), math.to_float(x).dtype, msg=backend.name)
self.assertEqual(DType(complex, 128), math.to_complex(x).dtype, msg=backend.name)
self.assertEqual(DType(int, 64), math.to_int64(x).dtype, msg=backend.name)
self.assertEqual(DType(int, 32), math.to_int32(x).dtype, msg=backend.name)
self.assertEqual(DType(float, 16), math.cast(x, DType(float, 16)).dtype, msg=backend.name)
self.assertEqual(DType(complex, 128), math.cast(x, DType(complex, 128)).dtype, msg=backend.name)
try:
math.cast(x, DType(float, 3))
self.fail(msg=backend.name)
except KeyError:
pass
def step(self, state, dt=1.0, potentials=(), obstacles=()):
if len(potentials) == 0:
potential = 0
else:
potential = math.zeros_like(math.real(
state.amplitude)) # for the moment, allow only real potentials
for pot in potentials:
potential = effect_applied(pot, potential, dt)
potential = potential.data
amplitude = state.amplitude.data
# Rotate by potential
rotation = math.exp(1j * math.to_complex(potential * dt))
amplitude = amplitude * rotation
# Move by rotating in Fourier space
amplitude_fft = math.fft(amplitude)
laplace = math.fftfreq(state.resolution, mode='square')
amplitude_fft *= math.exp(-1j * (2 * np.pi)**2 * math.to_complex(dt) *
laplace / (2 * state.mass))
amplitude = math.ifft(amplitude_fft)
obstacle_mask = union_mask([
obstacle.geometry for obstacle in obstacles
]).at(state.amplitude).data
amplitude *= 1 - obstacle_mask
normalized = False
symmetric = False
if not symmetric:
boundary_mask = math.zeros(
state.domain.centered_shape(1, batch_size=1)).data
boundary_mask[[slice(None)] + [
slice(self.margin, -self.margin)
for i in math.spatial_dimensions(boundary_mask)
] + [slice(None)]] = 1
amplitude *= boundary_mask
if len(obstacles) > 0 or not symmetric:
amplitude = normalize_probability(amplitude)
normalized = True
return state.copied_with(amplitude=amplitude)
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 normalize_probability(probability_amplitude):
p = math.to_float(abs(probability_amplitude)**2)
P = math.sum(p, math.spatial_dimensions(p), keepdims=True)
return probability_amplitude / math.to_complex(math.sqrt(P))