molify 0.0.1a0__py3-none-any.whl

molify/__init__.py

@@ -0,0 +1,41 @@
+from molify.ase2x import ase2networkx, ase2rdkit
+from molify.com import get_centers_of_mass
+from molify.compress import compress
+from molify.networkx2x import networkx2ase, networkx2rdkit
+from molify.pack import pack
+from molify.rdkit2x import rdkit2ase, rdkit2networkx
+from molify.smiles2x import smiles2atoms, smiles2conformers
+from molify.substructure import (
+    get_substructures,
+    iter_fragments,
+    match_substructure,
+    select_atoms_flat_unique,
+    select_atoms_grouped,
+    visualize_selected_molecules,
+)
+from molify.utils import unwrap_structures
+
+__all__ = [
+    "ase2rdkit",
+    "smiles2atoms",
+    "pack",
+    "smiles2conformers",
+    "compress",
+    "match_substructure",
+    "get_substructures",
+    "iter_fragments",
+    "select_atoms_flat_unique",
+    "visualize_selected_molecules",
+    "unwrap_structures",
+    "select_atoms_grouped",
+    #
+    "ase2networkx",
+    #
+    "rdkit2networkx",
+    "rdkit2ase",
+    #
+    "networkx2rdkit",
+    "networkx2ase",
+    #
+    "get_centers_of_mass",
+]

molify/ase2x.py

