def test_random_int(self):
for backend in BACKENDS:
with backend:
# 32 bits
a = math.random_uniform(instance(values=1000), low=-1, high=1, dtype=(int, 32))
self.assertEqual(a.dtype, DType(int, 32), msg=backend.name)
self.assertEqual(a.min, -1, msg=backend.name)
self.assertEqual(a.max, 0, msg=backend.name)
# 64 bits
a = math.random_uniform(instance(values=1000), low=-1, high=1, dtype=(int, 64))
self.assertEqual(a.dtype.kind, int, msg=backend.name) # Jax may downcast 64-bit to 32
self.assertEqual(a.min, -1, msg=backend.name)
self.assertEqual(a.max, 0, msg=backend.name)
def test_write_read_batch_batched_files(self):
DOMAIN = Domain(x=32, y=32, boundaries=CLOSED)
smoke = DOMAIN.scalar_grid(1) * math.random_uniform(count=2, config=3)
vel = DOMAIN.staggered_grid(2) * math.random_uniform(count=2, vel=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_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_write_read_batch_duplicate(self):
DOMAIN = Domain(x=32, y=32, boundaries=CLOSED)
smoke = DOMAIN.scalar_grid(1) * math.random_uniform(count=2)
vel = DOMAIN.staggered_grid(2) * math.random_uniform(count=2)
# write
scene = Scene.create(DIR, more=2)
scene.write({'smoke': smoke, 'vel': vel})
# read batch
smoke_ = scene.read('smoke')
vel_ = scene.read('vel')
self.assertEqual(4, smoke_.shape.batch.volume)
self.assertEqual(4, vel_.shape.batch.volume)
field.assert_close(smoke, smoke_)
field.assert_close(vel, vel_)
scene.remove()
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_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_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 distribute_points(density, particles_per_cell=1, distribution='uniform'):
"""
Distribute points according to the distribution specified in density.
:param density: binary tensor
:param particles_per_cell: integer
:param distribution: 'uniform' or 'center'
:return: tensor of shape (batch_size, point_count, rank)
"""
assert distribution in ('center', 'uniform')
index_array = []
batch_size = math.staticshape(density)[0] if math.staticshape(density)[0] is not None else 1
for batch in range(batch_size):
indices = math.where(density[batch, ..., 0] > 0)
indices = math.to_float(indices)
temp = []
for _ in range(particles_per_cell):
if distribution == 'center':
temp.append(indices + 0.5)
elif distribution == 'uniform':
temp.append(indices + math.random_uniform(math.shape(indices)))
index_array.append(math.concat(temp, axis=0))
try:
index_array = math.stack(index_array)
return index_array
except ValueError:
raise ValueError("all arrays in the batch must have the same number of active cells.")
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_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_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_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_convert_point_cloud(self):
loc = math.random_uniform(instance(points=4), channel(vector=2))
val = math.random_normal(instance(points=4), channel(vector=2))
points = PointCloud(Sphere(loc, radius=1), val)
for backend in BACKENDS:
converted = field.convert(points, backend)
self.assertEqual(converted.values.default_backend, backend)
self.assertEqual(converted.elements.center.default_backend,
backend)
self.assertEqual(converted.elements.radius.default_backend,
backend)
def test_unstack(self):
a = math.random_uniform(batch(b=10), spatial(x=4, y=3), channel(vector=2))
u = math.unstack(a, 'vector')
self.assertIsInstance(u, tuple)
self.assertEqual(len(u), 2)
math.assert_close(u, math.unstack(a, channel))
# Multiple dimensions
u = math.unstack(a, 'x,y')
self.assertIsInstance(u, tuple)
self.assertEqual(len(u), 12)
math.assert_close(u, math.unstack(a, spatial))
def test_grid_sample_backend_equality_2d_batched(self):
grid = math.random_normal(mybatch=10, y=10, x=7)
coords = math.random_uniform(mybatch=10, x=3, y=2) * (12, 9)
grid_ = math.tensor(grid.native('mybatch,x,y'), 'mybatch,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, coords, grid_, coords_ = math.tensors(
grid, coords, grid_, 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_optimize_u_net(self):
for lib in LIBRARIES:
net = lib.u_net(1, 1, levels=2)
optimizer = lib.adam(net, 1e-3)
def loss_function(x):
print("Running loss_function")
assert isinstance(x, math.Tensor)
pred = math.native_call(net, x)
return math.l2_loss(pred)
for i in range(2):
lib.update_weights(
net, optimizer, loss_function,
math.random_uniform(math.batch(batch=10),
math.spatial(x=8, y=8)))
def test_map_types(self):
def f(x, y):
assert x.shape.batch.names == ('batch', 'x', 'y')
assert x.shape.channel.names == ('vector', )
assert y.shape == x.shape
return x, y
for f_ in [
# math.map_types(f, 'x,y', batch),
# math.map_types(f, spatial('x,y'), batch),
math.map_types(f, spatial, batch),
]:
x = math.random_uniform(batch(batch=10), spatial(x=4, y=3),
channel(vector=2))
x_, y_ = f_(x, x)
assert x_.shape == x.shape
math.assert_close(x, x_)
def test_cast(self):
for backend in BACKENDS:
with backend:
x = math.random_uniform(dtype=DType(float, 64))
self.assertEqual(DType(float, 32), math.to_float(x).dtype, msg=backend.name)
self.assertEqual(DType(complex, 64), math.to_complex(x).dtype, msg=backend.name)
with math.precision(64):
self.assertEqual(DType(float, 64), math.to_float(x).dtype, msg=backend.name)
self.assertEqual(DType(complex, 128), math.to_complex(x).dtype, msg=backend.name)
self.assertEqual(DType(int, 64), math.to_int64(x).dtype, msg=backend.name)
self.assertEqual(DType(int, 32), math.to_int32(x).dtype, msg=backend.name)
self.assertEqual(DType(float, 16), math.cast(x, DType(float, 16)).dtype, msg=backend.name)
self.assertEqual(DType(complex, 128), math.cast(x, DType(complex, 128)).dtype, msg=backend.name)
try:
math.cast(x, DType(float, 3))
self.fail(msg=backend.name)
except KeyError:
pass
def test_reshaped_native(self):
a = math.random_uniform(channel(vector=2) & spatial(x=4, y=3))
nat = math.reshaped_native(a, ['batch', a.shape.spatial, 'vector'], force_expand=False)
self.assertEqual((1, 12, 2), nat.shape)
nat = math.reshaped_native(a, [batch(batch=10), a.shape.spatial, channel(vector=2)], force_expand=False)
self.assertEqual((1, 12, 2), nat.shape)
nat = math.reshaped_native(a, [batch(batch=10), a.shape.spatial, channel(vector=2)], force_expand=['x'])
self.assertEqual((1, 12, 2), nat.shape)
nat = math.reshaped_native(a, [batch(batch=10), a.shape.spatial, channel(vector=2)], force_expand=True)
self.assertEqual((10, 12, 2), nat.shape)
nat = math.reshaped_native(a, [batch(batch=10), a.shape.spatial, channel(vector=2)], force_expand=['batch'])
self.assertEqual((10, 12, 2), nat.shape)
nat = math.reshaped_native(a, [a.shape.spatial, channel(vector=2, v2=2)], force_expand=False)
self.assertEqual((12, 4), nat.shape)
try:
math.reshaped_native(a, [channel(vector=2, v2=2)], force_expand=False)
except AssertionError as err:
print(err)
pass
def test_optimize_dense_net(self):
for lib in LIBRARIES:
net = lib.dense_net(2,
3,
layers=[10],
batch_norm=True,
activation='Sigmoid')
optimizer = lib.adam(net)
def loss_function(x):
print("Running loss_function")
assert isinstance(x, math.Tensor)
pred = math.native_call(net, x)
return math.l2_loss(pred)
for i in range(2):
lib.update_weights(
net, optimizer, loss_function,
math.random_uniform(batch(batch=10), channel(vector=2)))
def _distribute_points(mask: math.Tensor, points_per_cell: int = 1, center: bool = False) -> math.Tensor:
"""
Generates points (either uniformly distributed or at the cell centers) according to the given tensor mask.
Args:
mask: Tensor with nonzero values at the indices where particles should get generated.
points_per_cell: Number of particles to generate at each marked index
center: Set points to cell centers. If False, points will be distributed using a uniform
distribution within each cell.
Returns:
A tensor containing the positions of the generated points.
"""
indices = math.to_float(math.nonzero(mask, list_dim=instance('points')))
temp = []
for _ in range(points_per_cell):
if center:
temp.append(indices + 0.5)
else:
temp.append(indices + (math.random_uniform(indices.shape)))
return math.concat(temp, dim=instance('points'))
def test_plot_point_cloud_3d(self):
points = math.random_uniform(instance(points=50), channel(vector=3))
cloud = PointCloud(Sphere(points, radius=.1),
bounds=Box(x=2, y=1, z=1))
self._test_plot(cloud)
def sample_uniform(self, *shape: math.Shape) -> Tensor:
uniform = math.random_uniform(self.shape.non_singleton, *shape,
math.channel(vector=self.spatial_rank))
return self.lower + uniform * self.size
def test_plot_point_cloud_vector_field_2d_bounded(self):
points = math.random_uniform(instance(points='a,b,c,d,e'),
channel(vector='x,y'))
velocity = PointCloud(Sphere(points, radius=.1), bounds=Box(x=1, y=1))
self._test_plot(velocity * (.05, 0))
def test_random_complex(self):
for backend in BACKENDS:
with backend:
a = math.random_uniform(instance(values=4), low=-1, high=0, dtype=(complex, 64))
self.assertEqual(a.dtype, DType(complex, 64), msg=backend.name)
math.assert_close(a.imag, 0, msg=backend.name)
