sawnergy 1.0.0__py3-none-any.whl

sawnergy/rin/rin_builder.py

@@ -0,0 +1,936 @@
+from __future__ import annotations
+
+# third-pary
+import numpy as np
+import threadpoolctl
+# built-in
+from pathlib import Path
+import logging
+from concurrent.futures import ThreadPoolExecutor
+import re
+import math
+import os
+# local
+from . import rin_util
+from .. import sawnergy_util
+
+# *----------------------------------------------------*
+#                        GLOBALS
+# *----------------------------------------------------*
+
+_logger = logging.getLogger(__name__) 
+
+# *----------------------------------------------------*
+#                        CLASSES
+# *----------------------------------------------------*
+
+class RINBuilder:
+    """Builds Residue Interaction Networks (RINs) from MD trajectories.
+
+    This class orchestrates running cpptraj to:
+    * compute per-frame, per-residue centers of mass (COMs),
+    * compute pairwise atomic non-bonded energies (electrostatics + van der Waals),
+    * project atomic interactions to residue-level interactions,
+    * post-process residue matrices (split, prune, remove self-interactions, symmetrize, normalize),
+    * package outputs into a compressed archive.
+
+    Args:
+        cpptraj_path (Path | str | None): Optional explicit path to the `cpptraj`
+            executable. If not provided, an attempt is made to locate it via
+            `rin_util.locate_cpptraj`.
+
+    Attributes:
+        cpptraj (Path): Resolved path to the `cpptraj` executable.
+    """
+
+    def __init__(self, cpptraj_path: Path | str | None = None):
+        _logger.debug("Initializing RINBuilder with cpptraj_path=%s", cpptraj_path)
+        if isinstance(cpptraj_path, str):
+            cpptraj_path = Path(cpptraj_path)
+
+        self.cpptraj = rin_util.locate_cpptraj(explicit=cpptraj_path, verify=True)
+        _logger.info("Using cpptraj at %s", self.cpptraj)
+
+    # ---------------------------------------------------------------------------------------------- #
+    #                                         CPPTRAJ HELPERS
+    # ---------------------------------------------------------------------------------------------- #
+
+    # NOTE: the pattern might be version specific
+    _elec_vdw_pattern = re.compile(r"""
+        ^\s*\[printdata\s+PW\[EMAP\]\s+square2d\s+noheader\]\s*\r?\n
+        ([0-9.eE+\-\s]+?)
+        ^\s*\[printdata\s+PW\[VMAP\]\s+square2d\s+noheader\]\s*\r?\n
+        ([0-9.eE+\-\s]+?)
+        (?=^\s*\[|^\s*TIME:|\Z)
+    """, re.MULTILINE | re.DOTALL | re.VERBOSE)
+
+    # NOTE: the pattern might be version specific
+    _com_block_pattern = lambda _, N: re.compile(rf"""
+        ^[^\n]*\bCOMZ{N}\b[^\n]*\n
+        ([0-9.eE+\-\s]+?)
+        (?=^\s*\[quit\]\s*$)
+    """, re.MULTILINE | re.DOTALL | re.VERBOSE)
+
+    _com_row_pattern = re.compile(r'^\s*\d+\s+(.+?)\s*$', re.MULTILINE)
+
+    def _get_number_frames(self,
+                        topology_file: str,
+                        trajectory_file: str,
+                        *,
+                        subprocess_env: dict | None = None,
+                        timeout: float | None = None) -> int:
+        """Return total number of frames in a trajectory.
+
+        Args:
+            topology_file (str): Path to topology (parm/prmtop) file.
+            trajectory_file (str): Path to trajectory file readable by cpptraj.
+            subprocess_env (dict | None): Optional environment overrides for the
+                cpptraj subprocess (e.g., thread settings).
+            timeout (float | None): Optional time limit (seconds) for the cpptraj call.
+
+        Returns:
+            int: Total number of frames.
+
+        Raises:
+            RuntimeError: If cpptraj output cannot be parsed into an integer.
+        """
+        _logger.debug("Requesting number of frames (topology=%s, trajectory=%s, timeout=%s)",
+                      topology_file, trajectory_file, timeout)
+        raw_out = rin_util.run_cpptraj(self.cpptraj,
+                                argv=["-p", topology_file, "-y", trajectory_file, "-tl"],
+                                env=subprocess_env,
+                                timeout=timeout)
+        _logger.debug("cpptraj -tl raw output: %r", raw_out)
+        out = raw_out.replace("Frames: ", "")
+        try:
+            frames = int(out)
+            _logger.info("Detected %d frames in trajectory %s", frames, trajectory_file)
+            return frames
+        except ValueError:
+            _logger.exception("Failed parsing frame count from cpptraj output: %r", out)
+            raise RuntimeError(f"Could not retrieve the number of frames from '{trajectory_file}' trajectory")
+    
+    def _get_atomic_composition_of_molecule(self,
+                                        topology_file: str,
+                                        trajectory_file: str,
+                                        molecule_id: int,
+                                        *,
+                                        subprocess_env: dict | None = None,
+                                        timeout: float | None = None) -> dict:
+        """Extract per-residue atomic composition for a molecule.
+
+        Runs a small cpptraj script that prints residue/atom membership and parses
+        it into a dictionary mapping residue IDs to sets of atom IDs.
+
+        Args:
+            topology_file (str): Path to topology (parm/prmtop) file.
+            trajectory_file (str): Path to trajectory file.
+            molecule_id (int): Molecule selector/ID used by cpptraj (e.g., `^1`).
+            subprocess_env (dict | None): Optional environment overrides for cpptraj.
+            timeout (float | None): Optional time limit (seconds).
+
+        Returns:
+            dict: Mapping ``{residue_id: set(atom_ids)}``.
+
+        Raises:
+            KeyError: If the requested `molecule_id` is not present in parsed output.
+        """
+        _logger.debug("Extracting atomic composition (molecule_id=%s)", molecule_id)
+        tmp_file: Path = sawnergy_util.temporary_file(prefix="mol_comp", suffix=".dat")
+        _logger.debug("Temporary composition file: %s", tmp_file)
+        try:
+            molecule_compositions_script = (self._load_data_from(topology_file, trajectory_file, 1, 1) + \
+                                       self._extract_molecule_compositions()) > str(tmp_file)
+            script = molecule_compositions_script.render()
+            _logger.debug("Running composition cpptraj script (len=%d chars)", len(script))
+            rin_util.run_cpptraj(self.cpptraj, script=script, env=subprocess_env, timeout=timeout)
+            hierarchy = rin_util.CpptrajMaskParser.hierarchize_molecular_composition(tmp_file)
+            if molecule_id not in hierarchy:
+                _logger.error("Molecule ID %s not found in composition hierarchy (available keys: %s)",
+                              molecule_id, list(hierarchy.keys())[:10])
+            comp = hierarchy[molecule_id]
+            _logger.info("Retrieved composition for molecule %s (residues=%d)", molecule_id, len(comp))
+            return comp
+        finally:
+            try:
+                tmp_file.unlink()
+                _logger.debug("Cleaned up temp file %s", tmp_file)
+            except OSError:
+                _logger.warning("Failed to remove temp file %s", tmp_file, exc_info=True)
+
+    def _calc_avg_atomic_interactions_in_frames(self,
+                                        frame_range: tuple[int, int],
+                                        topology_file: str,
+                                        trajectory_file: str,
+                                        molecule_id: int,
+                                        *,
+                                        subprocess_env: dict | None = None,
+                                        timeout: float | None = None) -> np.ndarray:
+        """Compute average atomic interaction matrix over a frame range.
+
+        Uses cpptraj `pairwise` to compute electrostatic (EMAP) and van der Waals
+        (VMAP) atomic interaction matrices, sums them, and returns the result.
+
+        Args:
+            frame_range (tuple[int, int]): Inclusive (start_frame, end_frame).
+            topology_file (str): Path to topology file.
+            trajectory_file (str): Path to trajectory file.
+            molecule_id (int): Molecule selector/ID for restricting computation.
+            subprocess_env (dict | None): Optional environment for cpptraj.
+            timeout (float | None): Optional time limit (seconds).
+
+        Returns:
+            np.ndarray: 2D array (n_atoms, n_atoms) of summed interactions. For the given
+            frame range, cpptraj's pairwise driver accumulates/averages internally and the
+            printed EMAP/VMAP “square2d” blocks correspond to the range specified.
+
+        Raises:
+            ValueError: If EMAP/VMAP blocks are not found, sizes mismatch, or the
+                block cannot be reshaped into a square matrix.
+        """
+        start_frame, end_frame = frame_range
+        _logger.debug("Calculating avg atomic interactions (frames=%s..%s, molecule_id=%s)",
+                      start_frame, end_frame, molecule_id)
+        interaction_calc_script = (
+            self._load_data_from(topology_file, trajectory_file, start_frame, end_frame)
+            + self._calc_nonbonded_energies_in_molecule(molecule_id)
+        ) > rin_util.PAIRWISE_STDOUT
+        script = interaction_calc_script.render()
+        _logger.debug("Running pairwise cpptraj script (len=%d chars)", len(script))
+        output = rin_util.run_cpptraj(self.cpptraj, script=script, env=subprocess_env, timeout=timeout)
+        _logger.debug("cpptraj pairwise output length: %d", len(output))
+
+        m = self._elec_vdw_pattern.search(output)
+        if not m:
+            _logger.error("EMAP/VMAP blocks not found in cpptraj output.")
+            raise ValueError("Could not find EMAP/VMAP blocks in cpptraj output. "
+                             "Potentially due to cpptraj version mismatch. "
+                             "The data retrieval is stable for CPPTRAJ of Version V6.18.1 (AmberTools)")
+        emap_txt, vmap_txt = m.group(1), m.group(2)
+
+        # Robust to wrapped lines: read all numbers, ignore line structure
+        emap_flat = np.fromstring(emap_txt, dtype=np.float32, sep=' ')
+        vmap_flat = np.fromstring(vmap_txt, dtype=np.float32, sep=' ')
+        _logger.debug("Parsed EMAP=%d values, VMAP=%d values", emap_flat.size, vmap_flat.size)
+
+        if emap_flat.size != vmap_flat.size:
+            _logger.error("Size mismatch EMAP(%d) vs VMAP(%d)", emap_flat.size, vmap_flat.size)
+            raise ValueError(f"EMAP and VMAP sizes differ: {emap_flat.size} vs {vmap_flat.size} "
+                             "Potentially due to cpptraj version mismatch. "
+                             "The data retrieval is stable for CPPTRAJ of Version V6.18.1 (AmberTools)")
+
+        n = int(round(math.sqrt(emap_flat.size)))
+        if n * n != emap_flat.size:
+            _logger.error("Non-square block: %d values (cannot form nxn)", emap_flat.size)
+            raise ValueError(f"Block is not square: {emap_flat.size} values (cannot reshape to nxn). "
+                             "Potentially due to cpptraj version mismatch. "
+                             "The data retrieval is stable for CPPTRAJ of Version V6.18.1 (AmberTools)")
+
+        elec_matrix = emap_flat.reshape(n, n)
+        vdw_matrix  = vmap_flat.reshape(n, n)
+        _logger.debug("Reshaped EMAP/VMAP to (%d, %d)", n, n)
+
+        interaction_matrix = (elec_matrix + vdw_matrix).astype(np.float32)
+        _logger.info("Computed interaction matrix shape: %s", interaction_matrix.shape)
+        return interaction_matrix
+
+    def _get_residue_COMs_per_frame(
+        self,
+        frame_range: tuple[int, int],
+        topology_file: str,
+        trajectory_file: str,
+        molecule_id: int,
+        number_residues: int,
+        *,
+        subprocess_env: dict | None = None,
+        timeout: float | None = None,
+    ) -> list[np.ndarray]:
+        """Compute per-residue COM coordinates for each frame.
+
+        Runs a cpptraj loop to compute ``vector COM<i>`` per residue and parses the
+        printed data.
+
+        Args:
+            frame_range (tuple[int, int]): Inclusive ``(start_frame, end_frame)`` (1-based).
+            topology_file (str): Path to topology file.
+            trajectory_file (str): Path to trajectory file.
+            molecule_id (int): Molecule selector/ID used by cpptraj for residue iteration.
+            number_residues (int): Expected residue count (used for validation).
+            subprocess_env (dict | None): Optional environment overrides for cpptraj.
+            timeout (float | None): Optional time limit (seconds) for cpptraj.
+
+        Returns:
+            list[np.ndarray]: A list of length ``n_frames`` where each element is a
+            ``(n_residues, 3)`` array of ``float32`` giving XYZ COM coordinates for that
+            frame. Element ``0`` corresponds to ``start_frame``, element ``-1`` to
+            ``end_frame``.
+
+        Raises:
+            ValueError: If ``frame_range`` is invalid (end < start).
+            RuntimeError: If the COM print block is missing/malformed or row sizes mismatch.
+        """
+        start_frame, end_frame = frame_range
+        _logger.debug("Getting COMs per frame (frames=%s..%s, residues=%d, molecule_id=%s)",
+                    start_frame, end_frame, number_residues, molecule_id)
+        if end_frame < start_frame:
+            _logger.error("Bad frame_range %s: end < start", frame_range)
+            raise ValueError(f"Bad frame_range {frame_range}: end < start")
+        number_frames = end_frame - start_frame + 1
+
+        # build and run cpptraj script
+        COM_script = (
+            self._load_data_from(topology_file, trajectory_file, start_frame, end_frame)
+            + self._compute_residue_COMs_in_molecule(molecule_id)
+        ) > rin_util.COM_STDOUT(molecule_id)
+        script_rendered = COM_script.render()
+        _logger.debug("Running COM cpptraj script (len=%d chars)", len(script_rendered))
+        output = rin_util.run_cpptraj(self.cpptraj, script=script_rendered,
+                                    env=subprocess_env, timeout=timeout)
+        _logger.debug("cpptraj COM output length: %d", len(output))
+
+        # extract COM block and per-frame rows
+        m = self._com_block_pattern(number_residues).search(output)
+        if not m:
+            _logger.error("COM print block not found in cpptraj output (expected COMZ%d header).",
+                        number_residues)
+            raise RuntimeError("Could not find COM print block in cpptraj output. "
+                             "Potentially due to cpptraj version mismatch. "
+                             "The data retrieval is stable for CPPTRAJ of Version V6.18.1 (AmberTools)")
+        block = m.group(1)
+        lines = self._com_row_pattern.findall(block)  # list[str], coords only (no frame #)
+        _logger.debug("Extracted %d COM rows (expected %d)", len(lines), number_frames)
+
+        if len(lines) != number_frames:
+            _logger.error("Frame row count mismatch: expected %d, got %d",
+                        number_frames, len(lines))
+            raise RuntimeError(f"Expected {number_frames} frame rows, got {len(lines)}. "
+                             "Potentially due to cpptraj version mismatch. "
+                             "The data retrieval is stable for CPPTRAJ of Version V6.18.1 (AmberTools)")
+
+        # parse and reshape to (n_residues, 3) per frame
+        rows = [np.fromstring(line, dtype=np.float32, sep=' ') for line in lines]
+        bad = [i for i, arr in enumerate(rows) if arr.size != number_residues * 3]
+        if bad:
+            _logger.error("Row(s) with wrong length detected (showing first few): %s", bad[:5])
+            raise RuntimeError(
+                f"Row(s) {bad[:5]} have wrong length; expected {number_residues*3} floats. "
+                "Potentially due to cpptraj version mismatch. "
+                "The data retrieval is stable for CPPTRAJ of Version V6.18.1 (AmberTools)"
+            )
+
+        coords: list[np.ndarray] = [row.reshape(3, number_residues).T for row in rows]
+        _logger.info("Built %d COM arrays of shape %s (one per frame)",
+                    len(coords), coords[0].shape)
+        return coords
+
+    # ---------------------------------------------------------------------------------------------- #
+    #                                       CPPTRAJ COMMANDS
+    # ---------------------------------------------------------------------------------------------- #
+
+    @staticmethod
+    def _load_data_from(topology_file: str,
+                       trajectory_file: str,
+                       start_frame: int,
+                       end_frame: int) -> rin_util.CpptrajScript:
+        """Create a cpptraj script that loads topology/trajectory and selects frames.
+
+        Args:
+            topology_file (str): Path to topology file.
+            trajectory_file (str): Path to trajectory file.
+            start_frame (int): First frame (1-based inclusive).
+            end_frame (int): Last frame (1-based inclusive).
+
+        Returns:
+            rin_util.CpptrajScript: Composable script object.
+        """
+        _logger.debug("Preparing data load (parm=%s, trajin=%s %s %s)",
+                      topology_file, trajectory_file, start_frame, end_frame)
+        return rin_util.CpptrajScript((f"parm {topology_file}",
+                                       f"trajin {trajectory_file} {start_frame} {end_frame}",
+                                       "noprogress silenceactions"))
+
+    @staticmethod
+    def _calc_nonbonded_energies_in_molecule(molecule_id: int) -> rin_util.CpptrajScript:
+        """Create a cpptraj command to compute pairwise non-bonded energies.
+
+        Args:
+            molecule_id (int): Molecule selector/ID for pairwise computation.
+
+        Returns:
+            rin_util.CpptrajScript: Script with `pairwise PW` command.
+        """
+        _logger.debug("Preparing pairwise command for molecule_id=%s", molecule_id)
+        return rin_util.CpptrajScript.from_cmd(f"pairwise PW ^{molecule_id} cuteelec 0.0 cutevdw 0.0")
+    
+    @staticmethod
+    def _extract_molecule_compositions() -> rin_util.CpptrajScript:
+        """Create a cpptraj command that emits residue/atom masks.
+
+        Returns:
+            rin_util.CpptrajScript: Script with `mask :*` command.
+        """
+        _logger.debug("Preparing mask extraction command")
+        return rin_util.CpptrajScript.from_cmd(f"mask :*")
+
+    @staticmethod
+    def _compute_residue_COMs_in_molecule(molecule_id: int):
+        """Create a cpptraj loop to compute per-residue COM vectors.
+
+        Args:
+            molecule_id (int): Molecule selector/ID whose residues are iterated.
+
+        Returns:
+            rin_util.CpptrajScript: Script that defines COM vectors (COM1, COM2, ...).
+        """
+        _logger.debug("Preparing COM vectors loop for molecule_id=%s", molecule_id)
+        return rin_util.CpptrajScript((
+            "autoimage",
+            "unwrap byres",
+            f"for residues R inmask ^{molecule_id}  i=1;i++",
+            "vector COM$i center $R",
+            "done"
+        ))
+
+    # ---------------------------------------------------------------------------------------------- #
+    #                                          POST-CPPTRAJ
+    # ---------------------------------------------------------------------------------------------- #
+
+    # -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= #
+    #  CONVERSION OF ATOMIC LEVEL INTERACTIONS INTO RESIDUE LEVEL
+    # -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= #
+    @staticmethod
+    def _compute_residue_membership_matrix(
+        res_to_atoms: dict[int, set[int]],
+        *,
+        dtype=np.float32
+    ) -> np.ndarray:
+        """Build a binary (n_atoms x n_residues) membership matrix.
+
+        Contiguously re-indexes arbitrary atom/residue IDs to 0..N-1 and sets
+        ``P[row(atom), col(residue)] = 1`` when the atom belongs to the residue.
+
+        Args:
+            res_to_atoms (dict[int, set[int]]): Mapping of residue IDs to sets of atom IDs.
+            dtype (np.dtype): Output dtype.
+
+        Returns:
+            np.ndarray: Membership matrix of shape (n_atoms, n_residues).
+        """
+        if not res_to_atoms:
+            _logger.info("Empty residue->atoms mapping; returning (0,0) matrix.")
+            return np.zeros((0, 0), dtype=dtype)
+
+        # ----- Build contiguous indices for residues (columns) -----
+        # Use numeric sort so indices are stable and predictable.
+        res_ids = sorted(res_to_atoms.keys())
+        res_to_col = {r: i for i, r in enumerate(res_ids)}
+        n_res = len(res_ids)
+
+        # ----- Build contiguous indices for atoms (rows) -----
+        # Union all atom IDs, then sort numerically.
+        atom_ids_set = set()
+        for r in res_ids:
+            atom_ids_set.update(res_to_atoms[r])
+        atom_ids = sorted(atom_ids_set)
+        atom_to_row = {a: i for i, a in enumerate(atom_ids)}
+        n_atoms = len(atom_ids)
+
+        _logger.debug("Membership dims: atoms=%d, residues=%d", n_atoms, n_res)
+
+        # ----- Fill membership matrix -----
+        P = np.zeros((n_atoms, n_res), dtype=dtype)
+        for r in res_ids:
+            c = res_to_col[r]
+            for a in res_to_atoms[r]:
+                P[atom_to_row[a], c] = 1.0
+
+        _logger.info("Built membership matrix with shape %s and density %.6f",
+                     P.shape, float(P.sum()) / (P.size if P.size else 1.0))
+        return P
+
+    @staticmethod
+    def _convert_atomic_to_residue_interactions(atomic_matrix: np.ndarray,
+                                                membership_matrix: np.ndarray) -> np.ndarray:
+        """Project atomic interaction matrix to residue space.
+
+        Computes ``R = Pᵀ @ A @ P`` where `A` is atomic (n_atoms x n_atoms) and
+        `P` is membership (n_atoms x n_residues).
+
+        Args:
+            atomic_matrix (np.ndarray): Atomic interaction matrix (n_atoms, n_atoms).
+            membership_matrix (np.ndarray): Membership matrix (n_atoms, n_residues).
+
+        Returns:
+            np.ndarray: Residue interaction matrix (n_residues, n_residues).
+        """
+        _logger.debug("Converting atomic->residue: atomic_matrix=%s, membership=%s",
+                      atomic_matrix.shape, membership_matrix.shape)
+        
+        if atomic_matrix.ndim != 2 or atomic_matrix.shape[0] != atomic_matrix.shape[1]:
+            raise ValueError(f"atomic_matrix must be square 2D; got shape {atomic_matrix.shape}")
+        if membership_matrix.ndim != 2 or membership_matrix.shape[0] != atomic_matrix.shape[0]:
+            raise ValueError(
+                f"Row count mismatch: atomic_matrix is {atomic_matrix.shape}, "
+                f"membership_matrix is {membership_matrix.shape}. Rows must match (#atoms)."
+            )
+
+        thread_count = os.cpu_count() or 1 
+        with threadpoolctl.threadpool_limits(limits=thread_count):
+            result = (membership_matrix.T @ atomic_matrix @ membership_matrix).astype(dtype=np.float32)
+        _logger.info("Residue interaction matrix shape: %s", result.shape)
+        return result
+
+    # -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= #
+    #  POST-PROCESSING OF RESIDUE LEVEL INTERACTIONS
+    # -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= #
+
+    def _split_into_attractive_repulsive(self, residue_matrix: np.ndarray) -> np.ndarray:
+        """Split residue interactions into attractive and repulsive channels.
+
+        Negative values go to the attractive channel (as positive magnitudes),
+        positive values go to the repulsive channel.
+
+        Args:
+            residue_matrix (np.ndarray): Residue interaction matrix (N, N).
+
+        Returns:
+            np.ndarray: Array of shape (2, N, N): [attractive, repulsive].
+        """
+        _logger.debug("Splitting matrix into attractive/repulsive channels; input shape=%s",
+                      residue_matrix.shape)
+        attr = np.where(residue_matrix <= 0, -residue_matrix, 0.0).astype(np.float32, copy=False)
+        rep  = np.where(residue_matrix >  0, residue_matrix, 0.0).astype(np.float32, copy=False)
+        out = np.stack([attr, rep], axis=0) # (2, N, N)
+        _logger.info("Two-channel matrix shape: %s", out.shape)
+        return out
+
+    def _prune_low_energies(self, two_channel_residue_matrix: np.ndarray, q: float) -> np.ndarray:
+        """Zero out values below a per-row quantile threshold.
+
+        Applies independently to attractive and repulsive channels.
+
+        Args:
+            two_channel_residue_matrix (np.ndarray): Array (2, N, N).
+            q (float): Quantile in (0, 1] used as threshold.
+
+        Returns:
+            np.ndarray: Pruned two-channel matrix.
+
+        Raises:
+            ValueError: If `q` is not in (0, 1].
+        """
+        _logger.debug("Pruning low energies with q=%s on matrix shape=%s", q, two_channel_residue_matrix.shape)
+        if not (0.0 < q <= 1.0):
+            _logger.error("Invalid pruning quantile q=%s", q)
+            raise ValueError(f"Invalid 'q' value. Expected a value between 0 and 1; received: {q}")
+        A = two_channel_residue_matrix[0]
+        R = two_channel_residue_matrix[1]
+        Ath = np.quantile(A, q, axis=1, keepdims=True)
+        Rth = np.quantile(R, q, axis=1, keepdims=True)
+        two_channel_residue_matrix[0] = np.where(A < Ath, 0.0, A)
+        two_channel_residue_matrix[1] = np.where(R < Rth, 0.0, R)
+        _logger.info("Pruning done at q=%s", q)
+        return two_channel_residue_matrix
+
+    def _remove_self_interactions(self, two_channel_residue_matrix: np.ndarray) -> np.ndarray:
+        """Zero the diagonal in both channels.
+
+        Args:
+            two_channel_residue_matrix (np.ndarray): Array (2, N, N).
+
+        Returns:
+            np.ndarray: Same array with zeroed diagonals.
+        """
+        _logger.debug("Zeroing self-interactions on shape=%s", two_channel_residue_matrix.shape)
+        np.fill_diagonal(two_channel_residue_matrix[0], 0.0); np.fill_diagonal(two_channel_residue_matrix[1], 0.0)
+        return two_channel_residue_matrix
+   
+    def _symmetrize(self, two_channel_residue_matrix: np.ndarray) -> np.ndarray:
+        """Symmetrize both channels via (M + Mᵀ)/2.
+
+        Args:
+            two_channel_residue_matrix (np.ndarray): Array (2, N, N).
+
+        Returns:
+            np.ndarray: Symmetrized two-channel matrix.
+        """
+        _logger.debug("Symmetrizing two-channel matrix shape=%s", two_channel_residue_matrix.shape)
+        A = two_channel_residue_matrix[0]
+        R = two_channel_residue_matrix[1]
+        two_channel_residue_matrix[0] = (A + A.T) * 0.5
+        two_channel_residue_matrix[1] = (R + R.T) * 0.5
+        _logger.info("Symmetrization complete")
+        return two_channel_residue_matrix
+
+    def _L1_normalize(self, two_channel_residue_matrix: np.ndarray) -> np.ndarray:
+        """Row-wise L1-normalization of both channels.
+
+        Each row is divided by its sum; zero rows remain zero.
+
+        Args:
+            two_channel_residue_matrix (np.ndarray): Array (2, N, N).
+
+        Returns:
+            np.ndarray: L1-normalized two-channel matrix.
+        
+        Note:
+            Row-wise normalization breaks symmetry because it converts energies into
+            per-row transition probabilities (rows sum to 1). Even if (i, j) == (j, i)
+            before normalization, differing row totals generally yield (i, j) != (j, i)
+            afterward.
+        """
+        _logger.debug("L1-normalizing two-channel matrix shape=%s", two_channel_residue_matrix.shape)
+        A = two_channel_residue_matrix[0]
+        R = two_channel_residue_matrix[1]
+        eps = 1e-12
+        Asum = A.sum(axis=1, keepdims=True)
+        Rsum = R.sum(axis=1, keepdims=True)
+        two_channel_residue_matrix[0] = np.divide(A, np.clip(Asum, eps, None),
+                                                out=np.zeros_like(A), where=Asum > 0)
+        two_channel_residue_matrix[1] = np.divide(R, np.clip(Rsum, eps, None),
+                                                out=np.zeros_like(R), where=Rsum > 0)
+        _logger.info("L1 normalization complete (zero-row counts: A=%d, R=%d)",
+                     int((Asum <= eps).sum()), int((Rsum <= eps).sum()))
+        return two_channel_residue_matrix
+
+    def _store_two_channel_array(
+        self,
+        arr: np.ndarray,
+        storage: sawnergy_util.ArrayStorage,
+        arrays_per_chunk: int,
+        attractive_dataset_name: str | None,
+        repulsive_dataset_name: str | None,
+    ) -> None:
+        """Persist a two-channel residue interaction array to storage.
+
+        This writes the two channels (index 0 → attractive, index 1 → repulsive)
+        into two separate dataset blocks inside the provided ``ArrayStorage``,
+        using the specified chunking policy.
+
+        Args:
+            arr: A 3-D array with shape ``(2, N, N)`` where the first axis indexes
+                the channels: ``arr[0]`` is the attractive channel and ``arr[1]`` is
+                the repulsive channel.
+            storage: An open ``ArrayStorage`` handle to write into.
+            arrays_per_chunk: Number of matrices per chunk along the leading axis
+                when writing into the Zarr arrays.
+            attractive_dataset_name: Dataset (block) name to store the attractive
+                channel under (if None, the dataset isn't persisted).
+            repulsive_dataset_name: Dataset (block) name to store the repulsive
+                channel under (if None, the dataset isn't persisted).
+
+        Returns:
+            None
+
+        Notes:
+            - If ``arr`` does not have the expected shape ``(2, N, N)``, a warning
+            is logged and the function still attempts to write ``arr[0]`` and
+            ``arr[1]`` as the two channels.
+            - Exceptions raised by the storage layer propagate to the caller.
+        """
+        if arr.ndim != 3 or arr.shape[0] != 2:
+            _logger.warning(
+                "Expected two-channel array with shape (2, N, N); got %s",
+                getattr(arr, "shape", None),
+            )
+
+        _logger.debug(
+            "Storing two-channel array: shape=%s, chunksize=%s, datasets=(%s, %s)",
+            arr.shape, arrays_per_chunk, attractive_dataset_name, repulsive_dataset_name
+        )
+
+        if attractive_dataset_name is not None:
+            storage.write(
+                these_arrays=[arr[0]],
+                to_block_named=attractive_dataset_name,
+                arrays_per_chunk=arrays_per_chunk
+            )
+
+        if repulsive_dataset_name is not None:
+            storage.write(
+                these_arrays=[arr[1]],
+                to_block_named=repulsive_dataset_name,
+                arrays_per_chunk=arrays_per_chunk
+            )
+
+        _logger.info(
+            "Stored attractive/repulsive arrays to '%s' / '%s'",
+            attractive_dataset_name, repulsive_dataset_name
+        )
+
+    # ---------------------------------------------------------------------------------------------- #
+    #                                           PUBLIC API
+    # ---------------------------------------------------------------------------------------------- #
+
+    def build_rin(
+        self,
+        topology_file: str,
+        trajectory_file: str,
+        molecule_of_interest: int,
+        frame_range: tuple[int, int] | None = None,
+        frame_batch_size: int = -1,
+        prune_low_energies_frac: float = 0.3,
+        output_path: str | Path | None = None,
+        keep_prenormalized_energies: bool = True,
+        *,
+        include_attractive: bool = True,
+        include_repulsive: bool = True,
+        parallel_cpptraj: bool = False,
+        simul_cpptraj_instances: int | None = None,
+        num_matrices_in_compressed_blocks: int = 10,
+        compression_level: int = 3,
+        cpptraj_run_time_limit: float | None = None
+    ) -> str:
+        """Build a Residue Interaction Network (RIN) archive from an MD trajectory.
+
+        High-level pipeline:
+        
+        1. Discover MD metadata (trajectory frame count; residue membership of the
+            target molecule).
+
+        2. For each frame batch:
+            
+            a) Run cpptraj `pairwise` on atoms → EMAP + VMAP → sum (atomic matrix).
+            
+            b) Project atomic → residue with ``R = Pᵀ @ A @ P``.
+            
+            c) Post-process residue matrix:
+                split into (attractive, repulsive) channels,
+                per-row quantile pruning,
+                remove self-interactions,
+                symmetrize.
+
+            d. Optionally store **pre-normalized energies** (attractive or repulsive or both, depending on `include_<kind>`).
+            
+            e. Row-wise L1 normalize (directed transition probabilities) and store.
+
+        3. Compute per-residue COM coordinates across requested frames and store.
+
+        4. Close and compress the temporary store into a zip (Zarr v3). Return path.
+
+        Args:
+            topology_file: Path to the topology (parm/prmtop) file.
+            trajectory_file: Path to a cpptraj-readable trajectory file.
+            molecule_of_interest: Molecule selector/ID used by cpptraj (e.g., ``1``
+                for ``^1``).
+            frame_range: 1-based inclusive ``(start, end)`` frames to process. If
+                ``None``, uses the full trajectory.
+            frame_batch_size: Number of frames per batch for pairwise calculations.
+                If ``<= 0``, processes all frames in a single batch.
+            prune_low_energies_frac: Per-row quantile ``q`` in ``(0, 1]`` used to
+                zero out small values independently in both channels.
+            output_path: Destination path (with or without ``.zip``). Defaults to
+                ``RIN_<timestamp>.zip`` in the current working directory.
+            keep_prenormalized_energies: If ``True``, stores the pre-normalized
+                attractive/repulsive matrices under ``ATTRACTIVE|REPULSIVE_energies``.
+            parallel_cpptraj: If ``True``, run multiple cpptraj frame batches in
+                parallel using threads (safe w.r.t. pickling).
+            simul_cpptraj_instances: Maximum concurrent cpptraj tasks (defaults to
+                ``os.cpu_count()`` when ``None``).
+            num_matrices_in_compressed_blocks: Number of matrices per chunk along
+                the leading axis when writing Zarr arrays.
+            compression_level: Blosc Zstd compression level for the final ZipStore.
+            cpptraj_run_time_limit: Optional timeout (seconds) for cpptraj calls.
+
+        Returns:
+            str: Path to the created ``.zip`` archive (Zarr v3).
+
+        Raises:
+            RuntimeError: Propagated from helper methods (e.g., cpptraj failures).
+            ValueError: Propagated from helper methods (e.g., bad frame ranges or
+                pruning quantile).
+
+        Notes:
+            * Row-wise L1 normalization produces **directed** transition
+            probabilities (rows sum to 1) and therefore breaks symmetry.
+            * All linear algebra runs in a single Python thread; BLAS may use
+            multiple threads internally. cpptraj parallelism is optional and uses
+            threads to avoid pickling constraints.
+        """
+        _logger.info(
+            "Building RIN (mol=%s, traj=%s, frame_range=%s, frame_batch_size=%s, "
+            "keep_abs=%s, parallel_cpptraj=%s, simul_instances=%s, comp_level=%s)",
+            molecule_of_interest, trajectory_file, frame_range, frame_batch_size,
+            keep_prenormalized_energies, parallel_cpptraj, simul_cpptraj_instances, compression_level
+        )
+
+        # ----------------------------------- MD META DATA -------------------------------------
+        total_frames = self._get_number_frames(
+            topology_file,
+            trajectory_file,
+            timeout=cpptraj_run_time_limit
+        )
+
+        molecule_composition = self._get_atomic_composition_of_molecule(
+            topology_file,
+            trajectory_file,
+            molecule_of_interest,
+            timeout=cpptraj_run_time_limit
+        )
+        number_residues = len(molecule_composition)
+        _logger.info("MD metadata: total_frames=%d, residues=%d", total_frames, number_residues)
+        # --------------------------------------------------------------------------------------
+
+        # --------------------- AUXILIARY VARIABLES' / TOOLS PREPARATION -----------------------
+        current_time = sawnergy_util.current_time()
+        attractive_transitions_name = "ATTRACTIVE_transitions"
+        repulsive_transitions_name  = "REPULSIVE_transitions"
+        attractive_energies_name    = "ATTRACTIVE_energies"
+        repulsive_energies_name     = "REPULSIVE_energies"
+        simul_cpptraj_instances = simul_cpptraj_instances or (os.cpu_count() or 1)
+        output_path = Path((output_path or (Path(os.getcwd()) /
+                        f"RIN_{current_time}"))).with_suffix(".zip")
+        _logger.debug("Output archive path: %s", output_path)
+
+        # -=- FRAMES OF THE MD SIMULATION -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
+        if frame_batch_size <= 0:
+            frame_batch_size = total_frames
+        if frame_range is None:
+            start_frame, end_frame = 1, total_frames
+        else:
+            start_frame, end_frame = frame_range
+        
+        if not (1 <= start_frame <= end_frame <= total_frames):
+            raise ValueError(f"frame_range must lie within [1, {total_frames}] and be ordered; got {frame_range}.")
+
+        frames = (
+            (s, min(s + frame_batch_size - 1, end_frame))
+            for s in range(start_frame, end_frame + 1, max(1, frame_batch_size))
+        )
+        _logger.debug("Frame selection: [%d..%d], batch_size=%d", start_frame, end_frame, frame_batch_size)
+        # -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=--=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
+
+        # -=- DATA PROCESSORS -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
+        frame_processor = sawnergy_util.elementwise_processor(
+            in_parallel=parallel_cpptraj,
+            Executor=ThreadPoolExecutor,
+            max_workers=simul_cpptraj_instances,
+            capture_output=True
+        )
+        matrix_processor = sawnergy_util.elementwise_processor(
+            in_parallel=False,   # <- BLAS handles lin. alg. parallelism &
+            capture_output=True  # the rest of the code is vectorized by default
+        )
+        # -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=--=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
+
+        # -=- ADJUST CPPTRAJ PARALLELISM -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
+        non_bonded_energies_subprocess_env = (
+            sawnergy_util.create_updated_subprocess_env(
+                OMP_NUM_THREADS=1,  # PREVENTING OVERSUBSCRIPTION
+                MKL_NUM_THREADS=1,
+                OPENBLAS_NUM_THREADS=1,
+                MKL_DYNAMIC=False
+            ) if parallel_cpptraj else None
+        )
+        _logger.debug("cpptraj parallel: %s (instances=%s)", parallel_cpptraj, simul_cpptraj_instances)
+        # -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
+
+        # create a membership matrix for atoms in residues
+        membership_matrix = self._compute_residue_membership_matrix(molecule_composition)
+        _logger.info("Membership matrix ready: shape=%s, nnz=%d",
+                    membership_matrix.shape, int(membership_matrix.sum()))
+        # --------------------------------------------------------------------------------------
+
+        # -------------------- INTERACTION DATA EXTRACTION AND PROCESSING ----------------------
+        pipeline = sawnergy_util.compose_steps(
+            (self._convert_atomic_to_residue_interactions, {"membership_matrix": membership_matrix}),
+            (self._split_into_attractive_repulsive, None),
+            (self._prune_low_energies, {"q": prune_low_energies_frac}),
+            (self._remove_self_interactions, None),
+            (self._symmetrize, None),
+        )
+        _logger.debug("Post-processing pipeline assembled")
+
+        with sawnergy_util.ArrayStorage.compress_and_cleanup(output_path, compression_level) as storage:
+            _logger.debug("Opened temporary store for writing")
+            com_avgs: list[np.ndarray] = []
+
+            for frame_batch in sawnergy_util.batches_of(
+                frames, batch_size=simul_cpptraj_instances if parallel_cpptraj else 1
+            ):
+                _logger.debug("Processing next frame batch")
+
+                atomic_matrices = frame_processor(
+                    frame_batch,
+                    self._calc_avg_atomic_interactions_in_frames,
+                    topology_file,
+                    trajectory_file,
+                    molecule_of_interest,
+                    subprocess_env=non_bonded_energies_subprocess_env,
+                    timeout=cpptraj_run_time_limit,
+                )
+
+                interaction_matrices = matrix_processor(atomic_matrices, pipeline)
+
+                if keep_prenormalized_energies:
+                    _logger.debug("Writing absolute energy channels")
+                    for arr in interaction_matrices:
+                        self._store_two_channel_array(
+                            arr,
+                            storage,
+                            num_matrices_in_compressed_blocks,
+                            attractive_energies_name if include_attractive else None,
+                            repulsive_energies_name if include_repulsive else None
+                        )
+
+                transition_matrices = matrix_processor(interaction_matrices, self._L1_normalize)
+
+                _logger.debug("Writing normalized transition channels")
+                for arr in transition_matrices:
+                    self._store_two_channel_array(
+                        arr,
+                        storage,
+                        num_matrices_in_compressed_blocks,
+                        attractive_transitions_name if include_attractive else None,
+                        repulsive_transitions_name if include_repulsive else None
+                    )
+
+                _logger.debug("Computing batch COMs for this frame batch (len=%d)", len(frame_batch))
+                com_lists_per_range = frame_processor(
+                    frame_batch,
+                    self._get_residue_COMs_per_frame,
+                    topology_file,
+                    trajectory_file,
+                    molecule_of_interest,
+                    number_residues,
+                    timeout=cpptraj_run_time_limit,
+                )
+
+                for i, com_frames in enumerate(com_lists_per_range):
+                    avg = np.stack(com_frames, axis=0).mean(axis=0).astype(np.float32, copy=False)
+                    com_avgs.append(avg)
+                    _logger.debug("Batch %d: COM avg shape=%s", i, avg.shape)
+
+            _logger.debug("Writing %d batch-averaged COM snapshots (chunk=%d)",
+                        len(com_avgs), num_matrices_in_compressed_blocks)
+
+            storage.write(
+                com_avgs,
+                to_block_named="COM",
+                arrays_per_chunk=num_matrices_in_compressed_blocks,
+            )
+
+            # ----------------------------- META-DATA --------------------------------
+            storage.add_attr("time_created", current_time)
+            storage.add_attr("com_name", "COM")
+            storage.add_attr("molecule_of_interest", molecule_of_interest)
+            storage.add_attr("frame_range", frame_range)
+            storage.add_attr("frame_batch_size", frame_batch_size)
+            storage.add_attr("prune_low_energies_frac", prune_low_energies_frac)
+            storage.add_attr("attractive_transitions_name", attractive_transitions_name if include_attractive else None)
+            storage.add_attr("repulsive_transitions_name", repulsive_transitions_name if include_repulsive else None)
+            storage.add_attr("attractive_energies_name",
+                            attractive_energies_name if include_attractive and keep_prenormalized_energies else None)
+            storage.add_attr("repulsive_energies_name",
+                            repulsive_energies_name if include_repulsive and keep_prenormalized_energies else None)
+            # ------------------------------------------------------------------------
+
+        _logger.info("RIN build complete -> %s", output_path)
+        return str(output_path)
+
+__all__ = [
+    "RINBuilder"
+]
+
+if __name__ == "__main__":
+    pass