prism-pruner 0.0.3__py3-none-any.whl

prism_pruner/__init__.py

@@ -0,0 +1 @@
+"""PRISM - PRuning Interface for Similar Molecules."""

prism_pruner/algebra.py

@@ -0,0 +1,163 @@
+"""Algebra utilities."""
+
+from typing import Sequence
+
+import numpy as np
+
+from prism_pruner.typing import Array1D_float, Array2D_float, Array3D_float
+
+
+def normalize(vec: Array1D_float) -> Array1D_float:
+    """Normalize a vector."""
+    return vec / np.linalg.norm(vec)
+
+
+def vec_angle(v1: Array1D_float, v2: Array1D_float) -> float:
+    """Return the planar angle defined by two 3D vectors."""
+    return float(
+        np.degrees(
+            np.arccos(
+                np.clip(
+                    np.dot(
+                        v1 / np.linalg.norm(v1),
+                        v2 / np.linalg.norm(v2),
+                    ),
+                    -1.0,
+                    1.0,
+                ),
+            )
+        )
+    )
+
+
+def dihedral(p: Array2D_float) -> float:
+    """
+    Find dihedral angle in degrees from 4 3D vecs.
+
+    Praxeolitic formula: 1 sqrt, 1 cross product.
+    """
+    p0, p1, p2, p3 = p
+
+    b0 = -1.0 * (p1 - p0)
+    b1 = p2 - p1
+    b2 = p3 - p2
+
+    # normalize b1 so that it does not influence magnitude of vector
+    # rejections that come next
+    b1 /= np.linalg.norm(b1)
+
+    # vector rejections
+    # v = projection of b0 onto plane perpendicular to b1
+    #   = b0 minus component that aligns with b1
+    # w = projection of b2 onto plane perpendicular to b1
+    #   = b2 minus component that aligns with b1
+    v = b0 - np.dot(b0, b1) * b1
+    w = b2 - np.dot(b2, b1) * b1
+
+    # angle between v and w in a plane is the torsion angle
+    # v and w may not be normalized but that's fine since tan is y/x
+    x = np.dot(v, w)
+    y = np.dot(np.cross(b1, v), w)
+
+    return float(np.degrees(np.arctan2(y, x)))
+
+
+def rot_mat_from_pointer(pointer: Array1D_float, angle: float) -> Array2D_float:
+    """
+    Get the rotation matrix from the rotation pivot using a quaternion.
+
+    :param pointer: 3D vector representing the rotation pivot
+    :param angle: rotation angle in degrees
+    :return rotation_matrix: matrix that applied to a point, rotates it along the pointer
+    """
+    assert pointer.shape[0] == 3
+
+    angle_2 = np.radians(angle) / 2
+    sin = np.sin(angle_2)
+    pointer = pointer / np.linalg.norm(pointer)
+    return quaternion_to_rotation_matrix(
+        [
+            sin * pointer[0],
+            sin * pointer[1],
+            sin * pointer[2],
+            np.cos(angle_2),
+        ]
+    )
+
+
+def quaternion_to_rotation_matrix(quat: Array1D_float | Sequence[float]) -> Array2D_float:
+    """
+    Convert a quaternion into a full three-dimensional rotation matrix.
+
+    This rotation matrix converts a point in the local reference frame to a
+    point in the global reference frame.
+
+    :param quat: 4-element array representing the quaternion (q0, q1, q2, q3)
+    :return: 3x3 element array representing the full 3D rotation matrix
+    """
+    # Extract the values from Q (adjusting for scalar last in input)
+    q1, q2, q3, q0 = quat
+
+    # First row of the rotation matrix
+    r00 = 2 * (q0 * q0 + q1 * q1) - 1
+    r01 = 2 * (q1 * q2 - q0 * q3)
+    r02 = 2 * (q1 * q3 + q0 * q2)
+
+    # Second row of the rotation matrix
+    r10 = 2 * (q1 * q2 + q0 * q3)
+    r11 = 2 * (q0 * q0 + q2 * q2) - 1
+    r12 = 2 * (q2 * q3 - q0 * q1)
+
+    # Third row of the rotation matrix
+    r20 = 2 * (q1 * q3 - q0 * q2)
+    r21 = 2 * (q2 * q3 + q0 * q1)
+    r22 = 2 * (q0 * q0 + q3 * q3) - 1
+
+    # 3x3 rotation matrix
+    return np.array([[r00, r01, r02], [r10, r11, r12], [r20, r21, r22]])
+
+
+def get_inertia_moments(coords: Array3D_float, masses: Array1D_float) -> Array1D_float:
+    """Compute the principal moments of inertia of a molecule.
+
+    Returns a length-3 array [I_x, I_y, I_z], sorted ascending.
+    """
+    # Shift to center of mass
+    com = np.sum(coords * masses[:, np.newaxis], axis=0) / np.sum(masses)
+    coords = coords - com
+
+    # Compute inertia tensor
+    norms_sq = np.einsum("ni,ni->n", coords, coords)
+    total = np.sum(masses * norms_sq)
+    I_matrix = total * np.eye(3) - np.einsum("n,ni,nj->ij", masses, coords, coords)
+
+    # Principal moments via symmetric eigendecomposition
+    moments, _ = np.linalg.eigh(I_matrix)
+
+    return np.sort(moments)
+
+
+def diagonalize(a: Array2D_float) -> Array2D_float:
+    """Build the diagonalized matrix."""
+    eigenvalues_of_a, eigenvectors_of_a = np.linalg.eig(a)
+    b = eigenvectors_of_a[:, np.abs(eigenvalues_of_a).argsort()]
+    return np.dot(np.linalg.inv(b), np.dot(a, b))  # type: ignore[no-any-return]
+
+
+def get_alignment_matrix(p: Array1D_float, q: Array1D_float) -> Array2D_float:
+    """
+    Build the rotation matrix that aligns vectors q to p (Kabsch algorithm).
+
+    Assumes centered vector sets (i.e. their mean is the origin).
+    """
+    # calculate the covariance matrix
+    cov_mat = p.T @ q
+
+    # Compute the SVD
+    v, _, w = np.linalg.svd(cov_mat)
+
+    # Ensure proper rotation (det = 1, not -1)
+    if np.linalg.det(v) * np.linalg.det(w) < 0.0:
+        v[:, -1] *= -1
+
+    return v @ w  # type: ignore[no-any-return]

