File size: 5,033 Bytes
24c2ab9 | 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """Position-Based Fluids in 3D (Macklin & Müller 2013).
Same skeleton as everything else in this project:
per-particle LOCAL PROJECTION (tied) -> GLOBAL RECONCILIATION
The local projection is the density constraint lambda — one rule shared
by every particle (the fluid's tied projector, roadmap §3). The global
step is the position-correction scatter + neighbor grid.
lambda_fn hook: the analytic lambda = -C / (sum|gradC|^2 + eps) can be
replaced by a learned rule (the fluid-token experiment, W7).
"""
import numpy as np
from scipy.spatial import cKDTree
class PBF3D:
def __init__(self, positions, bounds_min, bounds_max, h=0.1, rho0=1000.0,
dt=0.008, iters=4, gravity=(0.0, -9.81, 0.0),
scorr_k=1e-4, xsph_c=0.05, eps=100.0, lambda_fn=None):
self.x = np.asarray(positions, dtype=np.float64).copy()
self.v = np.zeros_like(self.x)
self.lo = np.asarray(bounds_min, dtype=np.float64)
self.hi = np.asarray(bounds_max, dtype=np.float64)
self.h, self.rho0, self.dt, self.iters = h, rho0, dt, iters
self.gravity = np.asarray(gravity, dtype=np.float64)
self.scorr_k, self.xsph_c, self.eps = scorr_k, xsph_c, eps
self.lambda_fn = lambda_fn
self.poly6_c = 315.0 / (64.0 * np.pi * h**9)
self.spiky_c = -45.0 / (np.pi * h**6)
self.w_dq = self.poly6_c * (h**2 - (0.3 * h)**2) ** 3 # scorr ref
# particle mass so a filled lattice sits at rest density: the most
# crowded particle of the initial configuration defines rho0
S = self._kernel_sums(self.x)
self.mass = rho0 / S.max()
self.rho = np.full(len(self.x), rho0)
def _kernel_sums(self, p):
tree = cKDTree(p)
pairs = tree.query_pairs(self.h, output_type="ndarray")
r2 = ((p[pairs[:, 0]] - p[pairs[:, 1]]) ** 2).sum(axis=1)
w = self.poly6_c * (self.h**2 - r2) ** 3
S = np.full(len(p), self.poly6_c * self.h**6) # self term W(0)
np.add.at(S, pairs[:, 0], w)
np.add.at(S, pairs[:, 1], w)
return S
def step(self, record=None):
"""One PBF step. record: optional list collecting (C, g2, lam)
rows from the final constraint iteration (corpus hook)."""
h, dt, rho0, m = self.h, self.dt, self.rho0, self.mass
v = self.v + dt * self.gravity
p = self.x + dt * v
np.clip(p, self.lo, self.hi, out=p)
tree = cKDTree(p)
pairs = tree.query_pairs(h, output_type="ndarray")
i, j = pairs[:, 0], pairs[:, 1]
N = len(p)
for it in range(self.iters):
d = p[i] - p[j]
r2 = (d * d).sum(axis=1)
r = np.sqrt(np.maximum(r2, 1e-12))
w = self.poly6_c * np.maximum(h**2 - r2, 0.0) ** 3
gw = (self.spiky_c * np.maximum(h - r, 0.0) ** 2 / r)[:, None] * d
gw *= m / rho0 # fold mass: gw is now a grad-C contribution
rho = np.full(N, m * self.poly6_c * h**6)
np.add.at(rho, i, m * w)
np.add.at(rho, j, m * w)
C = rho / rho0 - 1.0
sumg = np.zeros_like(p)
np.add.at(sumg, i, gw)
np.add.at(sumg, j, -gw)
g2 = (sumg * sumg).sum(axis=1)
gw2 = (gw * gw).sum(axis=1)
np.add.at(g2, i, gw2)
np.add.at(g2, j, gw2)
if self.lambda_fn is not None:
lam = self.lambda_fn(C, g2)
else:
lam = -C / (g2 + self.eps * (m / rho0) ** 2)
scorr = -self.scorr_k * (w * (m / rho0) / self.w_dq) ** 4
coef = (lam[i] + lam[j] + scorr)[:, None] * gw
dp = np.zeros_like(p)
np.add.at(dp, i, coef)
np.add.at(dp, j, -coef)
p += dp
np.clip(p, self.lo, self.hi, out=p)
if record is not None:
record.append(np.stack([C, g2, lam], axis=1))
self.v = (p - self.x) / dt
# XSPH viscosity on the final neighbor set
dv = self.v[j] - self.v[i]
wv = (self.poly6_c * np.maximum(
h**2 - ((p[i] - p[j]) ** 2).sum(axis=1), 0.0) ** 3
* m / rho0)[:, None] * dv
xs = np.zeros_like(p)
np.add.at(xs, i, wv)
np.add.at(xs, j, -wv)
self.v += self.xsph_c * xs
self.x = p
self.rho = rho
return C
def density_error(self):
"""Mean |rho/rho0 - 1| over particles at or above rest density
(free surfaces are legitimately under-dense in SPH kernels)."""
c = self.rho / self.rho0 - 1.0
return np.abs(c[c > -0.05]).mean()
def dam_break_block(lo, extent, spacing, jitter=1e-3, seed=0):
"""Lattice block of particles for dam-break initial conditions."""
rng = np.random.default_rng(seed)
axes = [np.arange(lo[k] + spacing / 2, lo[k] + extent[k], spacing)
for k in range(3)]
g = np.stack(np.meshgrid(*axes, indexing="ij"), axis=-1).reshape(-1, 3)
return g + rng.normal(scale=jitter, size=g.shape)
|