Make fluid pressure accessible
Fluid now has a non-struct dict solve_info.
Physics objects may update the dict during step()
IncompressibleFlow stores 'iterations' and 'pressure'.
@@ -23,6 +23,7 @@ class Fluid(DomainState):
def __init__(self, domain, density=0.0, velocity=0.0, buoyancy_factor=0.0, tags=('fluid', 'velocityfield'), name='fluid', **kwargs):
DomainState.__init__(self, **struct.kwargs(locals()))
+ self.solve_info = {}
def default_physics(self): return INCOMPRESSIBLE_FLOW
@@ -76,7 +77,7 @@ Supports obstacles, density effects, velocity effects, global gravity.
velocity = fluid.velocity
density = fluid.density
if self.make_input_divfree:
- velocity = divergence_free(velocity, fluid.domain, obstacles, pressure_solver=self.pressure_solver)
+ velocity, fluid.solve_info = divergence_free(velocity, fluid.domain, obstacles, pressure_solver=self.pressure_solver, return_info=True)
# --- Advection ---
density = advect.semi_lagrangian(density, velocity, dt=dt)
velocity = advect.semi_lagrangian(velocity, velocity, dt=dt)
@@ -90,7 +91,7 @@ Supports obstacles, density effects, velocity effects, global gravity.
velocity += buoyancy(fluid.density, gravity, fluid.buoyancy_factor) * dt
# --- Pressure solve ---
if self.make_output_divfree:
- velocity = divergence_free(velocity, fluid.domain, obstacles, pressure_solver=self.pressure_solver)
+ velocity, fluid.solve_info = divergence_free(velocity, fluid.domain, obstacles, pressure_solver=self.pressure_solver, return_info=True)
return fluid.copied_with(density=density, velocity=velocity, age=fluid.age + dt)
@@ -138,9 +139,10 @@ def solve_pressure(divergence, fluiddomain, pressure_solver=None):
return pressure, iteration
-def divergence_free(velocity, domain=None, obstacles=(), pressure_solver=None):
+def divergence_free(velocity, domain=None, obstacles=(), pressure_solver=None, return_info=False):
"""
Projects the given velocity field by solving for and subtracting the pressure.
+ :param return_info: if True, returns a dict holding information about the solve as a second object
:param velocity: StaggeredGrid
:param domain: Domain matching the velocity field, used for boundary conditions
:param obstacles: list of Obstacles
@@ -162,8 +164,8 @@ Projects the given velocity field by solving for and subtracting the pressure.
# --- Boundary Conditions, Pressure Solve ---
velocity = fluiddomain.with_hard_boundary_conditions(velocity)
divergence_field = velocity.divergence(physical_units=False)
- pressure, _ = solve_pressure(divergence_field, fluiddomain, pressure_solver=pressure_solver)
+ pressure, iterations = solve_pressure(divergence_field, fluiddomain, pressure_solver=pressure_solver)
pressure *= velocity.dx[0]
gradp = StaggeredGrid.gradient(pressure)
velocity -= fluiddomain.with_hard_boundary_conditions(gradp)
- return velocity
+ return velocity if not return_info else (velocity, {'pressure': pressure, 'iterations': iterations})
