v0.2.11 (#55)

* bump version

* empty molecule edge cases

* rename `unwrap_structures`

* update docstring

* add `scale` parameter

* add TODO

* simplify unwrap code

* remove tests that did nothing

* test cyclic molecules

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* address comments

* include node mapping for non-sequential data

* add test

* add one more test

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: PythonFZ

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rdkit2ase"
-version = "0.1.10"
+version = "0.1.11"
 description = "Interface between rdkit and ASE"
 readme = "README.md"
 license = "Apache-2.0"

rdkit2ase/__init__.py

@@ -5,6 +5,7 @@
 from rdkit2ase.rdkit2x import rdkit2ase, rdkit2networkx
 from rdkit2ase.smiles2x import smiles2atoms, smiles2conformers
 from rdkit2ase.substructure import get_substructures, iter_fragments, match_substructure
+from rdkit2ase.utils import unwrap_structures
 
 __all__ = [
     "ase2rdkit",
@@ -15,6 +16,7 @@
     "match_substructure",
     "get_substructures",
     "iter_fragments",
+    "unwrap_structures",
     #
     "ase2networkx",
     #

rdkit2ase/ase2x.py

@@ -12,8 +12,89 @@
     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=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:
+        i, j, d, s = vesin.ase_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"] = 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, suggestions: list[str] | None = None, pbc: bool = True
+    atoms: ase.Atoms,
+    suggestions: list[str] | None = None,
+    pbc: bool = True,
+    scale: float = 1.2,
 ) -> nx.Graph:
     """Convert an ASE Atoms object to a NetworkX graph with bonding information.
 
@@ -33,6 +114,9 @@ def ase2networkx(
         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
     -------
@@ -70,85 +154,25 @@ def ase2networkx(
     >>> 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"]
-        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=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
-
-    # 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=1.2))
-    pairwise_cutoffs = atom_radii[:, None] + atom_radii[None, :]
-
-    max_cutoff = np.max(pairwise_cutoffs)
-
-    if vesin is not None:
-        i, j, d, s = vesin.ase_neighbor_list(
-            "ijdS", atoms, cutoff=max_cutoff, self_interaction=False
+        return _create_graph_from_connectivity(
+            atoms, atoms.info["connectivity"], charges
         )
-    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:
-        # Keep only bonds where all shift vectors are zero (no periodic wrapping)
-        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
+    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
 
-    for i, atom in enumerate(atoms):
-        graph.nodes[i]["position"] = atom.position
-        graph.nodes[i]["atomic_number"] = 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
+    _add_node_properties(graph, atoms, charges, non_bonding_atomic_numbers)
 
     graph.graph["pbc"] = atoms.pbc
     graph.graph["cell"] = atoms.cell
@@ -188,6 +212,9 @@ def ase2rdkit(atoms: ase.Atoms, suggestions: list[str] | None = None) -> Chem.Mo
     >>> mol.GetNumAtoms()
     4
     """
+    if len(atoms) == 0:
+        return Chem.Mol()
+
     from rdkit2ase import ase2networkx, networkx2rdkit
 
     graph = ase2networkx(atoms, suggestions=suggestions)

rdkit2ase/bond_order.py

@@ -1,7 +1,11 @@
 import networkx as nx
 from networkx.algorithms import isomorphism
 
-from rdkit2ase.utils import rdkit_determine_bonds, suggestions2networkx, unwrap_molecule
+from rdkit2ase.utils import (
+    rdkit_determine_bonds,
+    suggestions2networkx,
+    unwrap_structures,
+)
 
 
 def sort_templates(graphs: list[nx.Graph]) -> list[nx.Graph]:
@@ -99,7 +103,7 @@ def update_bond_order_determine(graph: nx.Graph) -> None:
         if missing > 0:
             # Unwrapping could be made nicer, by utilizing the connectivity
             atoms = networkx2ase(subgraph)
-            atoms = unwrap_molecule(atoms)
+            atoms = unwrap_structures(atoms)
             rdkit_mol = rdkit_determine_bonds(atoms)
             rdkit_graph = rdkit2networkx(rdkit_mol)
             # update the bond order in the original graph

rdkit2ase/networkx2x.py

@@ -43,6 +43,9 @@ def networkx2ase(graph: nx.Graph) -> ase.Atoms:
     >>> 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])
@@ -58,7 +61,10 @@ def networkx2ase(graph: nx.Graph) -> ase.Atoms:
     connectivity = []
     for u, v, data in graph.edges(data=True):
         bond_order = data["bond_order"]
-        connectivity.append((u, v, 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
 

rdkit2ase/rdkit2x.py

@@ -33,6 +33,9 @@ def rdkit2ase(mol: Chem.Mol, seed: int = 42) -> ase.Atoms:
     >>> len(atoms)
     9
     """
+    if mol.GetNumAtoms() == 0:
+        return ase.Atoms()
+
     smiles = Chem.MolToSmiles(mol)
     mol = Chem.AddHs(mol)
     charges = [atom.GetFormalCharge() for atom in mol.GetAtoms()]