WP5: a segmented near field, because float32 cannot afford a scatter

tests/applications/test_svgd_vs_exact.py

@@ -12,6 +12,7 @@
 
 import jax
 import jax.numpy as jnp
+import numpy as np
 import pytest
 
 from yggdrax import DualTreeTraversalConfig
@@ -241,3 +242,62 @@ def test_svgd_phi_cutoff_bandwidths_is_the_c_times_h_convention():
     topo = build_svgd_topology(p, kernel_cutoff=6.0 * h, **kw)
     via_topo = svgd_phi_from_topology(p, sc, h, topo)
     assert float(jnp.max(jnp.abs(via_helper - via_topo))) < 1e-12
+
+
+# --- the two near-field accumulations --------------------------------------
+#
+# The near field can be summed by scattering each unordered pair's two
+# directions, or by a segmented reduction over the directed list. Which is
+# faster is a property of the dtype (float32's `.at[].add()` is 43x float64's
+# under index contention), so both exist. They must agree.
+
+
+@pytest.mark.parametrize("backend", ["radix", "leaf_kdtree"])
+def test_accumulations_agree(backend):
+    p = jax.random.normal(jax.random.PRNGKey(5), (1200, 3)) * 1.2
+    sc = p * 0.5
+    h = 0.6
+    topo = build_svgd_topology(
+        p, theta=0.4, leaf_size=16, backend=backend, traversal_config=_CFG
+    )
+    by_scatter = svgd_phi_from_topology(p, sc, h, topo, accumulate="scatter")
+    by_segment = svgd_phi_from_topology(p, sc, h, topo, accumulate="segment")
+    assert float(jnp.max(jnp.abs(by_scatter - by_segment))) < 1e-12
+
+    # "auto" is one of the two, never a third thing.
+    by_auto = svgd_phi_from_topology(p, sc, h, topo, accumulate="auto")
+    assert float(jnp.max(jnp.abs(by_auto - by_scatter))) < 1e-12
+
+
+def test_segment_accumulation_is_exact_at_theta_zero():
+    """The segmented path is a different summation order, not a different sum."""
+    p = jax.random.normal(jax.random.PRNGKey(6), (900, 3)) * 1.2
+    sc = p * 0.5
+    h = float(median_heuristic(p))
+    topo = build_svgd_topology(
+        p, theta=0.0, leaf_size=16, backend="radix", traversal_config=_CFG
+    )
+    ref = exact_phi(p, sc, h)
+    out = svgd_phi_from_topology(p, sc, h, topo, accumulate="segment")
+    assert float(jnp.linalg.norm(out - ref) / jnp.linalg.norm(ref)) < 1e-10
+
+
+def test_unknown_accumulation_is_rejected():
+    p = jax.random.normal(jax.random.PRNGKey(7), (200, 3))
+    topo = build_svgd_topology(
+        p, theta=0.4, leaf_size=16, backend="radix", traversal_config=_CFG
+    )
+    with pytest.raises(ValueError, match="accumulate must be"):
+        svgd_phi_from_topology(p, p * 0.5, 0.6, topo, accumulate="nonsense")
+
+
+def test_directed_pair_list_is_sorted_and_twice_the_halved_one():
+    """What makes the segmented reduction segmented."""
+    p = jax.random.normal(jax.random.PRNGKey(8), (1500, 3)) * 1.2
+    topo = build_svgd_topology(
+        p, theta=0.4, leaf_size=16, backend="radix", traversal_config=_CFG
+    )
+    directed = np.asarray(topo.near_dir_target)
+    assert directed.shape[0] == 2 * int(topo.near_target_row.shape[0])
+    assert directed.shape[0] == int(topo.num_near_leaf_pairs)
+    assert np.all(np.diff(directed) >= 0), "target rows must be non-decreasing"

yggdrax/applications/svgd/exact.py

@@ -13,31 +13,72 @@
 import jax.numpy as jnp
 from jaxtyping import Array, Float
 
-from yggdrax.applications.svgd.kernel import stein_pair_terms
-
 
 def exact_phi(
     particles: Float[Array, "n d"],
     scores: Float[Array, "n d"],
     h: float | Float[Array, ""],
+    *,
+    block_size: int | None = None,
 ) -> Float[Array, "n d"]:
     """Exact Stein update direction phi(x_i) for every particle, O(N^2).
 
+    The pair sum is contracted, not materialised. Writing it out,
+
+    .. math::
+
+        \\phi_i = \\frac{1}{N}\\Big[(K S)_i
+                 + \\big(x_i (K \\mathbf{1})_i - (K X)_i\\big) / h^2\\Big],
+
+    with :math:`K_{ij} = \\exp(-\\lVert x_i - x_j\\rVert^2 / 2h^2)`, so the whole
+    update is one kernel matrix and **two matmuls**. The obvious form builds an
+    ``(n, n, d)`` tensor of per-pair terms, which is *d* times the memory, runs
+    at elementwise rather than GEMM throughput, and cannot be evaluated at all
+    beyond N ~ 2e4 -- which mattered, because this is the baseline the tree
+    update is judged against, and a weak baseline flatters the tree.
+
     Args:
         particles: Particle positions, shape ``(n, d)``.
         scores: Target score at each particle, shape ``(n, d)``.
         h: Kernel bandwidth.
+        block_size: Targets per block. ``None`` does every target at once,
+            which costs an ``(n, n)`` kernel matrix; pass a block size to cap
+            that at ``(block_size, n)`` and reach large N.
 
     Returns:
         Update directions, shape ``(n, d)``.
     """
-    n = particles.shape[0]
-    # target axis 0 (i), source axis 1 (j): terms[i, j] contributes to phi[i].
-    x_t = particles[:, None, :]
-    x_s = particles[None, :, :]
-    s_s = scores[None, :, :]
-    terms = stein_pair_terms(x_t, x_s, s_s, h)  # (n, n, d)
-    return jnp.sum(terms, axis=1) / n
+    n, d = particles.shape
+    if block_size is None:
+        block_size = n
+    block = max(1, min(int(block_size), n))
+
+    def _block(x_t: Array) -> Array:
+        """Contribution of every source to one block of targets."""
+        # (B, n) kernel, then two contractions -- no (B, n, d) tensor is ever
+        # built. sum_j k_ij s_j is a matmul, and
+        # sum_j k_ij (x_i - x_j) = x_i * sum_j k_ij - sum_j k_ij x_j is another.
+        d2 = (
+            jnp.sum(x_t * x_t, axis=-1)[:, None]
+            - 2.0 * (x_t @ particles.T)
+            + jnp.sum(particles * particles, axis=-1)[None, :]
+        )
+        k = jnp.exp(-jnp.maximum(d2, 0.0) / (2.0 * h**2))  # (B, n)
+        attract = k @ scores
+        repulse = x_t * jnp.sum(k, axis=1)[:, None] - k @ particles
+        return attract + repulse / (h**2)
+
+    if block >= n:
+        return _block(particles) / n
+
+    pad = (-n) % block
+    padded = (
+        particles
+        if pad == 0
+        else jnp.concatenate([particles, jnp.zeros((pad, d), particles.dtype)])
+    )
+    out = jax.lax.map(_block, padded.reshape(-1, block, d))
+    return out.reshape(-1, d)[:n] / n
 
 
 def svgd_step(