def create_surface_mask(liquid_mask): """ Computes inner contours of the liquid_mask. A cell i is flagged 1 if liquid_mask[i] = 1 and it has a non-liquid neighbour. :param liquid_mask: binary tensor :return: tensor """ # When we create inner contour, we don't want the fluid-wall boundaries to show up as surface, so we should pad with symmetric edge values. mask = math.pad(liquid_mask, [[0, 0]] + [[1, 1]] * math.spatial_rank(liquid_mask) + [[0, 0]], "constant") dims = range(math.spatial_rank(mask)) bcs = math.zeros_like(liquid_mask) # Move in every possible direction to assure corners are properly set. directions = np.array( list(itertools.product(*np.tile((-1, 0, 1), (len(dims), 1))))) for d in directions: d_slice = tuple([(slice(2, None) if d[i] == -1 else slice(0, -2) if d[i] == 1 else slice(1, -1)) for i in dims]) center_slice = tuple([slice(1, -1) for _ in dims]) # Create inner contour of particles bc_d = math.maximum(mask[(slice(None),) + d_slice + (slice(None),)], mask[(slice(None),) + center_slice + (slice(None),)]) - \ mask[(slice(None),) + d_slice + (slice(None),)] bcs = math.maximum(bcs, bc_d) return bcs
def test_direct_initializers(self): np.testing.assert_equal(math.zeros([1, 16]), np.zeros([1, 16])) self.assertEqual(math.zeros([1, 16]).dtype, np.float32) np.testing.assert_equal(math.ones([1, 16, 1]), np.ones([1, 16, 1])) np.testing.assert_equal(math.zeros_like(math.ones([1, 16, 1])), np.zeros([1, 16, 1])) np.testing.assert_equal(math.randn([1, 4]).shape, [1, 4]) self.assertEqual(math.randn([1, 4]).dtype, np.float32)
def test_struct_initializers(self): obj = ([4], CenteredGrid([1, 4, 1], box[0:1], content_type=struct.shape), ([9], [8, 2])) z = math.zeros(obj) self.assertIsInstance(z, tuple) np.testing.assert_equal(z[0], np.zeros([4])) z2 = math.zeros_like(z) np.testing.assert_equal(math.shape(z)[0], math.shape(z2)[0])
def test_struct_initializers(self): bounds = box[0:1] # outside unsafe with struct.unsafe(): obj = ([4], CenteredGrid([1, 4, 1], bounds), ([9], [8, 2])) z = math.zeros(obj) self.assertIsInstance(z, tuple) np.testing.assert_equal(z[0], np.zeros([4])) z2 = math.zeros_like(z) np.testing.assert_equal(math.shape(z)[0], math.shape(z2)[0])
def _stagger_sample(self, box, resolution): """ Samples this field on a staggered grid. In addition to sampling, extrapolates the field using an occupancy mask generated from the points. :param box: physical dimensions of the grid :param resolution: grid resolution :return: StaggeredGrid """ resolution = np.array(resolution) valid_indices = math.to_int(math.floor(self.sample_points)) valid_indices = math.minimum(math.maximum(0, valid_indices), resolution - 1) # Correct format for math.scatter valid_indices = batch_indices(valid_indices) active_mask = math.scatter(self.sample_points, valid_indices, 1, math.concat([[valid_indices.shape[0]], resolution, [1]], axis=-1), duplicates_handling='any') mask = math.pad(active_mask, [[0, 0]] + [[1, 1]] * self.rank + [[0, 0]], "constant") if isinstance(self.data, (int, float, np.ndarray)): values = math.zeros_like(self.sample_points) + self.data else: values = self.data result = [] ones_1d = math.unstack(math.ones_like(values), axis=-1)[0] staggered_shape = [i + 1 for i in resolution] dx = box.size / resolution dims = range(len(resolution)) for d in dims: staggered_offset = math.stack([(0.5 * dx[i] * ones_1d if i == d else 0.0 * ones_1d) for i in dims], axis=-1) indices = math.to_int(math.floor(self.sample_points + staggered_offset)) valid_indices = math.maximum(0, math.minimum(indices, resolution)) valid_indices = batch_indices(valid_indices) values_d = math.expand_dims(math.unstack(values, axis=-1)[d], axis=-1) result.append(math.scatter(self.sample_points, valid_indices, values_d, [indices.shape[0]] + staggered_shape + [1], duplicates_handling=self.mode)) d_slice = tuple([(slice(0, -2) if i == d else slice(1,-1)) for i in dims]) u_slice = tuple([(slice(2, None) if i == d else slice(1,-1)) for i in dims]) active_mask = math.minimum(mask[(slice(None),) + d_slice + (slice(None),)], active_mask) active_mask = math.minimum(mask[(slice(None),) + u_slice + (slice(None),)], active_mask) staggered_tensor_prep = unstack_staggered_tensor(math.concat(result, axis=-1)) grid_values = StaggeredGrid(staggered_tensor_prep) # Fix values at boundary of liquids (using StaggeredGrid these might not receive a value, so we replace it with a value inside the liquid) grid_values, _ = extrapolate(grid_values, active_mask, voxel_distance=2) return grid_values
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 map_velocity_to_particles(previous_particle_velocity: PointCloud, velocity_grid: Grid, occupation_mask: Grid, previous_velocity_grid: Grid = None, viscosity: float = 0.) -> PointCloud: """ Maps result of velocity projection on grid back to particles. Provides option to choose between FLIP (particle velocities are updated by the change between projected and initial grid velocities) and PIC (particle velocities are replaced by the the projected velocities) method depending on the value of the `initial_v_field`. Args: previous_particle_velocity: PointCloud with particle positions as elements and their corresponding velocities as values velocity_grid: Divergence-free velocity grid occupation_mask: Binary Grid (same type as `velocity_grid`) indicating which cells hold particles previous_velocity_grid: Velocity field before projection and force update viscosity: If previous_velocity_grid is None, the particle-in-cell method (PIC) is applied. Otherwise this is the ratio between FLIP and PIC (0. for pure FLIP) Returns: PointCloud with particle positions as elements and updated particle velocities as values. """ viscosity = min(max(0., viscosity), 1.) if previous_velocity_grid is None: viscosity = 1. velocities = math.zeros_like(previous_particle_velocity.values) if viscosity > 0.: # --- PIC --- velocity_grid, _ = extrapolate_valid(velocity_grid, occupation_mask) velocities += viscosity * (velocity_grid @ previous_particle_velocity).values if viscosity < 1.: # --- FLIP --- v_change_field = velocity_grid - previous_velocity_grid v_change_field, _ = extrapolate_valid(v_change_field, occupation_mask) v_change = (v_change_field @ previous_particle_velocity).values velocities += (1 - viscosity) * (previous_particle_velocity.values + v_change) return previous_particle_velocity.with_values(velocities)
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
def data(self, data): if isinstance(data, (tuple, list, np.ndarray)): data = math.zeros_like(self.sample_points) + data return data