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)