[physics] Fix divergence removal in pressure make_incompressible()
@@ -36,21 +36,12 @@ def make_incompressible(velocity: Grid,
hard_bcs = field.stagger(accessible, math.minimum, input_velocity.extrapolation, type=type(velocity))
velocity = apply_boundary_conditions(velocity, obstacles)
div = divergence(velocity) * active
- if input_velocity.extrapolation in (math.extrapolation.ZERO, math.extrapolation.PERIODIC):
- div = _balance_divergence(div, active)
+ if not input_velocity.extrapolation.connects_to_outside:
+ solve = copy_with(solve, preprocess_y=lambda y: _balance_divergence(y, 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)
-
- # if input_velocity.extrapolation in (math.extrapolation.ZERO, math.extrapolation.PERIODIC):
- # def pressure_backward(_p, _p_, dp: CenteredGrid):
- # # re-generate active mask because value might not be accessible from forward pass (e.g. Jax jit)
- # active = CenteredGrid(HardGeometryMask(~union(*[obstacle.geometry for obstacle in obstacles])), resolution=dp.resolution, bounds=dp.bounds, extrapolation=extrapolation.NONE)
- # return _balance_divergence(dp, active),
- # pressure = math.custom_gradient(lambda p: p, pressure_backward)(pressure)
-
grad_pressure = field.spatial_gradient(pressure, input_velocity.extrapolation, type=type(velocity)) * hard_bcs
velocity = velocity - grad_pressure
return velocity, pressure
