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 cell_index(self, global_position):
local_position = self.box.global_to_local(
global_position) * self.resolution
position = math.to_int(local_position - 0.5)
position = math.maximum(0, position)
position = math.minimum(position, self.resolution - 1)
return position
def _frictionless_velocity_mask(self, velocity):
tensors = []
for axis in range(velocity.rank):
upper = self.accessible.padded([[0, 1] if ax == axis else [0, 0] for ax in range(self.rank)])
lower = self.accessible.padded([[1, 0] if ax == axis else [0, 0] for ax in range(self.rank)])
tensors.append(math.minimum(upper.data, lower.data))
return velocity.with_data(tensors)
def _grid_sample(self, box, resolution):
"""
Samples this field on a regular grid.
:param box: physical dimensions of the grid
:param resolution: grid resolution
:return: CenteredGrid
"""
sample_indices_nd = math.to_int(
math.round(box.global_to_local(self.sample_points) * resolution))
sample_indices_nd = math.minimum(
math.maximum(0, sample_indices_nd), resolution - 1
) # Snap outside points to edges, otherwise scatter raises an error
# Correct format for math.scatter
valid_indices = _batch_indices(sample_indices_nd)
shape = (math.shape(
self.data)[0], ) + tuple(resolution) + (self.data.shape[-1], )
scattered = math.scatter(self.sample_points,
valid_indices,
self.data,
shape,
duplicates_handling=self.mode)
return CenteredGrid(data=scattered,
box=box,
extrapolation='constant',
name=self.name + '_centered')
def _grid_sample(self, box, resolution):
"""
Samples this field on a regular grid.
:param box: physical dimensions of the grid
:param resolution: grid resolution
:return: CenteredGrid
"""
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)
scattered = math.scatter(self.sample_points, valid_indices, self.data, math.concat([[valid_indices.shape[0]], resolution, [1]], axis=-1), duplicates_handling=self.mode)
return CenteredGrid(data=scattered, box=box, extrapolation='constant', name=self.name+'_centered')
def minimum(f1: Field or Geometry or float, f2: Field or Geometry or float):
"""
Element-wise minimum.
One of the given fields needs to be an instance of `SampledField` and the the result will be sampled at the corresponding points.
If both are `SampledFields` but have different points, `f1` takes priority.
Args:
f1: `Field` or `Geometry` or constant.
f2: `Field` or `Geometry` or constant.
Returns:
`SampledField`
"""
f1, f2 = _auto_resample(f1, f2)
return f1.with_values(math.minimum(f1.values, f2.values))
def sparse_values(dimensions, extended_active_mask, extended_fluid_mask, sorting=None, periodic=False):
"""
Builds a sparse matrix such that when applied to a flattened pressure channel, it calculates the laplace
of that channel, taking into account obstacles and empty cells.
:param dimensions: valid simulation dimensions. Pressure channel should be of shape (batch size, dimensions..., 1)
:param extended_active_mask: Binary tensor with 2 more entries in every dimension than 'dimensions'.
:param extended_fluid_mask: Binary tensor with 2 more entries in every dimension than 'dimensions'.
:return: SciPy sparse matrix that acts as a laplace on a flattened pressure channel given obstacles and empty cells
"""
N = int(np.prod(dimensions))
d = len(dimensions)
dims = range(d)
values_list = []
diagonal_entries = 0 # diagonal matrix entries
gridpoints_linear = np.arange(N)
gridpoints = np.stack(np.unravel_index(gridpoints_linear, dimensions)) # d * (N^2) array mapping from linear to spatial frames
for dim in dims:
lower_active, self_active, upper_active = _dim_shifted(extended_active_mask, dim, (-1, 0, 1), diminish_others=(1, 1))
lower_accessible, upper_accessible = _dim_shifted(extended_fluid_mask, dim, (-1, 1), diminish_others=(1, 1))
stencil_upper = upper_active * self_active
stencil_lower = lower_active * self_active
stencil_center = - lower_accessible - upper_accessible
diagonal_entries += math.flatten(stencil_center)
dim_direction = math.expand_dims([1 if i == dim else 0 for i in range(d)], axis=-1)
# --- Stencil upper cells ---
upper_points, upper_idx = wrap_or_discard(gridpoints + dim_direction, dim, dimensions, periodic=collapsed_gather_nd(periodic, [dim, 1]))
values_list.append(math.gather(math.flatten(stencil_upper), upper_idx))
# --- Stencil lower cells ---
lower_points, lower_idx = wrap_or_discard(gridpoints - dim_direction, dim, dimensions, periodic=collapsed_gather_nd(periodic, [dim, 0]))
values_list.append(math.gather(math.flatten(stencil_lower), lower_idx))
values_list.insert(0, math.minimum(diagonal_entries, -1.))
values = math.concat(values_list, axis=0)
if sorting is not None:
values = math.gather(values, sorting)
return values
def sparse_pressure_matrix(dimensions, extended_active_mask, extended_fluid_mask, periodic=False):
"""
Builds a sparse matrix such that when applied to a flattened pressure channel, it calculates the laplace
of that channel, taking into account obstacles and empty cells.
:param dimensions: valid simulation dimensions. Pressure channel should be of shape (batch size, dimensions..., 1)
:param extended_active_mask: Binary tensor with 2 more entries in every dimension than 'dimensions'.
:param extended_fluid_mask: Binary tensor with 2 more entries in every dimension than 'dimensions'.
:return: SciPy sparse matrix that acts as a laplace on a flattened pressure channel given obstacles and empty cells
"""
N = int(np.prod(dimensions))
d = len(dimensions)
A = scipy.sparse.lil_matrix((N, N), dtype=np.float32)
dims = range(d)
diagonal_entries = np.zeros(N, extended_active_mask.dtype) # diagonal matrix entries
gridpoints_linear = np.arange(N)
gridpoints = np.stack(np.unravel_index(gridpoints_linear, dimensions)) # d * (N^2) array mapping from linear to spatial frames
for dim in dims:
lower_active, self_active, upper_active = _dim_shifted(extended_active_mask, dim, (-1, 0, 1), diminish_others=(1,1))
lower_accessible, upper_accessible = _dim_shifted(extended_fluid_mask, dim, (-1, 1), diminish_others=(1, 1))
stencil_upper = upper_active * self_active
stencil_lower = lower_active * self_active
stencil_center = - lower_accessible - upper_accessible
diagonal_entries += math.flatten(stencil_center)
dim_direction = math.expand_dims([1 if i == dim else 0 for i in range(d)], axis=-1)
# --- Stencil upper cells ---
upper_points, upper_idx = wrap_or_discard(gridpoints + dim_direction, dim, dimensions, periodic=collapsed_gather_nd(periodic, [dim, 1]))
A[gridpoints_linear[upper_idx], upper_points] = stencil_upper.flatten()[upper_idx]
# --- Stencil lower cells ---
lower_points, lower_idx = wrap_or_discard(gridpoints - dim_direction, dim, dimensions, periodic=collapsed_gather_nd(periodic, [dim, 0]))
A[gridpoints_linear[lower_idx], lower_points] = stencil_lower.flatten()[lower_idx]
A[gridpoints_linear, gridpoints_linear] = math.minimum(diagonal_entries, -1) # avoid 0, could lead to NaN
return scipy.sparse.csc_matrix(A)
def sample_at(self, points):
x = (points - self.center) / self.unit_distance
pot = math.sum(x**2, -1, keepdims=True) * self.data
if self.maximum_value is not None:
pot = math.minimum(pot, self.maximum_value)
return math.cast(pot, np.float32)
def sparse_pressure_matrix(dimensions, extended_active_mask,
extended_fluid_mask):
"""
Builds a sparse matrix such that when applied to a flattened pressure channel, it calculates the laplace
of that channel, taking into account obstacles and empty cells.
:param dimensions: valid simulation dimensions. Pressure channel should be of shape (batch size, dimensions..., 1)
:param extended_active_mask: Binary tensor with 2 more entries in every dimension than 'dimensions'.
:param extended_fluid_mask: Binary tensor with 2 more entries in every dimension than 'dimensions'.
:return: SciPy sparse matrix that acts as a laplace on a flattened pressure channel given obstacles and empty cells
"""
N = int(np.prod(dimensions))
d = len(dimensions)
A = scipy.sparse.lil_matrix((N, N), dtype=np.float32)
dims = range(d)
center_values = None # diagonal matrix entries
gridpoints_linear = np.arange(N)
gridpoints = np.stack(np.unravel_index(
gridpoints_linear,
dimensions)) # d * (N^2) array mapping from linear to spatial frames
for dim in dims:
upper_indices = tuple(
[slice(None)] +
[slice(2, None) if i == dim else slice(1, -1)
for i in dims] + [slice(None)])
center_indices = tuple(
[slice(None)] +
[slice(1, -1) if i == dim else slice(1, -1)
for i in dims] + [slice(None)])
lower_indices = tuple(
[slice(None)] +
[slice(0, -2) if i == dim else slice(1, -1)
for i in dims] + [slice(None)])
self_active = extended_active_mask[center_indices]
stencil_upper = extended_active_mask[upper_indices] * self_active
stencil_lower = extended_active_mask[lower_indices] * self_active
stencil_center = -extended_fluid_mask[
upper_indices] - extended_fluid_mask[lower_indices]
if center_values is None:
center_values = math.flatten(stencil_center)
else:
center_values = center_values + math.flatten(stencil_center)
# Find entries in matrix
dim_direction = np.zeros_like(gridpoints)
dim_direction[dim] = 1
# Upper frames
upper_indices = gridpoints + dim_direction
upper_in_range_inx = np.nonzero(upper_indices[dim] < dimensions[dim])
upper_indices_linear = np.ravel_multi_index(
upper_indices[:, upper_in_range_inx], dimensions)
A[gridpoints_linear[upper_in_range_inx],
upper_indices_linear] = stencil_upper.flatten()[upper_in_range_inx]
# Lower frames
lower_indices = gridpoints - dim_direction
lower_in_range_inx = np.nonzero(lower_indices[dim] >= 0)
lower_indices_linear = np.ravel_multi_index(
lower_indices[:, lower_in_range_inx], dimensions)
A[gridpoints_linear[lower_in_range_inx],
lower_indices_linear] = stencil_lower.flatten()[lower_in_range_inx]
A[gridpoints_linear, gridpoints_linear] = math.minimum(center_values, -1)
return scipy.sparse.csc_matrix(A)
def sparse_values(dimensions,
extended_active_mask,
extended_fluid_mask,
sorting=None):
"""
Builds a sparse matrix such that when applied to a flattened pressure channel, it calculates the laplace
of that channel, taking into account obstacles and empty cells.
:param dimensions: valid simulation dimensions. Pressure channel should be of shape (batch size, dimensions..., 1)
:param extended_active_mask: Binary tensor with 2 more entries in every dimension than 'dimensions'.
:param extended_fluid_mask: Binary tensor with 2 more entries in every dimension than 'dimensions'.
:return: SciPy sparse matrix that acts as a laplace on a flattened pressure channel given obstacles and empty cells
"""
N = int(np.prod(dimensions))
d = len(dimensions)
dims = range(d)
values_list = []
center_values = None # diagonal matrix entries
gridpoints_linear = np.arange(N)
gridpoints = np.stack(np.unravel_index(
gridpoints_linear,
dimensions)) # d * (N^2) array mapping from linear to spatial frames
for dim in dims:
upper_indices = tuple(
[slice(None)] +
[slice(2, None) if i == dim else slice(1, -1)
for i in dims] + [slice(None)])
center_indices = tuple(
[slice(None)] +
[slice(1, -1) if i == dim else slice(1, -1)
for i in dims] + [slice(None)])
lower_indices = tuple(
[slice(None)] +
[slice(0, -2) if i == dim else slice(1, -1)
for i in dims] + [slice(None)])
self_active = extended_active_mask[center_indices]
stencil_upper = extended_active_mask[upper_indices] * self_active
stencil_lower = extended_active_mask[lower_indices] * self_active
stencil_center = -extended_fluid_mask[
upper_indices] - extended_fluid_mask[lower_indices]
if center_values is None:
center_values = math.flatten(stencil_center)
else:
center_values = center_values + math.flatten(stencil_center)
dim_direction = np.zeros_like(gridpoints)
dim_direction[dim] = 1
# Upper frames
upper_indices = gridpoints + dim_direction
upper_in_range_inx = np.nonzero(
upper_indices[dim] < dimensions[dim])[0]
values_list.append(
math.gather(math.flatten(stencil_upper), upper_in_range_inx))
# Lower frames
lower_indices = gridpoints - dim_direction
lower_in_range_inx = np.nonzero(lower_indices[dim] >= 0)[0]
values_list.append(
math.gather(math.flatten(stencil_lower), lower_in_range_inx))
center_values = math.minimum(center_values, -1.)
values_list.insert(0, center_values)
values = math.concat(values_list, axis=0)
if sorting is not None:
values = math.gather(values, sorting)
return values
