Add Drift physics for passively advected Fields
@@ -5,6 +5,7 @@ from numbers import Number
import numpy as np
from phi import math, struct
+from phi.physics.field import Field
from .domain import Domain, DomainState
from .field import CenteredGrid, StaggeredGrid, advect, union_mask
@@ -20,7 +21,7 @@ class Fluid(DomainState):
A Fluid state consists of a density field (centered grid) and a velocity field (staggered grid).
"""
- def __init__(self, domain, density=0.0, velocity=0.0, buoyancy_factor=0.0, tags=('fluid', 'velocityfield'), name='fluid', **kwargs):
+ def __init__(self, domain, density=0.0, velocity=0.0, buoyancy_factor=0.0, tags=('fluid', 'velocityfield', 'velocity'), name='fluid', **kwargs):
DomainState.__init__(self, **struct.kwargs(locals()))
self.solve_info = {}
@@ -87,7 +88,7 @@ Supports obstacles, density effects, velocity effects, global gravity.
density = effect_applied(effect, density, dt)
for effect in velocity_effects:
velocity = effect_applied(effect, velocity, dt)
- velocity += buoyancy(fluid.density, gravity, fluid.buoyancy_factor).at(velocity) * dt
+ velocity += buoyancy(density, gravity, fluid.buoyancy_factor).at(velocity) * dt
# --- Pressure solve ---
if self.make_output_divfree:
velocity, fluid.solve_info = divergence_free(velocity, fluid.domain, obstacles, pressure_solver=self.pressure_solver, return_info=True)
@@ -97,6 +98,28 @@ Supports obstacles, density effects, velocity effects, global gravity.
INCOMPRESSIBLE_FLOW = IncompressibleFlow()
+class Drift(Physics):
+ """
+Passive advection with external velocity field.
+This Physics requires the world to contain a single velocity field or velocity-carrying state such as Fluid.
+
+This Physics can be applied to all built-in Fields.
+The fields will then be advected with the velocity field each time step.
+ """
+
+ def __init__(self, use_updated_velocity=False, conserve=True):
+ Physics.__init__(self, dependencies=[StateDependency('velocity', 'velocity', single_state=True, blocking=use_updated_velocity)])
+ self.conserve = conserve
+
+ def step(self, field, dt=1.0, velocity=None):
+ if not isinstance(velocity, Field):
+ velocity = velocity.velocity
+ advected = advect.advect(field, velocity, dt=dt).copied_with(age=field.age + dt)
+ if self.conserve and isinstance(field, (CenteredGrid, StaggeredGrid)) and np.all(~np.char.equal(struct.flatten(field.extrapolation), 'constant')): # If field has zero extrapolation, it cannot be conserved
+ advected = advected.normalized(field)
+ return advected
+
+
def buoyancy(density, gravity, buoyancy_factor):
"""
Computes the buoyancy force proportional to the density.
@@ -107,8 +130,7 @@ Computes the buoyancy force proportional to the density.
"""
if isinstance(gravity, (int, float)):
gravity = math.to_float(math.as_tensor([gravity] + ([0] * (density.rank - 1))))
- result = StaggeredGrid.from_scalar(density, -gravity * buoyancy_factor)
- return result
+ return density * (-gravity * buoyancy_factor)
def _is_div_free(velocity, is_div_free):
