From d9b82aa236a089507c8a7c86dff80bb7bca6a168 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 1 Jun 2026 19:07:54 +0000 Subject: [PATCH 1/3] Simplify Python RPCA: remove dead code, fix numba structure, vectorize hot loop - Move GPA accumulation kernel to a module-level @njit function instead of @jit on an instance method (numba cannot compile self); drop the ProcessPoolExecutor + wrapper + chunking machinery it required and let numba handle per-frame parallelism. Compute the reference mean once instead of every frame, and preallocate the coords array. - Delete dead helpers _calculate_rotation_matrix and _project_coordinates. - Collapse the O(n_res^2 * atoms^2) interaction-matrix loop into per-residue group vectors (G_i . G_j) for identical results in O(n_res^2). - Drop unused imports (align, cdist, ThreadPoolExecutor, multiprocessing, numba float64/int64) and no-op per-method OMP_NUM_THREADS env writes. - Replace the broken per-frame time scan in read_trajectory with a materialized times array + np.searchsorted. - Remove the invalid n_jobs kwarg passed to AnalysisBase and use zero-arg super() in the nested FastCovarianceAnalysis class. https://claude.ai/code/session_018SSc1GoJMfc7Ey8n9EKgrK --- numba_rpca_implementation.py | 316 +++++++++++------------------------ 1 file changed, 98 insertions(+), 218 deletions(-) diff --git a/numba_rpca_implementation.py b/numba_rpca_implementation.py index fa2f11f..a2f5f7b 100644 --- a/numba_rpca_implementation.py +++ b/numba_rpca_implementation.py @@ -13,15 +13,46 @@ import os import argparse import MDAnalysis as mda -from MDAnalysis.analysis import align from MDAnalysis.analysis.base import AnalysisBase from scipy import linalg import matplotlib.pyplot as plt -from scipy.spatial.distance import cdist import time -import multiprocessing -from numba import jit, prange, float64, int64 -from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor +from numba import njit, prange + + +@njit(parallel=True) +def _accumulate_rotated_coords(mobile_coords_array, reference_coords): + """Accumulate optimally-rotated coordinates for GPA (Numba-accelerated). + + Each frame is centered and superposed onto the centered reference via the + Kabsch/SVD rotation, then averaged. Numba handles the per-frame parallelism, + so callers do not need their own process pool. + """ + n_frames = mobile_coords_array.shape[0] + n_atoms = reference_coords.shape[0] + avg_coords = np.zeros((n_atoms, 3)) + + # Reference centering is constant across frames - compute it once. + ref_mean = np.sum(reference_coords, axis=0) / n_atoms + ref_centered = reference_coords - ref_mean + + for i in prange(n_frames): + mobile_coords = mobile_coords_array[i] + mobile_mean = np.sum(mobile_coords, axis=0) / n_atoms + mobile_centered = mobile_coords - mobile_mean + + # Optimal rotation via SVD of the correlation matrix. + correlation_matrix = np.dot(mobile_centered.T, ref_centered) + u, s, vh = np.linalg.svd(correlation_matrix) + + # Guard against reflections (improper rotations). + if np.sign(np.linalg.det(np.dot(vh.T, u.T))) < 0: + vh[-1] = -vh[-1] + rotation = np.dot(u, vh) + + avg_coords += np.dot(mobile_centered, rotation) + + return avg_coords / n_frames class RPCAAnalysis: @@ -52,23 +83,13 @@ def read_trajectory(self, trajectory_file, topology_file=None, start_time=0, end # Set trajectory start/end if specified if start_time > 0 or end_time > 0: - # Find frames corresponding to time range - times = universe.trajectory.time - start_frame = 0 - end_frame = len(universe.trajectory) - 1 - - if start_time > 0: - for i, t in enumerate(times): - if t >= start_time: - start_frame = i - break - - if end_time > 0: - for i, t in enumerate(times): - if t > end_time: - end_frame = i - 1 - break - + # Materialize per-frame times (trajectory.time is only the current frame), + # then locate the range with a single binary search per bound. + times = np.array([ts.time for ts in universe.trajectory]) + + start_frame = int(np.searchsorted(times, start_time, side='left')) if start_time > 0 else 0 + end_frame = int(np.searchsorted(times, end_time, side='right')) - 1 if end_time > 0 else len(times) - 1 + print(f"Using trajectory frames from {start_frame} to {end_frame}") universe.trajectory[start_frame:end_frame] @@ -78,45 +99,7 @@ def read_trajectory(self, trajectory_file, topology_file=None, start_time=0, end print(f"Error reading trajectory: {e}") return None - @jit(nopython=True, parallel=True) - def _accumulate_coords(self, mobile_coords_array, reference_coords, n_frames, n_atoms): - """Numba-accelerated function to accumulate coordinates for GPA""" - avg_coords = np.zeros((n_atoms, 3), dtype=np.float64) - - for i in prange(n_frames): - mobile_coords = mobile_coords_array[i] - - # Center coordinates - mobile_mean = np.mean(mobile_coords, axis=0) - ref_mean = np.mean(reference_coords, axis=0) - - mobile_centered = mobile_coords - mobile_mean - ref_centered = reference_coords - ref_mean - - # Calculate correlation matrix - correlation_matrix = np.dot(mobile_centered.T, ref_centered) - - # SVD - u, s, vh = np.linalg.svd(correlation_matrix) - - # Ensure proper rotation (not reflection) - d = np.sign(np.linalg.det(np.dot(vh.T, u.T))) - - # Construct rotation matrix - if d < 0: - vh[-1] = -vh[-1] # Flip last row of vh - - rotation = np.dot(u, vh) - - # Apply rotation - rotated_coords = np.dot(mobile_centered, rotation) - - # Accumulate for average - avg_coords += rotated_coords - - return avg_coords / n_frames - - def perform_gpa(self, universe, selection='protein', max_iterations=10, + def perform_gpa(self, universe, selection='protein', max_iterations=10, convergence=0.00001, ref_frame=0, n_jobs=None): """Perform Generalized Procrustes Analysis to find average structure @@ -132,117 +115,48 @@ def perform_gpa(self, universe, selection='protein', max_iterations=10, Average coordinates after GPA """ print("Performing Generalized Procrustes Analysis (GPA)...") - - if n_jobs is None: - n_jobs = multiprocessing.cpu_count() - + # Select atoms for analysis atoms = universe.select_atoms(selection) n_atoms = len(atoms) - - # Initialize with the first frame as reference + + # Initialize with the chosen frame as reference universe.trajectory[ref_frame] - reference_coords = atoms.positions.copy() - - # Pre-load all coordinates into memory for faster processing - # This trades memory for speed + reference_coords = atoms.positions.astype(np.float64) + + # Pre-load all coordinates into a preallocated array (trades memory for speed) print("Pre-loading trajectory coordinates...") - coords_array = [] - for ts in universe.trajectory: - coords_array.append(atoms.positions.copy()) - - coords_array = np.array(coords_array) - n_frames = len(coords_array) + n_frames = len(universe.trajectory) + coords_array = np.empty((n_frames, n_atoms, 3), dtype=np.float64) + for i, ts in enumerate(universe.trajectory): + coords_array[i] = atoms.positions print(f"Loaded {n_frames} frames") - - # Iterative GPA using optimized parallel processing + + # Iterative GPA - per-frame parallelism is handled inside the numba kernel. + avg_coords = reference_coords for iteration in range(max_iterations): print(f"GPA Iteration {iteration + 1}/{max_iterations}") - - # Process frames in parallel chunks - chunk_size = max(1, n_frames // n_jobs) - chunks = [coords_array[i:i+chunk_size] for i in range(0, n_frames, chunk_size)] - - avg_coords = np.zeros_like(reference_coords) - - with ProcessPoolExecutor(max_workers=n_jobs) as executor: - futures = [] - for chunk in chunks: - futures.append( - executor.submit( - self._accumulate_coords_wrapper, - chunk, reference_coords, len(chunk), n_atoms - ) - ) - - # Collect results - chunk_results = [future.result() for future in futures] - - # Weight by chunk size and combine - for i, result in enumerate(chunk_results): - weight = len(chunks[i]) / n_frames - avg_coords += result * weight - + + avg_coords = _accumulate_rotated_coords(coords_array, reference_coords) + # Calculate RMSD between old and new reference rmsd = np.sqrt(np.mean(np.sum((reference_coords - avg_coords)**2, axis=1))) print(f" RMSD between iterations: {rmsd:.6f}") - + # Check convergence if rmsd < convergence and iteration > 0: print(f"GPA converged after {iteration+1} iterations") break - + # Update reference for next iteration reference_coords = avg_coords.copy() - + return avg_coords - - def _accumulate_coords_wrapper(self, coords_chunk, reference_coords, chunk_size, n_atoms): - """Wrapper for the numba-accelerated function to handle Python objects""" - # Convert to numpy arrays of the right shape and type for numba - coords_array_np = np.array(coords_chunk, dtype=np.float64) - reference_coords_np = np.array(reference_coords, dtype=np.float64) - - return self._accumulate_coords(coords_array_np, reference_coords_np, chunk_size, n_atoms) - - @staticmethod - @jit(nopython=True) - def _calculate_rotation_matrix(mobile, reference): - """Calculate optimal rotation matrix using SVD (Numba-accelerated) - - Args: - mobile: Mobile coordinates (centered) - reference: Reference coordinates (centered) - - Returns: - Rotation matrix and RMSD - """ - # Compute correlation matrix - correlation_matrix = np.dot(mobile.T, reference) - - # Singular value decomposition - U, S, Vt = np.linalg.svd(correlation_matrix) - - # Ensure proper rotation (not reflection) - d = np.sign(np.linalg.det(np.dot(Vt.T, U.T))) - - # Construct rotation matrix - if d < 0: - S = np.diag([1, 1, d]) - rotation = np.dot(U, np.dot(S, Vt)) - else: - rotation = np.dot(U, Vt) - - # Calculate RMSD - rotated = np.dot(mobile, rotation) - rmsd = np.sqrt(np.mean(np.sum((rotated - reference)**2, axis=1))) - - return rotation, rmsd - + class FastCovarianceAnalysis(AnalysisBase): """Optimized covariance calculation using MDAnalysis analysis framework""" def __init__(self, atomgroup, reference_coords=None, **kwargs): - super(FastCovarianceAnalysis, self).__init__(atomgroup.universe.trajectory, **kwargs) + super().__init__(atomgroup.universe.trajectory, **kwargs) self.atomgroup = atomgroup self.n_atoms = len(atomgroup) self.n_dims = self.n_atoms * 3 @@ -305,24 +219,18 @@ def compute_covariance_matrix(self, universe, selection='protein', average_coord Covariance matrix and mean coordinates """ print("Computing covariance matrix...") - - # Set number of threads for BLAS operations - if n_jobs is None: - n_jobs = multiprocessing.cpu_count() - - os.environ['OMP_NUM_THREADS'] = str(n_jobs) - + # Select atoms for analysis atoms = universe.select_atoms(selection) - + # Initialize and run the optimized analysis - cov_analysis = self.FastCovarianceAnalysis(atoms, reference_coords=average_coords, n_jobs=n_jobs) + cov_analysis = self.FastCovarianceAnalysis(atoms, reference_coords=average_coords) cov_analysis.run() return cov_analysis.covariance_matrix, cov_analysis.mean_coords.flatten() @staticmethod - @jit(nopython=True) + @njit def _compute_kl_divergences(gevec, geigval, mean_diff, rank): """Compute KL divergences with Numba acceleration""" kl = np.zeros(rank, dtype=np.float64) @@ -426,24 +334,7 @@ def simultaneous_diagonalization(self, cov_a, cov_b, mean_a, mean_b): 'sum_kl': kl_sum, 'sum_kl_m': np.sum(kl_m) } - - @staticmethod - @jit(nopython=True, parallel=True) - def _project_coordinates(coords_array, mean_coords, gevec, first_vec, n_vecs): - """Optimized projection of coordinates onto eigenvectors""" - n_frames = coords_array.shape[0] - projections = np.zeros((n_frames, n_vecs), dtype=np.float64) - - for i in prange(n_frames): - coords = coords_array[i].flatten() - deviation = coords - mean_coords - - for j in range(n_vecs): - vec_idx = first_vec + j - projections[i, j] = np.dot(deviation, gevec[:, vec_idx]) - - return projections - + def project_trajectory(self, universe, selection, gevec, mean_coords, first_vec=0, last_vec=None, batch_size=1000, n_jobs=None): """Project trajectory onto generalized eigenvectors (optimized) @@ -462,18 +353,11 @@ def project_trajectory(self, universe, selection, gevec, mean_coords, first_vec= """ print("Projecting trajectory onto eigenvectors...") start_time = time.time() - - if n_jobs is None: - n_jobs = multiprocessing.cpu_count() - - # Set environment variables for optimal performance - os.environ['OMP_NUM_THREADS'] = str(n_jobs) - + # Select atoms for analysis atoms = universe.select_atoms(selection) n_atoms = len(atoms) - n_dims = n_atoms * 3 - + # Set number of eigenvectors to use if last_vec is None: last_vec = gevec.shape[1] - 1 @@ -533,43 +417,45 @@ def compute_residue_indices(self, atoms, residues): return residue_indices @staticmethod - @jit(nopython=True) + @njit def _compute_interaction_matrix(gevec_subset, kl_subset, residue_indices, n_res, n_atoms, n_vecs): - """Numba-accelerated calculation of interaction matrix""" + """Numba-accelerated calculation of interaction matrix. + + The contribution between residues i and j for one eigenvector is + sum_{a in i} sum_{b in j} (vec[a] . vec[b]) = G_i . G_j, where G_r is the + sum of the per-atom vectors over residue r. Collapsing the double atom + sum into per-residue group vectors turns the original + O(n_res^2 * atoms_per_res^2) loop into O(n_res^2), with identical results. + """ interaction_matrix = np.zeros((n_res, n_res), dtype=np.float64) atom_contrib = np.zeros(n_atoms, dtype=np.float64) - + # Pre-reshape gevec for better memory access patterns gevec_reshaped = gevec_subset.reshape(n_vecs, n_atoms, 3) - - # Loop over eigenvectors first for better cache locality + for k in range(n_vecs): vec = gevec_reshaped[k] weight = kl_subset[k] - - # Update per-residue contributions + + # Per-residue group vectors: G_r = sum of vec over the residue's atoms. + group = np.zeros((n_res, 3)) for i in range(n_res): idx_i = residue_indices[i] - + for a_idx in range(len(idx_i)): + group[i] += vec[idx_i[a_idx]] + + for i in range(n_res): for j in range(n_res): - idx_j = residue_indices[j] - - # Compute contribution between residue i and j - contrib = 0.0 - for a_idx in range(len(idx_i)): - a = idx_i[a_idx] - for b_idx in range(len(idx_j)): - b = idx_j[b_idx] - contrib += np.dot(vec[a], vec[b]) * weight - + contrib = weight * np.dot(group[i], group[j]) interaction_matrix[i, j] += contrib - - # Update per-atom contribution for diagonal elements + + # Spread the diagonal contribution evenly over the residue's atoms. if i == j: + idx_i = residue_indices[i] contrib_per_atom = contrib / len(idx_i) for a_idx in range(len(idx_i)): atom_contrib[idx_i[a_idx]] += contrib_per_atom - + return interaction_matrix, atom_contrib def compute_interaction_map(self, universe, selection, gevec, kl, first_vec=0, last_vec=None, n_jobs=None): @@ -589,21 +475,15 @@ def compute_interaction_map(self, universe, selection, gevec, kl, first_vec=0, l """ print("Computing interaction map...") start_time = time.time() - - if n_jobs is None: - n_jobs = multiprocessing.cpu_count() - - # Set environment variables for optimal performance - os.environ['OMP_NUM_THREADS'] = str(n_jobs) - + # Select atoms for analysis atoms = universe.select_atoms(selection) n_atoms = len(atoms) - + # Set number of eigenvectors to use if last_vec is None: last_vec = gevec.shape[1] - 1 - + n_vecs = last_vec - first_vec + 1 gevec_subset = gevec[:, first_vec:last_vec+1].copy() kl_subset = kl[first_vec:last_vec+1].copy() From 7d2e365f7e7b5f75031ae21464c30f653dbf6acb Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 1 Jun 2026 19:08:35 +0000 Subject: [PATCH 2/3] Add .gitignore for Python artifacts and analysis outputs https://claude.ai/code/session_018SSc1GoJMfc7Ey8n9EKgrK --- .gitignore | 11 +++++++++++ 1 file changed, 11 insertions(+) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..56b6f87 --- /dev/null +++ b/.gitignore @@ -0,0 +1,11 @@ +# Python build artifacts +__pycache__/ +*.py[cod] +*.egg-info/ +.pytest_cache/ + +# RPCA analysis outputs +*.npy +contributions.pdb +*.png +!figs/*.png From a2526f1b95af1b33aa3dc26b0cde3d662eb8cd44 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 1 Jun 2026 19:43:24 +0000 Subject: [PATCH 3/3] Add rpca PoC package: C sdiag core bound to NumPy via Cython, no GROMACS/MDA Proof-of-concept Python package that reuses the project's original C numerical core for relative PCA without any GROMACS or MDAnalysis dependency. - rpca_api.c/.h: thin buffer-in/buffer-out C entry point that calls Simultaneous_Diagonalization() from src/linear_algebra.c and adds the relative-PCA KL scoring/ordering glue. - _sdiag.pyx: zero-copy, nogil Cython binding exposing sdiag() on numpy arrays; returns geigval/gevec/kl/kl_m/acc_kl/rank. - setup.py links the package's C API against ../src/linear_algebra.c (single source of truth) and LAPACK/BLAS only. - Tests validate the defining property (G^T A G = I, G^T B G = diag), the KL formula and ordering, and cross-check eigenvalues against scipy.linalg.eigh. The built extension links only openblas/libm/libc/ libgfortran - verified no gromacs/mdanalysis. https://claude.ai/code/session_018SSc1GoJMfc7Ey8n9EKgrK --- pyrpca/.gitignore | 9 +++ pyrpca/README.md | 90 +++++++++++++++++++++++++ pyrpca/pyproject.toml | 14 ++++ pyrpca/setup.py | 48 +++++++++++++ pyrpca/src/rpca/__init__.py | 13 ++++ pyrpca/src/rpca/_csrc/rpca_api.c | 112 +++++++++++++++++++++++++++++++ pyrpca/src/rpca/_csrc/rpca_api.h | 60 +++++++++++++++++ pyrpca/src/rpca/_sdiag.pyx | 84 +++++++++++++++++++++++ pyrpca/tests/test_sdiag.py | 82 ++++++++++++++++++++++ 9 files changed, 512 insertions(+) create mode 100644 pyrpca/.gitignore create mode 100644 pyrpca/README.md create mode 100644 pyrpca/pyproject.toml create mode 100644 pyrpca/setup.py create mode 100644 pyrpca/src/rpca/__init__.py create mode 100644 pyrpca/src/rpca/_csrc/rpca_api.c create mode 100644 pyrpca/src/rpca/_csrc/rpca_api.h create mode 100644 pyrpca/src/rpca/_sdiag.pyx create mode 100644 pyrpca/tests/test_sdiag.py diff --git a/pyrpca/.gitignore b/pyrpca/.gitignore new file mode 100644 index 0000000..2dc5d2a --- /dev/null +++ b/pyrpca/.gitignore @@ -0,0 +1,9 @@ +build/ +*.so +*.o +# Cython-generated C +src/rpca/_sdiag.c +__pycache__/ +*.pyc +.pytest_cache/ +*.egg-info/ diff --git a/pyrpca/README.md b/pyrpca/README.md new file mode 100644 index 0000000..bb0acb8 --- /dev/null +++ b/pyrpca/README.md @@ -0,0 +1,90 @@ +# rpca (Python package) — proof of concept + +A Python/NumPy package that reuses the **original C numerical core** of this +project for speed, with **no dependency on GROMACS or MDAnalysis**. + +This is the first slice of a larger plan (see *Roadmap* below). It wires up the +highest-value, lowest-risk piece — the simultaneous diagonalization + relative +PCA Kullback–Leibler ranking — and proves the whole approach end to end. + +## Why this works + +The numerical core (`src/linear_algebra.c`) is pure C + LAPACK/BLAS: it does not +include a single GROMACS header. The GROMACS coupling in this project lives only +in (a) the four `main()` drivers (`gmx_*.c`, `SIMDAIG.c`) and (b) typedefs/PBC +helpers pulled in through some headers — none of which the linear-algebra core +touches. So the core can be compiled and called directly from Python. + +``` +Python (numpy) ──▶ Cython (_sdiag.pyx) ──▶ rpca_api.c ──▶ Simultaneous_Diagonalization() + (src/linear_algebra.c, LAPACK) +``` + +The built extension links **only** `libopenblas`, `libm`, `libc`, `libgfortran` +— verified with `ldd`. No GROMACS, no MDAnalysis. + +## Layout + +``` +pyrpca/ +├── pyproject.toml build metadata (setuptools + Cython + numpy) +├── setup.py compiles the extension against ../src/linear_algebra.c +├── src/rpca/ +│ ├── __init__.py +│ ├── _sdiag.pyx Cython binding (numpy <-> C, zero-copy, nogil) +│ └── _csrc/ +│ ├── rpca_api.h thin buffer-in/buffer-out C API +│ └── rpca_api.c calls the core + KL scoring/ordering glue +└── tests/test_sdiag.py validates G^T A G = I, G^T B G = diag, KL, vs SciPy +``` + +## Build & test + +Requires a C compiler, LAPACK/BLAS dev libraries, and Python build deps: + +```bash +# system (Debian/Ubuntu) +sudo apt-get install -y libopenblas-dev liblapack-dev +# python +pip install numpy cython pytest scipy + +cd pyrpca +python setup.py build_ext --inplace +PYTHONPATH=src python -m pytest tests/ -v +``` + +## Usage + +```python +import numpy as np +from rpca import sdiag + +res = sdiag(cov_a, cov_b, mean_a, mean_b) # symmetric (n, n) covariances + means +res["geigval"] # generalized eigenvalues (var_B / var_A per mode) +res["gevec"] # (n, rank) eigenvectors; column k is mode k +res["kl"] # per-mode KL divergence, descending +res["kl_m"] # mean-shift contribution to KL +res["acc_kl"] # cumulative % of total KL +res["rank"] # effective rank +``` + +## Roadmap + +1. **[done]** De-GROMACS the core for the sdiag path; Cython binding; tests. +2. Stream covariance + GPA over `.xtc` directly in C (reuse `covariance.c`, + `fitting.c` behind a small `compat.h` that replaces the GROMACS typedefs and + makes PBC optional). +3. Restore covariance-weighted Mahalanobis fitting (`CWfitting.c` + the BFGS + minimizer) — a step the earlier numba port dropped. +4. Pure-Python `io` (gro/pdb parsers, PDB B-factor writer) + thin `xtc` wrapper + over the bundled `xdrfile`, plus a minimal atom-selection layer. +5. High-level `RPCA` pipeline class + `argparse` CLI mirroring the old tools. +6. Switch packaging to `scikit-build-core`/CMake (or `meson-python`) and vendor + the C core into the package for distributable wheels. + +## Notes + +- For the PoC the build references `../src/linear_algebra.c` so there is a single + source of truth. Productionising (step 6) vendors a copy into the package. +- The KL formula matches the published RPCA definition (Ahmad et al., *JCTC* + 2019) and the existing Python reference implementation. diff --git a/pyrpca/pyproject.toml b/pyrpca/pyproject.toml new file mode 100644 index 0000000..a91fb81 --- /dev/null +++ b/pyrpca/pyproject.toml @@ -0,0 +1,14 @@ +[build-system] +requires = ["setuptools>=64", "wheel", "Cython>=3.0", "numpy>=1.23"] +build-backend = "setuptools.build_meta" + +[project] +name = "rpca" +version = "0.0.1" +description = "Relative Principal Components Analysis backed by the original C core (no GROMACS/MDAnalysis)" +readme = "README.md" +requires-python = ">=3.9" +dependencies = ["numpy>=1.23"] + +[project.optional-dependencies] +test = ["pytest", "scipy"] diff --git a/pyrpca/setup.py b/pyrpca/setup.py new file mode 100644 index 0000000..83b62bb --- /dev/null +++ b/pyrpca/setup.py @@ -0,0 +1,48 @@ +"""Build script for the rpca extension. + +The extension links the package's thin C API (rpca_api.c) against the existing +numerical core in the parent repository (src/linear_algebra.c), so there is a +single source of truth for the algorithms. Productionising this would vendor a +copy of the core into the package; for the proof of concept we reference it in +place. +""" +import os + +import numpy as np +from setuptools import Extension, setup +from Cython.Build import cythonize + +HERE = os.path.dirname(os.path.abspath(__file__)) +REPO = os.path.dirname(HERE) # parent RPCA repository + +CORE_SRC = os.path.join(REPO, "src", "linear_algebra.c") +CORE_INC = os.path.join(REPO, "include") +PKG_CSRC = os.path.join(HERE, "src", "rpca", "_csrc") + +# OpenBLAS bundles both CBLAS and the LAPACK routine (dsyevr_) the core needs; +# liblapack is listed as a fallback provider of the Fortran symbols. +blas_libs = ["openblas", "lapack", "m"] + +extensions = [ + Extension( + name="rpca._sdiag", + sources=[ + os.path.join(HERE, "src", "rpca", "_sdiag.pyx"), + os.path.join(PKG_CSRC, "rpca_api.c"), + CORE_SRC, + ], + include_dirs=[CORE_INC, PKG_CSRC, np.get_include()], + libraries=blas_libs, + extra_compile_args=["-O3"], + define_macros=[("NPY_NO_DEPRECATED_API", "NPY_1_7_API_VERSION")], + ) +] + +setup( + name="rpca", + version="0.0.1", + package_dir={"": "src"}, + packages=["rpca"], + ext_modules=cythonize(extensions, language_level=3), + zip_safe=False, +) diff --git a/pyrpca/src/rpca/__init__.py b/pyrpca/src/rpca/__init__.py new file mode 100644 index 0000000..fbfd390 --- /dev/null +++ b/pyrpca/src/rpca/__init__.py @@ -0,0 +1,13 @@ +"""rpca - Relative Principal Components Analysis backed by the original C core. + +This package reuses the proven C numerical routines (simultaneous +diagonalization, whitening, LAPACK eigendecomposition) from the parent project +and exposes them to Python/NumPy without any GROMACS or MDAnalysis dependency. + +Proof-of-concept stage: only the simultaneous-diagonalization step is wired up. +""" + +from ._sdiag import sdiag + +__all__ = ["sdiag"] +__version__ = "0.0.1" diff --git a/pyrpca/src/rpca/_csrc/rpca_api.c b/pyrpca/src/rpca/_csrc/rpca_api.c new file mode 100644 index 0000000..4550f5d --- /dev/null +++ b/pyrpca/src/rpca/_csrc/rpca_api.c @@ -0,0 +1,112 @@ +/* + * rpca_api.c - implementation of the thin C entry points. + * + * The heavy lifting (whitening transformation + LAPACK eigendecomposition) is + * done by Simultaneous_Diagonalization() in the existing core; this file only + * adds the lightweight KL scoring and ordering glue and marshals the result + * into caller-owned buffers. + */ +#include +#include +#include + +#include "linear_algebra.h" +#include "rpca_api.h" + +/* Indices that sort `key` in descending order. r is the rank of a covariance + * matrix (small relative to a trajectory), so an insertion sort is plenty. */ +static void argsort_desc(const double *key, int r, int *idx) +{ + int i, j, tmp; + for (i = 0; i < r; i++) idx[i] = i; + for (i = 1; i < r; i++) { + tmp = idx[i]; + j = i - 1; + while (j >= 0 && key[idx[j]] < key[tmp]) { + idx[j + 1] = idx[j]; + j--; + } + idx[j + 1] = tmp; + } +} + +int rpca_sdiag(int n, + const double *A, const double *B, + const double *mean_a, const double *mean_b, + int algo, int verbose, + double *geigval, double *gevec, + double *kl, double *kl_m, double *acc_kl, + int *rank_out) +{ + double *gv = NULL; /* core-allocated eigenvalues, length r */ + double *ge = NULL; /* core-allocated eigenvectors, column-major n x r */ + double *dav = NULL; /* mean shift (state B relative to state A) */ + double *kl_tmp = NULL, *klm_tmp = NULL; + int *order = NULL; + int r = 0, i, k, rc = 0; + + if (n <= 0 || A == NULL || B == NULL || rank_out == NULL) return 1; + + dav = (double *) malloc((size_t) n * sizeof(double)); + if (!dav) { rc = 2; goto cleanup; } + for (i = 0; i < n; i++) { + double ma = mean_a ? mean_a[i] : 0.0; + double mb = mean_b ? mean_b[i] : 0.0; + dav[i] = mb - ma; + } + + /* The core takes non-const pointers but copies A/B internally before any + * destructive LAPACK call, so the caller's matrices are left untouched. */ + if (algo == 1) { + Simultaneous_Diagonalization_subspacing(n, (double *) A, (double *) B, + &gv, &ge, &r, dav, verbose); + } else { + Simultaneous_Diagonalization(n, (double *) A, (double *) B, + &gv, &ge, &r, verbose); + } + if (r <= 0 || gv == NULL || ge == NULL) { rc = 3; goto cleanup; } + + kl_tmp = (double *) malloc((size_t) r * sizeof(double)); + klm_tmp = (double *) malloc((size_t) r * sizeof(double)); + order = (int *) malloc((size_t) r * sizeof(int)); + if (!kl_tmp || !klm_tmp || !order) { rc = 4; goto cleanup; } + + /* Per-mode KL divergence: + * kl_m = 0.5 * (g_k . (mean_b - mean_a))^2 + * kl = 0.5 * (geigval_k - log(geigval_k) - 1) + kl_m */ + for (k = 0; k < r; k++) { + const double *vk = ge + (size_t) k * n; /* column k */ + double proj = 0.0; + for (i = 0; i < n; i++) proj += vk[i] * dav[i]; + klm_tmp[k] = 0.5 * proj * proj; + kl_tmp[k] = 0.5 * (gv[k] - log(gv[k]) - 1.0) + klm_tmp[k]; + } + + argsort_desc(kl_tmp, r, order); + + double kl_sum = 0.0; + for (k = 0; k < r; k++) kl_sum += kl_tmp[k]; + + double run = 0.0; + for (k = 0; k < r; k++) { + int src = order[k]; + geigval[k] = gv[src]; + kl[k] = kl_tmp[src]; + kl_m[k] = klm_tmp[src]; + run += kl_tmp[src]; + acc_kl[k] = (kl_sum != 0.0) ? (run / kl_sum * 100.0) : 0.0; + memcpy(gevec + (size_t) k * n, ge + (size_t) src * n, + (size_t) n * sizeof(double)); + } + + *rank_out = r; + +cleanup: + free(dav); + free(kl_tmp); + free(klm_tmp); + free(order); + free(gv); /* allocated by the core with malloc */ + free(ge); + return rc; +} diff --git a/pyrpca/src/rpca/_csrc/rpca_api.h b/pyrpca/src/rpca/_csrc/rpca_api.h new file mode 100644 index 0000000..6489c60 --- /dev/null +++ b/pyrpca/src/rpca/_csrc/rpca_api.h @@ -0,0 +1,60 @@ +/* + * rpca_api.h - thin, dependency-free C entry points for the Python binding. + * + * These wrap the proven RPCA numerical core (src/linear_algebra.c) behind a + * buffer-in / buffer-out interface that is trivial to call from Cython/cffi/ + * ctypes: every output is written into a caller-allocated array, so the + * binding never has to free C-owned memory. + * + * No GROMACS, no MDAnalysis - only LAPACK/BLAS (via the core) and libm. + */ +#ifndef RPCA_API_H +#define RPCA_API_H + +#ifdef __cplusplus +extern "C" { +#endif + +/* + * Simultaneous diagonalization of two covariance matrices followed by the + * relative-PCA Kullback-Leibler ranking (Ahmad et al., JCTC 2019). + * + * Finds the generalized eigenbasis G such that G^T A G = I and + * G^T B G = diag(geigval), then scores each mode by the KL divergence of + * state B from state A and returns the modes ordered by descending KL. + * + * Inputs: + * n - dimension (= 3 * n_atoms) + * A, B - n*n symmetric covariance matrices (state A / state B). + * Row- or column-major is equivalent because they are symmetric. + * Neither is modified. + * mean_a, - length-n mean coordinate vectors. Either may be NULL (treated + * mean_b as zero), in which case the mean-shift term vanishes. + * algo - 0: standard, 1: mean-fluctuation subspacing. + * verbose - passed through to the core (0 = silent). + * + * Outputs (caller-allocated; size for the worst case r == n): + * geigval - length >= n; first r entries are the generalized eigenvalues. + * gevec - length >= n*n; first n*r entries are the eigenvectors in + * COLUMN-MAJOR n x r layout (mode k = gevec[k*n .. k*n+n-1]). + * kl - length >= n; per-mode KL divergence, descending. + * kl_m - length >= n; per-mode KL contribution from the mean shift. + * acc_kl - length >= n; cumulative percentage of total KL. + * rank_out - the effective rank r (<= n). Only the first r entries / columns + * of the output buffers are meaningful. + * + * Returns 0 on success, non-zero on error. + */ +int rpca_sdiag(int n, + const double *A, const double *B, + const double *mean_a, const double *mean_b, + int algo, int verbose, + double *geigval, double *gevec, + double *kl, double *kl_m, double *acc_kl, + int *rank_out); + +#ifdef __cplusplus +} +#endif + +#endif /* RPCA_API_H */ diff --git a/pyrpca/src/rpca/_sdiag.pyx b/pyrpca/src/rpca/_sdiag.pyx new file mode 100644 index 0000000..c0d35ad --- /dev/null +++ b/pyrpca/src/rpca/_sdiag.pyx @@ -0,0 +1,84 @@ +# cython: language_level=3 +"""Cython binding for the RPCA simultaneous-diagonalization core.""" + +import numpy as np +cimport numpy as cnp + +cdef extern from "rpca_api.h": + int rpca_sdiag(int n, + const double *A, const double *B, + const double *mean_a, const double *mean_b, + int algo, int verbose, + double *geigval, double *gevec, + double *kl, double *kl_m, double *acc_kl, + int *rank_out) nogil + +cnp.import_array() + + +def sdiag(cov_a, cov_b, mean_a=None, mean_b=None, int algo=0, int verbose=0): + """Simultaneous diagonalization + relative-PCA KL ranking of two states. + + Parameters + ---------- + cov_a, cov_b : (n, n) array_like + Symmetric covariance matrices of state A and state B. + mean_a, mean_b : (n,) array_like, optional + Mean coordinate vectors. If omitted the mean-shift term is zero. + algo : int + 0 = standard diagonalization, 1 = mean-fluctuation subspacing. + verbose : int + Verbosity passed to the C core. + + Returns + ------- + dict with keys ``geigval``, ``gevec`` (n x rank, column k = mode k), + ``kl``, ``kl_m``, ``acc_kl`` (all length ``rank``), ``rank``, + ``sum_kl`` and ``sum_kl_m``. + """ + cdef cnp.ndarray[double, ndim=2, mode="c"] A = np.ascontiguousarray(cov_a, dtype=np.float64) + cdef cnp.ndarray[double, ndim=2, mode="c"] B = np.ascontiguousarray(cov_b, dtype=np.float64) + cdef int n = A.shape[0] + + if A.shape[1] != n or B.shape[0] != n or B.shape[1] != n: + raise ValueError("cov_a and cov_b must be square matrices of the same size") + + cdef cnp.ndarray[double, ndim=1, mode="c"] ma + cdef cnp.ndarray[double, ndim=1, mode="c"] mb + ma = (np.zeros(n, dtype=np.float64) if mean_a is None + else np.ascontiguousarray(mean_a, dtype=np.float64).ravel()) + mb = (np.zeros(n, dtype=np.float64) if mean_b is None + else np.ascontiguousarray(mean_b, dtype=np.float64).ravel()) + if ma.shape[0] != n or mb.shape[0] != n: + raise ValueError("mean vectors must have length n") + + cdef cnp.ndarray[double, ndim=1, mode="c"] geigval = np.zeros(n, dtype=np.float64) + cdef cnp.ndarray[double, ndim=1, mode="c"] gevec = np.zeros(n * n, dtype=np.float64) + cdef cnp.ndarray[double, ndim=1, mode="c"] kl = np.zeros(n, dtype=np.float64) + cdef cnp.ndarray[double, ndim=1, mode="c"] kl_m = np.zeros(n, dtype=np.float64) + cdef cnp.ndarray[double, ndim=1, mode="c"] acc_kl = np.zeros(n, dtype=np.float64) + cdef int rank = 0 + cdef int status + + with nogil: + status = rpca_sdiag(n, &A[0, 0], &B[0, 0], &ma[0], &mb[0], + algo, verbose, + &geigval[0], &gevec[0], + &kl[0], &kl_m[0], &acc_kl[0], &rank) + if status != 0: + raise RuntimeError(f"rpca_sdiag failed with status {status}") + + r = rank + # Core writes column-major n x r: column k occupies gevec[k*n : (k+1)*n]. + gevec_2d = np.ascontiguousarray(gevec[:n * r].reshape((r, n)).T) + + return { + "geigval": geigval[:r].copy(), + "gevec": gevec_2d, + "kl": kl[:r].copy(), + "kl_m": kl_m[:r].copy(), + "acc_kl": acc_kl[:r].copy(), + "rank": r, + "sum_kl": float(kl[:r].sum()), + "sum_kl_m": float(kl_m[:r].sum()), + } diff --git a/pyrpca/tests/test_sdiag.py b/pyrpca/tests/test_sdiag.py new file mode 100644 index 0000000..dd59e45 --- /dev/null +++ b/pyrpca/tests/test_sdiag.py @@ -0,0 +1,82 @@ +"""Tests for rpca.sdiag. + +The defining property of simultaneous diagonalization gives a self-contained +ground truth (no second library required): for the generalized eigenbasis G, + + G^T A G = I and G^T B G = diag(geigval). + +We also check the KL bookkeeping and the descending ordering, and - when SciPy +is available - cross-check the eigenvalues against scipy.linalg.eigh. +""" +import numpy as np +import pytest + +from rpca import sdiag + + +def _random_spd(n, rng, scale=1.0): + """A random symmetric positive-definite n x n matrix.""" + M = rng.standard_normal((n, n)) + return scale * (M @ M.T) + n * np.eye(n) + + +def test_simultaneous_diagonalization_property(): + rng = np.random.default_rng(0) + n = 15 + A = _random_spd(n, rng) + B = _random_spd(n, rng, scale=2.0) + + res = sdiag(A, B) + G = res["gevec"] + r = res["rank"] + + assert G.shape == (n, r) + assert r == n # both inputs full rank + + GtAG = G.T @ A @ G + GtBG = G.T @ B @ G + + np.testing.assert_allclose(GtAG, np.eye(r), atol=1e-8) + # off-diagonal of G^T B G must vanish; diagonal equals the eigenvalues + np.testing.assert_allclose(GtBG - np.diag(np.diag(GtBG)), 0.0, atol=1e-8) + np.testing.assert_allclose(np.diag(GtBG), res["geigval"], atol=1e-8) + + +def test_kl_formula_and_ordering(): + rng = np.random.default_rng(1) + n = 12 + A = _random_spd(n, rng) + B = _random_spd(n, rng, scale=1.5) + mean_a = rng.standard_normal(n) + mean_b = rng.standard_normal(n) + + res = sdiag(A, B, mean_a, mean_b) + G, gv = res["gevec"], res["geigval"] + + # KL must be sorted descending + assert np.all(np.diff(res["kl"]) <= 1e-12) + + # reproduce the published per-mode KL from the returned pieces + proj = G.T @ (mean_b - mean_a) + kl_m_expected = 0.5 * proj ** 2 + kl_expected = 0.5 * (gv - np.log(gv) - 1.0) + kl_m_expected + + np.testing.assert_allclose(res["kl_m"], kl_m_expected, atol=1e-9) + np.testing.assert_allclose(res["kl"], kl_expected, atol=1e-9) + + # accumulated KL ends at 100% + np.testing.assert_allclose(res["acc_kl"][-1], 100.0, atol=1e-6) + assert res["sum_kl"] == pytest.approx(res["kl"].sum()) + + +def test_eigenvalues_match_scipy(): + scipy_linalg = pytest.importorskip("scipy.linalg") + rng = np.random.default_rng(2) + n = 10 + A = _random_spd(n, rng) + B = _random_spd(n, rng, scale=3.0) + + res = sdiag(A, B) + # generalized problem B v = lambda A v -> eigenvalues are var_B/var_A + ref = np.sort(scipy_linalg.eigh(B, A, eigvals_only=True)) + np.testing.assert_allclose(np.sort(res["geigval"]), ref, rtol=1e-6, atol=1e-8)