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 _expand_axes(data, points, collapse_dimensions=True): assert math.spatial_rank(data) >= 0 data = math.expand_dims(data, 1, math.spatial_rank(points) - math.spatial_rank(data)) if collapse_dimensions: return data else: points_axes = math.staticshape(points)[1:-1] data_axes = math.staticshape(data)[1:-1] for d_points, d_data in zip(points_axes, data_axes): assert d_points % d_data == 0 tilings = [1] + [d_points // d_data for d_points, d_data in zip(math.staticshape(points)[1:-1], math.staticshape(data)[1:-1])] + [1] data = math.tile(data, tilings) return data
def _mg_solve_forward(divergence, domain, pressure_guess, solvers): fluid_mask = domain.accessible_tensor(extend=1) active_mask = domain.active_tensor(extend=1) if active_mask is not None or fluid_mask is not None: if not np.all([s.supports_continuous_masks for s in solvers[:-1]]): logging.warning( "MultiscaleSolver solver: There are boundary conditions inside the domain but " "not all intermediate solvers support continuous masks") div_lvls = [divergence] act_lvls = [active_mask] fld_lvls = [fluid_mask] for grid_i in range(len(solvers) - 1): div_lvls.insert(0, math.downsample2x(div_lvls[0])) act_lvls.insert( 0, math.downsample2x(act_lvls[0]) if act_lvls[0] is not None else None) fld_lvls.insert( 0, math.downsample2x(fld_lvls[0]) if fld_lvls[0] is not None else None) if pressure_guess is not None: pressure_guess = math.downsample2x(pressure_guess) iter_list = [] for i, div in enumerate(div_lvls): pressure_guess, iteration = solvers[i].solve( div, FluidDomain(act_lvls[i], fld_lvls[i], boundaries), pressure_guess) iter_list.append(iteration) if pressure_guess.shape[1] < divergence.shape[1]: pressure_guess = math.upsample2x( pressure_guess) * 2**math.spatial_rank(divergence) return pressure_guess, iter_list
def _expand_axes(data, points, batch_size=1): assert math.spatial_rank(data) >= 0 data = math.expand_dims( data, 1, math.spatial_rank(points) - math.spatial_rank(data)) points_axes = math.staticshape(points)[1:-1] data_axes = math.staticshape(data)[1:-1] for d_points, d_data in zip(points_axes, data_axes): assert d_points % d_data == 0 tilings = [batch_size or 1] + [ d_points // d_data for d_points, d_data in zip( math.staticshape(points)[1:-1], math.staticshape(data)[1:-1]) ] + [1] data = math.tile(data, tilings) return data
def _rank(rank): if rank is None: return None if isinstance(rank, int): return rank if isinstance(rank, Geometry): return rank.rank else: return math.spatial_rank(rank)
def _rank(rank): if rank is None: return None elif isinstance(rank, int): pass elif isinstance(rank, Geometry): rank = rank.rank else: rank = math.spatial_rank(rank) return None if rank == 0 else rank
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 _weighted_sliced_laplace_nd(tensor, weights): if tensor.shape[-1] != 1: raise ValueError('Laplace operator requires a scalar channel as input') dims = range(math.spatial_rank(tensor)) components = [] for dimension in dims: lower_weights, center_weights, upper_weights = _dim_shifted( weights, dimension, (-1, 0, 1), diminish_others=(1, 1)) lower_values, center_values, upper_values = _dim_shifted( tensor, dimension, (-1, 0, 1), diminish_others=(1, 1)) diff = math.mul( upper_values, upper_weights * center_weights) + math.mul( lower_values, lower_weights * center_weights) + math.mul( center_values, -lower_weights - upper_weights) components.append(diff) return math.sum(components, 0)
def _weighted_sliced_laplace_nd(tensor, weights): if tensor.shape[-1] != 1: raise ValueError('Laplace operator requires a scalar channel as input') dims = range(math.spatial_rank(tensor)) components = [] for dimension in dims: center_slices = tuple([(slice(1, -1) if i == dimension else slice(1,-1)) for i in dims]) upper_slices = tuple([(slice(2, None) if i == dimension else slice(1,-1)) for i in dims]) lower_slices = tuple([(slice(-2) if i == dimension else slice(1,-1)) for i in dims]) lower_weights = weights[(slice(None),) + lower_slices + (slice(None),)] * weights[(slice(None),) + center_slices + (slice(None),)] upper_weights = weights[(slice(None),) + upper_slices + (slice(None),)] * weights[(slice(None),) + center_slices + (slice(None),)] center_weights = - lower_weights - upper_weights lower_values = tensor[(slice(None),) + lower_slices + (slice(None),)] upper_values = tensor[(slice(None),) + upper_slices + (slice(None),)] center_values = tensor[(slice(None),) + center_slices + (slice(None),)] diff = math.mul(upper_values, upper_weights) + \ math.mul(lower_values, lower_weights) + \ math.mul(center_values, center_weights) components.append(diff) return math.sum(components, 0)
def rank(self): return math.spatial_rank(self.data)
def staggered_grid(tensor, name='manta_staggered'): tensor = tensor[..., ::-1] # manta: xyz, phiflow: zyx assert math.staticshape(tensor)[-1] == math.spatial_rank(tensor) return StaggeredGrid(tensor, name=name)
def centered_grid(tensor, name='manta_centered', crop_valid=False): if crop_valid: tensor = tensor[(slice(None), ) + (slice(-1), ) * math.spatial_rank(tensor) + (slice(None), )] return CenteredGrid(tensor, name=name)
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