def test_multi_dim_tensor_from_numpy(self):
v = math.tensor(np.ones([1, 4, 3, 2]), batch('batch'), spatial('x,y'),
channel('vector'))
self.assertEqual((1, 4, 3, 2), v.shape.sizes)
v = math.tensor(np.ones([10, 4, 3, 2]), batch('batch'), spatial('x,y'),
channel('vector'))
self.assertEqual((10, 4, 3, 2), v.shape.sizes)
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_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_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 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 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_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_zeros_nonuniform(self):
nonuniform = shape_stack(batch('stack'),
batch(time=1) & spatial(x=3, y=3),
spatial(x=3, y=4), channel())
self.assertEqual(math.zeros(nonuniform).shape, nonuniform)
self.assertEqual(math.ones(nonuniform).shape, nonuniform)
self.assertEqual(math.random_normal(nonuniform).shape, nonuniform)
self.assertEqual(math.random_uniform(nonuniform).shape, nonuniform)
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_item_names(self):
s = spatial(x=4, y=3)
named = s.shape
self.assertEqual(named.get_item_names('dims'), ('x', 'y'))
shape = math.concat_shapes(batch(b=10), named)
self.assertEqual(shape.get_item_names('dims'), ('x', 'y'))
shape = math.merge_shapes(batch(b=10), named)
self.assertEqual(shape.get_item_names('dims'), ('x', 'y'))
c = channel(vector='r,g,b')
self.assertEqual(('r', 'g', 'b'), c.get_item_names('vector'))
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_tensor_from_tuple_of_tensor_like(self):
native = ([1, 2, 3], math.zeros(channel(vector=3)))
for backend in BACKENDS:
with backend:
tens = wrap(native, batch(stack=2), channel(vector=3))
self.assertEqual(math.NUMPY, math.choose_backend(tens))
self.assertEqual(
batch(stack=2) & channel(vector=3), tens.shape)
tens = tensor(native, batch(stack=2), channel(vector=3))
self.assertEqual(backend, math.choose_backend(tens))
self.assertEqual(
batch(stack=2) & channel(vector=3), tens.shape)
def test_plot_scalar_2d_batch(self):
self._test_plot(
CenteredGrid(Noise(batch(b=2)),
0,
x=64,
y=8,
bounds=Box(0, [1, 1])))
self._test_plot(CenteredGrid(Noise(batch(b=2)),
0,
x=64,
y=8,
bounds=Box(0, [1, 1])),
row_dims='b',
size=(2, 4))
def test_write_read_batch_matching(self):
smoke = CenteredGrid(1, extrapolation.BOUNDARY, x=32,
y=32) * math.random_uniform(batch(count=2))
vel = StaggeredGrid(2, 0, x=32, y=32) * math.random_uniform(
batch(count=2))
# write
scene = Scene.create(DIR, count=2)
scene.write({'smoke': smoke, 'vel': vel})
# read batch
smoke_ = scene.read('smoke')
vel_ = scene.read('vel')
field.assert_close(smoke, smoke_)
field.assert_close(vel, vel_)
scene.remove()
def test_tensor_from_tensor(self):
ref = math.stack([math.zeros(spatial(x=5)),
math.zeros(spatial(x=4))], batch('stack'))
for backend in BACKENDS:
with backend:
tens = math.tensor(ref, convert=False)
self.assertEqual(math.NUMPY, math.choose_backend(tens))
self.assertEqual(2, tens.shape.get_size('stack'))
self.assertEqual(('stack', 'x'), tens.shape.names)
tens = math.tensor(ref)
self.assertEqual(backend, math.choose_backend(tens))
self.assertEqual(backend, math.choose_backend(tens.stack[0]))
self.assertEqual(backend, math.choose_backend(tens.stack[1]))
tens = math.tensor(ref, batch('n1', 'n2'))
self.assertEqual(backend, math.choose_backend(tens))
def test_stacked_shapes(self):
t0 = math.ones(batch(batch=10) & spatial(x=4, y=3) & channel(vector=2))
for dim in t0.shape.names:
tensors = t0.unstack(dim)
stacked = math.stack(tensors, t0.shape[dim].with_sizes([None]))
self.assertEqual(set(t0.shape.names), set(stacked.shape.names))
self.assertEqual(t0.shape.volume, stacked.shape.volume)
def _test_make_incompressible(self, grid_type: type,
extrapolation: math.Extrapolation,
**batch_dims):
result = None
for i, backend in enumerate(BACKENDS):
with backend:
smoke = CenteredGrid(Sphere(
center=(math.random_uniform(batch(**batch_dims)) * 100,
10),
radius=5),
extrapolation,
x=16,
y=20,
bounds=Box[0:100, 0:100])
velocity = grid_type(0,
extrapolation,
x=16,
y=20,
bounds=Box[0:100, 0:100])
for _ in range(2):
velocity += smoke * (0, 0.1) @ velocity
velocity, _ = fluid.make_incompressible(velocity)
math.assert_close(divergence(velocity).values,
0,
abs_tolerance=2e-5)
if result is None:
result = velocity
else:
field.assert_close(
result,
abs_tolerance=1e-5,
msg=
f"Simulation with {backend} does not match {BACKENDS[:i]}"
)
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_slice_centered_grid(self):
g = CenteredGrid(Noise(batch(batch=10), channel(vector=2)), x=10, y=20)
s1 = g[{'vector': 0, 'batch': 1, 'x': 1}]
s2 = g.vector[0].batch[1].x[1]
self.assertIsInstance(s1, CenteredGrid)
self.assertEqual(s1.bounds, Box[1:2, 0:20])
field.assert_close(s1, s2)
def test_semi_collapsed(self):
scalar = math.ones(spatial(x=4, y=3))
scalar = CollapsedTensor(scalar, scalar.shape._expand(batch(batch=10)))
self.assertEqual((10, 4, 3), scalar.shape.sizes)
self.assertEqual(4, len(scalar.x.unstack()))
self.assertEqual(10, len(scalar.batch.unstack()))
self.assertEqual(0, scalar.y[0].batch[0].x[0].shape.rank)
def test_plot_vector_2d_batch(self):
self._test_plot(
CenteredGrid(Noise(batch(b=2), vector=2),
extrapolation.ZERO,
bounds=Box[0:1, 0:1],
x=10,
y=10))
def test_slice_staggered_grid_along_batch(self):
v = StaggeredGrid(Noise(batch(batch=10)), x=10, y=20)
b1 = v[{'batch': 1}]
b2 = v.batch[1]
b3 = field.unstack(v, 'batch')[1]
self.assertIsInstance(b1, StaggeredGrid)
field.assert_close(b1, b2, b3)
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_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_hessian(self):
def f(x, y):
return math.l1_loss(x**2 * y), x, y
eval_hessian = math.hessian(f,
wrt=[
0,
],
get_output=True,
get_gradient=True,
dim_suffixes=('1', '2'))
for backend in BACKENDS:
if backend.supports(Backend.hessian):
with backend:
x = math.tensor([(0.01, 1, 2)], channel('vector', 'v'))
y = math.tensor([1., 2.], batch('batch'))
(L, x, y), g, H, = eval_hessian(x, y)
math.assert_close(L, (5.0001, 10.0002), msg=backend.name)
math.assert_close(g.batch[0].vector[0], (0.02, 2, 4),
msg=backend.name)
math.assert_close(g.batch[1].vector[0], (0.04, 4, 8),
msg=backend.name)
math.assert_close(2,
H.v1[0].v2[0].batch[0],
H.v1[1].v2[1].batch[0],
H.v1[2].v2[2].batch[0],
msg=backend.name)
math.assert_close(4,
H.v1[0].v2[0].batch[1],
H.v1[1].v2[1].batch[1],
H.v1[2].v2[2].batch[1],
msg=backend.name)
def test_slice_staggered_grid_along_vector(self):
v = StaggeredGrid(Noise(batch(batch=10)), x=10, y=20)
x1 = v[{'vector': 0}]
x2 = v.vector[0]
x3 = v.vector['x']
x4 = field.unstack(v, 'vector')[0]
self.assertIsInstance(x1, CenteredGrid)
field.assert_close(x1, x2, x3, x4)
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_stacked_get(self):
t0 = math.ones(batch(batch=10) & spatial(x=4, y=3) & channel(vector=2))
tensors = t0.unstack('vector')
stacked = math.stack(tensors, channel('channel'))
self.assertEqual(tensors, stacked.channel.unstack())
assert tensors[0] is stacked.channel[0]
assert tensors[1] is stacked.channel[1:2].channel.unstack()[0]
self.assertEqual(4, len(stacked.x.unstack()))
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)
