prism-pruner 0.0.1__py3-none-any.whl

prism_pruner/__init__.py

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

prism_pruner/algebra.py

@@ -0,0 +1,190 @@
+"""Algebra utilities."""
+
+from typing import Sequence
+
+import numpy as np
+
+from prism_pruner.typing import Array1D_float, Array1D_int, Array2D_float, Array3D_float
+
+
+def norm(vec: Array1D_int) -> Array1D_int:
+    """Normalize a vector (3D only)."""
+    return vec / np.sqrt((vec[0] * vec[0] + vec[1] * vec[1] + vec[2] * vec[2]))  # type: ignore[no-any-return]
+
+
+def norm_of(vec: Array1D_int) -> float:
+    """Norm of a vector (3D only)."""
+    return float(np.sqrt((vec[0] * vec[0] + vec[1] * vec[1] + vec[2] * vec[2])))
+
+
+def vec_angle(v1: Array1D_int, v2: Array1D_int) -> float:
+    """Return the planar angle defined by two 3D vectors."""
+    return float(
+        np.degrees(
+            np.arccos(
+                np.clip(np.dot(norm(v1), 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 /= norm_of(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_int, 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 = 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 kronecker_delta(i: int, j: int) -> int:
+    """Kronecker delta."""
+    return int(i == j)
+
+
+def get_inertia_moments(coords: Array3D_float, masses: Array1D_float) -> Array1D_float:
+    """
+    Find the moments of inertia of the three principal axes.
+
+    :return: diagonal of the diagonalized inertia tensor, that is
+    a shape (3,) array with the moments of inertia along the main axes.
+    (I_x, I_y and largest I_z last)
+    """
+    coords -= center_of_mass(coords, masses)
+    inertia_moment_matrix = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
+
+    for i in range(3):
+        for j in range(3):
+            k = kronecker_delta(i, j)
+            inertia_moment_matrix[i][j] = sum(
+                [
+                    masses[n] * ((norm_of(coords[n]) ** 2) * k - coords[n][i] * coords[n][j])
+                    for n, _ in enumerate(coords)
+                ]
+            )
+
+    inertia_moment_matrix = diagonalize(inertia_moment_matrix)
+
+    return np.diag(inertia_moment_matrix)
+
+
+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 center_of_mass(coords: Array3D_float, masses: Array1D_float) -> Array1D_float:
+    """Find the center of mass for the atomic system."""
+    total_mass = sum([masses[i] for i in range(len(coords))])
+    w = np.array([0.0, 0.0, 0.0])
+    for i in range(len(coords)):
+        w += coords[i] * masses[i]
+    return w / total_mass  # type: ignore[no-any-return]
+
+
+def get_moi_deviation_vec(
+    coords1: Array2D_float, coords2: Array2D_float, masses: Array1D_float
+) -> Array1D_float:
+    """Determine the relative difference of the three principal axes moments of inertia."""
+    im_1 = get_inertia_moments(coords1, masses)
+    im_2 = get_inertia_moments(coords2, masses)
+
+    return np.abs(im_1 - im_2) / im_1
+
+
+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 = np.ascontiguousarray(p.T) @ q
+
+    # Compute the SVD
+    v, _, w = np.linalg.svd(cov_mat)
+
+    if (np.linalg.det(v) * np.linalg.det(w)) < 0.0:
+        v[:, -1] = -v[:, -1]
+
+    return np.dot(v, w)  # type: ignore[no-any-return]

prism_pruner/graph_manipulations.py

@@ -0,0 +1,194 @@
+"""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 scipy.spatial.distance import cdist
+
+from prism_pruner.algebra import dihedral, norm_of
+from prism_pruner.pt import pt
+from prism_pruner.typing import Array1D_bool, Array1D_int, Array2D_float
+
+
+@lru_cache()
+def d_min_bond(a1: int, a2: int, factor: float = 1.2) -> float:
+    """Return the bond distance between two atoms."""
+    return factor * (pt[a1].covalent_radius + pt[a2].covalent_radius)  # type: ignore [no-any-return]
+
+
+def graphize(
+    atoms: Array1D_int,
+    coords: Array2D_float,
+    mask: Array1D_bool | None = None,
+) -> Graph:
+    """
+    Return a NetworkX undirected graph of molecular connectivity.
+
+    :param atoms: atomic numbers
+    :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 norm_of(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)), "atomnos")
+
+    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.
+    """
+    element = graph.nodes[index]["atomnos"]
+
+    if element not in {6, 7, 8, 15, 16}:
+        return None
+
+    d: dict[int, dict[int, int | None]] = {
+        6: {2: 1, 3: 2, 4: 3},  # C - 2 neighbors means sp, 3 nb means sp2, 4 nb sp3
+        7: {2: 2, 3: None, 4: 3},  # N - 2 neighbors means sp2, 3 nb could mean sp3 or sp2, 4 nb sp3
+        8: {1: 2, 2: 3, 3: 3, 4: 3},  # O
+        15: {2: 2, 3: 3, 4: 3},  # P - like N
+        16: {2: 2, 3: 3, 4: 3},  # S
+    }
+    return d[element].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]["atomnos"] == 7:
+        nb = set(graph.neighbors(index))
+        nb_atoms = [graph.nodes[j]["atomnos"] for j in nb]
+
+        if mode != -1:
+            # Primary amides need to have 1H, secondary amides none
+            if nb_atoms.count(1) != (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]["atomnos"] == 6:
+                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 8 in {graph.nodes[i]["atomnos"] 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]["atomnos"] == 8:
+        if 1 in (nb := set(graph.neighbors(index))):
+            return False
+
+        for n in nb:
+            if graph.nodes[n]["atomnos"] == 6:
+                nb_nb = set(graph.neighbors(n))
+                if len(nb_nb) == 3:
+                    nb_nb_sym = [graph.nodes[i]["atomnos"] for i in nb_nb]
+                    if nb_nb_sym.count(8) > 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]["atomnos"] == 1 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)
+    ]