def distribute_points(density, particles_per_cell=1, distribution='uniform'): """ Distribute points according to the distribution specified in density. :param density: binary tensor :param particles_per_cell: integer :param distribution: 'uniform' or 'center' :return: tensor of shape (batch_size, point_count, rank) """ assert distribution in ('center', 'uniform') index_array = [] batch_size = math.staticshape(density)[0] if math.staticshape(density)[0] is not None else 1 for batch in range(batch_size): indices = math.where(density[batch, ..., 0] > 0) indices = math.to_float(indices) temp = [] for _ in range(particles_per_cell): if distribution == 'center': temp.append(indices + 0.5) elif distribution == 'uniform': temp.append(indices + math.random_uniform(math.shape(indices))) index_array.append(math.concat(temp, axis=0)) try: index_array = math.stack(index_array) return index_array except ValueError: raise ValueError("all arrays in the batch must have the same number of active cells.")
def push(self, positions: Tensor, outward: bool = True, shift_amount: float = 0) -> Tensor: loc_to_center = positions - self.center sgn_dist_from_surface = math.abs(loc_to_center) - self.half_size if outward: # --- get negative distances (particles are inside) towards the nearest boundary and add shift_amount --- distances_of_interest = (sgn_dist_from_surface == math.max( sgn_dist_from_surface, 'vector')) & (sgn_dist_from_surface < 0) shift = distances_of_interest * (sgn_dist_from_surface - shift_amount) else: shift = (sgn_dist_from_surface + shift_amount) * ( sgn_dist_from_surface > 0 ) # get positive distances (particles are outside) and add shift_amount shift = math.where( math.abs(shift) > math.abs(loc_to_center), math.abs(loc_to_center), shift) # ensure inward shift ends at center return positions + math.where(loc_to_center < 0, 1, -1) * shift
def effect_applied(effect, field, dt): effect_field = effect.field.at(field) if effect._mode == GROW: return field + math.mul(effect_field, dt) elif effect._mode == ADD: return field + effect_field elif effect._mode == FIX: assert effect.bounds is not None inside = mask([effect.bounds]).at(field) return math.where(inside, effect_field, field) else: raise ValueError('Invalid mode: %s' % effect.mode)
def extrapolate(input_field, valid_mask, voxel_distance=10): """ Create a signed distance field for the grid, where negative signs are fluid cells and positive signs are empty cells. The fluid surface is located at the points where the interpolated value is zero. Then extrapolate the input field into the air cells. :param domain: Domain that can create new Fields :param input_field: Field to be extrapolated :param valid_mask: One dimensional binary mask indicating where fluid is present :param voxel_distance: Optional maximal distance (in number of grid cells) where signed distance should still be calculated / how far should be extrapolated. :return: ext_field: a new Field with extrapolated values, s_distance: tensor containing signed distance field, depending only on the valid_mask """ ext_data = input_field.data dx = input_field.dx if isinstance(input_field, StaggeredGrid): ext_data = input_field.staggered_tensor() valid_mask = math.pad(valid_mask, [[0, 0]] + [[0, 1]] * input_field.rank + [[0, 0]], "constant") dims = range(input_field.rank) # Larger than voxel_distance to be safe. It could start extrapolating velocities from outside voxel_distance into the field. signs = -1 * (2 * valid_mask - 1) s_distance = 2.0 * (voxel_distance + 1) * signs surface_mask = create_surface_mask(valid_mask) # surface_mask == 1 doesn't output a tensor, just a scalar, but >= works. # Initialize the voxel_distance with 0 at the surface # Previously initialized with -0.5*dx, i.e. the cell is completely full (center is 0.5*dx inside the fluid surface). For stability and looks this was changed to 0 * dx, i.e. the cell is only half full. This way small changes to the SDF won't directly change neighbouring empty cells to fluid cells. s_distance = math.where((surface_mask >= 1), -0.0 * math.ones_like(s_distance), s_distance) directions = np.array( list(itertools.product(*np.tile((-1, 0, 1), (len(dims), 1))))) # First make a move in every positive direction (StaggeredGrid velocities there are correct, we want to extrapolate these) if isinstance(input_field, StaggeredGrid): for d in directions: if (d <= 0).all(): continue # Shift the field in direction d, compare new distances to old ones. d_slice = tuple([(slice(1, None) if d[i] == -1 else slice(0, -1) if d[i] == 1 else slice(None)) for i in dims]) d_field = math.pad( ext_data, [[0, 0]] + [([0, 1] if d[i] == -1 else [1, 0] if d[i] == 1 else [0, 0]) for i in dims] + [[0, 0]], "symmetric") d_field = d_field[(slice(None), ) + d_slice + (slice(None), )] d_dist = math.pad( s_distance, [[0, 0]] + [([0, 1] if d[i] == -1 else [1, 0] if d[i] == 1 else [0, 0]) for i in dims] + [[0, 0]], "symmetric") d_dist = d_dist[(slice(None), ) + d_slice + (slice(None), )] d_dist += np.sqrt((dx * d).dot(dx * d)) * signs if (d.dot(d) == 1) and (d >= 0).all(): # Pure axis direction (1,0,0), (0,1,0), (0,0,1) updates = (math.abs(d_dist) < math.abs(s_distance)) & (surface_mask <= 0) updates_velocity = updates & (signs > 0) ext_data = math.where( math.concat([(math.zeros_like(updates_velocity) if d[i] == 1 else updates_velocity) for i in dims], axis=-1), d_field, ext_data) s_distance = math.where(updates, d_dist, s_distance) else: # Mixed axis direction (1,1,0), (1,1,-1), etc. continue for _ in range(voxel_distance): buffered_distance = 1.0 * s_distance # Create a copy of current voxel_distance. This should not be necessary... for d in directions: if (d == 0).all(): continue # Shift the field in direction d, compare new distances to old ones. d_slice = tuple([(slice(1, None) if d[i] == -1 else slice(0, -1) if d[i] == 1 else slice(None)) for i in dims]) d_field = math.pad( ext_data, [[0, 0]] + [([0, 1] if d[i] == -1 else [1, 0] if d[i] == 1 else [0, 0]) for i in dims] + [[0, 0]], "symmetric") d_field = d_field[(slice(None), ) + d_slice + (slice(None), )] d_dist = math.pad( s_distance, [[0, 0]] + [([0, 1] if d[i] == -1 else [1, 0] if d[i] == 1 else [0, 0]) for i in dims] + [[0, 0]], "symmetric") d_dist = d_dist[(slice(None), ) + d_slice + (slice(None), )] d_dist += np.sqrt((dx * d).dot(dx * d)) * signs # We only want to update velocity that is outside of fluid updates = (math.abs(d_dist) < math.abs(buffered_distance)) & (surface_mask <= 0) updates_velocity = updates & (signs > 0) ext_data = math.where( math.concat([updates_velocity] * math.spatial_rank(ext_data), axis=-1), d_field, ext_data) buffered_distance = math.where(updates, d_dist, buffered_distance) s_distance = buffered_distance # Cut off inaccurate values distance_limit = -voxel_distance * (2 * valid_mask - 1) s_distance = math.where( math.abs(s_distance) < voxel_distance, s_distance, distance_limit) if isinstance(input_field, StaggeredGrid): ext_field = input_field.with_data(ext_data) stagger_slice = tuple([slice(0, -1) for i in dims]) s_distance = s_distance[(slice(None), ) + stagger_slice + (slice(None), )] else: ext_field = input_field.copied_with(data=ext_data) return ext_field, s_distance