refactor: read atoms.info['connectivity'] through one shared helper

pyproject.toml

@@ -15,6 +15,7 @@ dependencies = [
     "packmol>=21.1.2",
     "rdkit>=2024",
     "typer>=0.27.1",
+    "typing-extensions>=4.15.0",
 ]
 
 [dependency-groups]

src/molify/ase2x.py

@@ -6,6 +6,7 @@
 from rdkit import Chem
 
 from molify.constants import EdgeAttr, GraphAttr, NodeAttr
+from molify.utils import read_connectivity
 
 try:
     import vesin
@@ -131,6 +132,12 @@ def ase2networkx(
     networkx.Graph
         An undirected NetworkX graph with connectivity information.
 
+    Raises
+    ------
+    ValueError
+        For an invalid bond in ``atoms.info['connectivity']``, see
+        :func:`molify.utils.read_connectivity`.
+
     Notes
     -----
     The graph contains the following information:
@@ -172,13 +179,7 @@ def ase2networkx(
     charges = atoms.get_initial_charges()
 
     if GraphAttr.CONNECTIVITY in atoms.info:
-        connectivity = atoms.info[GraphAttr.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
-        ]
+        connectivity = read_connectivity(atoms)
         return _create_graph_from_connectivity(atoms, connectivity, charges)
 
     pairs = _compute_bonded_pairs(atoms, scale, pbc)

src/molify/compress.py

@@ -3,7 +3,7 @@
 from ase.cell import Cell
 
 from molify.constants import GraphAttr
-from molify.utils import calculate_box_dimensions
+from molify.utils import calculate_box_dimensions, fragment_indices
 
 
 def compress(
@@ -13,58 +13,36 @@ def compress(
 ) -> ase.Atoms:
     """Compress an ASE Atoms object to a target density.
 
-    Arguments
-    ---------
+    Parameters
+    ----------
     atoms : ase.Atoms
         The Atoms object to compress.
     density : float
         The target density in kg/m^3.
     freeze_molecules : bool
         If True, freeze the internal degrees of freedom of the molecules
         during compression, to prevent bond compression.
+
+    Raises
+    ------
+    ValueError
+        With ``freeze_molecules=True``, for a missing ``info['connectivity']``
+        or an invalid bond in it, see :func:`molify.utils.read_connectivity`.
     """
     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(GraphAttr.CONNECTIVITY)
-        if connectivity is None:
+        if atoms.info.get(GraphAttr.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)
+        fragments = fragment_indices(atoms)
+        new_cell = Cell.new(new_dimensions)
+        t_mat = np.linalg.solve(atoms.get_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
+        for fragment in fragments:
+            com = atoms[fragment].get_center_of_mass()
+            positions[fragment] += com @ t_mat - com
 
         atoms.set_positions(positions)
         atoms.set_cell(new_cell, scale_atoms=False)

src/molify/pack.py

@@ -10,7 +10,7 @@
 from rdkit import Chem
 
 from molify.constants import GraphAttr
-from molify.utils import calculate_box_dimensions
+from molify.utils import calculate_box_dimensions, read_connectivity
 
 log = logging.getLogger(__name__)
 
@@ -90,13 +90,29 @@ def _write_molecule_files(
             ase.io.write(filepath, atoms)
 
 
+def _combine_bonds(
+    selected_images: list[ase.Atoms],
+) -> list[tuple[int, int, float | None]] | None:
+    """Offset each molecule's bonds into the packed box.
+
+    Returns None when a molecule lacks ``info['connectivity']``.
+    """
+    if not all(GraphAttr.CONNECTIVITY in atoms.info for atoms in selected_images):
+        return None
+    bonds = []
+    offset = 0
+    for atoms in selected_images:
+        for i, j, order in read_connectivity(atoms):
+            bonds.append((i + offset, j + offset, order))
+        offset += len(atoms)
+    return bonds
+
+
 def _extract_atom_arrays(
     selected_images: list[ase.Atoms], packed_atoms: ase.Atoms
 ) -> ase.Atoms:
     """Extracts and adds relevant atom arrays (if present)
 
-    Add bonds from the input structures to the packed structure if available.
-
     Parameters
     ----------
     selected_images : list[ase.Atoms]
@@ -107,7 +123,7 @@ def _extract_atom_arrays(
     Returns
     -------
     ase.Atoms
-        The packed ASE Atoms object with the copied arrays and bonds.
+        The packed ASE Atoms object with the copied arrays.
     """
     array_keys = [
         "occupancy",
@@ -126,14 +142,6 @@ def _extract_atom_arrays(
             )
             packed_atoms.arrays[key] = concatenated_array
 
-    if all(GraphAttr.CONNECTIVITY in atoms.info for atoms in selected_images):
-        bonds = []
-        offset = 0
-        for atoms in selected_images:
-            for bond in atoms.info[GraphAttr.CONNECTIVITY]:
-                bonds.append((bond[0] + offset, bond[1] + offset, bond[2]))
-            offset += len(atoms)
-        packed_atoms.info[GraphAttr.CONNECTIVITY] = bonds
     charges = np.concatenate([atom.get_initial_charges() for atom in selected_images])
     if any(charge != 0 for charge in charges):
         packed_atoms.set_initial_charges(charges)
@@ -189,6 +197,12 @@ def pack(
     ase.Atoms
         An ASE Atoms object representing the packed system.
 
+    Raises
+    ------
+    ValueError
+        For an invalid bond in a molecule's ``info['connectivity']``, see
+        :func:`molify.utils.read_connectivity`.
+
     Example
     -------
     >>> from molify import pack, smiles2conformers
@@ -200,6 +214,7 @@ def pack(
     Atoms(symbols='C10H44O12', pbc=True, cell=[8.4, 8.4, 8.4])
     """
     selected_images = _select_conformers(data, counts, seed)
+    bonds = _combine_bonds(selected_images)
     cell = calculate_box_dimensions(images=selected_images, density=density)
 
     # Adjust cell dimensions according to ratio while keeping volume unchanged
@@ -240,4 +255,6 @@ def pack(
     packed_atoms.cell = cell
     packed_atoms.pbc = True
     packed_atoms = _extract_atom_arrays(selected_images, packed_atoms)
+    if bonds is not None:
+        packed_atoms.info[GraphAttr.CONNECTIVITY] = bonds
     return packed_atoms

src/molify/repeat.py

@@ -5,6 +5,7 @@
 from ase.geometry import find_mic, minkowski_reduce
 
 from molify.constants import GraphAttr, NodeAttr
+from molify.utils import read_connectivity
 
 
 def _normalize_rep(rep: int | Sequence[int]) -> tuple[int, ...]:
@@ -26,24 +27,10 @@ def _tile_bonds(
     atoms: ase.Atoms, reps: tuple[int, ...]
 ) -> list[tuple[int, int, float | None]]:
     n_atoms = len(atoms)
-    first, second, orders = [], [], []
-    for bond in atoms.info[GraphAttr.CONNECTIVITY]:
-        i, j = int(bond[0]), int(bond[1])
-        if not (0 <= i < n_atoms and 0 <= j < n_atoms):
-            raise ValueError(
-                f"bond ({i}, {j}) in atoms.info['connectivity'] needs atom "
-                f"indices in the range 0..{n_atoms - 1}"
-            )
-        if i == j:
-            raise ValueError(
-                f"bond ({i}, {j}) in atoms.info['connectivity'] links atom {i} to "
-                "itself"
-            )
-        first.append(i)
-        second.append(j)
-        orders.append(None if bond[2] is None else float(bond[2]))
-    i = np.array(first, dtype=int)
-    j = np.array(second, dtype=int)
+    bonds = read_connectivity(atoms)
+    i = np.array([bond[0] for bond in bonds], dtype=int)
+    j = np.array([bond[1] for bond in bonds], dtype=int)
+    orders = [bond[2] for bond in bonds]
 
     periodic = atoms.pbc & atoms.cell.any(1)
     d = atoms.positions[j] - atoms.positions[i]

src/molify/substructure.py

@@ -8,7 +8,7 @@
 
 from molify.ase2x import ase2rdkit
 from molify.constants import GraphAttr
-from molify.utils import find_connected_components
+from molify.utils import fragment_indices
 
 
 def match_substructure(  # noqa: C901
@@ -289,26 +289,26 @@ def iter_fragments(atoms: ase.Atoms) -> list[ase.Atoms]:
     ase.Atoms
         Each connected component (fragment) in the input structure.
 
+    Raises
+    ------
+    ValueError
+        For an invalid bond in ``atoms.info['connectivity']``, see
+        :func:`molify.utils.read_connectivity`.
+
     Examples
     --------
-    >>> from molify import smiles2atoms
-    >>> from rdkit.Chem import CombineMols
+    >>> from molify import iter_fragments, pack, smiles2conformers
     >>>
-    >>> # Create multi-fragment system
-    >>> ethanol = smiles2atoms("CCO")
-    >>> methanol = smiles2atoms("CO")
-    >>> combined = ethanol + methanol
+    >>> water = smiles2conformers("O", 1)
+    >>> ethanol = smiles2conformers("CCO", 1)
+    >>> box = pack([water, ethanol], [2, 1], density=800)
     >>>
-    >>> # Iterate over fragments
-    >>> fragments = list(iter_fragments(combined))
-    >>> len(fragments)
-    2
+    >>> [fragment.get_chemical_formula() for fragment in iter_fragments(box)]
+    ['H2O', 'H2O', 'C2H6O']
     """
     if GraphAttr.CONNECTIVITY in atoms.info:
-        # connectivity is a list of tuples (i, j, bond_type)
-        connectivity = atoms.info[GraphAttr.CONNECTIVITY]
-        for component in find_connected_components(connectivity):
-            yield atoms[list(component)]
+        for indices in fragment_indices(atoms):
+            yield atoms[indices]
     else:
         for molecule in separate(atoms):
             yield molecule