[physics] Update SPH kernels

phi/physics/sph.py

@@ -1,14 +1,15 @@
-from typing import Dict, Tuple, Any, Union
+from typing import Dict, Tuple, Any, Union, Sequence
 
 from phi import math
 from phi.field import Field
 from phi.math import Tensor, pairwise_distances, vec_length, Shape, non_channel, dual, where, PI
 from phi.geom import Geometry, Graph
-from phiml.math import channel, stack
+from phiml.math import channel, stack, vec, concat
 
 _DEFAULT_DESIRED_NEIGHBORS = {
     'quintic-spline': 34,
     'wendland-c2': 22,
+    'poly6': 30,
 }
 
 
@@ -20,7 +21,7 @@ def neighbor_graph(nodes: Geometry,
                    format='csr',
                    search_method='auto') -> Graph:
     """
-    Build a `phi.geom.Graph` based on proximity of `nodes`.
+    Build a `phi.geom.Graph` based on proximity of `nodes` and evaluates the kernel function.
 
     Args:
         nodes: Particles including obstacle particles as `Geometry` collection.
@@ -38,21 +39,21 @@ def neighbor_graph(nodes: Geometry,
     boundary = {} if boundary is None else boundary
     desired_neighbors = _DEFAULT_DESIRED_NEIGHBORS[kernel] if desired_neighbors is None else desired_neighbors
     # --- neighbor search ---
-    support = _optimal_support_radius(nodes.volume, desired_neighbors, nodes.spatial_rank)
+    support = _get_support_radius(nodes.volume, desired_neighbors, nodes.spatial_rank)
     # --- evaluate kernel and derivatives ---
     deltas = math.pairwise_differences(nodes.center, max_distance=support, format=format, method=search_method)
     distances = math.vec_length(deltas, eps=1e-5)
-    q = distances / support
-    edges = {'kernel': evaluate_kernel(q, nodes.spatial_rank, kernel) / support ** nodes.spatial_rank}
+    kernel = evaluate_kernel(deltas, distances, support, nodes.spatial_rank, kernel, derivative=(0, 1) if kernel_derivative else (0,))
     if kernel_derivative:
-        kernel_derivatives = evaluate_kernel(q, nodes.spatial_rank, kernel, derivative=1) * support ** (nodes.spatial_rank + 1) * deltas
-        for dim, val in dict(**kernel_derivatives.vector).items():
-            edges[dim] = val
-    edges = stack(edges, channel('vector'))
+        val, grad = kernel
+        val = vec(kernel=val)
+        edges = concat([val, grad], 'vector')
+    else:
+        edges = vec(kernel=kernel[0])
     return Graph(nodes, edges, boundary, deltas=deltas, distances=distances, bounding_distance=support)
 
 
-def _optimal_support_radius(volume: Tensor, desired_neighbors: float, spatial_rank: int) -> Tensor:  # volumeToSupport
+def _get_support_radius(volume: Tensor, desired_neighbors: float, spatial_rank: int) -> Tensor:  # volumeToSupport
     """
     Calculates the optimal kernel support radius so that on average `desired_neighbors` neighbors lie within reach of each particle.
 
