perf(swarm): vectorize IndexSwarmVariable proxy update (17-80x speedup)

src/underworld3/swarm.py

@@ -2558,6 +2558,22 @@ def _update_proxy_variables(self):
                 # n, d, b = kd_swarm.find_closest_point(self._meshLevelSetVars[0].coords)
                 d, n = kd_swarm.query(self._meshLevelSetVars[0].coords, k=1, sqr_dists=False)
 
+                # Which (particle, node) pairs are valid:
+                # - node is at same distance as the nearest node
+                # - node is within radius_s
+                is_nearest = np.isclose(n_distance, n_distance[:, [0]])
+                within_radius = n_distance < self.radius_s
+                valid = is_nearest & within_radius
+
+                # IDW weights (only for valid pairs; others zeroed)
+                weights = 1.0 / (n_distance + 1e-16)
+                weights[~valid] = 0.0
+
+                # Total weight per node — material independent, compute once
+                n_mesh_nodes = self._meshLevelSetVars[0].data.shape[0]
+                w = np.zeros(n_mesh_nodes)
+                np.add.at(w, n_indices, weights)
+
             for ii in range(self.indices):
                 meshVar = self._meshLevelSetVars[ii]
 
@@ -2571,23 +2587,15 @@ def _update_proxy_variables(self):
                 final_values = np.array(meshVar.data[:, 0], copy=True)
 
                 if not starved:
+                    # Material presence mask: (n_particles, 1) for broadcasting
+                    mat_mask = (self.data.flatten() == ii).astype(weights.dtype)[:, None]
+
+                    # Weighted sum per node
                     node_values = np.zeros(final_values.shape[0])
-                    w = np.zeros(final_values.shape[0])
-
-                    for i in range(self.swarm.local_size):
-                        tem = np.isclose(n_distance[i, :], n_distance[i, 0])
-                        dist = n_distance[i, tem]
-                        indices = n_indices[i, tem]
-                        tem = dist < self.radius_s
-                        dist = dist[tem]
-                        indices = indices[tem]
-                        for j, ind in enumerate(indices):
-                            node_values[ind] += (
-                                np.isclose(self.data[i], ii) / (1.0e-16 + dist[j])
-                            )[0]
-                            w[ind] += 1.0 / (1.0e-16 + dist[j])
-
-                    node_values[np.where(w > 0.0)[0]] /= w[np.where(w > 0.0)[0]]
+                    np.add.at(node_values, n_indices, weights * mat_mask)
+
+                    # Normalize
+                    node_values[w > 0] /= w[w > 0]
                     final_values = node_values
 
                     # if there is no material found,
@@ -2602,49 +2610,67 @@ def _update_proxy_variables(self):
                 # current values, i.e. the proxy is left unchanged there)
                 meshVar.data[:, 0] = final_values
         elif self.update_type == 1:
-            # NOTE: this branch performs data-dependent MeshVariable writes
-            # outside any access context, which is not parallel-safe
-            # independently of the starved-rank issue (pre-existing).
-            # The guard here only prevents the empty-rank KDTree crash.
-            if starved:
-                return
-            kd = uw.kdtree.KDTree(self.swarm._particle_coordinates.data)
-            n_distance, n_indices = kd.query(
-                self._meshLevelSetVars[0].coords, k=self.nnn, sqr_dists=False
-            )
+            if not starved:
+                kd = uw.kdtree.KDTree(self.swarm._particle_coordinates.data)
+                n_distance, n_indices = kd.query(
+                    self._meshLevelSetVars[0].coords, k=self.nnn, sqr_dists=False
+                )
+
+                # IDW weights and validity mask for all (node, particle) pairs
+                valid = n_distance < self.radius_s
+                a = 1.0 / (n_distance + 1e-16)
+                a[~valid] = 0.0
+
+                # Total weight per node (material-independent)
+                w = a.sum(axis=1)
+
+                # Handle boundary nodes: restrict to nnn_bc particles
+                if hasattr(self, 'ind_bc') and self.ind_bc is not None:
+                    bc_idx = np.array(list(self.ind_bc))
+                    bc_idx = bc_idx[bc_idx < a.shape[0]]
+                    if len(bc_idx) > 0:
+                        valid_bc = n_distance[bc_idx, :self.nnn_bc] < self.radius_s
+                        a_bc = 1.0 / (n_distance[bc_idx, :self.nnn_bc] + 1e-16)
+                        a_bc[~valid_bc] = 0.0
+                        w_bc = a_bc.sum(axis=1)
+
+                        a[bc_idx] = 0.0
+                        a[bc_idx, :self.nnn_bc] = a_bc
+                        w[bc_idx] = w_bc
 
             for ii in range(self.indices):
                 meshVar = self._meshLevelSetVars[ii]
