def fourier(field: GridType, diffusivity: float or math.Tensor, dt: float
or math.Tensor) -> FieldType:
"""
Exact diffusion of a periodic field in frequency space.
For non-periodic fields or non-constant diffusivity, use another diffusion function such as `explicit()`.
Args:
field:
diffusivity: Diffusion per time. `diffusion_amount = diffusivity * dt`
dt: Time interval. `diffusion_amount = diffusivity * dt`
Returns:
Diffused field of same type as `field`.
"""
if isinstance(field, ConstantField):
return field
assert isinstance(field,
Grid), "Cannot diffuse field of type '%s'" % type(field)
assert field.extrapolation == math.extrapolation.PERIODIC, "Fourier diffusion can only be applied to periodic fields."
amount = diffusivity * dt
k = math.fftfreq(field.resolution)
k2 = math.vec_squared(k)
fft_laplace = -(2 * math.PI)**2 * k2
diffuse_kernel = math.exp(fft_laplace * amount)
result_k = math.fft(field.values) * diffuse_kernel
result_values = math.real(math.ifft(result_k))
return field.with_(values=result_values)
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 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 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 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):
frequencies = math.fft(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
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 frequencies(self):
return self.with_data(
math.fftfreq(self.resolution, mode='vector') / self.dx)
