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