-                node_values = np.zeros((meshVar.data.shape[0],))
-                w = np.zeros((meshVar.data.shape[0],))
-                for i in range(meshVar.data.shape[0]):
-                    if i not in self.ind_bc:
-                        ind = np.where(n_distance[i, :] < self.radius_s)
-                        a = 1.0 / (n_distance[i, ind] + 1.0e-16)
-                        w[i] = np.sum(a)
-                        b = np.isclose(self.data[n_indices[i, ind]], ii)
-                        node_values[i] = np.sum(np.dot(a, b))
-                        if ind[0].size == 0:
-                            w[i] = 0
-                    else:
-                        ind = np.where(n_distance[i, : self.nnn_bc] < self.radius_s)
-                        a = 1.0 / (n_distance[i, : self.nnn_bc][ind] + 1.0e-16)
-                        w[i] = np.sum(a)
-                        b = np.isclose(self.data[n_indices[i, : self.nnn_bc][ind]], ii)
-                        node_values[i] = np.sum(np.dot(a, b))
-                        if ind[0].size == 0:
-                            w[i] = 0
-
-                node_values[np.where(w > 0.0)[0]] /= w[np.where(w > 0.0)[0]]
-                meshVar.data[:, 0] = node_values[...]
-
-                # if there is no material found,
-                # impose a near-neighbour hunt for a valid material and set that one
-                ind_w0 = np.where(w == 0.0)[0]
-                if len(ind_w0) > 0:
-                    ind_ = np.where(self.data[n_indices[ind_w0]] == ii)[0]
-                    if len(ind_) > 0:
-                        meshVar.data[ind_w0[ind_]] = 1.0
+
+                # MeshVariable reads/writes perform collective ghost
+                # synchronisation, so every rank must execute exactly the
+                # same read-then-write sequence per level set (same pattern
+                # as update_type=0 above). Starved ranks read their current
+                # proxy values and write them back unchanged; populated ranks
+                # compute and write new values.
+                final_values = np.array(meshVar.data[:, 0], copy=True)
+
+                if not starved:
+                    # Material presence at each (node, particle) pair
+                    # self.data has shape (n_particles, 1); flatten to (n_particles,)
+                    # so fancy indexing yields (n_nodes, nnn) — matching `a`.
+                    mat_present = (self.data.flatten()[n_indices] == ii).astype(a.dtype)
+
+                    # Weighted sum per node, then normalize
+                    node_values = (a * mat_present).sum(axis=1)
+                    node_values[w > 0] /= w[w > 0]
+                    final_values = node_values
+
+                    # if there is no material found,
+                    # impose a near-neighbour hunt for a valid material and set that one
+                    ind_w0 = np.where(w == 0.0)[0]
+                    if len(ind_w0) > 0:
+                        ind_ = np.where(self.data[n_indices[ind_w0]] == ii)[0]
+                        if len(ind_) > 0:
+                            final_values[ind_w0[ind_]] = 1.0
+
+                # single symmetric write (starved ranks write back their
+                # current values, i.e. the proxy is left unchanged there)
+                meshVar.data[:, 0] = final_values
         return
 
 

tests/benchmark_index_swarm_vectorized.py