prism_pruner/conformer_ensemble.py

@@ -0,0 +1,57 @@
+"""ConformerEnsemble class."""
+
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Self
+
+import numpy as np
+
+from prism_pruner.typing import Array1D_float, Array1D_str, Array2D_float, Array3D_float
+
+
+@dataclass
+class ConformerEnsemble:
+    """Class representing a conformer ensemble."""
+
+    coords: Array3D_float
+    atoms: Array1D_str
+    energies: Array1D_float = field(default_factory=lambda: np.array([]))
+
+    @classmethod
+    def from_xyz(cls, file: Path | str, read_energies: bool = False) -> Self:
+        """Generate ensemble from a multiple conformer xyz file."""
+        coords = []
+        atoms = []
+        energies = []
+        with Path(file).open() as f:
+            for num in f:
+                if read_energies:
+                    energy = next(re.finditer(r"-*\d+\.\d+", next(f))).group()
+                    energies.append(float(energy))
+                else:
+                    _comment = next(f)
+
+                conf_atoms = []
+                conf_coords = []
+                for _ in range(int(num)):
+                    atom, *xyz = next(f).split()
+                    conf_atoms.append(atom)
+                    conf_coords.append([float(x) for x in xyz])
+
+                atoms.append(conf_atoms)
+                coords.append(conf_coords)
+
+        return cls(coords=np.array(coords), atoms=np.array(atoms[0]), energies=np.array(energies))
+
+    def to_xyz(self, file: Path | str) -> None:
+        """Write ensemble to an xyz file."""
+
+        def to_xyz(coords: Array2D_float) -> str:
+            return "\n".join(
+                f"{atom} {x:15.8f} {y:15.8f} {z:15.8f}"
+                for atom, (x, y, z) in zip(self.atoms, coords, strict=True)
+            )
+
+        with Path(file).open("w") as f:
+            f.write("\n".join(map(to_xyz, self.coords)))

