def divergence_free(velocity,
domain=None,
obstacles=(),
pressure_solver=None,
return_info=False,
gradient='implicit'):
"""
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
:param pressure_solver: PressureSolver. Uses default solver if none provided.
:return: divergence-free velocity as StaggeredGrid
"""
assert isinstance(velocity, StaggeredGrid)
# --- Set up FluidDomain ---
if domain is None:
domain = Domain(velocity.resolution, OPEN)
obstacle_mask = mask(union([obstacle.geometry for obstacle in obstacles]),
antialias=False)
if obstacle_mask is not None:
obstacle_grid = obstacle_mask.at(
velocity.center_points).copied_with(extrapolation='constant')
active_mask = 1 - obstacle_grid
else:
active_mask = math.ones(
domain.centered_shape(name='active', extrapolation='constant'))
accessible_mask = active_mask.copied_with(
extrapolation=Material.accessible_extrapolation_mode(
domain.boundaries))
fluiddomain = FluidDomain(domain,
active=active_mask,
accessible=accessible_mask)
# --- Boundary Conditions, Pressure Solve ---
velocity = fluiddomain.with_hard_boundary_conditions(velocity)
for obstacle in obstacles:
if not obstacle.is_stationary:
obs_mask = mask(obstacle.geometry, antialias=True)
angular_velocity = AngularVelocity(
location=obstacle.geometry.center,
strength=obstacle.angular_velocity,
falloff=None)
velocity = ((1 - obs_mask) * velocity + obs_mask *
(angular_velocity + obstacle.velocity)).at(velocity)
divergence_field = velocity.divergence(physical_units=False)
pressure, iterations = poisson_solve(divergence_field,
fluiddomain,
solver=pressure_solver,
gradient=gradient)
pressure *= velocity.dx[0]
gradp = StaggeredGrid.gradient(pressure)
velocity -= fluiddomain.with_hard_boundary_conditions(gradp)
return velocity if not return_info else (velocity, {
'pressure': pressure,
'iterations': iterations,
'divergence': divergence_field
})
def effect_applied(effect, field, dt):
effect_field = effect.field.at(field)
if effect._mode == GROW:
return field + math.mul(effect_field, dt)
elif effect._mode == ADD:
return field + effect_field
elif effect._mode == FIX:
assert effect.bounds is not None
inside = mask([effect.bounds]).at(field)
return math.where(inside, effect_field, field)
else:
raise ValueError('Invalid mode: %s' % effect.mode)
effect_field = effect.field.at(field)
if effect._mode == GROW:
return field + math.mul(effect_field, dt)
elif effect._mode == ADD:
return field + effect_field
elif effect._mode == FIX:
assert effect.bounds is not None
inside = mask([effect.bounds]).at(field)
return math.where(inside, effect_field, field)
else:
raise ValueError('Invalid mode: %s' % effect.mode)
# pylint: disable-msg = invalid-name
Inflow = lambda geometry, rate=1.0, target='density': FieldEffect(
mask(geometry, antialias=True) * rate,
target,
GROW,
tags=('inflow', 'effect'))
Accelerator = lambda geometry, acceleration: FieldEffect(
mask(geometry, antialias=True) * acceleration, ('velocity', ),
GROW,
tags=('fan', 'effect'))
ConstantVelocity = lambda geometry, velocity: FieldEffect(
ConstantField(velocity),
bounds=geometry,
targets=('velocity', ),
mode=FIX,
tags=('effect', ))
HeatSource = lambda geometry, rate, name=None: FieldEffect(
mask(geometry, antialias=True) * rate, ('temperature', ), GROW, name=name)
def ColdSource(geometry, rate, name=None):
return FieldEffect(mask(geometry, antialias=True) * -rate,
('temperature', ),
GROW,
name=name)
def Accelerator(geometry, acceleration):
return FieldEffect(mask(geometry, antialias=True) * acceleration,
('velocity', ),
GROW,
tags=('fan', 'effect'))
def Inflow(geometry, rate=1.0, target='density'):
return FieldEffect(mask(geometry, antialias=True) * rate,
target,
GROW,
tags=('inflow', 'effect'))