@@ -72,50 +73,95 @@ def _optimal_support_radius(volume: Tensor, desired_neighbors: float, spatial_ra
         return (desired_neighbors / math.PI * .75 * volume) ** (1/3)  # N(h) = 4/3 𝜋 h³ / v
 
 
-def evaluate_kernel(q: Tensor, spatial_rank: int, kernel='wendland-c2', derivative=0):
+def expected_neighbors(volume: Tensor, support_radius, spatial_rank: int):
+    """
+    Given the average element volume and support radius, returns the average number of neighbors for a region filled with particles.
+
+    Args:
+        volume: Average particle volume.
+        support_radius: Other elements are considered neighbors if their center lies within a sphere of this radius around a particle.
+        spatial_rank: Spatial rank of the simulation.
+
+    Returns:
+        Number of expected neighbors.
+    """
+    if spatial_rank == 1:
+        return 2 * support_radius / volume
+    elif spatial_rank == 2:
+        return PI * support_radius**2 / volume
+    else:
+        return 4/3 * PI * support_radius**3 / volume
+
+
+
+def evaluate_kernel(delta, distance, h, spatial_rank: int, kernel: str, derivative=(0,)) -> Sequence[Tensor]:
     """
-    Compute the SPH kernel value at a normalized scalar distance `q` or a derivative of the kernel function.
+    Compute the SPH kernel value or a derivative of the kernel function.
+
+    For kernels that only depends on the squared distance, such as `poly6`, the `distance` variable is not used.
+    Instead, the squared distance is derived from `delta`.
 
     Args:
-        q: Normalized distance `phi.math.Tensor`. All values must lie between 0 and 1.
+        delta: Vectors to neighbors, i.e. position differences.
+        distance: Scalar distance to neighbors.
+        h: Support radius / smoothing length / maximum distance / cutoff.
         spatial_rank: Dimensionality of the simulation.
-        kernel: Which kernel to use, one of `'wendland-c2'`, `'quintic-spline'`.
-        derivative: Derivative order, `0` for kernel function, `1` for gradient.
+        kernel: Which kernel to use, one of `'quintic-spline'`, `'wendland-c2'`, `'poly6'`.
+        derivative: Ordered `tuple` of derivatives to compute, `0` for kernel function, `1` for gradient.
 
     Returns:
         `phi.math.Tensor`
     """
+    # SPH kernels must be divided by h^d for the kernel and h^(d+1) for the gradient
+    assert all(d <= 1 for d in derivative), f"Only derivative=0 and 1 are supported but got {derivative}"
+    assert tuple(sorted(derivative)) == derivative, f"Derivatives must be ordered but got {derivative}"
+    d = spatial_rank
+    result = []
     if kernel == 'quintic-spline':  # cutoff at q = 3 (d=3h)
-        alpha_d = {1: 1 / 120, 2: 7 / 478 / PI, 3: 1 / 120 / PI}[spatial_rank]
-        if not derivative:
+        q = 3 * distance / h
+        if 0 in derivative:
+            norm = 6 * 1 / 120 / h if d == 1 else 6 * 7 / 478 / PI / (h*h) if d == 2 else 6 * 1 / 120 / PI / (h*h*h)
             w1 = (3 - q) ** 5 - 6 * (2 - q) ** 5 + 15 * (1 - q) ** 5
             w2 = (3 - q) ** 5 - 6 * (2 - q) ** 5
             w3 = (3 - q) ** 5
-            fac = alpha_d
-        elif derivative == 1:
+            result.append(norm * where(q > 2, w3, where(q > 1, w2, w1)))
+        if 1 in derivative:
+            norm = 6 * -5 * 3 / 120 / (h*h) if d == 1 else 6 * -5 * 3 * 7 / 478 / PI / (h*h*h) if d == 2 else 6 * -5 * 3 / 120 / PI / h**4
             w1 = (3 - q) ** 4 - 6 * (2 - q) ** 4 + 15 * (1 - q) ** 4
             w2 = (3 - q) ** 4 - 6 * (2 - q) ** 4
             w3 = (3 - q) ** 4
-            fac = -5 * alpha_d
-        else:
-            raise NotImplementedError(f"{derivative}th derivative of {kernel} is not supported")
-        return fac * where(q > 2, w3, where(q > 1, w2, w1))
+            result.append(norm * where(q > 2, w3, where(q > 1, w2, w1)))
     elif kernel == 'wendland-c2':  # cutoff at q=2 (d=2h)
-        alpha_d = {2: 7 / 4 / PI, 3: 21 / 16 / PI}[spatial_rank]
-        if not derivative:
-            w = (1 - 0.5 * q) ** 4 * (2 * q + 1)
-            fac = alpha_d
-        elif derivative == 1:
-            w = (1 - 0.5 * q) ** 3 * q
-            fac = -5 * alpha_d
-        else:
-            raise NotImplementedError(f"{derivative}th derivative of {kernel} is not supported")
-        return fac * w
+        q = distance / h
+        if 0 in derivative:
+            norm = 3 / h if d == 1 else 7 / PI / (h*h) if d == 2 else 21 / 2 / PI / (h*h*h)
+            result.append(norm * (1-q) ** 4 * (4*q + 1))
+        if 1 in derivative:
+            norm = -20 * 3 / (h*h) if d == 1 else -20 * 7 / PI / (h*h*h) if d == 2 else -20 * 21 / 2 / PI / h**4
+            result.append(norm * q * (1-q)**3)
+    elif kernel == 'poly6':  # from Müller et al., Particle-based fluid simulation for interactive applications
+        diff = h**2 - math.vec_squared(delta)
+        if 0 in derivative:
+            norm = 35 / 16 / h**7 if d == 1 else 4 / PI / h**8 if d == 2 else 315 / 64 / PI / h**9
+            result.append(norm * diff ** 3)
+        if 1 in derivative:
+            norm = -6 * 35 / 16 / h**7 if d == 1 else -6 * 4 / PI / h**8 if d == 2 else -6 * 315 / 64 / PI / h**9
+            result.append(delta * norm * diff ** 2)
     else:
         raise ValueError(kernel)
+    return result
 
 
-def density(graph: Graph):
+def density(graph: Graph) -> Tensor:
+    """
+    Sum the kernel function over all neighbors within the support radius.
+
+    Args:
+        graph: `Graph` with `kernel` values stored in the edges.
+
+    Returns:
+        Relative density, i.e. not yet scaled by particle mass.
+    """
     return math.sum(graph.edges['kernel'], dual)
 
 

tests/commit/physics/test_sph.py

@@ -0,0 +1,34 @@
+from unittest import TestCase
+
+from phi import math
+from phi.physics import sph
+from phiml.math import spatial
+
+
+class TestSPH(TestCase):
+
+    def test_evaluate_kernel(self):
+        eps = 1e-8
+        with math.precision(64):
+            r = math.linspace(0, 1, spatial(x=100))
+            for kernel in ['quintic-spline', 'wendland-c2', 'poly6']:
+                val, grad = sph.evaluate_kernel(r, r, 1, 1, kernel, derivative=(0, 1))
+                h_val, = sph.evaluate_kernel(r + eps, r + eps, 1, 1, kernel, derivative=(0,))
+                fd_grad = (h_val - val) / eps
+                math.assert_close(fd_grad, grad, rel_tolerance=1e-5, abs_tolerance=1e-5)  # gradient correct
+                math.assert_close(val.x[-1], 0)  # W(1) = 0
+                math.assert_close(grad.x[-1], 0)  # dW(1) = 0
+                math.assert_close(1, math.sum(val / 100), abs_tolerance=0.01)  # integral = 1
+
+    def test_evaluate_kernel_scaled(self):
+        eps = 1e-8
+        with math.precision(64):
+            r = math.linspace(0, 10, spatial(x=100))
+            for kernel in ['quintic-spline', 'wendland-c2', 'poly6']:
+                val, grad = sph.evaluate_kernel(r, r, 10, 1, kernel, derivative=(0, 1))
+                h_val, = sph.evaluate_kernel(r + eps, r + eps, 10, 1, kernel, derivative=(0,))
+                fd_grad = (h_val - val) / eps
+                math.assert_close(fd_grad, grad, rel_tolerance=1e-5, abs_tolerance=1e-5)
+                math.assert_close(val.x[-1], 0)
+                math.assert_close(grad.x[-1], 0)
+                math.assert_close(1, math.sum(val / 10), abs_tolerance=0.01)