def test_write_read(self):
DOMAIN = Domain(x=32, y=32, boundaries=CLOSED)
smoke = DOMAIN.scalar_grid(1)
vel = DOMAIN.staggered_grid(2)
# 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_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 = 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_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_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_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_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_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_custom_spatial_dims(self):
domain = Domain(a=4, b=3)
grid = domain.scalar_grid(1)
self.assertEqual(math.shape(a=4, b=3), grid.shape)
grid = domain.staggered_grid(1)
self.assertEqual(math.shape(a=4, b=3, vector=2), grid.shape)
def test_domain_grid_memory_allocation(self):
domain = Domain(x=10000, y=10000, z=10000, w=10000)
grid = domain.scalar_grid()
self.assertEqual((10000, ) * 4, grid.shape.sizes)
sgrid = domain.staggered_grid()
self.assertEqual((10000, 10000, 10000, 10000, 4), sgrid.shape.sizes)
def make_incompressible(
velocity: StaggeredGrid,
domain: Domain,
obstacles: tuple or list or StaggeredGrid = (),
particles: PointCloud or None = None,
solve_params: math.LinearSolve = math.LinearSolve(),
pressure_guess: CenteredGrid = None
) -> Tuple[StaggeredGrid, CenteredGrid, math.Tensor, CenteredGrid,
StaggeredGrid]:
"""
Projects the given velocity field by solving for the pressure and subtracting its spatial_gradient.
Args:
velocity: Current velocity field as StaggeredGrid
obstacles: Sequence of `phi.physics.Obstacle` objects or binary StaggeredGrid marking through-flow cell faces
particles (Optional if occupation masks are provided): Pointcloud holding the current positions of the particles
domain (Optional if occupation masks are provided): Domain object
pressure_guess (Optional): Initial pressure guess as CenteredGrid
solve_params: Parameters for the pressure solve
Returns:
velocity: divergence-free velocity of type `type(velocity)`
pressure: solved pressure field, `CenteredGrid`
iterations: Number of iterations required to solve for the pressure
divergence: divergence field of input velocity, `CenteredGrid`
occupation_mask: StaggeredGrid
"""
points = particles.with_(values=math.wrap(1))
occupied_centered = points >> domain.grid()
occupied_staggered = points >> domain.staggered_grid()
if isinstance(obstacles, StaggeredGrid):
accessible = obstacles
else:
accessible = domain.accessible_mask(
union(*[obstacle.geometry for obstacle in obstacles]),
type=StaggeredGrid)
# --- Extrapolation is needed to exclude border divergence from the `occupied_centered` mask and thus
# from the pressure solve. If particles are randomly distributed, the `occupied_centered` mask
# could sometimes include the divergence at the borders (due to single particles right at the edge
# which temporarily deform the `occupied_centered` mask when moving into a new cell) which would then
# get compensated by the pressure. This is unwanted for falling liquids and therefore prevented by this
# extrapolation. ---
velocity_field, _ = extrapolate_valid(velocity * occupied_staggered,
occupied_staggered, 1)
velocity_field *= accessible # Enforces boundary conditions after extrapolation
div = field.divergence(
velocity_field
) * occupied_centered # Multiplication with `occupied_centered` excludes border divergence from pressure solve
def matrix_eq(p):
return field.where(
occupied_centered,
field.divergence(
field.spatial_gradient(p, type=StaggeredGrid) * accessible), p)
converged, pressure, iterations = field.solve(matrix_eq,
div,
pressure_guess
or domain.grid(),
solve_params=solve_params)
gradp = field.spatial_gradient(pressure,
type=type(velocity_field)) * accessible
return velocity_field - gradp, pressure, iterations, div, occupied_staggered
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()
