def test_pad_tensor(self): for backend in BACKENDS: with backend: a = math.meshgrid(x=4, y=3) # 0 p = math.pad(a, {'x': (1, 2), 'y': (0, 1)}, ZERO) self.assertEqual((7, 4, 2), p.shape.sizes) # dimension check math.assert_close(p.x[1:-2].y[:-1], a) # copy inner math.assert_close(p.x[0], 0) # 1 p = math.pad(a, {'x': (1, 2), 'y': (0, 1)}, ONE) self.assertEqual((7, 4, 2), p.shape.sizes) # dimension check math.assert_close(p.x[1:-2].y[:-1], a) # copy inner math.assert_close(p.x[0], 1) # periodic p = math.pad(a, {'x': (1, 2), 'y': (0, 1)}, PERIODIC) self.assertEqual((7, 4, 2), p.shape.sizes) # dimension check math.assert_close(p.x[1:-2].y[:-1], a) # copy inner math.assert_close(p.x[0].y[:-1], a.x[-1]) math.assert_close(p.x[-2:].y[:-1], a.x[:2]) # boundary p = math.pad(a, {'x': (1, 2), 'y': (0, 1)}, BOUNDARY) self.assertEqual((7, 4, 2), p.shape.sizes) # dimension check math.assert_close(p.x[1:-2].y[:-1], a) # copy inner math.assert_close(p.x[0].y[:-1], a.x[0]) math.assert_close(p.x[-2:].y[:-1], a.x[-1]) # mixed p = math.pad(a, {'x': (1, 2), 'y': (0, 1)}, combine_sides({'x': PERIODIC, 'y': (ONE, REFLECT)})) math.print(p) self.assertEqual((7, 4, 2), p.shape.sizes) # dimension check math.assert_close(p.x[1:-2].y[:-1], a) # copy inner math.assert_close(p.x[0].y[:-1], a.x[-1]) # periodic math.assert_close(p.x[-2:].y[:-1], a.x[:2]) # periodic
def at_faces(self, face_dimension_xyz): dims = range(self.spatial_rank) face_dimension_zyx = len( dims) - face_dimension_xyz - 1 # 0=Z, 1=Y, 2=X, etc. components = [] for d in dims: # z,y,x if d == face_dimension_zyx: components.append(self.staggered[..., len(dims) - d - 1]) else: # Interpolate other components vq = self.staggered[..., len(dims) - d - 1] t = vq for d2 in dims: # z,y,x slices1 = [(slice(1, None) if i == d2 else slice(None)) for i in dims] slices2 = [(slice(-1) if i == d2 else slice(None)) for i in dims] t = t[[slice(None)] + slices1] + t[[slice(None)] + slices2] if d2 == d: t = math.pad( t, [[0, 0]] + [([0, 1] if i == d2 else [0, 0]) for i in dims]) / 2 else: t = math.pad( t, [[0, 0]] + [([1, 0] if i == d2 else [0, 0]) for i in dims]) / 2 components.append(t) return math.stack(components[::-1], axis=-1)
def test_pad_collapsed(self): a = math.zeros(b=2, x=10, y=10, batch=10) p = math.pad(a, {'x': (1, 2)}, ZERO) self.assertIsInstance(p, CollapsedTensor) self.assertEqual((10, 2, 13, 10), p.shape.sizes) p = math.pad(a, {'x': (1, 2)}, PERIODIC) self.assertIsInstance(p, CollapsedTensor) self.assertEqual((10, 2, 13, 10), p.shape.sizes)
def stagger(field: CenteredGrid, face_function: Callable, extrapolation: math.extrapolation.Extrapolation, type: type = StaggeredGrid): """ Creates a new grid by evaluating `face_function` given two neighbouring cells. One layer of missing cells is inferred from the extrapolation. This method returns a Field of type `type` which must be either StaggeredGrid or CenteredGrid. When returning a StaggeredGrid, the new values are sampled at the faces of neighbouring cells. When returning a CenteredGrid, the new grid has the same resolution as `field`. Args: field: centered grid face_function: function mapping (value1: Tensor, value2: Tensor) -> center_value: Tensor extrapolation: extrapolation mode of the returned grid. Has no effect on the values. type: one of (StaggeredGrid, CenteredGrid) field: CenteredGrid: face_function: Callable: extrapolation: math.extrapolation.Extrapolation: type: type: (Default value = StaggeredGrid) Returns: grid of type matching the `type` argument """ all_lower = [] all_upper = [] if type == StaggeredGrid: for dim in field.shape.spatial.names: lo_valid, up_valid = extrapolation.valid_outer_faces(dim) width_lower = {dim: (int(lo_valid), int(up_valid) - 1)} width_upper = { dim: (int(lo_valid or up_valid) - 1, int(lo_valid and up_valid)) } all_lower.append( math.pad(field.values, width_lower, field.extrapolation)) all_upper.append( math.pad(field.values, width_upper, field.extrapolation)) all_upper = math.stack(all_upper, channel('vector')) all_lower = math.stack(all_lower, channel('vector')) values = face_function(all_lower, all_upper) result = StaggeredGrid(values, bounds=field.bounds, extrapolation=extrapolation) assert result.shape.spatial == field.shape.spatial return result elif type == CenteredGrid: left, right = math.shift(field.values, (-1, 1), padding=field.extrapolation, stack_dim=channel('vector')) values = face_function(left, right) return CenteredGrid(values, bounds=field.bounds, extrapolation=extrapolation) else: raise ValueError(type)
def linear_function(val): val = -val val *= 2 val = math.pad(val, {'x': (2, 0), 'y': (0, 1)}, extrapolation.PERIODIC) val = val.x[:-2].y[1:] + val.x[2:].y[:-1] val = math.pad(val, {'x': (0, 0), 'y': (0, 1)}, extrapolation.ZERO) val = math.pad(val, {'x': (2, 2), 'y': (0, 1)}, extrapolation.BOUNDARY) # sl = sl.vector[0] return val val = val.x[1:4].y[:2] return math.sum([val, sl], axis=0) - sl
def gradient(scalar_field, padding_mode='replicate'): assert isinstance(scalar_field, CenteredGrid) data = scalar_field.data if data.shape[-1] != 1: raise ValueError('input must be a scalar field') tensors = [] for dim in math.spatial_dimensions(data): upper = math.pad(data, [[0,1] if d == dim else [0,0] for d in math.all_dimensions(data)], padding_mode) lower = math.pad(data, [[1,0] if d == dim else [0,0] for d in math.all_dimensions(data)], padding_mode) tensors.append((upper - lower) / scalar_field.dx[dim - 1]) return StaggeredGrid(tensors, scalar_field.box, name='grad(%s)' % scalar_field.name, batch_size=scalar_field._batch_size)
def bake_extrapolation(grid: GridType) -> GridType: """ Pads `grid` with its current extrapolation. For `StaggeredGrid`s, the resulting grid will have a consistent shape, independent of the original extrapolation. Args: grid: `CenteredGrid` or `StaggeredGrid`. Returns: Padded grid with extrapolation `phi.math.extrapolation.NONE`. """ if grid.extrapolation == math.extrapolation.NONE: return grid if isinstance(grid, StaggeredGrid): values = grid.values.unstack('vector') padded = [] for dim, value in zip(grid.shape.spatial.names, values): lower, upper = grid.extrapolation.valid_outer_faces(dim) padded.append( math.pad(value, {dim: (0 if lower else 1, 0 if upper else 1)}, grid.extrapolation)) return StaggeredGrid(math.stack(padded, channel('vector')), bounds=grid.bounds, extrapolation=math.extrapolation.NONE) elif isinstance(grid, CenteredGrid): return pad(grid, 1).with_extrapolation(math.extrapolation.NONE) else: raise ValueError(f"Not a valid grid: {grid}")
def with_extrapolation(self, extrapolation: math.Extrapolation): if all( extrapolation.valid_outer_faces(dim) == self.extrapolation.valid_outer_faces(dim) for dim in self.resolution.names): return StaggeredGrid(self.values, extrapolation=extrapolation, bounds=self.bounds) else: values = [] for dim, component in zip(self.shape.spatial.names, self.values.unstack('vector')): old_lo, old_hi = [ int(v) for v in self.extrapolation.valid_outer_faces(dim) ] new_lo, new_hi = [ int(v) for v in extrapolation.valid_outer_faces(dim) ] widths = (new_lo - old_lo, new_hi - old_hi) values.append( math.pad(component, {dim: widths}, self.extrapolation)) values = math.stack(values, channel('vector')) return StaggeredGrid(values, extrapolation=extrapolation, bounds=self.bounds)
def create_surface_mask(liquid_mask): """ Computes inner contours of the liquid_mask. A cell i is flagged 1 if liquid_mask[i] = 1 and it has a non-liquid neighbour. :param liquid_mask: binary tensor :return: tensor """ # When we create inner contour, we don't want the fluid-wall boundaries to show up as surface, so we should pad with symmetric edge values. mask = math.pad(liquid_mask, [[0, 0]] + [[1, 1]] * math.spatial_rank(liquid_mask) + [[0, 0]], "constant") dims = range(math.spatial_rank(mask)) bcs = math.zeros_like(liquid_mask) # Move in every possible direction to assure corners are properly set. directions = np.array( list(itertools.product(*np.tile((-1, 0, 1), (len(dims), 1))))) for d in directions: d_slice = tuple([(slice(2, None) if d[i] == -1 else slice(0, -2) if d[i] == 1 else slice(1, -1)) for i in dims]) center_slice = tuple([slice(1, -1) for _ in dims]) # Create inner contour of particles bc_d = math.maximum(mask[(slice(None),) + d_slice + (slice(None),)], mask[(slice(None),) + center_slice + (slice(None),)]) - \ mask[(slice(None),) + d_slice + (slice(None),)] bcs = math.maximum(bcs, bc_d) return bcs
def padded(self, widths): data = math.pad(self.data, [[0, 0]] + widths + [[0, 0]], _pad_mode(self.extrapolation)) w_lower, w_upper = np.transpose(widths) box = AABox(self.box.lower - w_lower * self.dx, self.box.upper + w_upper * self.dx) return self.copied_with(data=data, box=box)
def _staggered_curl_3d(self): """ Calculates the curl operator on a staggered three-dimensional field. The resulting vector field is a staggered grid. If the velocities of the vector potential were sampled at the lower faces of a cube, the resulting velocities are sampled at the centers of the upper edges. :param vector_potential: three-dimensional vector potential :return: three-dimensional staggered vector field """ kernel = np.zeros((2, 2, 2, 3, 3), np.float32) derivative = np.array([-1, 1]) # x-component: dz/dy - dy/dz kernel[0, :, 0, 2, 0] = derivative kernel[:, 0, 0, 1, 0] = -derivative # y-component: dx/dz - dz/dx kernel[:, 0, 0, 0, 1] = derivative kernel[0, 0, :, 2, 1] = -derivative # z-component: dy/dx - dx/dy kernel[0, 0, :, 1, 2] = derivative kernel[0, :, 0, 0, 2] = -derivative vector_potential = math.pad(self.staggered, [[0, 0], [0, 1], [0, 1], [0, 1], [0, 0]], "SYMMETRIC") vector_field = math.conv(vector_potential, kernel, padding="VALID") return StaggeredGrid(vector_field)
def upsample2x(tensor, interpolation="LINEAR"): if interpolation.lower() != "linear": raise ValueError("Only linear interpolation supported") dims = range(spatial_rank(tensor)) vlen = tensor.shape[-1] spatial_dims = tensor.shape[1:-1] tensor = math.pad(tensor, [[0, 0]] + [[1, 1]] * spatial_rank(tensor) + [[0, 0]], "SYMMETRIC") for dim in dims: left_slices_1 = [(slice(2, None) if i == dim else slice(None)) for i in dims] left_slices_2 = [(slice(1, -1) if i == dim else slice(None)) for i in dims] right_slices_1 = [(slice(1, -1) if i == dim else slice(None)) for i in dims] right_slices_2 = [(slice(-2) if i == dim else slice(None)) for i in dims] left = 0.75 * tensor[[slice(None)] + left_slices_2 + [slice(None)]] + 0.25 * tensor[ [slice(None)] + left_slices_1 + [slice(None)]] right = 0.25 * tensor[[slice(None)] + right_slices_2 + [ slice(None) ]] + 0.75 * tensor[[slice(None)] + right_slices_1 + [slice(None)]] combined = math.stack([right, left], axis=2 + dim) tensor = math.reshape(combined, [-1] + [ spatial_dims[dim] * 2 if i == dim else tensor.shape[i + 1] for i in dims ] + [vlen]) return tensor
def linear_function(val): val = -val val *= 2 val = math.pad(val, { 'x': (2, 0), 'y': (0, 1) }, math.extrapolation.PERIODIC) val = val.x[:-2].y[1:] + val.x[2:].y[:-1] val = math.pad(val, { 'x': (0, 0), 'y': (0, 1) }, math.extrapolation.ZERO) val = math.pad(val, { 'x': (2, 2), 'y': (0, 1) }, math.extrapolation.BOUNDARY) return math.sum([val, val], dim='0') - val
def active_tensor(self, extend=0): """ Scalar channel encoding active cells as ones and inactive (open/obstacle) as zero. Active cells are those for which physical constants_dict such as pressure or velocity are calculated. :param extend: Extend the grid in all directions beyond the grid size specified by the domain """ return math.pad(self.active.data, [[0, 0]] + [[extend, extend]] * self.rank + [[0, 0]], "constant")
def padded(self, widths): extrapolation = self.extrapolation if isinstance( self.extrapolation, six.string_types) else ['constant'] + list( self.extrapolation) + ['constant'] data = math.pad(self.data, [[0, 0]] + widths + [[0, 0]], _pad_mode(extrapolation)) w_lower, w_upper = np.transpose(widths) box = AABox(self.box.lower - w_lower * self.dx, self.box.upper + w_upper * self.dx) return self.copied_with(data=data, box=box)
def padded(self, widths): if isinstance(widths, int): widths = [[widths, widths]] * self.rank data = math.pad(self.data, [[0, 0]] + widths + [[0, 0]], _pad_mode(self.extrapolation), constant_values=_pad_value(self.extrapolation_value)) w_lower, w_upper = np.transpose(widths) box = AABox(self.box.lower - w_lower * self.dx, self.box.upper + w_upper * self.dx) return self.copied_with(data=data, box=box)
def stack_staggered_components(data: Tensor) -> Tensor: padded = [] for dim, component in zip(data.shape.spatial.names, data.unstack('vector')): padded.append( math.pad( component, {d: (0, 1) for d in data.shape.spatial.without(dim).names}, mode=math.extrapolation.ZERO)) return math.channel_stack(padded, 'vector')
def accessible_tensor(self, extend=0): """ Scalar channel encoding cells that are accessible, i.e. not solid, as ones and obstacles as zero. :param extend: Extend the grid in all directions beyond the grid size specified by the domain """ pad_values = struct.map(lambda solid: int(not solid), Material.solid(self.domain.boundaries)) if isinstance(pad_values, (list, tuple)): pad_values = [0] + list(pad_values) + [0] result = math.pad(self.accessible.data, [[0,0]] + [[extend, extend]] * self.rank + [[0,0]], constant_values=pad_values) return result
def _staggered_curl_2d(self): kernel = np.zeros((2, 2, 1, 2), np.float32) derivative = np.array([-1, 1]) # x-component: dz/dy kernel[:, 0, 0, 0] = derivative # y-component: - dz/dx kernel[0, :, 0, 1] = -derivative scalar_potential = math.pad(self.staggered, [[0, 0], [0, 1], [0, 1], [0, 0]], "SYMMETRIC") vector_field = math.conv(scalar_potential, kernel, padding="VALID") return StaggeredGrid(vector_field)
def _central_diff_nd(field, dims): field = math.pad(field, [[0, 0]] + [[1, 1]] * spatial_rank(field) + [[0, 0]], "symmetric") df_dq = [] 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 = field[[slice(None)] + upper_slices + [0]] - field[[slice(None)] + lower_slices + [0]] df_dq.append(diff) return math.stack(df_dq[::-1], axis=-1)
def _forward_diff_nd(field, dims): df_dq = [] for dimension in dims: 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 = field[[slice(None)] + upper_slices] - field[[slice(None)] + lower_slices] padded = math.pad(diff, [[0, 0]] + [([0, 1] if i == dimension else [0, 0]) for i in dims]) df_dq.append(padded) return math.stack(df_dq[::-1], axis=-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 padded(self, widths): extrapolation = self.extrapolation if isinstance( self.extrapolation, six.string_types) else ['constant'] + list( self.extrapolation) + ['constant'] data = math.pad(self.data, [[0, 0]] + widths + [[0, 0]], _pad_mode(extrapolation)) w_lower, w_upper = np.transpose(widths) box = AABox(self.box.lower - w_lower * self.dx, self.box.upper + w_upper * self.dx) return CenteredGrid(data, box, extrapolation=self.extrapolation, name=self.name, batch_size=self._batch_size)
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 _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 laplace(tensor, weights=None, padding="symmetric"): if tensor.shape[-1] != 1: raise ValueError("Laplace operator requires a scalar field as input") rank = spatial_rank(tensor) if padding.lower() != "valid": tensor = math.pad(tensor, [[0, 0]] + [[1, 1]] * rank + [[0, 0]], padding) if weights is not None: return _weighted_sliced_laplace_nd(tensor, weights) if rank == 2: return _conv_laplace_2d(tensor) elif rank == 3: return _conv_laplace_3d(tensor) else: return _sliced_laplace_nd(tensor)
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 downsample2x(tensor, interpolation="LINEAR"): if interpolation.lower() != "linear": raise ValueError("Only linear interpolation supported") dims = range(spatial_rank(tensor)) tensor = math.pad(tensor, [[0, 0]] + [([0, 1] if (dim % 2) != 0 else [0, 0]) for dim in tensor.shape[1:-1]] + [[0, 0]], "SYMMETRIC") for dimension in dims: upper_slices = [(slice(1, None, 2) if i == dimension else slice(None)) for i in dims] lower_slices = [(slice(0, None, 2) if i == dimension else slice(None)) for i in dims] sum = tensor[[slice(None)] + upper_slices + [slice(None)]] + tensor[[slice(None)] + lower_slices + [slice(None)]] tensor = sum / 2 return tensor
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))