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