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 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