Skip to content

Performance: significant speedup and memory usage improvement for PCA().run() #5352

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jberg5 wants to merge 2 commits into
base: develop
Choose a base branch
from
+20 −6
Open

Performance: significant speedup and memory usage improvement for PCA().run() #5352

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
29e0e21
Compute covariance matrix once and use eigh for significant speedup
jberg5 Mar 29, 2026
372e8d1
docs
jberg5 Mar 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions package/AUTHORS
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,7 @@ Chronological list of authors
- Ayush Agarwal
- Parth Uppal
- Olivier Languin--Cattoën
- Jonathan Berg

External code
-------------
Expand Down
5 changes: 4 additions & 1 deletion package/CHANGELOG
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ The rules for this file:
??/??/?? IAlibay, orbeckst, marinegor, tylerjereddy, ljwoods2, marinegor,
spyke7, talagayev, tanii1125, BradyAJohnston, hejamu, jeremyleung521,
harshitgajjela-droid, kunjsinha, aygarwal, jauy123, Dreamstick9,
ollyfutur
ollyfutur, jberg

* 2.11.0

Expand Down Expand Up @@ -53,6 +53,9 @@ Enhancements
* Added new function `MDAnalysis.fetch.from_PDB` to download structure files from wwPDB
using `pooch` as optional dependency (Issue #4907, PR #4943)
* Added benchmarks for package.MDAnalysis.analysis.msd.EinsteinMSD (PR #5277)
* Improved PCA performance by computing covariance matrix via BLAS matrix
multiply instead of per-frame outer product accumulation, and using
symmetric eigendecomposition (eigh) instead of general eig (PR #5352)
* Improved performance of `AtomGroup.wrap()` with compounds (PR #5220)
* Adds support for parsing `.tpr` files produced by GROMACS 2026.0
* Enables parallelization for analysis.diffusionmap.DistanceMatrix
Expand Down
20 changes: 15 additions & 5 deletions package/MDAnalysis/analysis/pca.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -238,6 +238,11 @@ class PCA(AnalysisBase):
manually iterate through ``self._trajectory``, which would
incorrectly handle cases where the ``frame`` argument
was passed.
.. versionchanged:: 2.11.0
Covariance matrix is now computed via BLAS matrix multiply instead of
per-frame outer product accumulation, and uses symmetric
eigendecomposition (``eigh``) instead of general ``eig``, resulting
in significant performance improvements for large systems.
"""

_analysis_algorithm_is_parallelizable = False
Expand Down Expand Up @@ -290,6 +295,7 @@ def _prepare(self):
)
n_dim = self._n_atoms * 3
self.cov = np.zeros((n_dim, n_dim))
Copy link
Copy Markdown
Author

@jberg5 jberg5 Mar 29, 2026 •
edited
Loading

Choose a reason for hiding this comment

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

Technically we can remove this, but I wasn't sure if anyone out there would expect it to be accessible from _prepare onwards. No strong feelings either way, it's relatively cheap to initialize and maximizes backwards compatibility , but future readers might be confused by it.

self._frame_data = np.empty((self.n_frames, n_dim))
self._ref_atom_positions = self._reference.positions
self._ref_cog = self._reference.center_of_geometry()
self._ref_atom_positions -= self._ref_cog
Expand Down Expand Up @@ -329,14 +335,18 @@ def _single_frame(self):
else:
x = self._atoms.positions.ravel()
x -= self._xmean
self.cov += np.dot(x[:, np.newaxis], x[:, np.newaxis].T)
Copy link
Copy Markdown
Member

@orbeckst orbeckst Mar 31, 2026

Choose a reason for hiding this comment

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

One reason for building the covariance matrix incrementally is that it ensures bounds on memory usage. In your approach, memory consumption grows linearly with trajectory length because you're copying coordinates into memory and thus a longer trajectory will eventually fill up the memory. The test trajectory is very short. Real trajectories are easily 1000-100,000 times longer.

If there's really an advantage to multiplying huge matrices then perhaps a batched approach is a good middle ground.

Copy link
Copy Markdown
Member

@orbeckst orbeckst Mar 31, 2026

Choose a reason for hiding this comment

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

Some manual benchmarking:

import numpy as np
import MDAnalysis as mda; from MDAnalysis.tests import datafiles as data
u = mda.Universe(data.PSF, data.DCD)
x = u.atoms.positions.ravel()

shows that replacing the np.dot with transposes with np.outer will help already

>>> %timeit np.dot(x[:, np.newaxis], x[:, np.newaxis].T)
61.4 ms ± 466 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)

>>> %timeit np.outer(x, x)
35.4 ms ± 249 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)

However, the addition update is actually slow

%%timeit cov = np.zeros((len(x), len(x))); xx = np.outer(x,x)
    ...: cov += xx
    ...:
97.5 ms ± 3.08 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

so it makes sense to do the addition as part of matrix multiplication.

Batching is probably the way to go (but that makes streaming more difficult).

self._frame_data[self._frame_index] = x

def _conclude(self):
self.cov = self._frame_data.T @ self._frame_data
self.cov /= self.n_frames - 1
e_vals, e_vects = np.linalg.eig(self.cov)
sort_idx = np.argsort(e_vals)[::-1]
self._variance = e_vals[sort_idx]
self._p_components = e_vects[:, sort_idx]
del self._frame_data
# covariance matrix is symmetric and real-valued by
# construction, so we can use faster eigh
eigenvals, eigenvects = np.linalg.eigh(self.cov)
# eigh returns ascending order; reverse for descending
self._variance = eigenvals[::-1]
self._p_components = eigenvects[:, ::-1]
self._calculated = True
self.n_components = self._n_components

Expand Down
Loading