def tensor_as_field(t: Tensor): """ Interpret a `Tensor` as a `CenteredGrid` or `PointCloud` depending on its dimensions. Unlike the `CenteredGrid` constructor, this function will have the values sampled at integer points for each spatial dimension. Args: t: `Tensor` with either `spatial` or `instance` dimensions. Returns: `CenteredGrid` or `PointCloud` """ if spatial(t): assert not instance( t ), f"Cannot interpret tensor as Field because it has both spatial and instance dimensions: {t.shape}" bounds = Box(-0.5, math.wrap(spatial(t), channel('vector')) - 0.5) return CenteredGrid(t, 0, bounds=bounds) if instance(t): assert not spatial( t ), f"Cannot interpret tensor as Field because it has both spatial and instance dimensions: {t.shape}" assert 'vector' in t.shape, f"Cannot interpret tensor as PointCloud because it has not vector dimension." point_count = instance(t).volume bounds = data_bounds(t) radius = math.vec_length( bounds.size) / (1 + point_count**(1 / t.vector.size)) return PointCloud(Sphere(t, radius=radius))
def points(self, points: Tensor or Number or tuple or list, values: Tensor or Number = None, radius: Tensor or float or int or None = None, extrapolation: math.Extrapolation = math.extrapolation.ZERO, color: str or Tensor or tuple or list or None = None) -> PointCloud: """ Create a `phi.field.PointCloud` from the given `points`. The created field has no channel dimensions and all points carry the value `1`. Args: points: point locations in physical units values: (optional) values of the particles, defaults to 1. radius: (optional) size of the particles extrapolation: (optional) extrapolation to use, defaults to extrapolation.ZERO color: (optional) color used when plotting the points Returns: `phi.field.PointCloud` object """ extrapolation = extrapolation if isinstance(extrapolation, math.Extrapolation) else self.boundaries[extrapolation] if radius is None: radius = math.mean(self.bounds.size) * 0.005 # --- Parse points: tuple / list --- if isinstance(points, (tuple, list)): if len(points) == 0: # no points points = math.zeros(instance(points=0), channel(vector=1)) elif isinstance(points[0], Number): # single point points = math.tensor([points], instance('points'), channel('vector')) else: points = math.tensor(points, instance('points'), channel('vector')) elements = Sphere(points, radius) if values is None: values = math.tensor(1.) return PointCloud(elements, values, extrapolation, add_overlapping=False, bounds=self.bounds, color=color)
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 test_subshapes(self): s = batch(batch=10) & spatial(x=4, y=3) & channel(vector=2) & instance(points=1) self.assertEqual(batch(batch=10), s.batch) self.assertEqual(spatial(x=4, y=3), s.spatial) self.assertEqual(channel(vector=2), s.channel) self.assertEqual(instance(points=1), s.instance) self.assertEqual(batch(batch=10), batch(s)) self.assertEqual(spatial(x=4, y=3), spatial(s)) self.assertEqual(channel(vector=2), channel(s)) self.assertEqual(instance(points=1), instance(s))
def test_convert_point_cloud(self): loc = math.random_uniform(instance(points=4), channel(vector=2)) val = math.random_normal(instance(points=4), channel(vector=2)) points = PointCloud(Sphere(loc, radius=1), val) for backend in BACKENDS: converted = field.convert(points, backend) self.assertEqual(converted.values.default_backend, backend) self.assertEqual(converted.elements.center.default_backend, backend) self.assertEqual(converted.elements.radius.default_backend, backend)
def test_plot_point_cloud_2d_large(self): spheres = PointCloud( Sphere(wrap([(2, 4), (9, 8), (7, 8)], instance('points'), channel('vector')), radius=1)) cells = PointCloud( geom.pack_dims( CenteredGrid(0, 0, x=3, y=3, bounds=Box[4:6, 2:4]).elements, 'x,y', instance('points'))) cloud = field.stack([spheres, cells], instance('stack')) self._test_plot(cloud)
def test_grid_sample_gradient_1d(self): for backend in BACKENDS: if backend.supports(Backend.gradients): with backend: grid = math.tensor([0., 1, 2, 3], spatial('x')) coords = math.tensor([0.5, 1.5], instance('points')) with math.record_gradients(grid, coords): sampled = math.grid_sample(grid, coords, extrapolation.ZERO) loss = math.mean(math.l2_loss(sampled)) / 2 grad_grid, grad_coords = math.gradients(loss, grid, coords) math.assert_close(grad_grid, math.tensor([0.125, 0.5, 0.375, 0], spatial('x')), msg=backend) math.assert_close(grad_coords, math.tensor([0.25, 0.75], instance('points')), msg=backend)
def test_plot_point_cloud_2d(self): spheres = PointCloud( Sphere(wrap([(.2, .4), (.9, .8), (.7, .8)], instance('points'), channel('vector')), radius=.1)) cells = PointCloud( geom.pack_dims( CenteredGrid(0, 0, x=3, y=3, bounds=Box[.4:.6, .2:.4]).elements, 'x,y', instance('points'))) cloud = field.stack([spheres, cells], instance('stack')) self._test_plot(cloud)
def test_random_int(self): for backend in BACKENDS: with backend: # 32 bits a = math.random_uniform(instance(values=1000), low=-1, high=1, dtype=(int, 32)) self.assertEqual(a.dtype, DType(int, 32), msg=backend.name) self.assertEqual(a.min, -1, msg=backend.name) self.assertEqual(a.max, 0, msg=backend.name) # 64 bits a = math.random_uniform(instance(values=1000), low=-1, high=1, dtype=(int, 64)) self.assertEqual(a.dtype.kind, int, msg=backend.name) # Jax may downcast 64-bit to 32 self.assertEqual(a.min, -1, msg=backend.name) self.assertEqual(a.max, 0, msg=backend.name)
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 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 test_grid_sample(self): for backend in BACKENDS: with backend: grid = math.sum(math.meshgrid(x=[1, 2, 3], y=[0, 3]), 'vector') # 1 2 3 | 4 5 6 coords = math.tensor([(0, 0), (0.5, 0), (0, 0.5), (-2, -1)], instance('list'), channel('vector')) interp = math.grid_sample(grid, coords, extrapolation.ZERO) math.assert_close(interp, [1, 1.5, 2.5, 0], msg=backend.name)
def union(*geometries) -> Geometry: """ Union of the given geometries. A point lies inside the union if it lies within at least one of the geometries. Args: geometries: arbitrary geometries with same spatial dims. Arbitrary batch dims are allowed. *geometries: Returns: union Geometry """ if len(geometries) == 1 and isinstance(geometries[0], (tuple, list)): geometries = geometries[0] if len(geometries) == 0: return NO_GEOMETRY elif len(geometries) == 1: return geometries[0] elif all(type(g) == type(geometries[0]) for g in geometries): attrs = variable_attributes(geometries[0]) values = { a: math.stack([getattr(g, a) for g in geometries], math.instance('union')) for a in attrs } return copy_with(geometries[0], **values) else: base_geometries = () for geometry in geometries: base_geometries += geometry.geometries if isinstance( geometry, Union) else (geometry, ) return Union(base_geometries)
def test_plot_multiple(self): grid = CenteredGrid(Noise(batch(b=2)), 0, Box[0:1, 0:1], x=50, y=10) grid2 = CenteredGrid(grid, 0, Box[0:2, 0:1], x=20, y=50) points = wrap([(.2, .4), (.9, .8)], instance('points'), channel('vector')) cloud = PointCloud(Sphere(points, radius=0.1), bounds=Box(0, [1, 1])) titles = math.wrap([['b=0', 'b=0', 'points'], ['b=1', 'b=1', '']], spatial('rows,cols')) self._test_plot(grid, grid2, cloud, row_dims='b', title=titles)
def test_tensor_as_field(self): t = math.random_normal(spatial(x=4, y=3), channel(vector='x,y')) grid = field.tensor_as_field(t) self.assertIsInstance(grid, CenteredGrid) math.assert_close(grid.dx, 1) math.assert_close(grid.points.x[0].y[0], 0) t = math.random_normal(instance(points=5), channel(vector='x,y')) points = field.tensor_as_field(t) self.assertTrue((points.elements.bounding_radius() > 0).all) self.assertIsInstance(points, PointCloud)
def test_overlay(self): grid = CenteredGrid(Noise(), extrapolation.ZERO, x=64, y=8, bounds=Box(0, [1, 1])) points = wrap([(.2, .4), (.9, .8)], instance('points'), channel('vector')) cloud = PointCloud(Sphere(points, radius=.1)) self._test_plot(overlay(grid, grid * (0.1, 0.02), cloud), title='Overlay')
def test_infinite_cylinder(self): cylinder = geom.infinite_cylinder(x=.5, y=.5, radius=.5, inf_dim=math.spatial('z')) self.assertEqual(cylinder.spatial_rank, 3) cylinder = geom.infinite_cylinder(x=.5, y=.5, radius=.5, inf_dim='z') loc = math.wrap([(0, 0, 0), (.5, .5, 0), (1, 1, 0), (0, 0, 100), (.5, .5, 100)], math.instance('points'), math.channel(vector='x,y,z')) inside = math.wrap([False, True, False, False, True], math.instance('points')) math.assert_close(cylinder.lies_inside(loc), inside) loc = math.wrap([(0, 0, 0), (.5, 1, 0), (1, 1, 0), (0, 0, 100), (.5, 1, 100)], math.instance('points'), math.channel(vector='x,y,z')) corner_distance = math.sqrt(2) / 2 - .5 distance = math.wrap( [corner_distance, 0, corner_distance, corner_distance, 0], math.instance('points')) math.assert_close(cylinder.approximate_signed_distance(loc), distance)
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 test_grid_sample_gradient_2d(self): grads_grid = [] grads_coords = [] for backend in BACKENDS: if backend.supports(Backend.gradients): with backend: grid = math.tensor([[1., 2, 3], [1, 2, 3]], spatial('x,y')) coords = math.tensor([(0.5, 0.5), (1, 1.1), (-0.8, -0.5)], instance('points'), channel('vector')) with math.record_gradients(grid, coords): sampled = math.grid_sample(grid, coords, extrapolation.ZERO) loss = math.sum(sampled) / 3 grad_grid, grad_coords = math.gradients(loss, grid, coords) grads_grid.append(grad_grid) grads_coords.append(grad_coords) math.assert_close(*grads_grid) math.assert_close(*grads_coords)
def _plot_points(axis, data: PointCloud, **plt_args): x, y = [d.numpy() for d in data.points.vector.unstack_spatial('x,y')] color = [d.native() for d in data.color.points.unstack(len(x))] if isinstance(data.elements, Sphere): symbol = 'o' size = data.elements.bounding_radius().numpy() * 1.41 elif isinstance(data.elements, BaseBox): symbol = 's' size = math.mean(data.elements.bounding_half_extent(), 'vector').numpy() elif isinstance(data.elements, Point): symbol = 'x' size = 6 / _get_pixels_per_unit(axis.figure, axis) else: symbol = '*' size = data.elements.bounding_radius().numpy() size_px = size * _get_pixels_per_unit(axis.figure, axis) axis.scatter(x, y, marker=symbol, color=color, s=size_px ** 2, alpha=0.8) _annotate_points(axis, data.points, instance(data))
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')))
def test_plot_point_cloud_3d_points(self): self._test_plot( PointCloud( math.random_normal(instance(points=5), channel(vector='x,y,z'))))
def test_plot_spheres_2d(self): spheres = Sphere(wrap([(.2, .4), (.9, .8), (.7, .8)], instance('points'), channel('vector')), radius=.1) self._test_plot(spheres)
def test_join_dimensions(self): grid = math.random_normal(batch(batch=10) & spatial(x=4, y=3) & channel(vector=2)) points = math.pack_dims(grid, grid.shape.spatial, instance('points')) self.assertEqual(('batch', 'points', 'vector'), points.shape.names) self.assertEqual(grid.shape.volume, points.shape.volume) self.assertEqual(grid.shape.non_spatial, points.shape.non_instance)
def test_split_dimension(self): grid = math.random_normal(batch(batch=10) & spatial(x=4, y=3) & channel(vector=2)) points = math.pack_dims(grid, grid.shape.spatial, instance('points')) split = points.points.split(grid.shape.spatial) self.assertEqual(grid.shape, split.shape) math.assert_close(grid, split)
def test_vector_length(self): v = tensor([(0, 0), (1, 1), (-1, 0)], instance('values'), channel('vector')) le = math.vec_length(v) math.assert_close(le, [0, 1.41421356237, 1]) le = math.vec_length(v, eps=0.01) math.assert_close(le, [1e-1, 1.41421356237, 1])
def test_quantile(self): for backend in BACKENDS: if backend.name != "TensorFlow": # TensorFlow probability import problem with backend: t = math.tensor([(1, 2, 3, 4), (1, 2, 3, 4), (6, 7, 8, 9)], batch('batch'), instance('list')) q = math.quantile(t, 0.5) math.assert_close(q, (2.5, 2.5, 7.5), msg=backend.name) q = math.quantile(t, [0.5, 0.6]) math.assert_close(q, [(2.5, 2.5, 7.5), (2.8, 2.8, 7.8)], msg=backend.name)
def test_median(self): t = math.tensor([(1, 2, 3, 10), (0, 1, 3, 10)], batch('batch'), instance('list')) math.assert_close(math.median(t), [2.5, 2])
def test_sum_by_type(self): a = math.ones(spatial(x=3, y=4), batch(b=10), instance(i=2), channel(vector=2)) math.assert_close(math.sum(a, spatial), 12)
def test_random_complex(self): for backend in BACKENDS: with backend: a = math.random_uniform(instance(values=4), low=-1, high=0, dtype=(complex, 64)) self.assertEqual(a.dtype, DType(complex, 64), msg=backend.name) math.assert_close(a.imag, 0, msg=backend.name)