prism_pruner/graph_manipulations.py

@@ -0,0 +1,195 @@
+"""Graph manipulation utilities for molecular structures."""
+
+from functools import lru_cache
+
+import numpy as np
+from networkx import Graph, all_simple_paths, from_numpy_array, set_node_attributes
+from periodictable import elements
+from scipy.spatial.distance import cdist
+
+from prism_pruner.algebra import dihedral
+from prism_pruner.typing import Array1D_bool, Array1D_str, Array2D_float
+
+
+@lru_cache()
+def d_min_bond(a1: str, a2: str, factor: float = 1.2) -> float:
+    """Return the bond distance between two atoms."""
+    return factor * (elements.symbol(a1).covalent_radius + elements.symbol(a2).covalent_radius)  # type: ignore [no-any-return]
+
+
+def graphize(
+    atoms: Array1D_str,
+    coords: Array2D_float,
+    mask: Array1D_bool | None = None,
+) -> Graph:
+    """
+    Return a NetworkX undirected graph of molecular connectivity.
+
+    :param atoms: atomic symbols
+    :param coords: atomic coordinates as 3D vectors
+    :param mask: bool array, with False for atoms to be excluded in the bond evaluation
+    :return: connectivity graph
+    """
+    mask = np.array([True for _ in atoms], dtype=bool) if mask is None else mask
+    assert len(coords) == len(atoms)
+    assert len(coords) == len(mask)
+
+    matrix = np.zeros((len(coords), len(coords)))
+    for i, mask_i in enumerate(mask):
+        if not mask_i:
+            continue
+
+        for j, mask_j in enumerate(mask[i + 1 :], start=i + 1):
+            if not mask_j:
+                continue
+
+            if np.linalg.norm(coords[i] - coords[j]) < d_min_bond(atoms[i], atoms[j]):
+                matrix[i][j] = 1
+
+    graph = from_numpy_array(matrix)
+    set_node_attributes(graph, dict(enumerate(atoms)), "atoms")
+
+    return graph
+
+
+def get_sp_n(index: int, graph: Graph) -> int | None:
+    """
+    Get hybridization of selected atom.
+
+    Return n, that is the apex of sp^n hybridization for CONPS atoms.
+    This is just an assimilation to the carbon geometry in relation to sp^n:
+    - sp¹ is linear
+    - sp² is planar
+    - sp³ is tetraedral
+    This is mainly used to understand if a torsion is to be rotated or not.
+    """
+    atom = graph.nodes[index]["atoms"]
+
+    if atom not in {"C", "N", "O", "P", "S"}:
+        return None
+
+    # Relationship of number of neighbors to sp^n hybridization
+    d: dict[str, dict[int, int | None]] = {
+        "C": {2: 1, 3: 2, 4: 3},
+        "N": {2: 2, 3: None, 4: 3},  # 3 could mean sp3 or sp2
+        "O": {1: 2, 2: 3, 3: 3, 4: 3},
+        "P": {2: 2, 3: 3, 4: 3},
+        "S": {2: 2, 3: 3, 4: 3},
+    }
+    return d[atom].get(len(set(graph.neighbors(index))))
+
+
+def is_amide_n(index: int, graph: Graph, mode: int = -1) -> bool:
+    """
+    Assess if the atom is an amide-like nitrogen.
+
+    Note: carbamates and ureas are considered amides.
+
+    mode:
+    -1 - any amide
+    0 - primary amide (CONH2)
+    1 - secondary amide (CONHR)
+    2 - tertiary amide (CONR2)
+    """
+    # Must be a nitrogen atom
+    if graph.nodes[index]["atoms"] == "N":
+        nb = set(graph.neighbors(index))
+        nb_atoms = [graph.nodes[j]["atoms"] for j in nb]
+
+        if mode != -1:
+            # Primary amides need to have 1H, secondary amides none
+            if nb_atoms.count("H") != (2, 1, 0)[mode]:
+                return False
+
+        for n in nb:
+            # There must be at least one carbon atom next to N
+            if graph.nodes[n]["atoms"] == "C":
+                nb_nb = set(graph.neighbors(n))
+                # Bonded to three atoms
+                if len(nb_nb) == 3:
+                    # and at least one of them has to be an oxygen
+                    if "O" in {graph.nodes[i]["atoms"] for i in nb_nb}:
+                        return True
+    return False
+
+
+def is_ester_o(index: int, graph: Graph) -> bool:
+    """
+    Assess if the index is an ester-like oxygen.
+
+    Note: carbamates and carbonates return True, carboxylic acids return False.
+    """
+    if graph.nodes[index]["atoms"] == "O":
+        if "H" in (nb := set(graph.neighbors(index))):
+            return False
+
+        for n in nb:
+            if graph.nodes[n]["atoms"] == "C":
+                nb_nb = set(graph.neighbors(n))
+                if len(nb_nb) == 3:
+                    nb_nb_sym = [graph.nodes[i]["atoms"] for i in nb_nb]
+                    if nb_nb_sym.count("O") > 1:
+                        return True
+    return False
+
+
+def is_phenyl(coords: Array2D_float) -> bool:
+    """
+    Assess if the six atomic coords refer to a phenyl-like ring.
+
+    Note: quinones evaluate to True
+
+    :params coords: six coordinates of C/N atoms
+    :return: bool indicating if the six atoms look like part of a phenyl/naphtyl/pyridine
+             system, coordinates for the center of that ring
+    """
+    # if any atomic couple is more than 3 A away from each other, this is not a Ph
+    if np.max(cdist(coords, coords)) > 3:
+        return False
+
+    threshold_delta: float = 1 - np.cos(10 * np.pi / 180)
+    flat_delta: float = 1 - np.abs(np.cos(dihedral(coords[[0, 1, 2, 3]]) * np.pi / 180))
+
+    return flat_delta < threshold_delta
+
+
+def get_phenyl_ids(index: int, graph: Graph) -> list[int] | None:
+    """If index is part of a phenyl, return the six heavy atoms ids associated with the ring."""
+    for n in graph.neighbors(index):
+        for path in all_simple_paths(graph, source=index, target=n, cutoff=6):
+            if len(path) != 6 or any(graph.nodes[n]["atoms"] == "H" for n in path):
+                continue
+            if all(len(set(graph.neighbors(i))) == 3 for i in path):
+                return path  # type: ignore [no-any-return]
+
+    return None
+
+
+def find_paths(
+    graph: Graph,
+    u: int,
+    n: int,
+    exclude_set: set[int] | None = None,
+) -> list[list[int]]:
+    """
+    Find paths in graph.
+
+    Recursively find all paths of a NetworkX graph with length = n, starting from node u.
+
+    :param graph: NetworkX graph
+    :param u: starting node
+    :param n: path length
+    :param exclude_set: set of nodes to exclude from the paths
+    :return: list of paths (each path is a list of node indices)
+    """
+    exclude_set = (exclude_set or set()) | {u}
+
+    if n == 0:
+        return [[u]]
+
+    return [
+        [u, *path]
+        for neighbor in graph.neighbors(u)
+        if neighbor not in exclude_set
+        for path in find_paths(graph, neighbor, n - 1, exclude_set)
+    ]