@@ -0,0 +1,220 @@
+import ase
+import networkx as nx
+import numpy as np
+from ase.neighborlist import natural_cutoffs, neighbor_list
+from rdkit import Chem
+
+try:
+    import vesin
+except ImportError:
+    vesin = None
+
+
+def _create_graph_from_connectivity(
+    atoms: ase.Atoms, connectivity, charges
+) -> nx.Graph:
+    """Create NetworkX graph from explicit connectivity information."""
+    graph = nx.Graph()
+    graph.graph["pbc"] = atoms.pbc
+    graph.graph["cell"] = atoms.cell
+
+    for i, atom in enumerate(atoms):
+        graph.add_node(
+            i,
+            position=atom.position,
+            atomic_number=int(atom.number),
+            original_index=atom.index,
+            charge=charges[i],
+        )
+
+    for i, j, bond_order in connectivity:
+        graph.add_edge(i, j, bond_order=bond_order)
+    return graph
+
+
+def _compute_connectivity_matrix(atoms: ase.Atoms, scale: float, pbc: bool):
+    """Compute connectivity matrix from distance-based cutoffs."""
+    # non-bonding positive charged atoms / ions.
+    non_bonding_atomic_numbers = {3, 11, 19, 37, 55, 87}
+
+    atomic_numbers = atoms.get_atomic_numbers()
+    excluded_mask = np.isin(atomic_numbers, list(non_bonding_atomic_numbers))
+
+    atom_radii = np.array(natural_cutoffs(atoms, mult=scale))
+    pairwise_cutoffs = atom_radii[:, None] + atom_radii[None, :]
+    max_cutoff = np.max(pairwise_cutoffs)
+
+    if vesin is not None:
+        try:
+            i, j, d, s = vesin.ase_neighbor_list(
+                "ijdS", atoms, cutoff=max_cutoff, self_interaction=False
+            )
+        except Exception as e:
+            print(f"vesin failed with {e}, trying native ASE implementation")
+            i, j, d, s = neighbor_list(
+                "ijdS", atoms, cutoff=max_cutoff, self_interaction=False
+            )
+    else:
+        i, j, d, s = neighbor_list(
+            "ijdS", atoms, cutoff=max_cutoff, self_interaction=False
+        )
+
+    # If pbc=False, filter out bonds that cross periodic boundaries
+    if not pbc:
+        non_periodic_mask = np.all(s == 0, axis=1)
+        i = i[non_periodic_mask]
+        j = j[non_periodic_mask]
+        d = d[non_periodic_mask]
+
+    d_ij = np.full((len(atoms), len(atoms)), np.inf)
+    d_ij[i, j] = d
+    np.fill_diagonal(d_ij, 0.0)
+
+    # mask out non-bonding atoms
+    d_ij[excluded_mask, :] = np.inf
+    d_ij[:, excluded_mask] = np.inf
+
+    connectivity_matrix = np.zeros((len(atoms), len(atoms)), dtype=int)
+    np.fill_diagonal(d_ij, np.inf)
+    connectivity_matrix[d_ij <= pairwise_cutoffs] = 1
+
+    return connectivity_matrix, non_bonding_atomic_numbers
+
+
+def _add_node_properties(
+    graph: nx.Graph, atoms: ase.Atoms, charges, non_bonding_atomic_numbers
+):
+    """Add node properties to the graph."""
+    for i, atom in enumerate(atoms):
+        graph.nodes[i]["position"] = atom.position
+        graph.nodes[i]["atomic_number"] = int(atom.number)
+        graph.nodes[i]["original_index"] = atom.index
+        graph.nodes[i]["charge"] = float(charges[i])
+        if atom.number in non_bonding_atomic_numbers:
+            graph.nodes[i]["charge"] = 1.0
+
+
+def ase2networkx(
+    atoms: ase.Atoms,
+    pbc: bool = True,
+    scale: float = 1.2,
+) -> nx.Graph:
+    """Convert an ASE Atoms object to a NetworkX graph.
+
+    Determines which atoms are bonded (connectivity).
+    All edges will have bond_order=None unless atoms.info['connectivity']
+    already has bond orders.
+
+    Parameters
+    ----------
+    atoms : ase.Atoms
+        The ASE Atoms object to convert into a graph.
+    pbc : bool, optional
+        Whether to consider periodic boundary conditions when calculating
+        distances (default is True). If False, only connections within
+        the unit cell are considered.
+    scale : float, optional
+        Scaling factor for the covalent radii when determining bond cutoffs
+        (default is 1.2).
+
+    Returns
+    -------
+    networkx.Graph
+        An undirected NetworkX graph with connectivity information.
+
+    Notes
+    -----
+    The graph contains the following information:
+
+    - Nodes represent atoms with properties:
+        * position: Cartesian coordinates (numpy.ndarray)
+        * atomic_number: Element atomic number (int)
+        * original_index: Index in original Atoms object (int)
+        * charge: Formal charge (float)
+    - Edges represent bonds with:
+        * bond_order: Bond order (float or None if unknown)
+    - Graph properties include:
+        * pbc: Periodic boundary conditions
+        * cell: Unit cell vectors
+
+    Connectivity is determined by:
+
+    1. Using explicit connectivity if present in atoms.info
+    2. Otherwise using distance-based cutoffs (edges will have bond_order=None)
+
+    To get bond orders, pass the graph to networkx2rdkit().
+
+    Examples
+    --------
+    >>> from rdkit2ase import ase2networkx, smiles2atoms
+    >>> atoms = smiles2atoms(smiles="O")
+    >>> graph = ase2networkx(atoms)
+    >>> len(graph.nodes)
+    3
+    >>> len(graph.edges)
+    2
+    """
+    if len(atoms) == 0:
+        return nx.Graph()
+    charges = atoms.get_initial_charges()
+
+    if "connectivity" in atoms.info:
+        connectivity = atoms.info["connectivity"]
+        # ensure connectivity is list[tuple[int, int, float|None]] and
+        # does not contain np.generic
+        connectivity = [
+            (int(i), int(j), float(bond_order) if bond_order is not None else None)
+            for i, j, bond_order in connectivity
+        ]
+        return _create_graph_from_connectivity(atoms, connectivity, charges)
+
+    connectivity_matrix, non_bonding_atomic_numbers = _compute_connectivity_matrix(
+        atoms, scale, pbc
+    )
+
+    graph = nx.from_numpy_array(connectivity_matrix, edge_attr=None)
+    for u, v in graph.edges():
+        graph.edges[u, v]["bond_order"] = None
+
+    _add_node_properties(graph, atoms, charges, non_bonding_atomic_numbers)
+
+    graph.graph["pbc"] = atoms.pbc
+    graph.graph["cell"] = atoms.cell
+
+    return graph
+
+
+def ase2rdkit(atoms: ase.Atoms, suggestions: list[str] | None = None) -> Chem.Mol:
+    """Convert an ASE Atoms object to an RDKit molecule.
+
+    Convenience function that chains:
+    ase2networkx() → networkx2rdkit(suggestions=...)
+
+    Parameters
+    ----------
+    atoms : ase.Atoms
+        The ASE Atoms object to convert.
+    suggestions : list[str], optional
+        SMILES/SMARTS patterns for bond order determination.
+        Passed directly to networkx2rdkit().
+
+    Returns
+    -------
+    rdkit.Chem.Mol
+        The resulting RDKit molecule with bond orders determined.
+
+    Examples
+    --------
+    >>> from rdkit2ase import ase2rdkit, smiles2atoms
+    >>> atoms = smiles2atoms(smiles="C=O")
+    >>> mol = ase2rdkit(atoms)
+    >>> mol.GetNumAtoms()
+    4
+    """
+    if len(atoms) == 0:
+        return Chem.Mol()
+
+    from molify import ase2networkx, networkx2rdkit
+
+    graph = ase2networkx(atoms)
+    return networkx2rdkit(graph, suggestions=suggestions)