@@ -0,0 +1,95 @@
+"""
+Benchmark the vectorized IndexSwarmVariable._update_proxy_variables
+against the original loop-based reference.
+
+Uses a larger mesh to make the speedup measurable.
+Run with: python tests/benchmark_index_swarm_vectorized.py
+"""
+import numpy as np
+import underworld3 as uw
+import time
+
+from test_0115_index_swarm_vectorized import (
+    _reference_update_type0,
+    _reference_update_type1,
+)
+
+
+def benchmark(update_type=1, mesh_res=(40, 40), fill_param=5, nnn=10, radius=0.1):
+    uw.reset_default_model()
+    mesh = uw.meshing.StructuredQuadBox(
+        elementRes=mesh_res,
+        minCoords=(0.0, 0.0),
+        maxCoords=(10.0, 10.0),
+    )
+    swarm = uw.swarm.Swarm(mesh)
+    material = uw.swarm.IndexSwarmVariable(
+        "M_bench", swarm, indices=4, proxy_degree=1, proxy_continuous=True,
+        update_type=update_type, npoints=nnn, radius=radius,
+    )
+    swarm.populate(fill_param=fill_param)
+    with swarm.access():
+        # 4-layer material assignment
+        y = swarm.data[:, 1]
+        material.data[:, 0] = np.where(
+            y < 2.5, 0,
+            np.where(y < 5.0, 1,
+                     np.where(y < 7.5, 2, 3))
+        ).astype(int)
+
+    print(f"\n=== update_type={update_type}, mesh={mesh_res}, fill={fill_param}, "
+          f"nnn={nnn}, radius={radius} ===")
+    print(f"  Particles (local): {swarm.local_size}")
+    print(f"  Mesh nodes: {material._meshLevelSetVars[0].data.shape[0]}")
+    print(f"  Materials: {material.indices}")
+
+    # Benchmark vectorized (in-class)
+    material._proxy_stale = True
+    t0 = time.perf_counter()
+    material._update_proxy_if_stale()
+    t_vec = time.perf_counter() - t0
+    print(f"  Vectorized: {t_vec*1000:.2f} ms")
+
+    # Read vectorized result
+    vec = np.zeros((material._meshLevelSetVars[0].data.shape[0], material.indices))
+    for ii in range(material.indices):
+        vec[:, ii] = material._meshLevelSetVars[ii].data[:, 0]
+
+    # Benchmark reference (original loop)
+    t0 = time.perf_counter()
+    if update_type == 0:
+        ref = _reference_update_type0(swarm, material, nnn, radius)
+    else:
+        ref = _reference_update_type1(swarm, material, nnn, radius)
+    t_ref = time.perf_counter() - t0
+    print(f"  Reference:  {t_ref*1000:.2f} ms")
+
+    speedup = t_ref / t_vec if t_vec > 0 else float('inf')
+    print(f"  Speedup:    {speedup:.1f}x")
+
+    # Verify correctness
+    max_diff = np.max(np.abs(ref - vec))
+    print(f"  Max abs diff: {max_diff:.2e}")
+    assert np.allclose(ref, vec, atol=1e-10, rtol=1e-7), (
+        f"Vectorized output does not match reference! Max diff: {max_diff}"
+    )
+    print(f"  ✓ Output matches reference")
+    return speedup
+
+
+if __name__ == "__main__":
+    print("Benchmarking vectorized IndexSwarmVariable._update_proxy_variables")
+    print("=" * 70)
+
+    speedups = []
+    for res in [(10, 10), (20, 20), (40, 40), (60, 60)]:
+        for ut in [0, 1]:
+            s = benchmark(update_type=ut, mesh_res=res, nnn=10, radius=0.5)
+            speedups.append((ut, res, s))
+
+    print("\n" + "=" * 70)
+    print("Summary:")
+    print(f"{'update_type':<15} {'mesh':<15} {'speedup':<10}")
+    print("-" * 40)
+    for ut, res, s in speedups:
+        print(f"{ut:<15} {str(res):<15} {s:.1f}x")