def sm_resnet(initial, target, frames, training=True, trainable=True, reuse=tf.AUTO_REUSE):
frames = math.expand_dims(math.expand_dims(frames, -1), -1)
frames = tf.tile(frames, [1]+list(initial.shape[1:]))
y = tf.concat([initial, target, frames], axis=-1)
downres_padding = sum([2**i for i in range(5)])
y = tf.pad(y, [[0,0], [0,downres_padding], [0,0]])
resolutions = [ y ]
for i, filters in enumerate([4, 8, 16, 16, 16]):
y = tf.layers.conv1d(resolutions[0], filters, 2, strides=2, activation=tf.nn.relu, padding='valid', name='downconv_%d'%i, trainable=trainable, reuse=reuse)
for j in range(2):
y = residual_block_1d(y, filters, name='downrb_%d_%d' % (i,j), training=training, trainable=trainable, reuse=reuse)
resolutions.insert(0, y)
for j, nb_channels in enumerate([16, 16, 16]):
y = residual_block_1d(y, nb_channels, name='centerrb_%d' % j, training=training, trainable=trainable, reuse=reuse)
for i, resolution_data in enumerate(resolutions[1:]):
y = upsample2x(y)
res_in = resolution_data[:, 0:y.shape[1], :]
y = tf.concat([y, res_in], axis=-1)
if i < len(resolutions)-2:
y = tf.pad(y, [[0, 0], [0, 1], [0, 0]], mode='SYMMETRIC')
y = tf.layers.conv1d(y, 16, 2, 1, activation=tf.nn.relu, padding='valid', name='upconv_%d' % i, trainable=trainable, reuse=reuse)
for j, nb_channels in enumerate([16, 16]):
y = residual_block_1d(y, nb_channels, 2, name='uprb_%d_%d' % (i, j), training=training, trainable=trainable, reuse=reuse)
else:
# Last iteration
y = tf.pad(y, [[0,0], [0,1], [0,0]], mode='SYMMETRIC')
y = tf.layers.conv1d(y, 1, 2, 1, activation=None, padding='valid', name='upconv_%d'%i, trainable=trainable, reuse=reuse)
return y
def _convert_constant_to_data(value):
if isinstance(value, numbers.Number):
value = math.to_float(math.expand_dims(value))
if isinstance(value, (list, tuple)):
value = math.to_float(numpy.array(value))
if len(math.staticshape(value)) < 2:
value = math.expand_dims(value)
return value
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 gravity_tensor(gravity, rank):
if isinstance(gravity, Gravity):
gravity = gravity.gravity
if math.is_scalar(gravity):
return math.to_float(
math.expand_dims([gravity] + [0] * (rank - 1), 0, rank + 1))
else:
assert math.staticshape(gravity)[-1] == rank
return math.to_float(
math.expand_dims(gravity, 0,
rank + 2 - len(math.staticshape(gravity))))
def spatial_sum(tensor):
if isinstance(tensor, StaggeredGrid):
tensor = tensor.staggered
summed = math.sum(tensor, axis=math.dimrange(tensor))
for i in math.dimrange(tensor):
summed = math.expand_dims(summed, i)
return summed
def sample_at(self, points):
envelope = math.exp(-0.5 * math.sum(
(points - self.center)**2, axis=-1, keepdims=True) / self.size**2)
envelope = math.to_float(envelope)
wave = math.exp(1j * math.to_float(
math.expand_dims(np.dot(points, self.wave_vector), -1))) * envelope
return wave * self.data
def staggered_points(self, dimension):
idx_zyx = np.meshgrid(*[
np.arange(0.5, dim + 1.5, 1) if dim != dimension else np.arange(
0, dim + 1, 1) for dim in self.resolution
],
indexing="ij")
return math.expand_dims(math.stack(idx_zyx, axis=-1), 0)
def unstack_staggered_tensor(tensor):
tensors = math.unstack(tensor, -1)
result = []
for i, dim in enumerate(math.spatial_dimensions(tensor)):
slices = [slice(None, -1) if d != dim else slice(None) for d in math.spatial_dimensions(tensor)]
result.append(math.expand_dims(tensors[i][tuple([slice(None)]+slices)], -1))
return result
def sparse_indices(dimensions, periodic=False):
N = int(np.prod(dimensions))
d = len(dimensions)
dims = range(d)
gridpoints_linear = np.arange(N)
gridpoints = np.stack(np.unravel_index(
gridpoints_linear,
dimensions)) # d * (N^2) array mapping from linear to spatial frames
indices_list = [np.stack([gridpoints_linear] * 2, axis=-1)]
for dim in dims:
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]))
indices_list.append(
np.stack([gridpoints_linear[upper_idx], upper_points], axis=-1))
# --- Stencil lower cells ---
lower_points, lower_idx = wrap_or_discard(gridpoints - dim_direction,
dim,
dimensions,
periodic=collapsed_gather_nd(
periodic, [dim, 0]))
indices_list.append(
np.stack([gridpoints_linear[lower_idx], lower_points], axis=-1))
indices = np.concatenate(indices_list, axis=0)
# --- Sort indices ---
sorting = np.lexsort(np.transpose(indices)[:, ::-1])
sorted_indices = indices[sorting]
return sorted_indices, sorting
def total(self):
v_length = math.sqrt(
math.add(
[self.staggered[..., i]**2 for i in range(self.shape[-1])]))
total = math.sum(v_length, axis=range(1, self.spatial_rank + 1))
for i in range(self.spatial_rank + 1):
total = math.expand_dims(total, -1)
return total
def data(self, data):
assert math.is_tensor(data), data
rank = math.ndims(data)
assert rank <= 3
if rank < 3:
assert rank in (0, 1) # Scalar field / vector field
data = math.expand_dims(data, 0, 3 - rank)
return math.to_float(data)
def normalize(self):
v_length = math.sqrt(
math.add(
[self.staggered[..., i]**2 for i in range(self.shape[-1])]))
global_mean = math.mean(v_length, axis=range(1, self.spatial_rank + 1))
for i in range(self.spatial_rank + 1):
global_mean = math.expand_dims(global_mean, -1)
return StaggeredGrid(self.staggered / global_mean)
def data(self, data):
if data is None:
return None
if isinstance(data, (tuple, list)):
data = np.array(data) # numbers or objects
while math.ndims(data) < 2:
data = math.expand_dims(data)
return data
def gravity_tensor(gravity, rank):
if isinstance(gravity, Gravity):
gravity = gravity.gravity
if math.is_scalar(gravity):
gravity = gravity * GLOBAL_AXIS_ORDER.up_vector(rank)
assert math.staticshape(gravity)[-1] == rank
return math.to_float(
math.expand_dims(gravity, 0,
rank + 2 - len(math.staticshape(gravity))))
def unstack(self):
flags = propagate_flags_children(self.flags, self.rank, 1)
return [
CenteredGrid(math.expand_dims(component),
box=self.box,
name='%s[...,%d]' % (self.name, i),
flags=flags,
batch_size=self._batch_size)
for i, component in enumerate(math.unstack(self.data, -1))
]
def indices(self):
"""
Constructs a grid containing the index-location as components.
Each index denotes the location within the tensor starting from zero.
Indices are encoded as vectors in the index tensor.
:param dtype: a numpy data type (default float32)
:return: an index tensor of shape (1, spatial dimensions..., spatial rank)
"""
idx_zyx = np.meshgrid(*[range(dim) for dim in self.resolution], indexing="ij")
return math.expand_dims(np.stack(idx_zyx, axis=-1))
def data(self, data):
if data is None:
return None
if isinstance(data, (tuple, list)):
data = np.array(data) # numbers or objects
if self.content_type in (struct.shape, struct.staticshape):
assert math.ndims(data) == 1
else:
if math.ndims(data) < 2:
data = math.expand_dims(data, 0, number=2 - math.ndims(data))
return data
def getpoints(box, resolution):
idx_zyx = np.meshgrid(*[
np.linspace(0.5 / dim, 1 - 0.5 / dim, dim) for dim in resolution
],
indexing="ij")
local_coords = math.expand_dims(math.stack(idx_zyx, axis=-1),
0).astype(np.float32)
points = box.local_to_global(local_coords)
return CenteredGrid(points,
box,
name='grid_centers(%s, %s)' % (box, resolution),
flags=[SAMPLE_POINTS])
def _expand_axes(data, points, collapse_dimensions=True):
assert math.spatial_rank(data) >= 0
data = math.expand_dims(data, 1, math.spatial_rank(points) - math.spatial_rank(data))
if collapse_dimensions:
return data
else:
points_axes = math.staticshape(points)[1:-1]
data_axes = math.staticshape(data)[1:-1]
for d_points, d_data in zip(points_axes, data_axes):
assert d_points % d_data == 0
tilings = [1] + [d_points // d_data for d_points, d_data in zip(math.staticshape(points)[1:-1], math.staticshape(data)[1:-1])] + [1]
data = math.tile(data, tilings)
return data
def divergence(self):
dims = range(self.spatial_rank)
components = []
for dimension in dims:
comp = self.spatial_rank - dimension - 1
upper_slices = [(slice(1, None) if i == dimension else slice(-1))
for i in dims]
lower_slices = [(slice(-1) if i == dimension else slice(-1))
for i in dims]
diff = self.staggered[[slice(None)] + upper_slices + [comp]] - \
self.staggered[[slice(None)] + lower_slices + [comp]]
components.append(diff)
return math.expand_dims(math.add(components), -1)
def _central_divergence_nd(tensor):
rank = spatial_rank(tensor)
dims = range(rank)
components = []
tensor = math.pad(tensor, [[0, 0]] + [[1, 1]] * rank + [[0, 0]])
for dimension in dims:
upper_slices = [(slice(2, None) if i == dimension else slice(1, -1))
for i in dims]
lower_slices = [(slice(-2) if i == dimension else slice(1, -1))
for i in dims]
diff = tensor[[slice(None)] + upper_slices + [rank - dimension - 1]] - \
tensor[[slice(None)] + lower_slices + [rank - dimension - 1]]
components.append(diff)
return math.expand_dims(math.add(components), -1)
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 sample_at(self, points):
points_rank = math.spatial_rank(points)
src_rank = math.spatial_rank(self.location)
# --- Expand shapes to format (batch_size, points_dims..., src_dims..., channels) ---
points = math.expand_dims(points, axis=-2, number=src_rank)
src_points = math.expand_dims(self.location,
axis=-2,
number=points_rank)
src_strength = math.expand_dims(self.strength, axis=-1)
src_strength = math.batch_align(src_strength, 0, self.location)
src_strength = math.expand_dims(src_strength,
axis=-1,
number=points_rank)
src_axes = tuple(range(-2, -2 - src_rank, -1))
# --- Compute distances and falloff ---
distances = points - src_points
if self.falloff is not None:
raise NotImplementedError()
# distances_squared = math.sum(distances ** 2, axis=-1, keepdims=True)
# unit_distances = distances / math.sqrt(distances_squared)
# strength = src_strength * math.exp(-distances_squared)
else:
strength = src_strength
# --- Compute velocities ---
if math.staticshape(points)[-1] == 2: # Curl in 2D
dist_1, dist_2 = math.unstack(distances, axis=-1)
if GLOBAL_AXIS_ORDER.is_x_first:
velocity = strength * math.stack([-dist_2, dist_1], axis=-1)
else:
velocity = strength * math.stack([dist_2, -dist_1], axis=-1)
elif math.staticshape(points)[-1] == 3: # Curl in 3D
raise NotImplementedError('not yet implemented')
else:
raise AssertionError(
'Vector product not available in > 3 dimensions')
velocity = math.sum(velocity, axis=src_axes)
return velocity
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 _expand_axes(data, points, batch_size=1):
assert math.spatial_rank(data) >= 0
data = math.expand_dims(
data, 1,
math.spatial_rank(points) - math.spatial_rank(data))
points_axes = math.staticshape(points)[1:-1]
data_axes = math.staticshape(data)[1:-1]
for d_points, d_data in zip(points_axes, data_axes):
assert d_points % d_data == 0
tilings = [batch_size or 1] + [
d_points // d_data for d_points, d_data in zip(
math.staticshape(points)[1:-1],
math.staticshape(data)[1:-1])
] + [1]
data = math.tile(data, tilings)
return data
def _forward_divergence_nd(field):
rank = spatial_rank(field)
dims = range(rank)
components = []
for dimension in dims:
vq = field[..., rank - dimension - 1]
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]
diff = vq[[slice(None)] + upper_slices] - vq[[slice(None)] +
lower_slices]
padded = math.pad(diff,
[[0, 0]] + [([0, 1] if i == dimension else [0, 0])
for i in dims])
components.append(padded)
return math.expand_dims(math.add(components), -1)
def wrap_or_discard(points, check_bounds_dim, dimensions, periodic=False):
"""
Handles points that lie outside the domain by either discarding them or wrapping them, depending on periodic.
:param points: grid indices, typically of shape (dimensions, cell_count)
:param check_bounds_dim: int
:param dimensions: domain resolution
:param periodic: if False: discard indices outside domain, if True: wrap indices outside domain
:return:
"""
if not periodic:
upper_in_range_inx = np.nonzero((points[check_bounds_dim] < dimensions[check_bounds_dim]) & (points[check_bounds_dim] >= 0))[0]
new_points = points[:, upper_in_range_inx] # discard points outside domain
else:
upper_in_range_inx = slice(None)
new_points = points % math.expand_dims(dimensions, -1) # wrap points
linear_points = np.ravel_multi_index(new_points, dimensions)
return linear_points, upper_in_range_inx
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 location(self, loc):
loc = math.to_float(loc)
assert math.staticshape(loc)[-1] in (2, 3)
if math.ndims(loc) < 2:
loc = math.expand_dims(loc, axis=0, number=2 - math.ndims(loc))
return loc