molify/bond_order.py

@@ -0,0 +1,110 @@
+import networkx as nx
+from networkx.algorithms import isomorphism
+
+from molify.utils import (
+    rdkit_determine_bonds,
+    suggestions2networkx,
+    unwrap_structures,
+)
+
+
+def sort_templates(graphs: list[nx.Graph]) -> list[nx.Graph]:
+    return sorted(
+        graphs, key=lambda g: (g.number_of_nodes(), g.number_of_edges()), reverse=True
+    )
+
+
+def has_bond_order(graph: nx.Graph) -> bool:
+    return all(x[2]["bond_order"] is not None for x in graph.edges(data=True))
+
+
+def update_graph_data(
+    working_graph: nx.Graph, graph: nx.Graph, mol_graph: nx.Graph
+) -> bool:
+    graph_matcher = isomorphism.GraphMatcher(
+        working_graph,
+        mol_graph,
+        node_match=lambda n1, n2: n1["atomic_number"] == n2["atomic_number"],
+    )
+    match = next(graph_matcher.subgraph_isomorphisms_iter(), None)
+    if match is None:
+        return False
+
+    # Remove matched nodes from future matching
+    working_graph.remove_nodes_from(match.keys())
+
+    # Transfer bond order from template
+    inv_match = {v: k for k, v in match.items()}
+    for u, v, data in mol_graph.edges(data=True):
+        if "bond_order" in data:
+            if data["bond_order"] is None:
+                raise ValueError(
+                    "Cannot update bond order with None value. "
+                    "Please check the input templates."
+                )
+            u_g = inv_match[u]
+            v_g = inv_match[v]
+            if graph.has_edge(u_g, v_g):
+                graph[u_g][v_g]["bond_order"] = data["bond_order"]
+
+    # Assign charges for matched nodes
+    for graph_node, template_node in match.items():
+        if "charge" in mol_graph.nodes[template_node]:
+            graph.nodes[graph_node]["charge"] = mol_graph.nodes[template_node]["charge"]
+    return True
+
+
+def update_bond_order_from_suggestions(graph, suggestions: list[nx.Graph]) -> None:
+    if has_bond_order(graph):
+        return
+
+    working_graph = graph.copy()
+    suggestions = sort_templates(suggestions)
+
+    for mol_graph in suggestions:
+        while True:
+            updated = update_graph_data(working_graph, graph, mol_graph)
+            if not updated:
+                break
+
+
+def update_bond_order(graph: nx.Graph, suggestions: list[str] | None = None) -> None:
+    if not has_bond_order(graph):
+        if suggestions is not None:
+            suggestion_graphs = suggestions2networkx(suggestions)
+            for component in nx.connected_components(graph):
+                subgraph = graph.subgraph(component)
+                update_bond_order_from_suggestions(subgraph, suggestion_graphs)
+                for u, v, data in subgraph.edges(data=True):
+                    graph[u][v].update(data)
+
+        update_bond_order_determine(graph)
+
+
+def update_bond_order_determine(graph: nx.Graph) -> None:
+    """
+    Update the bond order in the graph based on the suggestions.
+    If the graph already has bond order, it will not be updated.
+    """
+    from molify.networkx2x import networkx2ase
+    from molify.rdkit2x import rdkit2networkx
+
+    if has_bond_order(graph):
+        return
+
+    for component in nx.connected_components(graph):
+        subgraph = graph.subgraph(component)
+        # add pbc and cell information to the subgraph
+        subgraph.graph.update(graph.graph)
+        missing = sum(
+            data.get("bond_order") is None for u, v, data in subgraph.edges(data=True)
+        )
+        # if any bond information is missing, we update the bond order
+        if missing > 0:
+            # Unwrapping could be made nicer, by utilizing the connectivity
+            atoms = networkx2ase(subgraph)
+            atoms = unwrap_structures(atoms)
+            rdkit_mol = rdkit_determine_bonds(atoms)
+            rdkit_graph = rdkit2networkx(rdkit_mol)
+            # update the bond order in the original graph
+            update_bond_order_from_suggestions(graph, [rdkit_graph])

