Skip to content

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

Open
bknight1 wants to merge 1 commit into
underworldcode:developmentfrom
bknight1:perf/vectorize-index-swarm-proxy
Open

perf(swarm): vectorize IndexSwarmVariable proxy update (17-80x speedup)#374
bknight1 wants to merge 1 commit into
underworldcode:developmentfrom
bknight1:perf/vectorize-index-swarm-proxy

Conversation

@bknight1

Copy link
Copy Markdown
Member

Summary

Vectorize IndexSwarmVariable._update_proxy_variables() by replacing the per-particle and per-node Python loops with numpy vectorized operations using np.isclose, np.add.at, and broadcasting.

Performance

update_type Mesh (60×60, 130k particles, 4 materials) Speedup
0 (particle→node) 13.1 s → 0.17 s 77×
1 (node→particle) 0.34 s → 0.02 s 17×

Correctness

  • Output is numerically identical to original (max abs diff: 4.44e-16)
  • 26 parametrized verification tests across update_type, radius, and nnn values
  • 3 integration tests (fractions sum to one, interface location)
  • All existing IndexSwarmVariable tests pass

Changes

  • src/underworld3/swarm.py: Vectorize both update_type=0 and update_type=1 in _update_proxy_variables. Key optimizations:
    • w (total weight per node) computed once outside the material loop
    • np.add.at for scatter-add accumulation on mesh nodes
    • Boundary node handling pre-computed outside the material loop
    • self.data[n_indices] flattened to avoid trailing dimension mismatch
  • tests/test_0115_index_swarm_vectorized.py: New test file with reference implementations + 26 parametrized tests
  • tests/benchmark_index_swarm_vectorized.py: New benchmark for measuring speedup

Copilot AI review requested due to automatic review settings July 20, 2026 01:28
@bknight1
bknight1 requested a review from lmoresi as a code owner July 20, 2026 01:28

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Note

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

This PR accelerates IndexSwarmVariable._update_proxy_variables() by replacing Python loops with NumPy vectorized operations while adding verification tests and a benchmark to validate correctness and measure speedups.

Changes:

  • Vectorized proxy update paths for update_type=0 and update_type=1 using np.isclose, broadcasting, and np.add.at.
  • Added a new pytest module with reference (loop-based) implementations and parametrized comparisons.
  • Added a standalone benchmark script to quantify performance gains and assert numerical equivalence.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File Description
src/underworld3/swarm.py Reworks proxy update logic to use vectorized IDW accumulation and shared weight computations.
tests/test_0115_index_swarm_vectorized.py Adds reference implementations + parametrized tests to validate vectorized behavior.
tests/benchmark_index_swarm_vectorized.py Adds a runnable benchmark comparing vectorized vs reference implementations.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread src/underworld3/swarm.py
Comment on lines +2640 to +2652
# 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
Comment thread src/underworld3/swarm.py
Comment on lines +2662 to 2665
# Weighted sum per node, then normalize
node_values = (a * mat_present).sum(axis=1)
node_values[w > 0] /= w[w > 0]
meshVar.data[:, 0] = node_values[...]
Comment on lines +88 to +92
ind = np.where(n_distance[i, :] < radius_s)
a = 1.0 / (n_distance[i, ind] + 1.0e-16)
w[i] = np.sum(a)
b = np.isclose(material.data[n_indices[i, ind]], ii)
node_values[i] = np.sum(np.dot(a, b))
Replace the per-particle and per-node Python loops in
_update_proxy_variables with vectorized numpy operations:

- update_type=0 (particle->node): use np.isclose for validity
  masks and np.add.at for scatter-add accumulation. w (total
  weight per node) computed once outside the material loop
  instead of redundantly per material. ~77x speedup.

- update_type=1 (node->particle): use boolean masks for
  radius filtering and (a * mat_present).sum() for weighted
  sums. Boundary node handling pre-computed outside the
  material loop. ~17x speedup.

Output is numerically identical (max diff 4.44e-16).
@bknight1
bknight1 force-pushed the perf/vectorize-index-swarm-proxy branch from 87c72c2 to 3129c01 Compare July 21, 2026 07:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants