def laplace(p):
grad = gradient(p, type(velocity))
grad *= hard_bcs
grad = grad.with_(
extrapolation=domain.boundaries['near_vector_extrapolation'])
div = divergence(grad)
lap = where(active, div, p)
return lap
def step_gradient_2d(params, plasma, phi, dt=0):
"""time gradient of model"""
# Diffusion function
def diffuse(arr, N, dx):
for i in range(N):
arr = field.laplace(arr) # math.fourier_laplace(arr, dx)
return arr
# Calculate Gradients
dx_p, dy_p = field.gradient(phi).unstack('gradient')
# Get difference
diff = (phi - plasma.density)
# Step 2.1: New Omega.
nu = (-1)**(params['N'] + 1) * params['nu']
o = (params['c1'] * diff)
if params['arakawa_coeff']:
o += -params[
'arakawa_coeff'] * math._nd._periodic_2d_arakawa_poisson_bracket(
phi.values, plasma.omega.values, plasma.dx)
if nu and params['N']:
o += nu * diffuse(plasma.omega, params['N'], plasma.dx)
# Step 2.2: New Density.
n = (params['c1'] * diff)
if params['arakawa_coeff']:
n += -params[
'arakawa_coeff'] * math._nd._periodic_2d_arakawa_poisson_bracket(
phi.values, plasma.density.values, plasma.dx)
if params['kappa_coeff']:
n += -params['kappa_coeff'] * dy_p
if nu:
n += nu * diffuse(plasma.density, params['N'], plasma.dx)
return Namespace(
density=n,
omega=o,
phi=phi, # NOTE: NOT A GRADIENT
age=plasma.age + dt,
dx=plasma.dx)
def make_incompressible(velocity: Grid,
domain: Domain,
obstacles: tuple or list = (),
solve_params: math.LinearSolve = math.LinearSolve(
None, 1e-3),
pressure_guess: CenteredGrid = None):
"""
Projects the given velocity field by solving for the pressure and subtracting its gradient.
This method is similar to :func:`field.divergence_free()` but differs in how the boundary conditions are specified.
Args:
velocity: Vector field sampled on a grid
domain: Used to specify boundary conditions
obstacles: List of Obstacles to specify boundary conditions inside the domain (Default value = ())
pressure_guess: Initial guess for the pressure solve
solve_params: Parameters for the pressure solve
Returns:
velocity: divergence-free velocity of type `type(velocity)`
pressure: solved pressure field, `CenteredGrid`
iterations: Number of iterations required to solve for the pressure
divergence: divergence field of input velocity, `CenteredGrid`
"""
input_velocity = velocity
active = domain.grid(
HardGeometryMask(~union(*[obstacle.geometry
for obstacle in obstacles])),
extrapolation=domain.boundaries['active_extrapolation'])
accessible = domain.grid(
active, extrapolation=domain.boundaries['accessible_extrapolation'])
hard_bcs = field.stagger(accessible,
math.minimum,
domain.boundaries['accessible_extrapolation'],
type=type(velocity))
velocity = layer_obstacle_velocities(velocity * hard_bcs, obstacles).with_(
extrapolation=domain.boundaries['near_vector_extrapolation'])
div = divergence(velocity)
if domain.boundaries[
'near_vector_extrapolation'] == math.extrapolation.BOUNDARY:
div -= field.mean(div)
# Solve pressure
def laplace(p):
grad = gradient(p, type(velocity))
grad *= hard_bcs
grad = grad.with_(
extrapolation=domain.boundaries['near_vector_extrapolation'])
div = divergence(grad)
lap = where(active, div, p)
return lap
pressure_guess = pressure_guess if pressure_guess is not None else domain.scalar_grid(
0)
converged, pressure, iterations = field.solve(laplace,
y=div,
x0=pressure_guess,
solve_params=solve_params,
constants=[active, hard_bcs])
if math.all_available(converged) and not math.all(converged):
raise AssertionError(
f"pressure solve did not converge after {iterations} iterations\nResult: {pressure.values}"
)
# Subtract grad pressure
gradp = field.gradient(pressure, type=type(velocity)) * hard_bcs
velocity = (velocity -
gradp).with_(extrapolation=input_velocity.extrapolation)
return velocity, pressure, iterations, div
def test_gradient(self):
domain = Domain(x=4, y=3)
phi = domain.grid() * (1, 2)
grad = field.gradient(phi, stack_dim='gradient')
self.assertEqual(('spatial', 'spatial', 'channel', 'channel'),
grad.shape.types)
def make_incompressible(
velocity: StaggeredGrid,
domain: Domain,
obstacles: tuple or list or StaggeredGrid = (),
particles: PointCloud or None = None,
solve_params: math.LinearSolve = math.LinearSolve(),
pressure_guess: CenteredGrid = None
) -> Tuple[StaggeredGrid, CenteredGrid, math.Tensor, CenteredGrid,
StaggeredGrid]:
"""
Projects the given velocity field by solving for the pressure and subtracting its gradient.
Args:
velocity: Current velocity field as StaggeredGrid
obstacles: Sequence of `phi.physics.Obstacle` objects or binary StaggeredGrid marking through-flow cell faces
particles (Optional if occupation masks are provided): Pointcloud holding the current positions of the particles
domain (Optional if occupation masks are provided): Domain object
pressure_guess (Optional): Initial pressure guess as CenteredGrid
solve_params: Parameters for the pressure solve
Returns:
velocity: divergence-free velocity of type `type(velocity)`
pressure: solved pressure field, `CenteredGrid`
iterations: Number of iterations required to solve for the pressure
divergence: divergence field of input velocity, `CenteredGrid`
occupation_mask: StaggeredGrid
"""
points = particles.with_(values=math.wrap(1))
occupied_centered = points >> domain.grid()
occupied_staggered = points >> domain.staggered_grid()
if isinstance(obstacles, StaggeredGrid):
accessible = obstacles
else:
accessible = domain.accessible_mask(
union(*[obstacle.geometry for obstacle in obstacles]),
type=StaggeredGrid)
# --- Extrapolation is needed to exclude border divergence from the `occupied_centered` mask and thus
# from the pressure solve. If particles are randomly distributed, the `occupied_centered` mask
# could sometimes include the divergence at the borders (due to single particles right at the edge
# which temporarily deform the `occupied_centered` mask when moving into a new cell) which would then
# get compensated by the pressure. This is unwanted for falling liquids and therefore prevented by this
# extrapolation. ---
velocity_field, _ = extrapolate_valid(velocity * occupied_staggered,
occupied_staggered, 1)
velocity_field *= accessible # Enforces boundary conditions after extrapolation
div = field.divergence(
velocity_field
) * occupied_centered # Multiplication with `occupied_centered` excludes border divergence from pressure solve
def matrix_eq(p):
return field.where(
occupied_centered,
field.divergence(
field.gradient(p, type=StaggeredGrid) * accessible), p)
converged, pressure, iterations = field.solve(matrix_eq,
div,
pressure_guess
or domain.grid(),
solve_params=solve_params)
gradp = field.gradient(pressure, type=type(velocity_field)) * accessible
return velocity_field - gradp, pressure, iterations, div, occupied_staggered
def matrix_eq(p):
return field.where(
occupied_centered,
field.divergence(
field.gradient(p, type=StaggeredGrid) * accessible), p)