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