def lies_inside(self, location):
center = math.batch_align(self.center, 1, location)
radius = math.batch_align(self.radius, 0, location)
distance_squared = math.sum((location - center)**2,
axis=-1,
keepdims=True)
return distance_squared <= radius**2
def value_at(self, location):
center = math.batch_align(self.center, 1, location)
radius = math.batch_align(self.radius, 0, location)
distance_squared = math.sum((location - center)**2,
axis=-1,
keepdims=True)
bool_inside = distance_squared <= radius**2
return math.to_float(bool_inside)
def sample_at(self, x):
phase_offset = math.batch_align_scalar(self.phase_offset, 0, x)
k = math.batch_align(self.k, 1, x)
data = math.batch_align(self.data, 1, x)
spatial_phase = math.sum(k * x, -1, keepdims=True)
result = math.sin(
math.to_float(spatial_phase + phase_offset)) * math.to_float(data)
return result
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 sample_at(self, x):
phase_offset = math.batch_align_scalar(self.phase_offset, 0, x)
k = math.batch_align(self.k, 1, x)
data = math.batch_align_scalar(self.data, 0, x)
spatial_phase = math.sum(k * x, -1, keepdims=True)
wave = math.sin(spatial_phase + phase_offset) * data
return math.cast(wave, self.dtype)
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.
:param location: float tensor of shape (batch_size, ..., rank)
:return: float tensor of shape (*location.shape[:-1], 1).
"""
center = math.batch_align(self.center, 1, location)
radius = math.batch_align(self.radius, 0, location)
distance_squared = math.sum((location - center)**2,
axis=-1,
keepdims=True)
distance_squared = math.maximum(
distance_squared,
radius * 1e-2) # Prevent infinite gradient at sphere center
distance = math.sqrt(distance_squared)
return distance - radius
def global_to_child(self, location):
""" Inverse transform """
delta = location - self.center
if math.staticshape(location)[-1] == 2:
angle = -math.batch_align(self.angle, 0, location)
sin = math.sin(angle)
cos = math.cos(angle)
y, x = math.unstack(delta, axis=-1)
if GLOBAL_AXIS_ORDER.is_x_first:
x, y = y, x
rot_x = cos * x - sin * y
rot_y = sin * x + cos * y
rotated = math.stack([rot_y, rot_x], axis=-1)
elif math.staticshape(location)[-1] == 3:
angle = math.batch_align(self.angle, 1, location)
raise NotImplementedError(
'not yet implemented') # ToDo apply angle
else:
raise NotImplementedError('Rotation only supported in 2D and 3D')
final = rotated + self.center
return final
def approximate_signed_distance(self, location):
"""
Computes the signed L-infinity norm (manhattan distance) from the location to the nearest side of the box.
For an outside location `l` with the closest surface point `s`, the distance is `max(abs(l - s))`.
For inside locations it is `-max(abs(l - s))`.
:param location: float tensor of shape (batch_size, ..., rank)
:return: float tensor of shape (*location.shape[:-1], 1).
"""
lower, upper = math.batch_align([self.lower, self.upper], 1, location)
center = 0.5 * (lower + upper)
extent = upper - lower
distance = math.abs(location - center) - extent * 0.5
return math.max(distance, axis=-1, keepdims=True)
def sample_at(self, points):
points_rank = math.spatial_rank(points)
src_rank = math.spatial_rank(self.location)
# --- Expand shapes to format (batch_size, points_dims..., src_dims..., channels) ---
points = math.expand_dims(points, axis=-2, number=src_rank)
src_points = math.expand_dims(self.location,
axis=-2,
number=points_rank)
src_strength = math.expand_dims(self.strength, axis=-1)
src_strength = math.batch_align(src_strength, 0, self.location)
src_strength = math.expand_dims(src_strength,
axis=-1,
number=points_rank)
src_axes = tuple(range(-2, -2 - src_rank, -1))
# --- Compute distances and falloff ---
distances = points - src_points
if self.falloff is not None:
raise NotImplementedError()
# distances_squared = math.sum(distances ** 2, axis=-1, keepdims=True)
# unit_distances = distances / math.sqrt(distances_squared)
# strength = src_strength * math.exp(-distances_squared)
else:
strength = src_strength
# --- Compute velocities ---
if math.staticshape(points)[-1] == 2: # Curl in 2D
dist_1, dist_2 = math.unstack(distances, axis=-1)
if GLOBAL_AXIS_ORDER.is_x_first:
velocity = strength * math.stack([-dist_2, dist_1], axis=-1)
else:
velocity = strength * math.stack([dist_2, -dist_1], axis=-1)
elif math.staticshape(points)[-1] == 3: # Curl in 3D
raise NotImplementedError('not yet implemented')
else:
raise AssertionError(
'Vector product not available in > 3 dimensions')
velocity = math.sum(velocity, axis=src_axes)
return velocity
def lies_inside(self, location):
lower, upper = math.batch_align([self.lower, self.upper], 1, location)
bool_inside = (location >= lower) & (location <= upper)
return math.all(bool_inside, axis=-1, keepdims=True)
def local_to_global(self, local_position):
size, lower = math.batch_align([self.size, self.lower], 1,
local_position)
return local_position * size + lower
def global_to_local(self, global_position):
size, lower = math.batch_align([self.size, self.lower], 1,
global_position)
return (global_position - lower) / size
def value_at(self, global_position):
lower, upper = math.batch_align([self.lower, self.upper], 1,
global_position)
bool_inside = (global_position >= lower) & (global_position <= upper)
bool_inside = math.all(bool_inside, axis=-1, keepdims=True)
return math.to_float(bool_inside)
