def _test_advection(adv):
s = CenteredGrid(Noise(), x=4, y=3)
v = CenteredGrid(Noise(vector=2), x=4, y=3)
field.assert_close(s, adv(s, v, 0), adv(s, v * 0, 1))
sv = StaggeredGrid(Noise(), x=4, y=3)
field.assert_close(s, adv(s, sv, 0), adv(s, sv * 0, 1))
field.assert_close(sv, adv(sv, sv, 0), adv(sv, sv * 0, 1))
def test_zero_vector_grid(self):
for data in [(0, 0),
Noise(vector='x,y'),
Noise(vector=2), lambda x: x]:
grid = CenteredGrid(data, 0, x=4, y=3)
self.assertEqual(('x', 'y'), grid.values.vector.item_names)
self.assertEqual(('x', 'y'), grid.dx.vector.item_names)
def _test_advection(adv):
domain = Domain(x=4, y=3)
s = domain.scalar_grid(Noise())
v = domain.vector_grid(Noise(vector=2))
field.assert_close(s, adv(s, v, 0), adv(s, v * 0, 1))
sv = domain.staggered_grid(Noise())
field.assert_close(s, adv(s, sv, 0), adv(s, sv * 0, 1))
field.assert_close(sv, adv(sv, sv, 0), adv(sv, sv * 0, 1))
def test_runge_kutta_4(self):
domain = Domain(x=4, y=3)
points = domain.distribute_points(domain.bounds, points_per_cell=2)
v = CenteredGrid(Noise(vector=2), x=4, y=3)
field.assert_close(points, advect.runge_kutta_4(points, v, 0),
advect.runge_kutta_4(points, v * 0, 0))
sv = StaggeredGrid(Noise(), x=4, y=3)
field.assert_close(points, advect.runge_kutta_4(points, sv, 0),
advect.runge_kutta_4(points, sv * 0, 0))
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_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_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_plot_vector_grid_2d(self):
self._test_plot(
CenteredGrid(Noise(vector=2),
extrapolation.ZERO,
x=64,
y=8,
bounds=Box(0, [1, 1])) * 0.1)
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_diffuse_centered_batched(self):
grid = CenteredGrid(Noise(batch=2, vector=2),
extrapolation.PERIODIC,
x=4,
y=3)
diffuse.explicit(grid, 1, 1, substeps=10)
diffuse.implicit(grid, 1, 1, order=2)
diffuse.fourier(grid, 1, 1)
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_plot_multiple(self):
grid = CenteredGrid(Noise(batch(b=2)), 0, Box[0:1, 0:1], x=50, y=10)
grid2 = CenteredGrid(grid, 0, Box[0:2, 0:1], x=20, y=50)
points = wrap([(.2, .4), (.9, .8)], instance('points'),
channel('vector'))
cloud = PointCloud(Sphere(points, radius=0.1), bounds=Box(0, [1, 1]))
titles = math.wrap([['b=0', 'b=0', 'points'], ['b=1', 'b=1', '']],
spatial('rows,cols'))
self._test_plot(grid, grid2, cloud, row_dims='b', title=titles)
def test_diffuse_staggered_batched(self):
for diffusivity in [1, 0.5, math.wrap([1., 0.5], batch('batch'))]:
grid = StaggeredGrid(Noise(batch(batch=2), vector=2),
extrapolation.PERIODIC,
x=4,
y=3)
diffuse.explicit(grid, diffusivity, 1, substeps=10)
diffuse.implicit(grid, diffusivity, 1, order=2)
diffuse.fourier(grid, diffusivity, 1)
grid = StaggeredGrid(Noise(batch(batch=2), vector=2),
extrapolation.ZERO,
x=4,
y=3)
diffuse.explicit(grid, diffusivity, 1, substeps=10)
# diffuse.implicit(grid, diffusivity, 1, order=2) # not yet supported
grid = StaggeredGrid(Noise(batch(batch=2), vector=2),
extrapolation.BOUNDARY,
x=4,
y=3)
diffuse.explicit(grid, diffusivity, 1, substeps=10)
def test_overlay(self):
grid = CenteredGrid(Noise(),
extrapolation.ZERO,
x=64,
y=8,
bounds=Box(0, [1, 1]))
points = wrap([(.2, .4), (.9, .8)], instance('points'),
channel('vector'))
cloud = PointCloud(Sphere(points, radius=.1))
self._test_plot(overlay(grid, grid * (0.1, 0.02), cloud),
title='Overlay')
def test_staggered_grid_with_extrapolation(self):
grid = StaggeredGrid(Noise(vector='x,y'),
extrapolation.BOUNDARY,
x=20,
y=10)
grid_0 = grid.with_extrapolation(extrapolation.ZERO)
self.assertEqual(grid.resolution, grid_0.resolution)
grid_ = grid_0.with_extrapolation(extrapolation.BOUNDARY)
self.assertEqual(grid.resolution, grid_.resolution)
math.assert_close(grid_.values.vector['x'].x[0], 0)
math.assert_close(grid_.values.vector['x'].x[-1], 0)
math.assert_close(grid_.values.vector['y'].y[0], 0)
math.assert_close(grid_.values.vector['y'].y[-1], 0)
def test_make_incompressible_gradients_equal_tf_torch(self):
DOMAIN = Domain(x=16, y=16, boundaries=OPEN, bounds=Box[0:100, 0:100]) # TODO CLOSED solve fails because div is not subtracted from dx
velocity0 = DOMAIN.staggered_grid(Noise(vector=2))
grads = []
for backend in [TF_BACKEND, TORCH_BACKEND]:
with backend:
velocity = param = velocity0.with_(values=math.tensor(velocity0.values))
with math.record_gradients(param.values):
solve = math.LinearSolve()
velocity, _, _, _ = fluid.make_incompressible(velocity, DOMAIN, solve_params=solve)
loss = field.l2_loss(velocity)
assert math.isfinite(loss)
grad = math.gradients(loss, param.values)
assert math.all(math.isfinite(grad))
grads.append(grad)
math.assert_close(*grads, abs_tolerance=1e-5)
def test_make_incompressible_gradients_equal_tf_torch(self):
velocity0 = StaggeredGrid(Noise(),
ZERO,
x=16,
y=16,
bounds=Box[0:100, 0:100])
grads = []
for backend in BACKENDS:
if backend.supports(Backend.record_gradients):
with backend:
velocity = param = velocity0.with_values(
math.tensor(velocity0.values))
with math.record_gradients(param.values):
velocity, _ = fluid.make_incompressible(velocity)
loss = field.l2_loss(velocity)
assert math.isfinite(loss).all
grad = math.gradients(loss, param.values)
assert math.isfinite(grad).all
grads.append(grad)
math.assert_close(*grads, abs_tolerance=1e-5)
def test_staggered_grid_sizes_by_extrapolation(self):
s = spatial(x=20, y=10)
for initializer in [0, Noise(vector=2), (0, 1), Sphere((0, 0), 1)]:
g_const = StaggeredGrid(initializer,
extrapolation.ZERO,
resolution=s)
self.assertEqual(g_const.shape,
spatial(x=20, y=10) & channel(vector=2))
self.assertEqual(g_const.values.vector[0].shape, spatial(x=19,
y=10))
g_periodic = StaggeredGrid(initializer,
extrapolation.PERIODIC,
resolution=s)
self.assertEqual(g_periodic.shape,
spatial(x=20, y=10) & channel(vector=2))
self.assertEqual(g_periodic.values.vector[0].shape,
spatial(x=20, y=10))
g_boundary = StaggeredGrid(initializer,
extrapolation.BOUNDARY,
resolution=s)
self.assertEqual(g_boundary.shape,
spatial(x=20, y=10) & channel(vector=2))
self.assertEqual(g_boundary.values.vector[0].shape,
spatial(x=21, y=10))
def test_plot_scalar_tensor_2d(self):
self._test_plot(
CenteredGrid(Noise(), 0, x=64, y=8, bounds=Box(0, [1, 1])).values)
def test_plot_staggered_grid_2d(self):
self._test_plot(
StaggeredGrid(
Noise(), extrapolation.ZERO, x=16, y=10, bounds=Box(0, [1, 1]))
* 0.1)
def test_zero_staggered_grid(self):
for data in [(0, 0), 0, Noise(), lambda x: x]:
grid = StaggeredGrid(data, 0, x=4, y=3)
self.assertEqual(('x', 'y'), grid.values.vector.item_names)
self.assertEqual(('x', 'y'), grid.dx.vector.item_names)
def test_domain_grid_from_tensor(self):
domain = Domain(x=4, y=3)
grid = domain.vector_grid(Noise(vector=2))
grid2 = domain.vector_grid(grid.values)
math.assert_close(grid.values, grid2.values)
def test_domain_grid_from_field(self):
large_grid = Domain(x=4, y=3).grid(Noise())
small_grid = Domain(x=3, y=2).grid(large_grid)
math.assert_close(large_grid.values.x[:-1].y[:-1], small_grid.values)
