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_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_runge_kutta_4(self):
domain = Domain(x=4, y=3)
points = domain.distribute_points(domain.bounds, points_per_cell=2)
v = domain.vector_grid(Noise(vector=2))
field.assert_close(points, advect.runge_kutta_4(points, v, 0),
advect.runge_kutta_4(points, v * 0, 0))
sv = domain.staggered_grid(Noise())
field.assert_close(points, advect.runge_kutta_4(points, sv, 0),
advect.runge_kutta_4(points, sv * 0, 0))
def test_trace_function(self):
def f(x: StaggeredGrid, y: CenteredGrid):
return x + (y >> x)
ft = field.trace_function(f)
domain = Domain(x=4, y=3)
x = domain.staggered_grid(1)
y = domain.vector_grid(1)
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_gradient_function(self):
def f(x: StaggeredGrid, y: CenteredGrid):
pred = x + (y >> x)
loss = field.l2_loss(pred)
return loss
domain = Domain(x=4, y=3)
x = domain.staggered_grid(1)
y = domain.vector_grid(1)
with torch.TORCH_BACKEND:
dx, = field.gradient_function(f)(x, y)
self.assertIsInstance(dx, StaggeredGrid)
loss, dx, dy = field.gradient_function(f, (0, 1),
get_output=True)(x, y)
self.assertIsInstance(loss, math.Tensor)
self.assertIsInstance(dx, StaggeredGrid)
self.assertIsInstance(dy, CenteredGrid)
def test_gradient_function(self):
def f(x: StaggeredGrid, y: CenteredGrid):
pred = x + (y >> x)
loss = field.l2_loss(pred)
return loss
domain = Domain(x=4, y=3)
x = domain.staggered_grid(1)
y = domain.vector_grid(1)
for backend in BACKENDS:
if backend.supports(Backend.gradients):
with backend:
dx, = field.functional_gradient(f)(x, y)
self.assertIsInstance(dx, StaggeredGrid)
loss, dx, dy = field.functional_gradient(f, (0, 1), get_output=True)(x, y)
self.assertIsInstance(loss, math.Tensor)
self.assertIsInstance(dx, StaggeredGrid)
self.assertIsInstance(dy, CenteredGrid)
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_sample_at(self):
DOMAIN = Domain(x=4, y=3)
field = AngularVelocity([0, 0])
self.assertEqual(math.shape(vector=2), field.shape.channel)
field >> DOMAIN.vector_grid()
field >> DOMAIN.staggered_grid()