SmokePhysics now takes single Gravity
Adds default gravity argument for convenience.
@@ -31,16 +31,15 @@ class SmokePhysics(Physics):
def __init__(self, pressure_solver=None, make_input_divfree=False, make_output_divfree=True):
Physics.__init__(self, [StateDependency('obstacles', 'obstacle'),
- StateDependency('gravity', 'gravity'),
+ StateDependency('gravity', 'gravity', single_state=True),
StateDependency('density_effects', 'density_effect', blocking=True),
StateDependency('velocity_effects', 'velocity_effect', blocking=True)])
self.pressure_solver = pressure_solver
self.make_input_divfree = make_input_divfree
self.make_output_divfree = make_output_divfree
- def step(self, smoke, dt=1.0, obstacles=(), gravity=(), density_effects=(), velocity_effects=(), **dependent_states):
- assert len(dependent_states) == 0
- gravity = gravity_tensor(sum(gravity), smoke.rank)
+ def step(self, smoke, dt=1.0, obstacles=(), gravity=Gravity(), density_effects=(), velocity_effects=()):
+ gravity = gravity_tensor(gravity, smoke.rank)
velocity = smoke.velocity
density = smoke.density
obstacle_mask = union_mask([obstacle.geometry for obstacle in obstacles])
@@ -49,13 +48,13 @@ class SmokePhysics(Physics):
# --- Advection ---
density = advect.semi_lagrangian(density, velocity, dt=dt)
velocity = advect.semi_lagrangian(velocity, velocity, dt=dt)
- # --- Density effects ---
+ # --- Effects ---
for effect in density_effects:
density = effect_applied(effect, density, dt)
- # --- velocity effects
for effect in velocity_effects:
velocity = effect_applied(effect, velocity, dt)
velocity += buoyancy(smoke.density, gravity, smoke.buoyancy_factor) * dt
+ # --- Pressure solve ---
if self.make_output_divfree:
velocity = divergence_free(velocity, smoke.domain, obstacle_mask, pressure_solver=self.pressure_solver)
return smoke.copied_with(density=density, velocity=velocity, age=smoke.age + dt)
