Skip to content

perf(structure): vectorize RDF distance calculation and add automatic cell list heuristic - #972

Open
kokhlo wants to merge 1 commit into
biotite-dev:mainfrom
kokhlo:issue-307
Open

kokhlo wants to merge 1 commit into
biotite-dev:mainfrom
kokhlo:issue-307

Conversation

@kokhlo

@kokhlo kokhlo commented Oct 22, 2026

Copy link
Copy Markdown

Summary

rdf() evaluated the displacement for each center individually in a Python loop, even though each iteration only selected rows from the cell list query result. This dominated the runtime for all system sizes (see benchmarks below).

This PR

  1. replaces the per-center loop with a single vectorized displacement() evaluation over the (center, atom) pairs from CellList.get_atoms_in_cells(..., result_format=CellList.Result.PAIRS). For periodic=True the periodic image duplicates in the raw pairs are removed after mapping the image indices back to the original atoms.
  2. adds the use_cell_list parameter that [WIP] Optional brute force approach for RDF calculation #315 introduced for the old pure-Python implementation, including the automatic choice proposed in Improve distance calculation speed for small systems  #307: a work estimate compares the brute-force cost n_centers * n_atoms with the cell list cost n_atoms + factor * n_centers * n_atoms * visited_fraction, where the visited fraction is estimated from the ratio of the swept cell volume ((3 * interval_max)**3) to the system volume. For few centers or small systems the brute-force approach evaluates the full broadcast instead.
  3. keeps the exact prior code path whenever the interval exceeds the smallest box inradius with periodic=True: in that regime minimum-image displacements and cell-list preselection are not guaranteed to be interchangeable for strongly sheared triclinic boxes, so the automatic mode conservatively uses the cell list.

This supersedes #315 (written against the pre-Rust implementation) and closes #307.

Correctness

The distance evaluation per pair is unchanged — displacement() is called with the same coordinates, the same box and the same minimum-image convention as before; only the batching changed, and np.histogram counting is order-independent.

Verified against main with a test set of 9 cases (single/multiple centers, periodic/non-periodic, multi-model stacks, selection, adversarial small box with interval beyond the inradius, strongly sheared triclinic boxes): output bins and g_r are exactly identical (np.array_equal) to main for all cases, for both explicitly chosen paths and the automatic choice.

New tests in tests/structure/test_rdf.py:

  • test_rdf_use_cell_list_equivalence: cell list and brute-force paths give identical results (periodic and non-periodic).
  • test_rdf_small_system_heuristic: the automatic choice matches an explicitly chosen path for a single center.

Benchmarks

macOS (M-series laptop), Python 3.13, NumPy 2.x, median of 7 runs, bins=100, interval=(0, 10), fixed seed:

system atoms centers main this PR speedup
waterbox 1 851 1 0.14 ms 0.15 ms 1.0x
waterbox 1 851 10 0.47 ms 0.53 ms ~1x
waterbox 1 851 100 4.54 ms 4.21 ms 1.1x
waterbox 1 851 500 22.9 ms 19.3 ms 1.2x
1gya 1 976 10 0.41 ms 0.37 ms 1.1x
1gya 1 976 100 3.39 ms 2.81 ms 1.2x
1gya 1 976 500 17.0 ms 17.0 ms 1.0x
synthetic sparse box 30 000 10 1.31 ms 1.95 ms* —
synthetic sparse box 30 000 100 5.92 ms 6.81 ms* —
synthetic sparse box 30 000 500 37.1 ms 26.2 ms 1.4x

*The synthetic sparse-box rows include the heuristic overhead: the crossover from brute-force to cell list sits between 10 and 100 centers for this density; the per-component cost at 100 centers is 5.4 ms (cell list) vs 73 ms (brute force), i.e. the heuristic picks the cheaper path with an order of magnitude margin.

The interesting small-system regime (few centers, use_cell_list=False explicitly) is 1.5-2x faster than main:

case main (loop over cell list) this PR (brute force)
waterbox, 1 center 0.14 ms 0.11 ms
30k sparse, 1 center 0.86 ms 0.55 ms
30k sparse, 10 centers 1.31 ms 0.6 ms

Added benchmarks/structure/benchmark_rdf.py (few/many centers, small/large/sparse systems) for CodSpeed tracking.

Testing

  • pytest tests/structure/test_rdf.py — 9 passed (7 existing + 2 new)
  • pytest tests/structure/ --ignore=tests/structure/io --ignore=tests/structure/test_connect.py --ignore=tests/structure/test_info.py: failure set identical to main (306 pre-existing failures/errors from missing optional data, 0 regressions)
  • ruff check / ruff format --check on touched files: clean
  • pyright on rdf.py: same pre-existing findings as main, no new ones

Fixes #307
Supersedes #315

… cell list heuristic

The per-center Python loop in rdf() dominated the runtime, even for
large systems where the cell list approach should pay off.
Replace it with a single vectorized displacement evaluation over the
pairs returned by the cell list query.

For small workloads, i.e. few centers or few atoms, the cell list
overhead does not pay off. Add an optional use_cell_list parameter,
that allows to explicitly choose the approach, or let it be chosen
automatically from a work estimate based on the number of centers, the
number of atoms and the fraction of the system volume that the cell
list queries visit.

Fixes biotite-dev#307

This branch has not been deployed

No deployments
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.

Improve distance calculation speed for small systems

1 participant