molify/com.py

@@ -0,0 +1,75 @@
+import ase
+import networkx as nx
+
+
+def get_centers_of_mass(atoms: ase.Atoms, unwrap: bool = True, **kwargs) -> ase.Atoms:
+    """Compute the center of mass for each molecule in an ASE Atoms object.
+
+    Parameters
+    ----------
+    atoms : ase.Atoms
+        The ASE Atoms object containing the molecular structures.
+    unwrap : bool, optional
+        If True, unwrap the structures before computing the center of mass.
+    **kwargs : dict
+        Additional keyword arguments to pass to the `ase2networkx` function.
+
+    Returns
+    -------
+    ase.Atoms
+        An ASE Atoms object containing the centers of mass of each molecule,
+        with the masses of the molecules as the `masses` attribute.
+        The species will start at 1 and increment
+        for each unique molecule type.
+
+
+    Example
+    -------
+    >>> from rdkit2ase import smiles2conformers, pack, get_centers_of_mass
+    >>> ethanol = smiles2conformers("CCO", numConfs=10)
+    >>> box = pack([ethanol], [10], density=786)
+    >>> centers = get_centers_of_mass(box)
+    >>> print(centers.get_positions().shape)
+    (10, 3)
+    """
+    from molify import ase2networkx, networkx2ase, unwrap_structures
+
+    if unwrap:
+        graph = ase2networkx(atoms, **kwargs)
+        atoms = unwrap_structures(atoms)  # Ensure structures are unwrapped
+    graph = ase2networkx(atoms, **kwargs)
+
+    centers = []
+    masses = []
+
+    counts = {}
+
+    for component in nx.connected_components(graph):
+        subgraph = graph.subgraph(component)
+        molecule = networkx2ase(subgraph)
+        center_of_mass = molecule.get_center_of_mass()
+        centers.append(center_of_mass)
+        masses.append(molecule.get_masses().sum())
+
+        for key in counts:
+            g = counts[key]["graph"]
+            if nx.is_isomorphic(subgraph, g):
+                counts[key]["count"] += 1
+                break
+        else:
+            counts[len(counts)] = {
+                "graph": subgraph,
+                "count": 1,
+            }
+
+    numbers = []
+    for i, (key, value) in enumerate(counts.items()):
+        numbers.extend([i + 1] * value["count"])
+
+    return ase.Atoms(
+        positions=centers,
+        masses=masses,
+        cell=atoms.get_cell(),
+        pbc=atoms.pbc,
+        numbers=numbers,
+    )

molify/compress.py

