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_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_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_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_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_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_trace_function(self):
def f(x: StaggeredGrid, y: CenteredGrid):
return x + (y @ x)
ft = field.jit_compile(f)
x = StaggeredGrid(1, x=4, y=3)
y = CenteredGrid((1, 1), x=4, y=3)
res_f = f(x, y)
res_ft = ft(x, y)
self.assertEqual(res_f.shape, res_ft.shape)
field.assert_close(res_f, res_ft)
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_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_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_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_gradient_function(self):
def f(x: StaggeredGrid, y: CenteredGrid):
pred = x + (y @ x)
loss = field.l2_loss(pred)
return loss
grad = field.functional_gradient(f, get_output=False)
fgrad = field.functional_gradient(f, (0, 1), get_output=True)
x = StaggeredGrid(1, x=4, y=3)
y = CenteredGrid((1, 1), x=4, y=3)
for backend in BACKENDS:
if backend.supports(Backend.gradients):
with backend:
dx, = grad(x, y)
self.assertIsInstance(dx, StaggeredGrid)
loss, (dx, dy) = fgrad(x, y)
self.assertIsInstance(loss, math.Tensor)
self.assertIsInstance(dx, StaggeredGrid)
self.assertIsInstance(dy, CenteredGrid)
def test_write_read(self):
smoke = CenteredGrid(1, extrapolation.BOUNDARY, x=32, y=32)
vel = StaggeredGrid(2, 0, x=32, y=32)
# write
scene = Scene.create(DIR)
scene.write(smoke=smoke, vel=vel)
self.assertEqual(1, len(scene.frames))
self.assertEqual(1, len(scene.complete_frames))
self.assertEqual(2, len(scene.fieldnames))
# read single
smoke_ = scene.read('smoke')
vel_ = scene.read('vel')
field.assert_close(smoke, smoke_)
field.assert_close(vel, vel_)
self.assertEqual(smoke.extrapolation, smoke_.extrapolation)
self.assertEqual(vel.extrapolation, vel_.extrapolation)
# read multiple
smoke__, vel__ = scene.read(['smoke', 'vel']) # deprecated
field.assert_close(smoke, smoke__)
field.assert_close(vel, vel__)
smoke__, vel__ = scene.read('smoke', 'vel')
field.assert_close(smoke, smoke__)
field.assert_close(vel, vel__)
scene.remove()
def test_floor(self):
grid = StaggeredGrid(2.8, x=4, y=3)
field.assert_close(field.floor(grid), 2)
def test_downsample_staggered_2d(self):
grid = StaggeredGrid(1, extrapolation.BOUNDARY, x=32, y=40)
downsampled = field.downsample2x(grid)
self.assertEqual(set(spatial(x=16, y=20) & channel(vector=2)),
set(downsampled.shape))
def test_isfinite(self):
grid = StaggeredGrid(1, x=4, y=3)
field.assert_close(field.isfinite(grid), True)
def test_sqrt(self):
grid = StaggeredGrid(2, x=4, y=3)
field.assert_close(field.sqrt(grid), numpy.sqrt(2))
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_iter_dim(self):
slices = tuple(StaggeredGrid(0, x=4, y=3).vector)
self.assertEqual(2, len(slices))
self.assertFalse(slices[0].shape.non_spatial)
self.assertEqual(('x', 'y'), slices[0].bounds.size.vector.item_names)
def test_sample_at(self):
field = AngularVelocity([0, 0])
self.assertEqual(channel(vector=2), field.shape.channel)
field @ CenteredGrid(0, x=4, y=3)
field @ StaggeredGrid(0, x=4, y=3)
def test_create_grid_int_resolution(self):
g = CenteredGrid(0, 0, Box(x=4, y=3), resolution=10)
self.assertEqual(g.shape, spatial(x=10, y=10))
g = StaggeredGrid(0, 0, Box(x=4, y=3), resolution=10)
self.assertEqual(spatial(g), spatial(x=10, y=10))
def test_ceil(self):
grid = StaggeredGrid(1.1, x=4, y=3)
field.assert_close(field.ceil(grid), 2)
def test_real(self):
grid = StaggeredGrid(1, x=4, y=3)
field.assert_close(field.real(grid), grid)
def test_round_(self):
grid = StaggeredGrid(1.7, x=4, y=3)
field.assert_close(field.round(grid), 2)
def test_sign(self):
grid = StaggeredGrid(0.5, x=4, y=3)
field.assert_close(field.sign(grid), 1)
def test_imag(self):
grid = StaggeredGrid(1, x=4, y=3)
field.assert_close(field.imag(grid), 0)
def test_cos(self):
grid = StaggeredGrid(0, x=4, y=3)
field.assert_close(field.cos(grid), 1)
def test_curl_2d_staggered_to_centered(self):
velocity = StaggeredGrid((2, 0), extrapolation.BOUNDARY, x=2, y=2)
curl = field.curl(velocity)
self.assertEqual(spatial(x=3, y=3), curl.resolution)
def test_infinite_cylinder_to_grid(self):
cylinder = geom.infinite_cylinder(x=2, y=1.5, radius=.8, inf_dim='z')
StaggeredGrid(cylinder, 0, x=4, y=3, z=2)
def test_abs(self):
grid = StaggeredGrid(-1, x=4, y=3)
field.assert_close(field.abs(grid), abs(grid), 1)
