def make_incompressible(velocity: StaggeredGrid,
domain: Domain,
particles: PointCloud,
obstacles: tuple or list or StaggeredGrid = (),
solve=math.Solve('auto', 1e-5, 0, gradient_solve=math.Solve('auto', 1e-5, 1e-5))):
"""
Projects the given velocity field by solving for the pressure and subtracting its spatial_gradient.
Args:
velocity: Current velocity field as StaggeredGrid
domain: Domain object
particles: `PointCloud` holding the current positions of the particles
obstacles: Sequence of `phi.physics.Obstacle` objects or binary StaggeredGrid marking through-flow cell faces
solve: Parameters for the pressure solve_linear
Returns:
velocity: divergence-free velocity of type `type(velocity)`
pressure: solved pressure field, `CenteredGrid`
iterations: Number of iterations required to solve_linear for the pressure
divergence: divergence field of input velocity, `CenteredGrid`
occupation_mask: StaggeredGrid
"""
points = particles.with_values(math.tensor(1., convert=True))
occupied_centered = points @ domain.scalar_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_linear. 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). This 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_linear
@field.jit_compile_linear
def matrix_eq(p):
return field.where(occupied_centered, field.divergence(field.spatial_gradient(p, type=StaggeredGrid) * accessible), p)
if solve.x0 is None:
solve = copy_with(solve, x0=domain.scalar_grid())
pressure = field.solve_linear(matrix_eq, div, solve)
def pressure_backward(_p, _p_, dp):
return dp * occupied_centered.values,
add_mask_in_gradient = math.custom_gradient(lambda p: p, pressure_backward)
pressure = pressure.with_values(add_mask_in_gradient(pressure.values))
gradp = field.spatial_gradient(pressure, type=type(velocity_field)) * accessible
return velocity_field - gradp, pressure, occupied_staggered
def test_linear_solve_matrix_tape(self):
y = CenteredGrid(1, extrapolation.ZERO, x=3) * (1, 2)
x0 = CenteredGrid(0, extrapolation.ZERO, x=3)
for method in ['CG', 'CG-adaptive', 'auto']:
solve = math.Solve(method, 0, 1e-3, x0=x0, max_iterations=100)
with math.SolveTape() as solves:
x = field.solve_linear(math.jit_compile_linear(field.laplace),
y, solve)
math.assert_close(x.values, [[-1.5, -2, -1.5], [-3, -4, -3]],
abs_tolerance=1e-3)
assert len(solves) == 1
assert solves[0] == solves[solve]
math.assert_close(solves[solve].residual.values,
0,
abs_tolerance=1e-3)
assert math.close(solves[solve].iterations, 2) or math.close(
solves[solve].iterations, -1)
with math.SolveTape(record_trajectories=True) as solves:
x = field.solve_linear(math.jit_compile_linear(field.laplace),
y, solve)
math.assert_close(x.values, [[-1.5, -2, -1.5], [-3, -4, -3]],
abs_tolerance=1e-3)
assert solves[solve].x.trajectory.size == 3
math.assert_close(solves[solve].residual.trajectory[-1].values,
0,
abs_tolerance=1e-3)
def solve(y, method):
print(f"Tracing {method} with {backend}...")
solve = math.Solve(method, 0, 1e-3, x0=x0, max_iterations=100)
with SolveTape() as solves:
x = field.solve_linear(math.jit_compile_linear(field.laplace),
y, solve)
return x
def test_linear_solve_matrix_batched(
self): # TODO also test batched matrix
y = CenteredGrid(1, extrapolation.ZERO, x=3) * (1, 2)
x0 = CenteredGrid(0, extrapolation.ZERO, x=3)
for method in ['CG', 'CG-adaptive', 'auto']:
solve = math.Solve(method, 0, 1e-3, x0=x0, max_iterations=100)
x = field.solve_linear(math.jit_compile_linear(field.laplace), y,
solve)
math.assert_close(x.values, [[-1.5, -2, -1.5], [-3, -4, -3]],
abs_tolerance=1e-3)
def test_solve_linear_matrix_dirichlet(self):
for backend in BACKENDS:
with backend:
y = CenteredGrid(1, extrapolation.ONE, x=3)
x0 = CenteredGrid(0, extrapolation.ONE, x=3)
solve = math.Solve('CG', 0, 1e-3, x0=x0, max_iterations=100)
x_ref = field.solve_linear(field.laplace, y, solve)
x_jit = field.solve_linear(
math.jit_compile_linear(field.laplace), y, solve)
math.assert_close(x_ref.values,
x_jit.values, [-0.5, -1, -0.5],
abs_tolerance=1e-3,
msg=backend)
def test_solve_linear_matrix(self):
for backend in BACKENDS:
with backend:
y = CenteredGrid(1, extrapolation.ZERO, x=3)
x0 = CenteredGrid(0, extrapolation.ZERO, x=3)
for method in ['CG', 'CG-adaptive', 'auto']:
solve = math.Solve(method,
0,
1e-3,
x0=x0,
max_iterations=100)
x = field.solve_linear(
math.jit_compile_linear(field.laplace), y, solve)
math.assert_close(x.values, [-1.5, -2, -1.5],
abs_tolerance=1e-3,
msg=backend)
def test_solve_linear_function_batched(self):
y = CenteredGrid(1, extrapolation.ZERO, x=3) * (1, 2)
x0 = CenteredGrid(0, extrapolation.ZERO, x=3)
for method in ['CG', 'CG-adaptive', 'auto']:
solve = math.Solve(method, 0, 1e-3, x0=x0, max_iterations=100)
x = field.solve_linear(math.jit_compile_linear(field.laplace), y,
solve)
math.assert_close(x.values,
math.wrap([[-1.5, -2, -1.5], [-3, -4, -3]],
channel('vector'), spatial('x')),
abs_tolerance=1e-3)
with math.SolveTape() as solves:
x = field.solve_linear(math.jit_compile_linear(field.laplace),
y, solve)
math.assert_close(x.values,
math.wrap([[-1.5, -2, -1.5], [-3, -4, -3]],
channel('vector'), spatial('x')),
abs_tolerance=1e-3)
assert len(solves) == 1
assert solves[0] == solves[solve]
math.assert_close(solves[solve].residual.values,
0,
abs_tolerance=1e-3)
def solve(y, method):
solve = math.Solve(method, 0, 1e-3, x0=x0, max_iterations=100)
return field.solve_linear(math.jit_compile_linear(field.laplace),
y, solve)
def test_minimize(self):
def loss(x, y):
return math.l2_loss(x - 1) + math.l2_loss(y + 1)
for backend in BACKENDS:
if backend.supports(Backend.functional_gradient):
with backend:
x0 = tensor([0, 0, 0], spatial('x')), tensor([-1, -1, -1],
spatial('y'))
x, y = math.minimize(
loss, math.Solve('L-BFGS-B', 0, 1e-3, x0=x0))
math.assert_close(x,
1,
abs_tolerance=1e-3,
msg=backend.name)
math.assert_close(y,
-1,
abs_tolerance=1e-3,
msg=backend.name)
x0 = tensor([[0, 0, 0], [1, 1, 1]], batch('batch'),
spatial('x')), tensor(
[[0, 0, 0], [-1, -1, -1]], batch('batch'),
spatial('y'))
x, y = math.minimize(
loss, math.Solve('L-BFGS-B', 0, 1e-3, x0=x0))
math.assert_close(x,
1,
abs_tolerance=1e-3,
msg=backend.name)
math.assert_close(y,
-1,
abs_tolerance=1e-3,
msg=backend.name)
with math.SolveTape() as solves:
x, y = math.minimize(
loss, math.Solve('L-BFGS-B', 0, 1e-3, x0=x0))
math.assert_close(x,
1,
abs_tolerance=1e-3,
msg=backend.name)
math.assert_close(y,
-1,
abs_tolerance=1e-3,
msg=backend.name)
math.assert_close(solves[0].residual,
0,
abs_tolerance=1e-4)
assert (solves[0].iterations <= (4, 0)).all
assert (solves[0].function_evaluations <= (30, 1)).all
with math.SolveTape(
record_trajectories=True) as trajectories:
x, y = math.minimize(
loss, math.Solve('L-BFGS-B', 0, 1e-3, x0=x0))
math.assert_close(x,
1,
abs_tolerance=1e-3,
msg=backend.name)
math.assert_close(y,
-1,
abs_tolerance=1e-3,
msg=backend.name)
math.assert_close(trajectories[0].residual.trajectory[-1],
0,
abs_tolerance=1e-4)
assert (
trajectories[0].iterations == solves[0].iterations).all
assert trajectories[
0].residual.trajectory.size == trajectories[0].x[
0].trajectory.size
assert trajectories[0].residual.trajectory.size > 1
def make_incompressible(
velocity: GridType,
obstacles: tuple or list = (),
solve=math.Solve('auto',
1e-5,
1e-5,
gradient_solve=math.Solve('auto', 1e-5, 1e-5))
) -> Tuple[GridType, CenteredGrid]:
"""
Projects the given velocity field by solving for the pressure and subtracting its spatial_gradient.
This method is similar to :func:`field.divergence_free()` but differs in how the boundary conditions are specified.
Args:
velocity: Vector field sampled on a grid
obstacles: List of Obstacles to specify boundary conditions inside the domain (Default value = ())
solve: Parameters for the pressure solve as.
Returns:
velocity: divergence-free velocity of type `type(velocity)`
pressure: solved pressure field, `CenteredGrid`
"""
assert isinstance(
obstacles,
(tuple,
list)), f"obstacles must be a tuple or list but got {type(obstacles)}"
input_velocity = velocity
accessible_extrapolation = _accessible_extrapolation(
input_velocity.extrapolation)
active = CenteredGrid(
HardGeometryMask(~union(*[obstacle.geometry
for obstacle in obstacles])),
resolution=velocity.resolution,
bounds=velocity.bounds,
extrapolation=extrapolation.NONE)
accessible = active.with_extrapolation(accessible_extrapolation)
hard_bcs = field.stagger(accessible,
math.minimum,
input_velocity.extrapolation,
type=type(velocity))
velocity = apply_boundary_conditions(velocity, obstacles)
div = divergence(velocity) * active
if not input_velocity.extrapolation.connects_to_outside:
assert solve.preprocess_y is None, "fluid.make_incompressible() does not support custom preprocessing"
solve = copy_with(solve,
preprocess_y=_balance_divergence,
preprocess_y_args=(active, ))
if solve.x0 is None:
pressure_extrapolation = _pressure_extrapolation(
input_velocity.extrapolation)
solve = copy_with(solve,
x0=CenteredGrid(
0,
resolution=div.resolution,
bounds=div.bounds,
extrapolation=pressure_extrapolation))
pressure = math.solve_linear(masked_laplace,
f_args=[hard_bcs, active],
y=div,
solve=solve)
grad_pressure = field.spatial_gradient(
pressure, input_velocity.extrapolation, type=type(velocity)) * hard_bcs
velocity = velocity - grad_pressure
return velocity, pressure