From c884ebc62cac02d73b0e816264f5de5b64438db6 Mon Sep 17 00:00:00 2001 From: Akshit Boora Date: Tue, 18 Aug 2026 16:23:47 +0530 Subject: [PATCH] ENH: add subselection kwarg to AlignTraj (Issue #5380) Add a new 'subselection' keyword argument to AlignTraj.__init__ that accepts a selection string or AtomGroup. When provided, only those atoms are written to the output trajectory (enabling a performance/memory trade-off), while the superposition/RMSD fit is still computed using the atoms defined by 'select'. Also add a module-level guide in MDAnalysis.analysis.align documenting three usage patterns for aligning sub-systems with different memory characteristics. Changes: - package/MDAnalysis/analysis/align.py: add subselection param + docs - testsuite/MDAnalysisTests/analysis/test_align.py: 4 new unit tests - package/CHANGELOG: Enhancements + Documentation entries - package/AUTHORS: add Akshit Boora Closes #5380 --- .gitignore | 2 + package/AUTHORS | 1 + package/CHANGELOG | 12 ++- package/MDAnalysis/analysis/align.py | 102 +++++++++++++++++- .../MDAnalysisTests/analysis/test_align.py | 67 ++++++++++++ 5 files changed, 180 insertions(+), 4 deletions(-) diff --git a/.gitignore b/.gitignore index ed6eebde889..be10b5985f9 100644 --- a/.gitignore +++ b/.gitignore @@ -53,3 +53,5 @@ benchmarks/results .idea .vscode *.lock +.venv/ +.pytest_cache/ diff --git a/package/AUTHORS b/package/AUTHORS index ed82321a9e8..391111132be 100644 --- a/package/AUTHORS +++ b/package/AUTHORS @@ -284,6 +284,7 @@ Chronological list of authors - Sai Udayagiri - Apoorva Verma - Aryaman Chaudhri + - Akshit Boora External code ------------- diff --git a/package/CHANGELOG b/package/CHANGELOG index 3e7c3d309d3..da557a96378 100644 --- a/package/CHANGELOG +++ b/package/CHANGELOG @@ -18,7 +18,7 @@ The rules for this file: spyke7, talagayev, tanii1125, BradyAJohnston, hejamu, jeremyleung521, harshitgajjela-droid, kunjsinha, aygarwal, jauy123, Dreamstick9, ollyfutur, Amarendra22, charity-g, ParthUppal523, apoorva-01, RMeli, - raulloiscuns, Aryaman-Chaudhri + raulloiscuns, Aryaman-Chaudhri, akshitboora * 2.11.0 @@ -103,6 +103,10 @@ Enhancements * Adds support for parsing `.tpr` files produced by GROMACS 2026.0 * Enables parallelization for analysis.diffusionmap.DistanceMatrix (Issue #4679, PR #4745) + * `AlignTraj` now accepts a `subselection` keyword argument (a selection + string or AtomGroup) that restricts which atoms are written to the output + trajectory, enabling a performance/memory trade-off when only a subset + of atoms is of interest (Issue #5380) Changes * The msd.py inside analysis is changed, and ProgressBar is implemented inside @@ -110,6 +114,12 @@ Changes Deprecations +Documentation + * Added a guide on "Performance and memory trade-offs when aligning + sub-systems" to `MDAnalysis.analysis.align` documenting how to use + `AlignTraj` with `subselection` and in-memory trajectories to avoid + reading/writing full-system coordinates (Issue #5380) + 10/18/25 IAlibay, orbeckst, BHM-Bob, TRY-ER, Abdulrahman-PROG, pbuslaev, yuxuanzhuang, yuyuan871111, tanishy7777, tulga-rdn, Gareth-elliott, diff --git a/package/MDAnalysis/analysis/align.py b/package/MDAnalysis/analysis/align.py index 3dc4f84fb1d..145aca37a8a 100644 --- a/package/MDAnalysis/analysis/align.py +++ b/package/MDAnalysis/analysis/align.py @@ -155,6 +155,69 @@ (See the documentation of the functions for this advanced usage.) +Performance and memory trade-offs when aligning sub-systems +----------------------------------------------------------- + +When working with large molecular systems (such as a protein solvated in +thousands of water molecules), you are often only interested in aligning +and analyzing a specific sub-system or domain (e.g. a solute protein or a +binding domain). Understanding the trade-offs between memory consumption, +disk I/O, and CPU performance is key to selecting the right workflow: + +1. **Lazy disk iteration vs in-memory representation (`in_memory`)**: + - By default (``in_memory=False``), MDAnalysis loads frames lazily + one-by-one from disk into memory. This uses minimal constant RAM, making + it suitable for very large trajectories or memory-constrained + environments, at the cost of disk I/O when reading and writing. + - Setting ``in_memory=True`` transfers the trajectory coordinates to an + in-memory numpy array + (:class:`~MDAnalysis.coordinates.memory.MemoryReader`), enabling fast + in-place transformations without writing to disk. However, this + requires approximately :math:`N_{\\text{frames}} \\times N_{\\text{atoms}} \\times 3 \\times 4\\text{ bytes}` + (or 8 bytes for double precision) of RAM. For a system with 500,000 + atoms and 10,000 frames, holding the entire trajectory in RAM would + require ~60 GB. + +2. **Aligning only a sub-system with `subselection`**: + By default, :class:`AlignTraj` computes the optimal rotation/translation + matrix using the atoms in `select`, but applies the transformation and + writes out coordinates for all atoms in the universe (:attr:`mobile.atoms`). + If you only need the aligned coordinates of a sub-domain (e.g. ``protein``), + use the `subselection` keyword:: + + >>> aligner = align.AlignTraj( + ... trj, ref, select="protein and name CA", subselection="protein", + ... filename="protein_aligned.dcd" + ... ) # doctest: +SKIP + >>> aligner.run() # doctest: +SKIP + + This calculates the superposition on the C-alpha atoms, but only transforms + and writes the ``protein`` atoms to disk, drastically reducing output file + size and disk write overhead. + +3. **In-memory alignment of sub-systems without writing files**: + If you want the speed benefits of an in-memory trajectory without the huge + RAM footprint of solvent atoms, extract the sub-system into an in-memory + Universe using :func:`~MDAnalysis.core.universe.Merge` and + :func:`~MDAnalysis.analysis.base.AnalysisFromFunction` + (see :ref:`creating-in-memory-trajectory-label`):: + + >>> from MDAnalysis.coordinates.memory import MemoryReader + >>> from MDAnalysis.analysis.base import AnalysisFromFunction + >>> + >>> domain_a = trj.select_atoms("protein") + >>> coords_a = AnalysisFromFunction( + ... lambda ag: ag.positions.copy(), domain_a + ... ).run().results['timeseries'] # doctest: +SKIP + >>> u_sub = mda.Merge(domain_a) # doctest: +SKIP + >>> u_sub.load_new(coords_a, format=MemoryReader) # doctest: +SKIP + >>> + >>> # Now align u_sub in-place in memory with a minimal RAM footprint + >>> aligner = align.AlignTraj( + ... u_sub, u_sub, select="name CA", in_memory=True + ... ).run() # doctest: +SKIP + + Functions and Classes --------------------- @@ -720,6 +783,7 @@ def __init__( strict=False, force=True, in_memory=False, + subselection=None, writer_kwargs=None, **kwargs, ): @@ -758,6 +822,18 @@ def __init__( performance substantially in some cases. In this case, no file is written out (`filename` and `prefix` are ignored) and only the coordinates of `mobile` are *changed in memory*. + subselection : str or AtomGroup or None (optional) + Apply the transformation and write out only this selection. + + ``None`` [default] + Apply to and write ``mobile.universe.atoms`` (i.e., all atoms + in the context of `mobile`). + *selection-string* + Apply to and write ``mobile.select_atoms(selection-string)``. + :class:`~MDAnalysis.core.groups.AtomGroup` + Apply to and write the arbitrary group of atoms. + + .. versionadded:: 2.11.0 verbose : bool (optional) Set logger to show more information and show detailed progress of the calculation if set to ``True``; the default is ``False``. @@ -792,8 +868,8 @@ def __init__( Notes ----- - If set to ``verbose=False``, it is recommended to wrap the statement - in a ``try ... finally`` to guarantee restoring of the log level in - the case of an exception. + in a ``try ... finally`` to guarantee restoring of the log level in + the case of an exception. - The ``in_memory`` option changes the `mobile` universe to an in-memory representation (see :mod:`MDAnalysis.coordinates.memory`) for the remainder of the Python session. If ``mobile.trajectory`` is @@ -826,6 +902,10 @@ def __init__( .. versionchanged:: 2.8.0 Added ``writer_kwargs`` kwarg dict to pass to the writer + .. versionadded:: 2.11.0 + Added ``subselection`` keyword to allow applying and writing the + transformation to a subset of atoms. + """ select = rms.process_selection(select) self.ref_atoms = reference.select_atoms(*select["reference"]) @@ -856,7 +936,23 @@ def __init__( logging.disable(logging.WARN) # store reference to mobile atoms - self.mobile = mobile.atoms + if subselection is None: + self.mobile = ( + mobile.universe.atoms + if hasattr(mobile, "universe") + else mobile.atoms + ) + elif isinstance(subselection, str): + self.mobile = mobile.select_atoms(subselection) + else: + try: + self.mobile = subselection.atoms + except AttributeError: + err = ( + "subselection must be a selection string, an " + "AtomGroup or Universe or None" + ) + raise TypeError(err) from None self.filename = filename diff --git a/testsuite/MDAnalysisTests/analysis/test_align.py b/testsuite/MDAnalysisTests/analysis/test_align.py index 5ac553c16c5..06cb7268824 100644 --- a/testsuite/MDAnalysisTests/analysis/test_align.py +++ b/testsuite/MDAnalysisTests/analysis/test_align.py @@ -456,6 +456,73 @@ def test_AlignTraj_writer_kwargs(self, universe, reference, tmpdir): ).run() assert_equal(aligner._writer.precision, 2) + def test_AlignTraj_subselection_str(self, universe, reference, tmpdir): + outfile = str(tmpdir.join("sub_align.dcd")) + reference.trajectory[-1] + sub_sel = "resid 1-50" + sub_atoms = universe.select_atoms(sub_sel) + + aligner = align.AlignTraj( + universe, + reference, + select="name CA and resid 1-50", + subselection=sub_sel, + filename=outfile, + ).run() + + # Output trajectory should only have atoms from subselection + sub_u = mda.Merge(sub_atoms) + sub_u.load_new(outfile) + assert sub_u.atoms.n_atoms == sub_atoms.n_atoms + assert len(sub_u.trajectory) == len(universe.trajectory) + + def test_AlignTraj_subselection_atomgroup( + self, universe, reference, tmpdir + ): + outfile = str(tmpdir.join("sub_ag_align.dcd")) + reference.trajectory[-1] + sub_atoms = universe.select_atoms("resid 1-50") + + align.AlignTraj( + universe, + reference, + select="name CA and resid 1-50", + subselection=sub_atoms, + filename=outfile, + ).run() + + sub_u = mda.Merge(sub_atoms) + sub_u.load_new(outfile) + assert sub_u.atoms.n_atoms == sub_atoms.n_atoms + + def test_AlignTraj_subselection_in_memory(self, universe, reference): + reference.trajectory[-1] + sub_sel = "resid 1-50" + sub_atoms = universe.select_atoms(sub_sel) + orig_positions = sub_atoms.positions.copy() + + aligner = align.AlignTraj( + universe, + reference, + select="name CA and resid 1-50", + subselection=sub_sel, + in_memory=True, + ).run() + + assert aligner.filename is None + assert not np.allclose(sub_atoms.positions, orig_positions) + + def test_AlignTraj_subselection_invalid(self, universe, reference): + with pytest.raises( + TypeError, match="subselection must be a selection string" + ): + align.AlignTraj( + universe, + reference, + select="name CA", + subselection=12345, + ) + def _assert_rmsd(self, reference, fitted, frame, desired, weights=None): fitted.trajectory[frame] rmsd = rms.rmsd(