ENH: add subselection keyword argument to AlignTraj (Issue #5380)

.gitignore

@@ -53,3 +53,5 @@ benchmarks/results
 .idea
 .vscode
 *.lock
+.venv/
+.pytest_cache/

package/AUTHORS

@@ -284,6 +284,7 @@ Chronological list of authors
   - Sai Udayagiri
   - Apoorva Verma
   - Aryaman Chaudhri
+  - Akshit Boora
 
 External code
 -------------

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,13 +103,23 @@ 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
    _conclude_simple and _conclude_fft functions instead of tqdm (Issue #5144, PR #5153)
 
 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,

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
 

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(