File size: 2,917 Bytes
4450790 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
import numpy as np
from scipy.ndimage import map_coordinates, spline_filter
from scipy.sparse.linalg import factorized
from .numerical import difference, operator
class Fluid:
def __init__(self, shape, *quantities, pressure_order=1, advect_order=3):
self.shape = shape
self.dimensions = len(shape)
# Prototyping is simplified by dynamically
# creating advected quantities as needed.
self.quantities = quantities
for q in quantities:
setattr(self, q, np.zeros(shape))
self.indices = np.indices(shape)
self.velocity = np.zeros((self.dimensions, *shape))
laplacian = operator(shape, difference(2, pressure_order))
self.pressure_solver = factorized(laplacian)
self.advect_order = advect_order
def step(self):
# Advection is computed backwards in time as described in Stable Fluids.
advection_map = self.indices - self.velocity
# SciPy's spline filter introduces checkerboard divergence.
# A linear blend of the filtered and unfiltered fields based
# on some value epsilon eliminates this error.
def advect(field, filter_epsilon=10e-2, mode='constant'):
filtered = spline_filter(field, order=self.advect_order, mode=mode)
field = filtered * (1 - filter_epsilon) + field * filter_epsilon
return map_coordinates(field, advection_map, prefilter=False, order=self.advect_order, mode=mode)
# Apply advection to each axis of the
# velocity field and each user-defined quantity.
for d in range(self.dimensions):
self.velocity[d] = advect(self.velocity[d])
for q in self.quantities:
setattr(self, q, advect(getattr(self, q)))
# Compute the jacobian at each point in the
# velocity field to extract curl and divergence.
jacobian_shape = (self.dimensions,) * 2
partials = tuple(np.gradient(d) for d in self.velocity)
jacobian = np.stack(partials).reshape(*jacobian_shape, *self.shape)
divergence = jacobian.trace()
# If this curl calculation is extended to 3D, the y-axis value must be negated.
# This corresponds to the coefficients of the levi-civita symbol in that dimension.
# Higher dimensions do not have a vector -> scalar, or vector -> vector,
# correspondence between velocity and curl due to differing isomorphisms
# between exterior powers in dimensions != 2 or 3 respectively.
curl_mask = np.triu(np.ones(jacobian_shape, dtype=bool), k=1)
curl = (jacobian[curl_mask] - jacobian[curl_mask.T]).squeeze()
# Apply the pressure correction to the fluid's velocity field.
pressure = self.pressure_solver(divergence.flatten()).reshape(self.shape)
self.velocity -= np.gradient(pressure)
return divergence, curl, pressure |