@@ -0,0 +1,73 @@
+import ase
+import numpy as np
+from ase.cell import Cell
+
+from molify.utils import calculate_box_dimensions
+
+
+def compress(
+    atoms: ase.Atoms,
+    density: float,
+    freeze_molecules: bool = False,
+) -> ase.Atoms:
+    """Compress an ASE Atoms object to a target density.
+
+    Arguments
+    ---------
+    atoms : ase.Atoms
+        The Atoms object to compress.
+    density : float
+        The target density in g/cm^3.
+    freeze_molecules : bool
+        If True, freeze the internal degrees of freedom of the molecules
+        during compression, to prevent bond compression.
+    """
+    atoms = atoms.copy()
+    new_dimensions = np.array(calculate_box_dimensions([atoms], density))
+
+    if freeze_molecules:
+        new_cell = Cell.new(new_dimensions)
+        old_cell = atoms.get_cell()
+        connectivity = atoms.info.get("connectivity")
+        if connectivity is None:
+            raise ValueError("No connectivity info found for freeze_molecules=True")
+
+        # Build molecular fragments
+        from collections import defaultdict
+
+        groups = defaultdict(set)
+        parent = {}
+
+        def find(x):
+            while parent.get(x, x) != x:
+                x = parent[x]
+            return x
+
+        def union(x, y):
+            rx, ry = find(x), find(y)
+            if rx != ry:
+                parent[ry] = rx
+
+        for i, j, _ in connectivity:
+            union(i, j)
+
+        for idx in range(len(atoms)):
+            root = find(idx)
+            groups[root].add(idx)
+
+        t_mat = np.linalg.solve(old_cell.T, new_cell.T)
+
+        positions = atoms.get_positions()
+        for group in groups.values():
+            group = list(group)
+            com = atoms[group].get_center_of_mass()
+            com_new = com @ t_mat
+            shift = com_new - com
+            positions[group] += shift
+
+        atoms.set_positions(positions)
+        atoms.set_cell(new_cell, scale_atoms=False)
+    else:
+        atoms.set_cell(new_dimensions, scale_atoms=True)
+
+    return atoms

molify/networkx2x.py

