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 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 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 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_solve_diverge(self):
y = math.ones(spatial(x=2)) * (1, 2)
x0 = math.zeros(spatial(x=2))
for method in ['CG']:
solve = Solve(method, 0, 1e-3, x0=x0, max_iterations=100)
try:
field.solve_linear(math.jit_compile_linear(math.laplace), y,
solve)
assert False
except Diverged:
pass
with math.SolveTape(record_trajectories=True) as solves:
try:
field.solve_linear(math.jit_compile_linear(math.laplace),
y, solve) # impossible
assert False
except Diverged:
pass
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_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 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 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)