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_direct_initializers(self):
np.testing.assert_equal(math.zeros([1, 16]), np.zeros([1, 16]))
self.assertEqual(math.zeros([1, 16]).dtype, np.float32)
np.testing.assert_equal(math.ones([1, 16, 1]), np.ones([1, 16, 1]))
np.testing.assert_equal(math.zeros_like(math.ones([1, 16, 1])), np.zeros([1, 16, 1]))
np.testing.assert_equal(math.randn([1, 4]).shape, [1, 4])
self.assertEqual(math.randn([1, 4]).dtype, np.float32)
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_Dict(self):
d1 = math.Dict(a=1, b=math.ones(), c=math.ones(spatial(x=3)))
math.assert_close(d1 * 2, d1 + d1, 2 * d1, 2 / d1)
math.assert_close(0 + d1, d1, d1 - 0, abs(d1), round(d1))
math.assert_close(-d1, 0 - d1)
math.assert_close(d1 // 2, d1 * 0, d1 % 1)
math.assert_close(d1 / 2, d1 * 0.5, 0.5 * d1)
math.assert_close(math.sin(d1 * 0), d1 * 0)
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_collapsed_op2(self):
# Collapsed + Collapsed
a = math.zeros(vector=4)
b = math.ones(batch=3)
c = a + b
self.assertIsInstance(c, CollapsedTensor)
self.assertEqual(c.shape.volume, 12)
self.assertEqual(c._inner.shape.volume, 1)
# Collapsed + Native
n = math.ones(vector=3) + (0, 1, 2)
math.assert_close(n, (1, 2, 3))
def test_trace_function(self):
def f(x: math.Tensor, y: math.Tensor):
return x + y
for backend in BACKENDS:
with backend:
ft = math.jit_compile(f)
args1 = math.ones(x=2), math.ones(y=2)
args2 = math.ones(x=3), math.ones(y=3)
res1 = ft(*args1)
self.assertEqual(math.shape(x=2, y=2), res1.shape)
math.assert_close(res1, 2)
res2 = ft(*args2)
self.assertEqual(math.shape(x=3, y=3), res2.shape)
math.assert_close(res2, 2)
def points(self,
points: Tensor or Number or tuple or list,
radius: Tensor or float or int or None = None,
extrapolation: math.Extrapolation = None,
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
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 or math.extrapolation.ZERO
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(points=0, vector=1)
elif isinstance(points[0], Number): # single point
points = math.wrap([points], 'points, vector')
else:
points = math.wrap(points, 'points, vector')
elements = Sphere(points, radius)
return PointCloud(elements, math.ones(), extrapolation, add_overlapping=False, bounds=self.bounds, color=color)
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_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_zeros_nonuniform(self):
nonuniform = shape_stack('stack', BATCH_DIM, shape(time=1, x=3, y=3),
shape(x=3, y=4), shape())
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_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 divergence_free(velocity, domain=None, obstacles=(), pressure_solver=None, return_info=False):
"""
Projects the given velocity field by solving for and subtracting the pressure.
:param return_info: if True, returns a dict holding information about the solve as a second object
:param velocity: StaggeredGrid
:param domain: Domain matching the velocity field, used for boundary conditions
:param obstacles: list of Obstacles
:param pressure_solver: PressureSolver. Uses default solver if none provided.
:return: divergence-free velocity as StaggeredGrid
"""
assert isinstance(velocity, StaggeredGrid)
# --- Set up FluidDomain ---
if domain is None:
domain = Domain(velocity.resolution, OPEN)
obstacle_mask = union_mask([obstacle.geometry for obstacle in obstacles])
if obstacle_mask is not None:
obstacle_grid = obstacle_mask.at(velocity.center_points, collapse_dimensions=False).copied_with(extrapolation='constant')
active_mask = 1 - obstacle_grid
else:
active_mask = math.ones(domain.centered_shape(name='active', extrapolation='constant'))
accessible_mask = active_mask.copied_with(extrapolation=Material.accessible_extrapolation_mode(domain.boundaries))
fluiddomain = FluidDomain(domain, active=active_mask, accessible=accessible_mask)
# --- Boundary Conditions, Pressure Solve ---
velocity = fluiddomain.with_hard_boundary_conditions(velocity)
divergence_field = velocity.divergence(physical_units=False)
pressure, iterations = solve_pressure(divergence_field, fluiddomain, pressure_solver=pressure_solver)
pressure *= velocity.dx[0]
gradp = StaggeredGrid.gradient(pressure)
velocity -= fluiddomain.with_hard_boundary_conditions(gradp)
return velocity if not return_info else (velocity, {'pressure': pressure, 'iterations': iterations})
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_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_semi_collapsed(self):
scalar = math.ones(x=4, y=3)
scalar = CollapsedTensor(scalar, scalar.shape.expand(10, 'batch', BATCH_DIM))
self.assertEqual('(batch=10, x=4, y=3)', repr(scalar.shape))
self.assertEqual(4, len(scalar.x.unstack()))
self.assertEqual(10, len(scalar.batch.unstack()))
self.assertEqual('()', repr(scalar.y[0].batch[0].x[0].shape))
def test_stacked_shapes(self):
t0 = math.ones(batch=10, x=4, y=3, vector=2)
for dim in t0.shape.names:
tensors = t0.unstack(dim)
stacked = TensorStack(tensors, dim, t0.shape.get_type(dim))
self.assertEqual(set(t0.shape.names), set(stacked.shape.names))
self.assertEqual(t0.shape.volume, stacked.shape.volume)
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_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_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_custom_gradient_scalar(self):
def f(x):
return x
def grad(_x, _y, df):
return df * 0,
for backend in BACKENDS:
if backend.supports(Backend.functional_gradient):
with backend:
normal_gradient, = math.functional_gradient(
f, get_output=False)(math.ones())
math.assert_close(normal_gradient, 1)
f_custom_grad = math.custom_gradient(f, grad)
custom_gradient, = math.functional_gradient(
f_custom_grad, get_output=False)(math.ones())
math.assert_close(custom_gradient, 0)
def test_stacked_get(self):
t0 = math.ones(batch=10, x=4, y=3, vector=2)
tensors = t0.unstack('vector')
stacked = TensorStack(tensors, 'channel', CHANNEL_DIM)
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_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_pad(self):
test_in_func_out = [
(math.zeros(spatial(x=3, y=4, z=5,
a=1)), lambda tensor: ConstantExtrapolation(0).
pad(tensor, dict(x=[1, 1], y=[1, 0], z=[0, 1], a=[0, 0])),
math.zeros(spatial(x=5, y=5, z=6, a=1))),
(math.ones(spatial(x=3, y=4, z=5,
a=1)), lambda tensor: ConstantExtrapolation(1).
pad(tensor, dict(x=[1, 1], y=[1, 0], z=[0, 1], a=[0, 0])),
math.ones(spatial(x=5, y=5, z=6, a=1))),
(-math.ones(spatial(x=3, y=4, z=5, a=1)),
lambda tensor: ConstantExtrapolation(-1).pad(
tensor, dict(x=[1, 1], y=[1, 0], z=[0, 1], a=[0, 0])),
-math.ones(spatial(x=5, y=5, z=6, a=1))),
]
for val_in, func, val_out in test_in_func_out:
math.assert_close(val_out, func(val_in))
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_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_trace_function(self):
def f(x: math.Tensor, y: math.Tensor):
return x + y
for backend in [
math.NUMPY_BACKEND, tf.TF_BACKEND, torch.TORCH_BACKEND
]:
with backend:
ft = math.trace_function(f)
args1 = math.ones(x=2), math.ones(y=2)
args2 = math.ones(x=3), math.ones(y=3)
res1 = ft(*args1)
self.assertEqual(math.shape(x=2, y=2), res1.shape)
math.assert_close(res1, 2)
res2 = ft(*args2)
self.assertEqual(math.shape(x=3, y=3), res2.shape)
math.assert_close(res2, 2)
def test_dimension_types(self):
v = math.ones(batch=10, x=4, y=3, vector=2)
self.assertEqual(v.x.index, 1)
self.assertEqual(v.x.name, 'x')
self.assertTrue(v.x.is_spatial)
self.assertTrue(v.batch.is_batch)
b = v.x.as_batch()
self.assertTrue(b.x.is_batch)
def test_pad(self):
test_in_func_out = [
(math.zeros(x=3, y=4, z=5, a=1),
lambda tensor: ConstantExtrapolation(0).pad(tensor, dict(x=[1, 1], y=[1, 0], z=[0, 1], a=[0, 0])),
math.zeros(x=5, y=5, z=6, a=1)),
(math.ones(x=3, y=4, z=5, a=1),
lambda tensor: ConstantExtrapolation(1).pad(tensor, dict(x=[1, 1], y=[1, 0], z=[0, 1], a=[0, 0])),
math.ones(x=5, y=5, z=6, a=1)),
(-math.ones(x=3, y=4, z=5, a=1),
lambda tensor: ConstantExtrapolation(-1).pad(tensor, dict(x=[1, 1], y=[1, 0], z=[0, 1], a=[0, 0])),
- math.ones(x=5, y=5, z=6, a=1)),
]
for val_in, func, val_out in test_in_func_out:
try:
math.assert_close(val_out, func(val_in))
# TypeError('__bool__ should return bool, returned NotImplementedType')
# self.assertEqual(val_out, func(val_in))
except Exception as e:
raise BaseException(AssertionError(e, val_in, func, val_out))