def multi_advect(self, fields, interpolation="LINEAR", dt=1):
assert isinstance(
fields,
(list, tuple)), "first parameter must be either a tuple or list"
inputs_lists = []
coords_lists = []
value_generators = []
for field in fields:
if isinstance(field, StaggeredGrid):
i, c, v = self._mac_block_advection(field.staggered, dt)
else:
i, c, v = self._centered_block_advection(field, dt)
inputs_lists.append(i)
coords_lists.append(c)
value_generators.append(v)
inputs = math.concat(sum(inputs_lists, []), 0)
coords = math.concat(sum(coords_lists, []), 0)
all_advected = math.resample(inputs,
coords,
interpolation=interpolation,
boundary="REPLICATE")
all_advected = math.reshape(all_advected, [self.spatial_rank, -1] +
list(all_advected.shape[1:]))
all_advected = math.unstack(all_advected)
results = []
abs_i = 0
for i in range(len(inputs_lists)):
n = len(inputs_lists[0])
assigned_advected = all_advected[abs_i:abs_i + n]
results.append(value_generators[i](assigned_advected))
abs_i += n
return results
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 distribute_points(density, particles_per_cell=1, distribution='uniform'):
"""
Distribute points according to the distribution specified in density.
:param density: binary tensor
:param particles_per_cell: integer
:param distribution: 'uniform' or 'center'
:return: tensor of shape (batch_size, point_count, rank)
"""
assert distribution in ('center', 'uniform')
index_array = []
batch_size = math.staticshape(density)[0] if math.staticshape(density)[0] is not None else 1
for batch in range(batch_size):
indices = math.where(density[batch, ..., 0] > 0)
indices = math.to_float(indices)
temp = []
for _ in range(particles_per_cell):
if distribution == 'center':
temp.append(indices + 0.5)
elif distribution == 'uniform':
temp.append(indices + math.random_uniform(math.shape(indices)))
index_array.append(math.concat(temp, axis=0))
try:
index_array = math.stack(index_array)
return index_array
except ValueError:
raise ValueError("all arrays in the batch must have the same number of active cells.")
def concat(geometries: tuple or list,
dim: str,
sizes: tuple or list or None = None):
"""
Concatenates multiple geometries of the same type.
Args:
geometries: sequence of `phi.geom.Geometry` objects of the same type
sizes: implicit
dim: dimension to concatenate
Returns:
New `phi.geom.Geometry` object
"""
if all(isinstance(g, type(geometries[0])) for g in geometries):
characteristics = [g.__characteristics__() for g in geometries]
if sizes is not None:
characteristics = [{
key: math.expand(val, dim, size)
for key, val in c.items()
} for c, size in zip(characteristics, sizes)]
new_attributes = {}
for key in characteristics[0].keys():
concatenated = math.concat([c[key] for c in characteristics], dim)
new_attributes[key] = concatenated
return geometries[0].__with__(**new_attributes)
else:
raise NotImplementedError()
def _generate_examples():
# --- Example 1 ---
ex1 = np.tile(np.linspace(1, 0, 5), [4, 1])
ex1 = math.expand_dims(math.expand_dims(ex1, -1), 0) - math.mean(ex1)
# --- Example 2 ---
ex2 = np.zeros([1, 4, 5, 1])
ex2[0, :, 2, 0] = 1
ex2 -= math.mean(ex2)
# --- Stack examples to batch ---
return math.concat([ex1, ex2], axis=0)
def concat(*fields: SampledField, dim: str):
assert all(isinstance(f, SampledField) for f in fields)
assert all(isinstance(f, type(fields[0])) for f in fields)
if any(f.extrapolation != fields[0].extrapolation for f in fields):
raise NotImplementedError("Concatenating extrapolations not supported")
if isinstance(fields[0], Grid):
values = math.concat([f.values for f in fields], dim=dim)
return fields[0].with_(values=values)
elif isinstance(fields[0], PointCloud):
elements = geom.concat([f.elements for f in fields],
dim,
sizes=[f.shape.get_size(dim) for f in fields])
values = math.concat([
math.expand(f.values, dim, f.shape.get_size(dim)) for f in fields
], dim)
colors = math.concat(
[math.expand(f.color, dim, f.shape.get_size(dim)) for f in fields],
dim)
return fields[0].with_(elements=elements, values=values, color=colors)
raise NotImplementedError(type(fields[0]))
def at(self, other_field):
if isinstance(other_field, StaggeredGrid) and other_field.box == self.box and np.allclose(other_field.resolution, self.resolution):
return self
try:
points = other_field.points
resampled = [centeredgrid.at(other_field) for centeredgrid in self.data]
data = math.concat([field.data for field in resampled], -1)
return other_field.copied_with(data=data, flags=propagate_flags_resample(self, other_field.flags, other_field.rank))
except IncompatibleFieldTypes:
return broadcast_at(self, other_field)
except StaggeredSamplePoints:
return broadcast_at(self, other_field)
def at(self, other_field, collapse_dimensions=True, force_optimization=False, return_self_if_compatible=False):
if isinstance(other_field, StaggeredGrid) and other_field.box == self.box:
return self
try:
points = other_field.points
resampled = [centeredgrid.at(other_field) for centeredgrid in self.data]
data = math.concat([field.data for field in resampled], -1)
return other_field.copied_with(data=data, flags=propagate_flags_resample(self, other_field.flags, other_field.rank))
except IncompatibleFieldTypes:
return broadcast_at(self, other_field)
except StaggeredSamplePoints:
return broadcast_at(self, other_field)
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 concat(fields: List[SampledFieldType], dim: Shape) -> SampledFieldType:
"""
Concatenates the given `SampledField`s along `dim`.
See Also:
`stack()`.
Args:
fields: List of matching `SampledField` instances.
dim: Concatenation dimension as `Shape`. Size is ignored.
Returns:
`SampledField` matching concatenated fields.
"""
assert all(isinstance(f, SampledField) for f in fields)
assert all(isinstance(f, type(fields[0])) for f in fields)
if any(f.extrapolation != fields[0].extrapolation for f in fields):
raise NotImplementedError("Concatenating extrapolations not supported")
if isinstance(fields[0], Grid):
values = math.concat([f.values for f in fields], dim)
return fields[0].with_values(values)
elif isinstance(fields[0], PointCloud):
elements = geom.concat([f.elements for f in fields],
dim,
sizes=[f.shape.get_size(dim) for f in fields])
values = math.concat(
[math.expand(f.values, f.shape.only(dim)) for f in fields], dim)
colors = math.concat(
[math.expand(f.color, f.shape.only(dim)) for f in fields], dim)
return PointCloud(elements=elements,
values=values,
color=colors,
extrapolation=fields[0].extrapolation,
add_overlapping=fields[0]._add_overlapping,
bounds=fields[0]._bounds)
raise NotImplementedError(type(fields[0]))
def batch_indices(indices):
"""
Reshapes the indices such that, aside from indices, they also contain batch number.
For example the entry (32, 40) as coordinates of batch 2 will become (2, 32, 40).
Transform shape (b, p, d) to (b, p, d+1) where batch size is b, number of particles is p and number of dimensions is d.
"""
batch_size = indices.shape[0]
out_spatial_rank = len(indices.shape) - 2
out_spatial_size = math.shape(indices)[1:-1]
batch_range = math.DYNAMIC_BACKEND.choose_backend(indices).range(batch_size)
batch_ids = math.reshape(batch_range, [batch_size] + [1] * out_spatial_rank)
tile_shape = math.pad(out_spatial_size, [[1,0]], constant_values=1)
batch_ids = math.expand_dims(math.tile(batch_ids, tile_shape), axis=-1)
return math.concat((batch_ids, indices), axis=-1)
def from_scalar(scalar_field, axis_forces, padding_mode="constant"):
if scalar_field.shape[-1] != 1:
raise ValueError("Resample requires a scalar field as input")
rank = spatial_rank(scalar_field)
dims = range(rank)
df_dq = []
for dimension in dims: # z,y,x
padded_field = math.pad(scalar_field, [[0, 0]] +
[[1, 1] if i == dimension else [0, 1]
for i in dims] + [[0, 0]], padding_mode)
upper_slices = [(slice(1, None) if i == dimension else slice(None))
for i in dims]
lower_slices = [(slice(-1) if i == dimension else slice(None))
for i in dims]
neighbour_sum = padded_field[[slice(None)] + upper_slices + [slice(None)]] + \
padded_field[[slice(None)] + lower_slices + [slice(None)]]
df_dq.append(axis_forces[dimension] * neighbour_sum * 0.5 / rank)
return StaggeredGrid(math.concat(df_dq[::-1], axis=-1))
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 gradient(scalar_field, padding="symmetric"):
if scalar_field.shape[-1] != 1:
raise ValueError("Gradient requires a scalar field as input")
rank = spatial_rank(scalar_field)
dims = range(rank)
field = math.pad(scalar_field, [[0, 0]] + [[1, 1]] * rank + [[0, 0]],
mode=padding)
df_dq = []
for dimension in dims:
upper_slices = [
(slice(1, None) if i == dimension else slice(1, None))
for i in dims
]
lower_slices = [(slice(-1) if i == dimension else slice(1, None))
for i in dims]
diff = field[[slice(None)] + upper_slices] - field[[slice(None)] +
lower_slices]
df_dq.append(diff)
return StaggeredGrid(math.concat(df_dq[::-1], axis=-1))
def _advect_mac(self, field_mac, dt, interpolation):
# resample in each dimension
idx = indices_tensor(self.staggered)
advected_component_fields = []
dims = range(len(self.staggered.shape) - 2)
for d in dims: # z,y,x
velocity_at_staggered_points = self.at_faces(len(dims) - d -
1)[..., ::-1]
sample_coords = idx - velocity_at_staggered_points * dt
d_comp = len(dims) - d - 1
advected = math.resample(field_mac[..., d_comp:d_comp + 1],
sample_coords,
interpolation=interpolation,
boundary="REPLICATE")
advected_component_fields.append(advected)
all_advected = math.concat(advected_component_fields[::-1], axis=-1)
return all_advected
def _distribute_points(mask: math.Tensor, points_per_cell: int = 1, center: bool = False) -> math.Tensor:
"""
Generates points (either uniformly distributed or at the cell centers) according to the given tensor mask.
Args:
mask: Tensor with nonzero values at the indices where particles should get generated.
points_per_cell: Number of particles to generate at each marked index
center: Set points to cell centers. If False, points will be distributed using a uniform
distribution within each cell.
Returns:
A tensor containing the positions of the generated points.
"""
indices = math.to_float(math.nonzero(mask, list_dim=instance('points')))
temp = []
for _ in range(points_per_cell):
if center:
temp.append(indices + 0.5)
else:
temp.append(indices + (math.random_uniform(indices.shape)))
return math.concat(temp, dim=instance('points'))
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 test_nonzero(self):
c = math.concat([math.zeros(b=3, a=2), math.ones(a=2, b=4)], 'b')
nz = math.nonzero(c)
self.assertEqual(nz.shape.nonzero, 8)
self.assertEqual(nz.shape.vector, 2)
def test_concat(self):
c = math.concat([math.zeros(b=3, a=2), math.ones(a=2, b=4)], 'b')
self.assertEqual(2, c.shape.a)
self.assertEqual(7, c.shape.b)
math.assert_close(c.b[:3], 0)
math.assert_close(c.b[3:], 1)
def test_nonzero(self):
c = math.concat([math.zeros(spatial(b=3, a=2)), math.ones(spatial(a=2, b=4))], spatial('b'))
nz = math.nonzero(c)
self.assertEqual(nz.shape.get_size('nonzero'), 8)
self.assertEqual(nz.shape.get_size('vector'), 2)
def post_advection(advected_list):
return StaggeredGrid(math.concat(advected_list[::-1], axis=-1))
def sample_at(self, points):
return math.concat([component.sample_at(points) for component in self.data], axis=-1)
def u_net(domain, input_fields, output_field, levels=2, filters=16, blocks_per_level=2, skip_combine='concat', training=False, trainable=True, reuse=None):
"""
Restrictions:
- 2D only
- Domain resolution must be multiple of 2**levels
:param skip_combine: 'concat'
:param blocks_per_level: number of residual blocks per level
:param filters: Number of convolutional filters
:type filters: int or tuple or list
:param levels: number of additional resolution levels, equals number of downsampling / upsampling operations
:param domain: the u-net is executed on this domain.
:type domain: Domain
:param input_fields: list of Fields to be passed to the network as input
:param output_field: determines sample points of the result
:param training: whether the network is executed in training or inference mode
:param trainable: whether the weights of the network are trainable
:param reuse: whether to reuse weights from previous unet calls
:return: Field sampled like output_field
"""
assert isinstance(domain, Domain)
assert isinstance(output_field, Field)
net_inputs = []
for input_field in input_fields:
assert isinstance(input_field, Field)
resampled = input_field.at(domain)
net_inputs.append(resampled)
y = CenteredGrid.sample(math.concat([math.to_float(grid.data) for grid in net_inputs], axis=-1), domain)
# --- Execute network ---
pad_width = sum([2 ** i for i in range(levels)])
y = y.padded([[0, pad_width]] * domain.rank)
resolutions = [y]
for level in range(levels):
level_filters = filters if isinstance(filters, int) else filters[level]
y = conv_layer(resolutions[0], level_filters, 2, strides=2, activation=tf.nn.relu, padding='valid', name='down_convolution_%d' % level, trainable=trainable, reuse=reuse)
for i in range(blocks_per_level):
y = residual_block(y, level_filters, name='down_res_block_%d_%d' % (level, i), training=training, trainable=trainable, reuse=reuse)
resolutions.insert(0, y)
y = resolutions.pop(0)
assert np.all(y.box.size == domain.box.size)
for level in range(levels):
y = math.upsample2x(y)
res_in = resolutions.pop(0)
res_in = res_in.at(y) # No resampling required, simply shaving off the top rows
if skip_combine == 'concat':
y = y.with_data(math.concat([y.data, res_in.data], axis=-1))
else:
raise NotImplementedError()
y = y + res_in
y = y.padded([[0, 1], [0, 1]])
if resolutions:
level_filters = filters if isinstance(filters, int) else reversed(filters)[level]
y = conv_layer(y, level_filters, kernel_size=2, activation=tf.nn.relu, padding='valid', name='up_convolution_%d' % level, trainable=trainable, reuse=reuse)
for i in range(blocks_per_level):
y = residual_block(y, level_filters, name='up_res_block_%d_%d' % (level, i), training=training, trainable=trainable, reuse=reuse)
else: # Last iteration
y = conv_layer(y, output_field.component_count, kernel_size=2, activation=None, padding='valid', name='up_convolution_%d' % level, trainable=trainable, reuse=reuse)
result = y.at(output_field)
return result
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
def sample_at(self, points, collapse_dimensions=True):
return math.concat(
[component.sample_at(points) for component in self.data], axis=-1)
def stack_staggered_components(tensors):
for i, tensor in enumerate(tensors):
paddings = [[0, 1] if d != i else [0, 0] for d in range(len(tensors))]
tensors[i] = math.pad(tensor, [[0, 0]] + paddings + [[0, 0]])
return math.concat(tensors, -1)
def test_concat(self):
c = math.concat([math.zeros(spatial(b=3, a=2)), math.ones(spatial(a=2, b=4))], spatial('b'))
self.assertEqual(2, c.shape.get_size('a'))
self.assertEqual(7, c.shape.get_size('b'))
math.assert_close(c.b[:3], 0)
math.assert_close(c.b[3:], 1)
