sawnergy 1.0.0__py3-none-any.whl

sawnergy/rin/rin_util.py

@@ -0,0 +1,391 @@
+from __future__ import annotations
+
+# built-in
+import logging
+from dataclasses import dataclass, field
+import os, shutil, subprocess
+from pathlib import Path
+import re
+# local
+from .. import sawnergy_util
+
+# *----------------------------------------------------*
+#                        GLOBALS
+# *----------------------------------------------------*
+
+_logger = logging.getLogger(__name__)
+PAIRWISE_STDOUT: CpptrajScript
+COM_STDOUT: CpptrajScript
+
+# *----------------------------------------------------*
+#                        CLASSES
+# *----------------------------------------------------*
+
+class CpptrajNotFound(RuntimeError):
+    """Error raised when a functional `cpptraj` executable cannot be found.
+
+    The exception message lists the candidate paths that were tried and gives a
+    brief hint on how to make `cpptraj` discoverable (install AmberTools, add
+    to PATH, or set the CPPTRAJ environment variable).
+    """
+    def __init__(self, candidates: list[Path]):
+        """Initialize the exception with the candidate paths.
+
+        Args:
+            candidates: Ordered list of filesystem paths that were checked for a
+                working `cpptraj` executable.
+        """
+        msg = (
+            "Could not locate a working `cpptraj` executable.\n"
+            f"Tried the following locations:\n" +
+            "\n".join(f"  - {p}" for p in candidates) +
+            "\nEnsure that AmberTools is installed and `cpptraj` is on your PATH, "
+            "or set the CPPTRAJ environment variable to its location."
+        )
+        super().__init__(msg)
+
+@dataclass(frozen=True)
+class CpptrajScript:
+    """Immutable builder for composing cpptraj input scripts.
+
+    Instances hold a tuple of command strings. You can:
+    - Append a command with `+ "cmd"`.
+    - Concatenate two scripts with `+ other_script`.
+    - Redirect the last command to a file with the overloaded `>` operator.
+    - Render the final script text with `render()`, which ensures a trailing
+      newline and injects a `run` command if one is not already present.
+
+    Attributes:
+        commands: Ordered tuple of cpptraj command lines (without trailing
+            newlines).
+    """
+    commands: tuple[str] = field(default_factory=tuple)
+
+    @classmethod
+    def from_cmd(cls, cmd: str) -> CpptrajScript:
+        """Create a script containing a single command.
+
+        Args:
+            cmd: A single cpptraj command line (no trailing newline required).
+
+        Returns:
+            CpptrajScript: A new script with exactly one command.
+        """
+        return cls((cmd,))
+
+    def __add__(self, other: str | CpptrajScript) -> CpptrajScript:
+        """Concatenate a command or another script.
+
+        Args:
+            other: Either a command string to append as a new line, or another
+                `CpptrajScript` whose commands will be appended in order.
+
+        Returns:
+            CpptrajScript: A new script with `other` appended.
+
+        Raises:
+            TypeError: If `other` is not a `str` or `CpptrajScript`.
+        """
+        if isinstance(other, str):
+            return CpptrajScript(self.commands + (other,))
+        elif isinstance(other, CpptrajScript):
+            return CpptrajScript(self.commands + other.commands)
+        else:
+            return NotImplemented
+
+    def __gt__(self, file_name: str | CpptrajScript) -> CpptrajScript: # >
+        """Overload `>` to add an `out <file>` target to the last command.
+
+        If `file_name` is a string, append `out <file_name>` to the last command.
+        If `file_name` is a `CpptrajScript`, treat this as concatenation (same
+        effect as `self + file_name`).
+
+        Args:
+            file_name: Output filename to attach to the last command, or another
+                script to concatenate.
+
+        Returns:
+            CpptrajScript: A new script with modified/concatenated commands.
+        """
+        if isinstance(file_name, CpptrajScript):
+            return self + file_name
+        else:
+            save_to = (self.commands[-1] + f" out {file_name}",)
+            return CpptrajScript(self.commands[:-1] + save_to)
+
+    def render(self) -> str:
+        """Render the script to text, auto-inserting `run` if missing.
+
+        Returns:
+            str: The full script text joined by newlines. If no `run` appears in
+            `commands`, a `run` line (plus a trailing blank line) is added.
+        """
+        commands = self.commands + ("",) if "run" in self.commands else self.commands + ("run", "")
+        return "\n".join(commands)
+
+PAIRWISE_STDOUT = CpptrajScript((
+                  "run",
+                  "printdata PW[EMAP] square2d noheader",
+                  "printdata PW[VMAP] square2d noheader"
+                ))
+
+COM_STDOUT = lambda mol_id: CpptrajScript((
+            "run",
+            f"for residues R inmask ^{mol_id}  i=1;i++",
+            "dataset legend $R COM$i",
+            "dataset vectorcoord X COM$i name COMX$i",
+            "dataset vectorcoord Y COM$i name COMY$i",
+            "dataset vectorcoord Z COM$i name COMZ$i",
+            "done",
+            "printdata COMX* COMY* COMZ*"
+           ))
+
+# *----------------------------------------------------*
+#                       FUNCTIONS
+# *----------------------------------------------------*
+
+# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= #
+#  WRAPPERS AND HELPERS FOR THE CPPTRAJ EXECUTABLE
+# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= #
+
+def locate_cpptraj(explicit: Path | None = None, verify: bool = True) -> str:
+    """Locate a working `cpptraj` executable.
+
+    This function attempts to resolve the path to the `cpptraj` binary used in
+    AmberTools. It searches for `cpptraj` in the following order:
+
+    1. An explicitly provided path.
+    2. The `CPPTRAJ` environment variable.
+    3. System PATH (via `shutil.which`).
+    4. The `AMBERHOME/bin` directory.
+    5. The `CONDA_PREFIX/bin` directory.
+
+    Each candidate path is checked for existence and executability. If
+    `verify=True`, the function also probes the binary with the `-h` flag
+    to ensure it responds correctly.
+
+    Args:
+        explicit (Path | None): An explicit path to the `cpptraj` executable.
+            If provided, this is the first candidate tested.
+        verify (bool): If True, run `cpptraj -h` to confirm functionality
+            of the executable. If False, only existence and executability
+            are checked. Defaults to True.
+
+    Returns:
+        str: The absolute path to a verified `cpptraj` executable.
+
+    Raises:
+        CpptrajNotFound: If no functional `cpptraj` instance can be located.
+        subprocess.TimeoutExpired: If the `cpptraj -h` verification command
+            exceeds the timeout limit.
+    """
+    _logger.info("Attempting to locate a `cpptraj` executable")
+
+    if explicit is not None: _logger.info(f"An explicit path was provided: {explicit.resolve()}")
+    else: _logger.info("No explicit path was provided")
+
+    candidates = []
+    if explicit: candidates.append(Path(explicit))
+
+    if os.getenv("CPPTRAJ"): candidates.append(Path(os.getenv("CPPTRAJ")))
+
+    for name in ("cpptraj", "cpptraj.exe"):
+        exe = shutil.which(name)
+        if exe: candidates.append(Path(exe))
+
+    if os.getenv("AMBERHOME"): candidates.append(Path(os.getenv("AMBERHOME")) / "bin" / "cpptraj")
+    if os.getenv("CONDA_PREFIX"): candidates.append(Path(os.getenv("CONDA_PREFIX")) / "bin" / "cpptraj")
+
+    _logger.info(f"Checking the following paths for cpptraj presence: {candidates}")
+    for p in candidates:
+        if p and p.exists() and os.access(p, os.X_OK):
+            _logger.info(f"Found a `cpptraj` instance at {p}")
+
+            if not verify:
+                _logger.info(f"No verification was prompted. Returning the path {p}")
+                return str(p.resolve())
+
+            _logger.info("Attempting to verify that it works")
+            try:
+                # cpptraj -h prints a help message
+                proc = subprocess.run([str(p), "-h"], capture_output=True, text=True, timeout=15) # 15 sec timeout
+            except subprocess.TimeoutExpired:
+                _logger.warning(f"The instance at {p} hung during verification (timeout). Skipping.")
+                continue
+
+            if proc.returncode in (0, 1):
+                _logger.info(f"The instance is functional. Returning the path {p}")
+                return str(p.resolve())
+            else:
+                _logger.warning(f"The instance is not functional: {proc.stderr}")
+    
+    _logger.error(f"No functional `cpptraj` instance was found")
+    raise CpptrajNotFound(candidates)
+
+def run_cpptraj(cpptraj: str,
+                script: str | None = None,
+                argv: list[str] | None = None,
+                timeout: float | None = None,
+                *,
+                env: dict | None = None):
+    """Run `cpptraj` and return its standard output.
+
+    If `script` text is provided, it is sent to cpptraj via STDIN. A trailing
+    `quit` line is appended automatically if the script does not already end
+    with one. Alternatively, you can pass command-line arguments via `argv`
+    (e.g., `["-i", "script.in"]`) and leave `script=None`.
+
+    Args:
+        cpptraj: Path to the `cpptraj` executable.
+        script: Complete cpptraj script to feed on STDIN. If not `None` and not
+            already terminated by `quit`, the function appends `quit\\n`.
+        argv: Additional command-line arguments to pass to `cpptraj`.
+        timeout: Maximum wall time in seconds for the subprocess.
+        env: Environment variables for the child process. Values must be strings.
+
+    Returns:
+        str: Captured `stdout` produced by `cpptraj`.
+
+    Raises:
+        subprocess.CalledProcessError: If cpptraj exits with a non-zero status.
+        Exception: For unexpected errors during subprocess execution.
+    """
+    if script is not None:
+        if not script.rstrip().lower().endswith("quit"):
+            script = script + "quit\n"
+
+    args = [cpptraj] + (argv or [])
+    try:
+        _logger.debug(f"Running cpptraj command: {script} with args: {args}")
+        proc = subprocess.run(
+            args,
+            input=script,
+            text=True,
+            capture_output=True,
+            check=True,
+            timeout=timeout,
+            env=env
+        )
+        return proc.stdout
+    except subprocess.CalledProcessError as e:
+        stderr = (e.stderr or "").strip()
+        _logger.error(f"Cpptraj execution failed (code {e.returncode}). Stderr:\n{stderr}")
+        raise
+    except Exception as e:
+        _logger.error(f"Unexpected error while running cpptraj: {e}")
+        raise
+
+# *----------------------------------------------------*
+#                 CPPTRAJ OUTPUT PARSERS
+# *----------------------------------------------------*
+
+# MD ITEMS (atoms, residues, molecules)
+class CpptrajMaskParser:
+    """Namespace container for cpptraj mask table parsing helpers."""
+    __slots__ = ()  # no instances allowed
+
+    # ---------- REGEX ----------
+    _spaces_pattern = re.compile(r"\s+")
+    _items_pattern = re.compile(r"\[(\w+)\]")  # captures fields of interest: [AtNum], [Rnum], [Mnum]
+
+    # --------- HELPERS ---------
+    @staticmethod
+    def _id2item_map(header: str) -> dict[str, int]:
+        """Map column names found in bracket tokens in `header` to their 0-based indices."""
+        cols = CpptrajMaskParser._items_pattern.findall(header)
+        return {name: i for i, name in enumerate(cols)}
+
+    @staticmethod
+    def _collapse_spaces(s: str) -> str:
+        """Normalize all whitespace runs to a single space and strip ends."""
+        return CpptrajMaskParser._spaces_pattern.sub(" ", s).strip()
+
+    @staticmethod
+    def _get_row_items(row: str, header_map: dict[str, int]) -> tuple[int, int, int]:
+        """Extract molecule/residue/atom IDs from a data row using the header map."""
+        items = CpptrajMaskParser._collapse_spaces(row).split()
+        try:
+            return (
+                int(items[header_map["Mnum"]]),
+                int(items[header_map["Rnum"]]),
+                int(items[header_map["AtNum"]]),
+            )
+        except KeyError as ke:
+            raise ValueError(f"Required column missing in header: {ke}") from ke
+        except IndexError as ie:
+            raise ValueError(f"Row has fewer fields than expected: {row!r}") from ie
+
+    # --------- PUBLIC ----------
+    @staticmethod
+    def hierarchize_molecular_composition(mol_compositions_file: str) -> dict[int, dict[int, set[int]]]:
+        """
+        Build {molecule_id: {residue_id: {atom_id, ...}, ...}} from a cpptraj mask table.
+
+        Assumes the file's header line contains bracketed column labels (e.g., [AtNum], [Rnum], [Mnum]).
+
+        Args:
+            mol_compositions_file: Path to a text file produced by cpptraj that
+                lists atoms with bracketed header tokens identifying molecule,
+                residue, and atom indices.
+
+        Returns:
+            dict[int, dict[int, set[int]]]: Nested mapping from molecule ID to
+            residue ID to the set of atom IDs.
+
+        Raises:
+            RuntimeError: If the input file is empty.
+            ValueError: If required columns are missing or a row is malformed.
+        """
+        lines = sawnergy_util.read_lines(mol_compositions_file, skip_header=False)
+        if not lines:
+            raise RuntimeError(f"0 lines were read from {mol_compositions_file}")
+
+        header = lines[0]
+        header_map = CpptrajMaskParser._id2item_map(header)
+
+        required = {"Mnum", "Rnum", "AtNum"}
+        missing = required.difference(header_map)
+        if missing:
+            raise ValueError(f"Missing required columns in header: {sorted(missing)}")
+
+        hierarchy: dict[int, dict[int, set[int]]] = {}
+
+        for line in lines[1:]:
+            if not line.strip():
+                continue
+            molecule_id, residue_id, atom_id = CpptrajMaskParser._get_row_items(line, header_map)
+
+            residues = hierarchy.setdefault(molecule_id, {})
+            atoms = residues.setdefault(residue_id, set())
+            atoms.add(atom_id)
+
+        return hierarchy
+
+# CENTER OF THE MASS
+def com_parser(line: str) -> str:
+    """Parse a cpptraj `center of mass` output line into CSV format.
+
+    The input line is expected to contain seven whitespace-separated fields:
+    `frame x y z vx vy vz` (velocity fields ignored here). The function emits
+    a CSV string with the first four values: `frame,x,y,z\\n`.
+
+    Args:
+        line: A single line from cpptraj's COM output.
+
+    Returns:
+        str: A CSV-formatted line containing `frame,x,y,z` and a trailing newline.
+
+    Raises:
+        ValueError: If the input line does not contain at least four fields.
+    """
+    frame, x, y, z, _, _, _= line.split()
+    return f"{frame},{x},{y},{z}\n"
+
+
+__all__ = [
+    "run_cpptraj",
+    "CpptrajScript"
+]
+
+if __name__ == "__main__":
+    pass