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_box_product(self):
a = Box(x=4)
b = Box(y=3).shifted(math.wrap(1))
ab = a * b
self.assertEqual(2, ab.spatial_rank)
math.assert_close(ab.size, (4, 3))
math.assert_close(ab.lower, (0, 1))
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 test_domain_grid_from_function(self):
grid = Domain(x=4, y=3).scalar_grid(lambda x: math.sum(x**2, 'vector'))
math.assert_close(grid.values.x[0].y[0], 0.5)
self.assertEqual(grid.shape.volume, 12)
grid = Domain(
x=4, y=3).scalar_grid(lambda x: math.ones(x.shape.non_channel))
math.assert_close(grid.values, 1)
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_batch_independence(self):
def simulate(centers):
world = World()
fluid = world.add(Fluid(Domain(x=5,
y=4,
boundaries=CLOSED,
bounds=Box(0, [40, 32])),
buoyancy_factor=0.1,
batch_size=centers.shape[0]),
physics=IncompressibleFlow())
world.add(Inflow(Sphere(center=centers, radius=3), rate=0.2))
world.add(
Fan(Sphere(center=centers, radius=5), acceleration=[1.0, 0]))
world.step(dt=1.5)
world.step(dt=1.5)
world.step(dt=1.5)
assert not math.close(fluid.density.values, 0)
print()
return fluid.density.values.batch[0], fluid.velocity.values.batch[
0]
d1, v1 = simulate(math.tensor([[5, 16], [5, 4]], names='batch,vector'))
d2, v2 = simulate(math.tensor([[5, 16], [5, 16]],
names='batch,vector'))
math.assert_close(d1, d2)
math.assert_close(v1, v2)
def test_closest_grid_values_1d(self):
grid = math.tensor([0, 1, 2, 3], names='x')
coords = math.tensor([[0.1], [1.9], [0.5], [3.1]], names='x,vector')
closest = math.closest_grid_values(grid, coords, extrapolation.ZERO)
math.assert_close(
closest,
math.tensor([(0, 1), (1, 2), (0, 1), (3, 0)], names='x,closest_x'))
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_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 test_collapsed(self):
scalar = math.zeros(spatial(x=4, y=3))
math.assert_close(scalar, 0)
self.assertEqual((4, 3), scalar.shape.sizes)
self.assertEqual(4, scalar.y[0].shape.size)
self.assertEqual(0, scalar.y[0].x[0].shape.rank)
self.assertEqual(3, len(scalar.y.unstack()))
def test_slice_by_item_name(self):
t = math.tensor(spatial(x=4, y=3))
math.assert_close(t.dims['x'], 4)
math.assert_close(t.dims['y'], 3)
math.assert_close(t.dims['y,x'], (3, 4))
math.assert_close(t.dims[('y', 'x')], (3, 4))
math.assert_close(t.dims[spatial('x,y')], (4, 3))
def test_collapsed(self):
scalar = math.zeros(x=4, y=3)
math.assert_close(scalar, 0)
self.assertEqual('(x=4, y=3)', repr(scalar.shape))
self.assertEqual('(x=4)', repr(scalar.y[0].shape))
self.assertEqual('()', repr(scalar.y[0].x[0].shape))
self.assertEqual(3, len(scalar.y.unstack()))
def test_constant_diffusion(self):
DOMAIN = Domain(x=4, y=3, boundaries=PERIODIC)
grid = DOMAIN.scalar_grid(1)
explicit = diffuse.explicit(grid, 1, 1, substeps=10)
implicit = diffuse.implicit(grid, 1, 1, order=2)
fourier = diffuse.fourier(grid, 1, 1)
math.assert_close(grid.values, explicit.values, implicit.values, fourier.values)
def test_constant_diffusion(self):
grid = CenteredGrid(1, extrapolation.PERIODIC, x=4, y=3)
explicit = diffuse.explicit(grid, 1, 1, substeps=10)
implicit = diffuse.implicit(grid, 1, 1, order=2)
fourier = diffuse.fourier(grid, 1, 1)
math.assert_close(grid.values, explicit.values, implicit.values,
fourier.values)
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_linear_operator(self):
GLOBAL_AXIS_ORDER.x_last()
direct = math.random_normal(batch=3, x=4, y=3) # , vector=2
op = lin_placeholder(direct)
def linear_function(val):
val = -val
val *= 2
val = math.pad(val, {'x': (2, 0), 'y': (0, 1)}, extrapolation.PERIODIC)
val = val.x[:-2].y[1:] + val.x[2:].y[:-1]
val = math.pad(val, {'x': (0, 0), 'y': (0, 1)}, extrapolation.ZERO)
val = math.pad(val, {'x': (2, 2), 'y': (0, 1)}, extrapolation.BOUNDARY)
# sl = sl.vector[0]
return val
val = val.x[1:4].y[:2]
return math.sum([val, sl], axis=0) - sl
functions = [
linear_function,
lambda val: math.gradient(val, difference='forward', padding=extrapolation.ZERO, dims='x').gradient[0],
lambda val: math.gradient(val, difference='backward', padding=extrapolation.PERIODIC, dims='x').gradient[0],
lambda val: math.gradient(val, difference='central', padding=extrapolation.BOUNDARY, dims='x').gradient[0],
]
for f in functions:
direct_result = f(direct)
# print(direct_result.batch[0], 'Direct result')
op_result = f(op)
# print(op_result.build_sparse_coordinate_matrix().todense())
self.assertIsInstance(op_result, ShiftLinOp)
op_result = NativeTensor(op_result.native(), op_result.shape)
# print(op_result.batch[0], 'Placeholder result')
math.assert_close(direct_result, op_result)
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_fourier_laplace_2d_periodic(self):
"""test for convergence of the laplace operator"""
test_params = {
'size': [16, 32, 40],
'L': [1, 2,
3], # NOTE: Cannot test with less than 1 full wavelength
}
test_cases = [
dict(zip(test_params, v)) for v in product(*test_params.values())
]
for params in test_cases:
vec = math.meshgrid(x=params['size'], y=params['size'])
sine_field = math.prod(
math.sin(2 * PI * params['L'] * vec / params['size'] + 1),
'vector')
sin_lap_ref = -2 * (
2 * PI * params['L'] / params['size']
)**2 * sine_field # leading 2 from from x-y cross terms
sin_lap = math.fourier_laplace(sine_field, 1)
try:
math.assert_close(sin_lap,
sin_lap_ref,
rel_tolerance=0,
abs_tolerance=1e-5)
except BaseException as e:
abs_error = math.abs(sin_lap - sin_lap_ref)
max_abs_error = math.max(abs_error)
max_rel_error = math.max(math.abs(abs_error / sin_lap_ref))
variation_str = "\n".join([
f"max_absolute_error: {max_abs_error}",
f"max_relative_error: {max_rel_error}",
])
print(f"{variation_str}\n{params}")
raise AssertionError(e, f"{variation_str}\n{params}")
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 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_stacked_native(self):
t0 = math.ones(batch=10, x=4, y=3, vector=2)
tensors = t0.unstack('vector')
stacked = TensorStack(tensors, 'vector2', CHANNEL_DIM)
math.assert_close(stacked, t0)
self.assertEqual((10, 4, 3, 2), stacked.native().shape)
self.assertEqual((4, 3, 2, 10), stacked.native(order=('x', 'y', 'vector2', 'batch')).shape)
self.assertEqual((2, 10, 3, 4), stacked.native(order=('vector2', 'batch', 'y', 'x')).shape) # this should re-stack since only the stacked dimension position is different
def test_make_incompressible_np_equal_torch(self):
import sys
print(sys.getrecursionlimit())
with math.NUMPY_BACKEND:
v_np = self._test_make_incompressible(StaggeredGrid)
with TORCH_BACKEND:
v_to = self._test_make_incompressible(StaggeredGrid)
math.assert_close(v_np, v_to, abs_tolerance=1e-5)
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_upsample2x(self):
meshgrid = math.meshgrid(x=(0, 1, 2, 3), y=(0, -1, -2))
double_size = math.upsample2x(meshgrid, extrapolation.BOUNDARY)
same_size = math.downsample2x(double_size)
math.print(meshgrid, 'Normal size')
math.print(double_size, 'Double size')
math.print(same_size, 'Same size')
math.assert_close(meshgrid.x[1:-1].y[1:-1], same_size.x[1:-1].y[1:-1])
def test_ifft(self):
dimensions = 'xyz'
for backend in BACKENDS:
with backend:
for d in range(1, len(dimensions) + 1):
x = math.random_normal(spatial(**{dim: 6 for dim in dimensions[:d]})) + math.tensor((0, 1), batch('batch'))
k = math.fft(x)
x_ = math.ifft(k)
math.assert_close(x, x_, abs_tolerance=1e-5, msg=backend.name)
def test_jit_compile_with_native(self):
@math.jit_compile
def scalar_mul(x, fac=1):
return x * fac
for backend in BACKENDS:
with backend:
x = backend.ones([3, 2])
math.assert_close(scalar_mul(x, fac=1), 1, msg=backend)
def _test_make_incompressible_batched(self, grid_type):
DOMAIN = Domain(x=16, y=16, boundaries=CLOSED, bounds=Box[0:100, 0:100])
smoke = DOMAIN.scalar_grid(Sphere(center=(math.random_uniform(batch=2) * 100, 10), radius=5))
velocity = DOMAIN.vector_grid(0, grid_type)
for _ in range(2):
velocity += smoke * (0, 0.1) >> velocity
velocity, pressure, _, _ = fluid.make_incompressible(velocity, DOMAIN)
math.assert_close(divergence(velocity).values, 0, abs_tolerance=2e-5)
return velocity.values
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_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_fft_dims(self):
for backend in BACKENDS:
with backend:
x = math.random_normal(batch(x=8, y=6, z=4))
k3 = math.fft(x, 'x,y,z')
k = x
for dim in 'xyz':
k = math.fft(k, dim)
math.assert_close(k, k3, abs_tolerance=1e-5, msg=backend.name)
