def test_extrapolate_valid_3D_3x3x3_2(self):
valid = tensor([[[0, 0, 0],
[0, 0, 0],
[0, 0, 0]],
[[0, 0, 1],
[0, 0, 0],
[1, 0, 0]],
[[0, 0, 0],
[0, 0, 0],
[0, 0, 0]]], spatial('x, y, z'))
values = tensor([[[0, 0, 0],
[0, 0, 0],
[0, 0, 0]],
[[1, 0, 4],
[0, 0, 0],
[2, 0, 0]],
[[0, 0, 0],
[0, 0, 0],
[0, 0, 0]]], spatial('x, y, z'))
expected_valid = math.ones(spatial(x=3, y=3, z=3))
expected_values = tensor([[[3, 4, 4],
[2, 3, 4],
[2, 2, 3]],
[[3, 4, 4],
[2, 3, 4],
[2, 2, 3]],
[[3, 4, 4],
[2, 3, 4],
[2, 2, 3]]], spatial('x, y, z'))
extrapolated_values, extrapolated_valid = math.extrapolate_valid_values(values, valid, 2)
math.assert_close(extrapolated_values, expected_values)
math.assert_close(extrapolated_valid, expected_valid)
def troubleshoot_torch():
from phi import math
try:
import torch
except ImportError:
return "Not installed."
try:
from phi import torch as phi_torch
except BaseException as err:
return f"Installed ({torch.__version__}) but not available due to internal error: {err}"
try:
gpu_count = len(phi_torch.TORCH.list_devices('GPU'))
except BaseException as err:
return f"Installed ({torch.__version__}) but device initialization failed with error: {err}"
with phi_torch.TORCH:
try:
math.assert_close(
math.ones(batch(batch=8) & spatial(x=64)) +
math.ones(batch(batch=8) & spatial(x=64)), 2)
except BaseException as err:
return f"Installed ({torch.__version__}) but tests failed with error: {err}"
if torch.__version__.startswith('1.10.'):
return f"Installed ({torch.__version__}), {gpu_count} GPUs available. This version has known bugs with JIT compilation. Recommended: 1.11+ or 1.8.2 LTS"
if torch.__version__.startswith('1.9.'):
return f"Installed ({torch.__version__}), {gpu_count} GPUs available. You may encounter problems importing torch.fft. Recommended: 1.11+ or 1.8.2 LTS"
return f"Installed ({torch.__version__}), {gpu_count} GPUs available."
def test_boolean_mask_dim_missing(self):
for backend in BACKENDS:
with backend:
x = math.random_uniform(spatial(x=2))
mask = math.tensor([True, False, True, True], batch('selection'))
selected = math.boolean_mask(x, 'selection', mask)
self.assertEqual(set(spatial(x=2) & batch(selection=3)), set(selected.shape), msg=backend.name)
def test_jit_compile(self):
@math.jit_compile
def scalar_mul(x, fac=1):
return x * fac
for backend in BACKENDS:
with backend:
x = math.ones(spatial(x=4))
trace_count_0 = len(scalar_mul.traces)
math.assert_close(scalar_mul(x, fac=1), 1, msg=backend)
if backend.supports(Backend.jit_compile):
self.assertEqual(len(scalar_mul.traces), trace_count_0 + 1)
math.assert_close(scalar_mul(x, fac=1), 1, msg=backend)
if backend.supports(Backend.jit_compile):
self.assertEqual(len(scalar_mul.traces), trace_count_0 + 1)
math.assert_close(scalar_mul(x, fac=2), 2, msg=backend)
if backend.supports(Backend.jit_compile):
self.assertEqual(len(scalar_mul.traces), trace_count_0 + 2)
math.assert_close(scalar_mul(math.zeros(spatial(x=4)), fac=2),
0,
msg=backend)
if backend.supports(Backend.jit_compile):
self.assertEqual(len(scalar_mul.traces), trace_count_0 + 2)
math.assert_close(scalar_mul(math.zeros(spatial(y=4)), fac=2),
0,
msg=backend)
if backend.supports(Backend.jit_compile):
self.assertEqual(len(scalar_mul.traces), trace_count_0 + 3)
def test_indexing(self):
s = batch(batch=10) & spatial(x=4, y=3) & channel(vector=2)
self.assertEqual(batch(batch=10), s[0:1])
self.assertEqual(batch(batch=10), s[[0]])
self.assertEqual(spatial(x=4, y=3), s[1:3])
self.assertEqual(spatial(x=4), s['x'])
self.assertEqual(spatial(x=4, y=3), s['x, y'])
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_boolean_mask_1d(self):
for backend in BACKENDS:
with backend:
x = math.range(spatial('range'), 4)
mask = math.tensor([True, False, True, False], spatial('range'))
math.assert_close(math.boolean_mask(x, 'range', mask), [0, 2], msg=backend.name)
math.assert_close(x.range[mask], [0, 2], msg=backend.name)
def test_np_speed_sum(self):
print()
np1, np2 = rnpv(64), rnpv(256)
t1 = math.tensor(np1, batch('batch'), spatial('x, y'), channel('vector'))
t2 = math.tensor(np2, batch('batch'), spatial('x, y'), channel('vector'))
_assert_equally_fast(lambda: np.sum(np1), lambda: math.sum(t1, dim=t1.shape), n=10000)
_assert_equally_fast(lambda: np.sum(np2), lambda: math.sum(t2, dim=t1.shape), n=10000)
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 test_downsample2x(self):
meshgrid = math.meshgrid(x=(0, 1, 2, 3), y=(0, -1, -2))
half_size = math.downsample2x(meshgrid, extrapolation.BOUNDARY)
math.print(meshgrid, 'Full size')
math.print(half_size, 'Half size')
math.assert_close(half_size.vector[0], wrap([[0.5, 2.5], [0.5, 2.5]], spatial('y,x')))
math.assert_close(half_size.vector[1], wrap([[-0.5, -0.5], [-2, -2]], spatial('y,x')))
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_dot_matrix(self):
for backend in BACKENDS:
with backend:
a = math.ones(spatial(x=2, a=4) & batch(batch=10))
b = math.ones(spatial(y=3, b=4))
dot = math.dot(a, 'a', b, 'b')
self.assertEqual(set(spatial(x=2, y=3) & batch(batch=10)), set(dot.shape), msg=backend.name)
math.assert_close(dot, 4, msg=backend.name)
def test_dot_vector(self):
for backend in BACKENDS:
with backend:
a = math.ones(spatial(a=4))
b = math.ones(spatial(b=4))
dot = math.dot(a, 'a', b, 'b')
self.assertEqual(0, dot.rank, msg=backend.name)
math.assert_close(dot, 4, a.a * b.b, msg=backend.name)
math.assert_close(math.dot(a, 'a', a, 'a'), 4, msg=backend.name)
def test_extrapolate_valid_3x3_sanity(self):
values = tensor([[0, 0, 0],
[0, 1, 0],
[0, 0, 0]], spatial('x, y'))
valid = values
extrapolated_values, extrapolated_valid = math.extrapolate_valid_values(values, valid)
expected_values = math.ones(spatial(x=3, y=3))
expected_valid = extrapolated_values
math.assert_close(extrapolated_values, expected_values)
math.assert_close(extrapolated_valid, expected_valid)
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_np_speed_op2(self):
print()
np1, np2 = rnpv(64), rnpv(64)
t1 = math.tensor(np1, batch('batch'), spatial('x, y'), channel('vector'))
t2 = math.tensor(np2, batch('batch'), spatial('x, y'), channel('vector'))
_assert_equally_fast(lambda: np1 + np2, lambda: t1 + t2, n=10000)
np1, np2 = rnpv(256), rnpv(256)
t1 = math.tensor(np1, batch('batch'), spatial('x, y'), channel('vector'))
t2 = math.tensor(np2, batch('batch'), spatial('x, y'), channel('vector'))
_assert_equally_fast(lambda: np1 + np2, lambda: t1 + t2, n=1000)
def test_stack(self):
stacked = shape_stack(batch('stack'), batch(time=1) & spatial(x=3, y=3), spatial(x=3, y=4), EMPTY_SHAPE)
print(stacked)
self.assertEqual(('stack', 'time', 'x', 'y'), stacked.names)
self.assertEqual(3, stacked.get_size('stack'))
self.assertEqual(1, stacked.get_size('time'))
math.assert_close((3, 3, 1), stacked.get_size('x'))
math.assert_close((3, 4, 1), stacked.get_size('y'))
print(stacked.shape)
self.assertEqual(('stack', 'dims'), stacked.shape.names)
self.assertEqual(12, stacked.shape.volume)
def test_boolean_mask_batched(self):
for backend in BACKENDS:
with backend:
x = math.expand(math.range(spatial('x'), 4), batch(batch=2)) * math.tensor([1, -1])
mask = math.tensor([[True, False, True, False], [False, True, False, False]], batch('batch'), spatial('x'))
selected = math.boolean_mask(x, 'x', mask)
expected_0 = math.tensor([(0, -0), (2, -2)], spatial('x'), channel('vector'))
expected_1 = math.tensor([(1, -1)], spatial('x'), channel('vector'))
math.assert_close(selected.batch[0], expected_0, msg=backend.name)
math.assert_close(selected.batch[1], expected_1, msg=backend.name)
math.assert_close(selected, x.x[mask], msg=backend.name)
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_boolean_mask_semi_batched(self):
for backend in BACKENDS:
with backend:
x = math.range(spatial('x'), 4)
mask = math.tensor([[True, False, True, False], [False, True, False, False]], batch('batch'), spatial('x'))
selected = math.boolean_mask(x, 'x', mask)
self.assertEqual(3, selected.shape.volume, msg=backend.name)
def test_dimension_types(self):
v = math.ones(batch(batch=10) & spatial(x=4, y=3) & channel(vector=2))
self.assertEqual(v.x.index, 1)
self.assertEqual(v.x.name, 'x')
self.assertEqual(('batch', 'spatial', 'spatial', 'channel'), v.shape.types)
b = v.x.as_batch()
self.assertEqual(('batch', 'batch', 'spatial', 'channel'), b.shape.types)
def test_meshgrid_names(self):
shape = spatial(x=2) & channel(vector='x')
indices = list(shape.meshgrid(names=True))
self.assertEqual([
dict(x=0, vector='x'),
dict(x=1, vector='x'),
], indices)
def test_cos(self):
for backend in BACKENDS:
with backend:
math.assert_close(math.cos(math.zeros(spatial(x=4))), 1, abs_tolerance=1e-6, msg=backend.name)
math.assert_close(math.cos(math.tensor(math.PI / 2)), 0, abs_tolerance=1e-6, msg=backend.name)
math.assert_close(math.cos(math.tensor(math.PI)), -1, abs_tolerance=1e-6, msg=backend.name)
math.assert_close(math.cos(math.tensor(math.PI * 3 / 2)), 0, abs_tolerance=1e-6, msg=backend.name)
def upsample_apply(params, inputs, **kwargs):
x = math.wrap(
inputs, math.batch('batch'),
*[math.spatial(f'{i}') for i in range(len(inputs.shape) - 2)],
math.channel('vector'))
x = math.upsample2x(x)
return x.native(x.shape)
def test_grid_sample_backend_equality_2d_batched(self):
grid = math.random_normal(batch(mybatch=10) & spatial(y=10, x=7))
coords = math.random_uniform(batch(mybatch=10) & spatial(x=3, y=2)) * (12, 9)
grid_ = math.tensor(grid.native('mybatch,x,y'), batch('mybatch'), spatial('x,y'))
coords_ = coords.vector.flip()
for extrap in (extrapolation.ZERO, extrapolation.ONE, extrapolation.BOUNDARY, extrapolation.PERIODIC):
sampled = []
for backend in BACKENDS:
with backend:
grid = math.tensor(grid)
coords = math.tensor(coords)
grid_ = math.tensor(grid_)
coords_ = math.tensor(coords_)
sampled.append(math.grid_sample(grid, coords, extrap))
sampled.append(math.grid_sample(grid_, coords_, extrap))
math.assert_close(*sampled, abs_tolerance=1e-5)
def test_extrapolate_valid_3x3(self):
valid = tensor([[0, 0, 0],
[0, 0, 1],
[1, 0, 0]], spatial('x, y'))
values = tensor([[1, 0, 2],
[0, 0, 4],
[2, 0, 0]], spatial('x, y'))
expected_valid = tensor([[0, 1, 1],
[1, 1, 1],
[1, 1, 1]], spatial('x, y'))
expected_values = tensor([[1, 4, 4],
[2, 3, 4],
[2, 3, 4]], spatial('x, y'))
extrapolated_values, extrapolated_valid = math.extrapolate_valid_values(values, valid)
math.assert_close(extrapolated_values, expected_values)
math.assert_close(extrapolated_valid, expected_valid)
def test_animate(self):
values = math.random_uniform(batch(time=3), spatial(x=32, y=32))
anim = plot(values,
animate='time',
show_color_bar=False,
frame_time=100,
lib='matplotlib')
anim.to_html5_video()
def test_grid_constant_extrapolation(self):
grid = CenteredGrid(math.random_uniform(spatial(x=50, y=10)), 0.,
Box[0:1, 0:1])
self.assertEqual(grid.extrapolation, extrapolation.ZERO)
grid = CenteredGrid(0, 0, Box[0:1, 0:1], x=50, y=10)
self.assertEqual(grid.extrapolation, extrapolation.ZERO)
grid = StaggeredGrid(0, 0, Box[0:1, 0:1], x=50, y=10)
self.assertEqual(grid.extrapolation, extrapolation.ZERO)
def test_dot_batched_vector(self):
for backend in BACKENDS:
with backend:
a = math.ones(batch(batch=10) & spatial(a=4))
b = math.ones(batch(batch=10) & spatial(b=4))
dot = math.dot(a, 'a', b, 'b')
self.assertEqual(batch(batch=10), dot.shape, msg=backend.name)
math.assert_close(dot, 4, a.a * b.b, msg=backend.name)
dot = math.dot(a, 'a', a, 'a')
self.assertEqual(batch(batch=10), dot.shape, msg=backend.name)
math.assert_close(dot, 4, a.a * a.a, msg=backend.name)
# more dimensions
a = math.ones(batch(batch=10) & spatial(a=4, x=2))
b = math.ones(batch(batch=10) & spatial(y=3, b=4))
dot = math.dot(a, 'a', b, 'b')
self.assertEqual(set(spatial(x=2, y=3) & batch(batch=10)), set(dot.shape), msg=backend.name)
math.assert_close(dot, 4, msg=backend.name)
def test_pad_collapsed(self):
a = math.zeros(spatial(b=2, x=10, y=10) & batch(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)
