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)
