def test_scatter_update_1d_batched(self): for backend in BACKENDS: with backend: # Only base batched base = math.zeros(spatial(x=4)) + math.tensor([0, 1]) indices = math.wrap([1, 2], instance('points')) values = math.wrap([11, 12], instance('points')) updated = math.scatter(base, indices, values, mode='update', outside_handling='undefined') math.assert_close(updated, math.tensor([(0, 1), (11, 11), (12, 12), (0, 1)], spatial('x'), channel('vector')), msg=backend.name) # Only values batched base = math.ones(spatial(x=4)) indices = math.wrap([1, 2], instance('points')) values = math.wrap([[11, 12], [-11, -12]], batch('batch'), instance('points')) updated = math.scatter(base, indices, values, mode='update', outside_handling='undefined') math.assert_close(updated, math.tensor([[1, 11, 12, 1], [1, -11, -12, 1]], batch('batch'), spatial('x')), msg=backend.name) # Only indices batched base = math.ones(spatial(x=4)) indices = math.wrap([[0, 1], [1, 2]], batch('batch'), instance('points')) values = math.wrap([11, 12], instance('points')) updated = math.scatter(base, indices, values, mode='update', outside_handling='undefined') math.assert_close(updated, math.tensor([[11, 12, 1, 1], [1, 11, 12, 1]], batch('batch'), spatial('x')), msg=backend.name) # Everything batched base = math.zeros(spatial(x=4)) + math.tensor([0, 1], batch('batch')) indices = math.wrap([[0, 1], [1, 2]], batch('batch'), instance('points')) values = math.wrap([[11, 12], [-11, -12]], batch('batch'), instance('points')) updated = math.scatter(base, indices, values, mode='update', outside_handling='undefined') math.assert_close(updated, math.tensor([[11, 12, 0, 0], [1, -11, -12, 1]], batch('batch'), spatial('x')), msg=backend.name)
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 test_scatter_1d(self): for backend in BACKENDS: with backend: base = math.ones(spatial(x=4)) indices = math.wrap([1, 2], instance('points')) values = math.wrap([11, 12], instance('points')) updated = math.scatter(base, indices, values, mode='update', outside_handling='undefined') math.assert_close(updated, [1, 11, 12, 1]) updated = math.scatter(base, indices, values, mode='add', outside_handling='undefined') math.assert_close(updated, [1, 12, 13, 1]) # with vector dim indices = math.expand(indices, channel(vector=1)) updated = math.scatter(base, indices, values, mode='update', outside_handling='undefined') math.assert_close(updated, [1, 11, 12, 1])
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 """ sample_indices_nd = math.to_int( math.round(box.global_to_local(self.sample_points) * resolution)) sample_indices_nd = math.minimum( math.maximum(0, sample_indices_nd), resolution - 1 ) # Snap outside points to edges, otherwise scatter raises an error # Correct format for math.scatter valid_indices = _batch_indices(sample_indices_nd) shape = (math.shape( self.data)[0], ) + tuple(resolution) + (self.data.shape[-1], ) scattered = math.scatter(self.sample_points, valid_indices, self.data, shape, duplicates_handling=self.mode) return CenteredGrid(data=scattered, box=box, extrapolation='constant', name=self.name + '_centered')
def test_scatter_add_2d(self): for backend in BACKENDS: with backend: base = math.ones(spatial(x=3, y=2)) indices = math.wrap([(0, 0), (0, 0), (0, 1), (2, 1)], instance('points'), channel('vector')) values = math.wrap([11, 11, 12, 13], instance('points')) updated = math.scatter(base, indices, values, mode='add', outside_handling='undefined') math.assert_close(updated, math.tensor([[23, 1, 1], [13, 1, 14]], spatial('y,x')))
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 grid_scatter(self, bounds: Box, resolution: math.Shape, outside_handling: str): """ Approximately samples this field on a regular grid using math.scatter(). Args: outside_handling: `str` passed to `phi.math.scatter()`. bounds: physical dimensions of the grid resolution: grid resolution Returns: CenteredGrid """ closest_index = bounds.global_to_local(self.points) * resolution - 0.5 mode = 'add' if self._add_overlapping else 'mean' base = math.zeros(resolution) if isinstance(self.extrapolation, math.extrapolation.ConstantExtrapolation): base += self.extrapolation.value scattered = math.scatter(base, closest_index, self.values, mode=mode, outside_handling=outside_handling) return scattered
def _grid_scatter(self, box: Box, resolution: math.Shape): """ Approximately samples this field on a regular grid using math.scatter(). Args: box: physical dimensions of the grid resolution: grid resolution box: Box: resolution: math.Shape: Returns: CenteredGrid """ closest_index = math.to_int(math.round(box.global_to_local(self.points) * resolution - 0.5)) if self._add_overlapping: duplicates_handling = 'add' else: duplicates_handling = 'mean' scattered = math.scatter(closest_index, self.values, resolution, duplicates_handling=duplicates_handling, outside_handling='discard', scatter_dims=('points',)) return scattered
def test_scatter_2d_discard(self): base = math.ones(spatial(x=3, y=2)) indices = math.wrap([(-1, 0), (0, 1), (3, 1)], instance('points'), channel('vector')) values = math.wrap([11, 12, 13], instance('points')) updated = math.scatter(base, indices, values, mode='update', outside_handling='discard') math.assert_close(updated, math.tensor([[1, 1, 1], [12, 1, 1]], spatial('y,x')))