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
