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)