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 show(self, caller_is_main):
while True:
for i in range(self.fig_count):
self.axes[i].clear()
v = self.app.get_field(self.shown_fields[i])
if isinstance(v, StaggeredGrid):
v = v.at_centers()
v = math.vec_squared(v.values).native('y, x')
self.axes[i].imshow(v, origin='lower')
self.axes[i].set_title(self.shown_fields[i])
plt.draw()
plt.pause(0.01)
def approximate_signed_distance(self, location):
"""
Computes the exact distance from location to the closest point on the sphere.
Very close to the sphere center, the distance takes a constant value.
Args:
location: float tensor of shape (batch_size, ..., rank)
Returns:
float tensor of shape (*location.shape[:-1], 1).
"""
distance_squared = math.vec_squared(location - self.center)
distance_squared = math.maximum(
distance_squared,
self.radius * 1e-2) # Prevent infinite gradient at sphere center
distance = math.sqrt(distance_squared)
return distance - self.radius
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 vec_squared(field: SampledField):
""" See `phi.math.vec_squared()` """
if isinstance(field, StaggeredGrid):
field = field.at_centers()
return field.with_values(math.vec_squared(field.values))