@@ -0,0 +1,187 @@
+import ase
+import networkx as nx
+import numpy as np
+from rdkit import Chem
+
+from molify.utils import bond_type_from_order
+
+
+def networkx2ase(graph: nx.Graph) -> ase.Atoms:
+    """Convert a NetworkX graph to an ASE Atoms object.
+
+    Parameters
+    ----------
+    graph : networkx.Graph
+        The molecular graph to convert. Node attributes must include:
+            * position (numpy.ndarray): Cartesian coordinates
+            * atomic_number (int): Element atomic number
+            * charge (float, optional): Formal charge
+        Edge attributes must include:
+            * bond_order (float): Bond order
+        Optional graph attributes:
+            * pbc (bool): Periodic boundary conditions
+            * cell (numpy.ndarray): Unit cell vectors
+
+    Returns
+    -------
+    ase.Atoms
+        The resulting Atoms object with:
+            * Atomic positions and numbers
+            * Initial charges if present in graph
+            * Connectivity information stored in atoms.info
+            * PBC and cell if present in graph
+
+    Examples
+    --------
+    >>> import networkx as nx
+    >>> from rdkit2ase import networkx2ase
+    >>> graph = nx.Graph()
+    >>> graph.add_node(0, position=[0,0,0], atomic_number=1, charge=0)
+    >>> graph.add_node(1, position=[1,0,0], atomic_number=1, charge=0)
+    >>> graph.add_edge(0, 1, bond_order=1.0)
+    >>> atoms = networkx2ase(graph)
+    >>> len(atoms)
+    2
+    """
+    # Create mapping from original node indices to new sequential indices
+    node_mapping = {node: i for i, node in enumerate(graph.nodes)}
+
+    positions = np.array([graph.nodes[n]["position"] for n in graph.nodes])
+    numbers = np.array([graph.nodes[n]["atomic_number"] for n in graph.nodes])
+    charges = np.array([graph.nodes[n]["charge"] for n in graph.nodes])
+
+    atoms = ase.Atoms(
+        positions=positions,
+        numbers=numbers,
+        charges=charges,
+        pbc=graph.graph.get("pbc", False),
+        cell=graph.graph.get("cell", None),
+    )
+
+    connectivity = []
+    for u, v, data in graph.edges(data=True):
+        bond_order = data["bond_order"]
+        # Map original node indices to new sequential indices
+        new_u = node_mapping[u]
+        new_v = node_mapping[v]
+        connectivity.append((new_u, new_v, bond_order))
+
+    atoms.info["connectivity"] = connectivity
+
+    return atoms
+
+
+def networkx2rdkit(graph: nx.Graph, suggestions: list[str] | None = None) -> Chem.Mol:
+    """Convert a NetworkX graph to an RDKit molecule.
+
+    Automatically determines bond orders for edges with bond_order=None.
+
+    Parameters
+    ----------
+    graph : networkx.Graph
+        The molecular graph to convert
+
+        Node attributes:
+            * atomic_number (int): Element atomic number
+            * charge (float, optional): Formal charge
+
+        Edge attributes:
+            * bond_order (float or None): Bond order. If None, will be determined
+              automatically.
+
+    suggestions : list[str], optional
+        Optional SMILES/SMARTS patterns to help determine bond orders.
+        Used when edges have bond_order=None.
+
+    Returns
+    -------
+    rdkit.Chem.Mol
+        The resulting RDKit molecule with:
+            * Atoms and bonds from the graph
+            * Formal charges if specified
+            * All bond orders determined
+            * Sanitized molecular structure
+
+    Raises
+    ------
+    ValueError
+        If nodes are missing atomic_number attribute, or if bond order
+        determination fails for any edge.
+
+    Examples
+    --------
+    >>> import networkx as nx
+    >>> from rdkit2ase import networkx2rdkit
+    >>> graph = nx.Graph()
+    >>> graph.add_node(0, atomic_number=6, charge=0)
+    >>> graph.add_node(1, atomic_number=8, charge=0)
+    >>> graph.add_edge(0, 1, bond_order=2.0)
+    >>> mol = networkx2rdkit(graph)
+    >>> mol.GetNumAtoms()
+    2
+
+    >>> # With automatic bond order determination
+    >>> graph2 = nx.Graph()
+    >>> graph2.add_node(0, atomic_number=6, charge=0, position=[0, 0, 0])
+    >>> graph2.add_node(1, atomic_number=8, charge=0, position=[1.2, 0, 0])
+    >>> graph2.add_edge(0, 1, bond_order=None)  # Will be determined
+    >>> mol2 = networkx2rdkit(graph2)
+    """
+    from molify.bond_order import update_bond_order
+
+    # Check if any edges have bond_order=None and determine them if needed
+    has_none_bond_orders = any(
+        data.get("bond_order") is None for u, v, data in graph.edges(data=True)
+    )
+
+    if has_none_bond_orders:
+        # Make a copy to avoid modifying the input graph
+        graph = graph.copy()
+        update_bond_order(graph, suggestions)
+
+        # Verify all bond orders are now determined
+        for u, v, data in graph.edges(data=True):
+            bond_order = data.get("bond_order")
+            if bond_order is None:
+                raise ValueError(
+                    f"Failed to determine bond order for edge ({u}, {v}).\n"
+                    f"This can happen with:\n"
+                    f"  - Unusual bonding situations\n"
+                    f"  - Disconnected atoms\n"
+                    f"  - Invalid atomic coordinates\n\n"
+                    f"Try:\n"
+                    f"  1. Check atom positions are reasonable\n"
+                    f"  2. Provide SMILES suggestions for complex molecules\n"
+                    f"  3. Manually set connectivity with bond orders"
+                )
+
+    mol = Chem.RWMol()
+    nx_to_rdkit_atom_map = {}
+
+    for node_id, attributes in graph.nodes(data=True):
+        atomic_number = attributes.get("atomic_number")
+        charge = attributes.get("charge", 0)
+
+        if atomic_number is None:
+            raise ValueError(f"Node {node_id} is missing 'atomic_number' attribute.")
+
+        atom = Chem.Atom(int(atomic_number))
+        atom.SetFormalCharge(int(charge))
+        idx = mol.AddAtom(atom)
+        nx_to_rdkit_atom_map[node_id] = idx
+
+    for u, v, data in graph.edges(data=True):
+        bond_order = data.get("bond_order")
+
+        if bond_order is None:
+            raise ValueError(f"Edge ({u}, {v}) is missing 'bond_order' attribute.")
+
+        mol.AddBond(
+            nx_to_rdkit_atom_map[u],
+            nx_to_rdkit_atom_map[v],
+            bond_type_from_order(bond_order),
+        )
+
+    Chem.SanitizeMol(mol)
+
+    return mol.GetMol()