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 nonzero(field: SampledField):
indices = math.nonzero(field.values, list_dim=instance('points'))
elements = field.elements[indices]
return PointCloud(elements, values=math.tensor(1.), extrapolation=math.extrapolation.ZERO, add_overlapping=False, bounds=field.bounds, color=None)
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_nonzero_batched(self):
grid = math.tensor([[(0, 1)], [(0, 0)]], 'batch,x,y')
nz = math.nonzero(grid, list_dim='nonzero', index_dim='vector')
self.assertEqual(('batch', 'nonzero', 'vector'), nz.shape.names)
self.assertEqual(1, nz.batch[0].shape.nonzero)
self.assertEqual(0, nz.batch[1].shape.nonzero)
def test_nonzero_batched(self):
grid = math.tensor([[(0, 1)], [(0, 0)]], batch('batch'), spatial('x,y'))
nz = math.nonzero(grid)
self.assertEqual(('batch', 'nonzero', 'vector'), nz.shape.names)
self.assertEqual(1, nz.batch[0].shape.get_size('nonzero'))
self.assertEqual(0, nz.batch[1].shape.get_size('nonzero'))
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)