prism-pruner 0.0.1__py3-none-any.whl

prism_pruner/pruner.py

@@ -0,0 +1,571 @@
+"""PRISM - PRuning Interface for Similar Molecules."""
+
+from copy import deepcopy
+from dataclasses import dataclass, field
+from time import perf_counter
+from typing import Any, Callable, Sequence, TypeVar
+
+import numpy as np
+from networkx import Graph, connected_components
+from scipy.spatial.distance import cdist
+
+from prism_pruner.algebra import get_moi_deviation_vec
+from prism_pruner.pt import pt
+from prism_pruner.rmsd import rmsd_and_max
+from prism_pruner.torsion_module import (
+    get_angles,
+    get_hydrogen_bonds,
+    get_torsions,
+    is_nondummy,
+    rotationally_corrected_rmsd_and_max,
+)
+from prism_pruner.typing import (
+    Array1D_bool,
+    Array1D_float,
+    Array1D_int,
+    Array2D_float,
+    Array2D_int,
+    Array3D_float,
+)
+from prism_pruner.utils import flatten, get_double_bonds_indices, time_to_string
+
+__version__ = "1.0.0"
+
+
+@dataclass
+class PrunerConfig:
+    """Configuration dataclass for Pruner."""
+
+    structures: Array3D_float
+
+    # Optional parameters that get initialized
+    energies: Array1D_float = field(default_factory=lambda: np.array([]))
+    ewin: float = field(default=0.0)
+    debugfunction: Callable[[str], None] | None = field(default=None)
+
+    # Computed fields
+    calls: int = field(default=0, init=False)
+    cache_calls: int = field(default=0, init=False)
+    cache: set[tuple[int, int]] = field(default_factory=lambda: set(), init=False)
+
+    def __post_init__(self) -> None:
+        """Validate inputs and initialize computed fields."""
+        self.mask = np.ones(shape=(self.structures.shape[0],), dtype=np.bool_)
+
+        if len(self.energies) != 0:
+            assert self.ewin > 0.0, (
+                "If you provide energies, please also provide an appropriate energy window ewin."
+            )
+
+        # Set defaults for optional parameters
+        if len(self.energies) == 0:
+            self.energies = np.zeros(self.structures.shape[0])
+
+        if self.ewin == 0.0:
+            self.ewin = 1.0
+
+    def evaluate_sim(self, *args: Any, **kwargs: Any) -> bool:
+        """Stub method - override in subclasses as needed."""
+        raise NotImplementedError
+
+
+PrunerConfigType = TypeVar("PrunerConfigType", bound=PrunerConfig)
+
+
+@dataclass
+class RMSDRotCorrPrunerConfig(PrunerConfig):
+    """Configuration dataclass for Pruner."""
+
+    atomnos: Array1D_int = field(kw_only=True)
+    max_rmsd: float = field(kw_only=True)
+    max_dev: float = field(kw_only=True)
+    angles: Sequence[Sequence[int]] = field(kw_only=True)
+    torsions: Array2D_int = field(kw_only=True)
+    graph: Graph = field(kw_only=True)
+    heavy_atoms_only: bool = True
+
+    def evaluate_sim(self, coord1: Array2D_float, coord2: Array2D_float) -> bool:
+        """Return if the structures are similar."""
+        rmsd, max_dev = rotationally_corrected_rmsd_and_max(
+            coord1,
+            coord2,
+            atomnos=self.atomnos,
+            torsions=self.torsions,
+            graph=self.graph,
+            angles=self.angles,
+            debugfunction=self.debugfunction,
+            heavy_atoms_only=self.heavy_atoms_only,
+        )
+
+        if rmsd > self.max_rmsd:
+            return False
+
+        if max_dev > self.max_dev:
+            return False
+
+        return True
+
+
+@dataclass
+class RMSDPrunerConfig(PrunerConfig):
+    """Configuration dataclass for Pruner."""
+
+    atomnos: Array1D_int = field(kw_only=True)
+    max_rmsd: float = field(kw_only=True)
+    max_dev: float = field(kw_only=True)
+    heavy_atoms_only: bool = True
+
+    def evaluate_sim(self, coord1: Array2D_float, coord2: Array2D_float) -> bool:
+        """Return if the structures are similar."""
+        if self.heavy_atoms_only:
+            mask = self.atomnos != 1
+        else:
+            mask = np.ones(self.structures[0].shape[0], dtype=bool)
+
+        rmsd, max_dev = rmsd_and_max(
+            coord1[mask],
+            coord2[mask],
+            center=True,
+        )
+
+        if rmsd > self.max_rmsd:
+            return False
+
+        if max_dev > self.max_dev:
+            return False
+
+        return True
+
+
+@dataclass
+class MOIPrunerConfig(PrunerConfig):
+    """Configuration dataclass for Pruner."""
+
+    masses: Array1D_float = field(kw_only=True)
+    max_dev: float = 0.01
+
+    def evaluate_sim(self, coord1: Array2D_float, coord2: Array2D_float) -> bool:
+        """Return if the structures are similar."""
+        dev_vec = get_moi_deviation_vec(
+            coord1,
+            coord2,
+            masses=self.masses,
+        )
+
+        return bool((dev_vec < self.max_dev).all())
+
+
+def _main_compute_subrow(
+    prunerconfig: PrunerConfigType,
+    ref: Array2D_float,
+    structures: Array3D_float,
+    in_mask: Array1D_bool,
+    first_abs_index: int,
+) -> bool:
+    """Evaluate the similarity of a subrow of the similarity matrix.
+
+    Return True if ref is similar to any
+    structure in structures, returning at the first instance of a match.
+    Ignores structures that are False (0) in in_mask and does not perform
+    the comparison if the energy difference between the structures is less
+    than self.ewin. Saves dissimilar structural pairs (i.e. that evaluate to
+    False (0)) by adding them to self.cache, avoiding redundant calcaulations.
+    """
+    # iterate over target structures
+    for i, structure in enumerate(structures):
+        # only compare active structures
+        if in_mask[i]:
+            # check if we have performed this comparison already:
+            # if so, we already know that those two structures are not similar,
+            # since the in_mask attribute is not False for ref nor for i
+            i1 = first_abs_index
+            i2 = first_abs_index + 1 + i
+            hash_value = (i1, i2)
+
+            prunerconfig.calls += 1
+            if hash_value in prunerconfig.cache:
+                prunerconfig.cache_calls += 1
+
+            # if we have not computed the value before, check if the two
+            # structures have close enough energy before running the comparison
+            elif np.abs(prunerconfig.energies[i1] - prunerconfig.energies[i2]) < prunerconfig.ewin:
+                # function will return True if the structures are similar,
+                # and will stop iterating on this row, returning
+                if prunerconfig.evaluate_sim(ref, structure):
+                    return True
+
+            # if structures are not similar, add the result to the
+            # cache, because they will return here,
+            # while similar structures are discarded and won't come back
+            prunerconfig.cache.add(hash_value)
+
+    return False
+
+
+def _main_compute_row(
+    prunerconfig: PrunerConfigType,
+    structures: Array3D_float,
+    in_mask: Array1D_bool,
+    first_abs_index: int,
+) -> Array1D_bool:
+    """Evaluate the similarity of a row of the similarity matrix.
+
+    For a given set of structures, check if each is similar
+    to any other after itself. Return a boolean mask to slice
+    the array, only retaining the structures that are dissimilar.
+    The inner subrow function caches computed non-similar pairs.
+
+    """
+    # initialize the result container
+    out_mask = np.ones(shape=in_mask.shape, dtype=np.bool_)
+
+    # loop over the structures
+    for i, ref in enumerate(structures):
+        # only check for similarity if the structure is active
+        if in_mask[i]:
+            # reject structure i if it is similar to any other after itself
+            similar = _main_compute_subrow(
+                prunerconfig,
+                ref,
+                structures[i + 1 :],
+                in_mask[i + 1 :],
+                first_abs_index=first_abs_index + i,
+            )
+            out_mask[i] = not similar
+
+        else:
+            out_mask[i] = 0
+
+    return out_mask
+
+
+def _main_compute_group(
+    prunerconfig: PrunerConfigType,
+    structures: Array2D_float,
+    in_mask: Array1D_bool,
+    k: int,
+) -> Array1D_bool:
+    """Evaluate the similarity of each chunk of the similarity matrix.
+
+    Acts individually on k chunks of the structures array,
+    returning the updated mask.
+    """
+    # initialize final result container
+    out_mask = np.ones(shape=structures.shape[0], dtype=np.bool_)
+
+    # calculate the size of each chunk
+    chunksize = int(len(structures) // k)
+
+    # iterate over chunks (multithreading here?)
+    for chunk in range(int(k)):
+        first = chunk * chunksize
+        if chunk == k - 1:
+            last = len(structures)
+        else:
+            last = chunksize * (chunk + 1)
+
+        # get the structure chunk
+        structures_chunk = structures[first:last]
+
+        # compare structures within that chunk and save results to the out_mask
+        out_mask[first:last] = _main_compute_row(
+            prunerconfig,
+            structures_chunk,
+            in_mask[first:last],
+            first_abs_index=first,
+        )
+    return out_mask
+
+
+def prune(prunerconfig: PrunerConfigType) -> tuple[Array2D_float, Array1D_bool]:
+    """Perform the similarity pruning.
+
+    Remove similar structures by repeatedly grouping them into k
+    subgroups and removing similar ones. A cache is present to avoid
+    repeating RMSD computations.
+
+    Similarity occurs for structures with both rmsd < self.max_rmsd and
+    maximum absolute atomic deviation < self.max_dev.
+
+    Sets the self.structures and the corresponding self.mask attributes.
+    """
+    start_t = perf_counter()
+
+    structures = deepcopy(prunerconfig.structures)
+
+    # initialize the output mask
+    out_mask = np.ones(shape=prunerconfig.structures.shape[0], dtype=np.bool_)
+    prunerconfig.cache = set()
+
+    # split the structure array in subgroups and prune them internally
+    for k in (
+        5e5,
+        2e5,
+        1e5,
+        5e4,
+        2e4,
+        1e4,
+        5000,
+        2000,
+        1000,
+        500,
+        200,
+        100,
+        50,
+        20,
+        10,
+        5,
+        2,
+        1,
+    ):
+        # choose only k values such that every subgroup
+        # has on average at least twenty active structures in it
+        if k == 1 or 20 * k < np.count_nonzero(out_mask):
+            before = np.count_nonzero(out_mask)
+
+            start_t_k = perf_counter()
+
+            # compute similarities and get back the out_mask
+            # and the pairings to be added to cache
+            out_mask = _main_compute_group(
+                prunerconfig,
+                structures,
+                out_mask,
+                k=int(k),
+            )
+
+            after = np.count_nonzero(out_mask)
+            newly_discarded = before - after
+
+            if prunerconfig.debugfunction is not None:
+                elapsed = start_t_k - perf_counter()
+                prunerconfig.debugfunction(
+                    f"DEBUG: {prunerconfig.__class__.__name__} - k={k}, rejected {newly_discarded} "
+                    + f"(keeping {after}/{len(out_mask)}), in {time_to_string(elapsed)}"
+                )
+
+    del prunerconfig.cache
+
+    if prunerconfig.debugfunction is not None:
+        elapsed = start_t - perf_counter()
+        prunerconfig.debugfunction(
+            f"DEBUG: {prunerconfig.__class__.__name__} - keeping "
+            + f"{after}/{len(out_mask)} "
+            + f"({time_to_string(elapsed)})"
+        )
+
+        fraction = 0 if prunerconfig.calls == 0 else prunerconfig.cache_calls / prunerconfig.calls
+        prunerconfig.debugfunction(
+            f"DEBUG: {prunerconfig.__class__.__name__} - Used cached data "
+            + f"{prunerconfig.cache_calls}/{prunerconfig.calls} times, "
+            + f"{100 * fraction:.2f}% of total calls"
+        )
+
+    return prunerconfig.structures[out_mask], out_mask
+
+
+def prune_by_rmsd(
+    structures: Array3D_float,
+    atomnos: Array1D_int,
+    max_rmsd: float = 0.25,
+    max_dev: float | None = None,
+    debugfunction: Callable[[str], None] | None = None,
+) -> tuple[Array3D_float, Array1D_bool]:
+    """Remove duplicate structures using a heavy-atom RMSD metric.
+
+    Remove similar structures by repeatedly grouping them into k
+    subgroups and removing similar ones. A cache is present to avoid
+    repeating RMSD computations.
+
+    Similarity occurs for structures with both RMSD < max_rmsd and
+    maximum deviation < max_dev. max_dev by default is 2 * max_rmsd.
+    """
+    # set default max_dev if not provided
+    max_dev = max_dev or 2 * max_rmsd
+
+    # set up PrunerConfig dataclass
+    prunerconfig = RMSDPrunerConfig(
+        structures=structures,
+        atomnos=atomnos,
+        max_rmsd=max_rmsd,
+        max_dev=max_dev,
+        debugfunction=debugfunction,
+    )
+
+    # run the pruning
+    return prune(prunerconfig)
+
+
+def prune_by_rmsd_rot_corr(
+    structures: Array3D_float,
+    atomnos: Array1D_int,
+    graph: Graph,
+    max_rmsd: float = 0.25,
+    max_dev: float | None = None,
+    logfunction: Callable[[str], None] | None = None,
+    debugfunction: Callable[[str], None] | None = None,
+) -> tuple[Array3D_float, Array1D_bool]:
+    """Remove duplicates using a heavy-atom RMSD metric, corrected for degenerate torsions.
+
+    Remove similar structures by repeatedly grouping them into k
+    subgroups and removing similar ones. A cache is present to avoid
+    repeating RMSD computations.
+
+    Similarity occurs for structures with both RMSD < max_rmsd and
+    maximum deviation < max_dev. max_dev by default is 2 * max_rmsd.
+
+    The RMSD and maximum deviation metrics used are the lowest ones
+    of all the degenerate rotamers of the input structure.
+    """
+    # center structures
+    structures = np.array([s - s.mean(axis=0) for s in structures])
+    ref = structures[0]
+
+    # get the number of molecular fragments
+    subgraphs = list(connected_components(graph))
+
+    # if they are more than two, give up on pruning by rot corr rmsd
+    if len(subgraphs) > 2:
+        return structures, np.ones(structures.shape[0], dtype=bool)
+
+    # if they are two, we can add a fictitious bond between the closest
+    # atoms on the two molecular fragment in the provided graph, and
+    # then removing it before returning
+    if len(subgraphs) == 2:
+        subgraphs = [list(vals) for vals in connected_components(graph)]
+        all_dists_array = cdist(ref[list(subgraphs[0])], ref[list(subgraphs[1])])
+        min_d = np.min(all_dists_array)
+        s1, s2 = np.where(all_dists_array == min_d)
+        i1, i2 = subgraphs[0][s1[0]], subgraphs[1][s2[0]]
+        graph.add_edge(i1, i2)
+
+        if debugfunction is not None:
+            debugfunction(
+                f"DEBUG: prune_by_rmsd_rot_corr - temporarily added "
+                f"edge {i1}-{i2} to the graph (will be removed before returning)"
+            )
+
+    # set default max_dev if not provided
+    max_dev = max_dev or 2 * max_rmsd
+
+    # add hydrogen bonds to molecular graph
+    hydrogen_bonds = get_hydrogen_bonds(ref, atomnos, graph)
+    for hb in hydrogen_bonds:
+        graph.add_edge(*hb)
+
+    # keep an unraveled set of atoms in hbs
+    flat_hbs = set(flatten(hydrogen_bonds))
+
+    # get all rotable bonds in the molecule, including dummy rotations
+    torsions = get_torsions(
+        graph,
+        hydrogen_bonds=hydrogen_bonds,
+        double_bonds=get_double_bonds_indices(ref, atomnos),
+        keepdummy=True,
+        mode="symmetry",
+    )
+
+    # only keep dummy rotations (checking both directions)
+    torsions = [
+        t
+        for t in torsions
+        if not (is_nondummy(t.i2, t.i3, graph) and (is_nondummy(t.i3, t.i2, graph)))
+    ]
+
+    # since we only compute RMSD based on heavy atoms, discard
+    # quadruplets that involve hydrogen atom as termini, unless
+    # they are involved in hydrogen bonding
+    torsions = [
+        t
+        for t in torsions
+        if (1 not in [atomnos[i] for i in t.torsion])
+        or (t.torsion[0] in flat_hbs or t.torsion[3] in flat_hbs)
+    ]
+
+    # get torsions angles
+    angles = [get_angles(t, graph) for t in torsions]
+
+    # Used specific directionality of torsions so that we always
+    # rotate the dummy portion (the one attached to the last index)
+    torsions_ids = np.asarray(
+        [
+            list(t.torsion) if is_nondummy(t.i2, t.i3, graph) else list(reversed(t.torsion))
+            for t in torsions
+        ]
+    )
+
+    # Set up final mask and cache
+    final_mask = np.ones(structures.shape[0], dtype=bool)
+
+    # Halt the run if there are too many structures or no subsymmetrical bonds
+    if len(torsions_ids) == 0:
+        if debugfunction is not None:
+            debugfunction(
+                "DEBUG: prune_by_rmsd_rot_corr - No subsymmetrical torsions found: skipping "
+                "symmetry-corrected RMSD pruning"
+            )
+
+        return structures[final_mask], final_mask
+
+    # Print out torsion information
+    if logfunction is not None:
+        logfunction("\n >> Dihedrals considered for rotamer corrections:")
+        for i, (torsion, angle) in enumerate(zip(torsions_ids, angles, strict=False)):
+            logfunction(
+                " {:2s} - {:21s} : {}{}{}{} : {}-fold".format(
+                    str(i + 1),
+                    str(torsion),
+                    pt[atomnos[torsion[0]]].symbol,
+                    pt[atomnos[torsion[1]]].symbol,
+                    pt[atomnos[torsion[2]]].symbol,
+                    pt[atomnos[torsion[3]]].symbol,
+                    len(angle),
+                )
+            )
+        logfunction("\n")
+
+    # Initialize PrunerConfig
+    prunerconfig = RMSDRotCorrPrunerConfig(
+        structures=structures,
+        atomnos=atomnos,
+        graph=graph,
+        torsions=torsions_ids,
+        debugfunction=debugfunction,
+        angles=angles,
+        max_rmsd=max_rmsd,
+        max_dev=max_dev,
+    )
+
+    # run pruning
+    structures_out, mask = prune(prunerconfig)
+
+    # remove the extra bond in the molecular graph
+    if len(subgraphs) == 2:
+        graph.remove_edge(i1, i2)
+
+    return structures_out, mask
+
+
+def prune_by_moment_of_inertia(
+    structures: Array3D_float,
+    atomnos: Array1D_int,
+    max_deviation: float = 1e-2,
+    debugfunction: Callable[[str], None] | None = None,
+) -> tuple[Array3D_float, Array1D_bool]:
+    """Remove duplicate structures using a moments of inertia-based metric.
+
+    Remove duplicate structures (enantiomeric or rotameric) based on the
+    moments of inertia on the principal axes. If all three MOI
+    deviate less than max_deviation percent from another structure,
+    they are classified as rotamers or enantiomers and therefore only one
+    of them is kept (i.e. max_deviation = 0.1 is 10% relative deviation).
+    """
+    # set up PrunerConfig dataclass
+    prunerconfig = MOIPrunerConfig(
+        structures=structures,
+        debugfunction=debugfunction,
+        max_dev=max_deviation,
+        masses=np.array([pt[a].mass for a in atomnos]),
+    )
+
+    return prune(prunerconfig)

prism_pruner/pt.py

@@ -0,0 +1,12 @@
+"""PRISM - PRuning Interface for Similar Molecules."""
+
+from periodictable import core, covalent_radius, mass
+
+for pt_n in range(5):
+    try:
+        pt = core.PeriodicTable(table=f"H={pt_n + 1}")
+        covalent_radius.init(pt)
+        mass.init(pt)
+    except ValueError:
+        continue
+    break

prism_pruner/rmsd.py

@@ -0,0 +1,39 @@
+"""PRISM - PRuning Interface for Similar Molecules."""
+
+import numpy as np
+
+from prism_pruner.algebra import get_alignment_matrix, norm_of
+from prism_pruner.typing import Array2D_float
+
+
+def rmsd_and_max(
+    p: Array2D_float,
+    q: Array2D_float,
+    center: bool = False,
+) -> tuple[float, float]:
+    """Return RMSD and max deviation.
+
+    Return a tuple with the RMSD between p and q
+    and the maximum deviation of their positions.
+    """
+    if center:
+        p -= p.mean(axis=0)
+        q -= q.mean(axis=0)
+
+    # get alignment matrix
+    rot_mat = get_alignment_matrix(p, q)
+
+    # Apply it to p
+    p = np.ascontiguousarray(p) @ rot_mat
+
+    # Calculate deviations
+    diff = p - q
+
+    # Calculate RMSD
+    rmsd = np.sqrt((diff * diff).sum() / len(diff))
+
+    # # Calculate max deviation
+    # max_delta = np.linalg.norm(diff, axis=1).max()
+    max_delta = max([norm_of(v) for v in diff])
+
+    return rmsd, max_delta