mod: define less internals with frozen atoms

pysisyphus/intcoords/RedundantCoords.py

@@ -60,6 +60,7 @@ def __init__(
         constrain_prims=None,
         freeze_atoms=None,
         freeze_atoms_exclude=False,
+        internals_with_frozen=False,
         define_for=None,
         bonds_only=False,
         check_bends=True,
@@ -92,6 +93,7 @@ def __init__(
             freeze_atoms = list()
         self.freeze_atoms = np.array(freeze_atoms, dtype=int)
         self.freeze_atoms_exclude = freeze_atoms_exclude
+        self.internals_with_frozen = internals_with_frozen
         self.define_for = define_for
         # Constrain primitives
         if constrain_prims is None:
@@ -542,6 +544,7 @@ def set_primitive_indices(
             hybrid=self.hybrid,
             hbond_angles=self.hbond_angles,
             freeze_atoms=self.freeze_atoms if self.freeze_atoms_exclude else None,
+            internals_with_frozen=self.internals_with_frozen,
             define_for=self.define_for,
             rm_for_frag=self.rm_for_frag,
             logger=self.logger,

pysisyphus/intcoords/setup.py

@@ -642,6 +642,7 @@ def setup_redundant(
     hbond_angles=False,
     freeze_atoms=None,
     define_for=None,
+    internals_with_frozen=False,
     rm_for_frag: Optional[set] = None,
     logger=None,
 ):
@@ -654,23 +655,27 @@ def setup_redundant(
     if rm_for_frag is None:
         rm_for_frag = set()
 
-    rm_for_frag
-
     log(
         logger,
         f"Detecting primitive internals for {len(atoms)} atoms.\n"
         f"Excluding {len(freeze_atoms)} frozen atoms from the internal coordinate setup.",
     )
-    # By default all atomes are used to generate coordinates
+
+    # Mask array. By default all atomes are used to generate internal coordinates.
     use_atoms = np.ones_like(atoms, dtype=bool)
     # Only use atoms in 'define_for' to generate internal coordinates
     if define_for:
         use_atoms[:] = False  # Disable/mask all others
         use_atoms[define_for] = True
-    else:
+    # If not explicitly enabled, don't form internal coordinates containing frozen atoms.
+    # With 'internals_with_frozen', the bonds will be filtered for bonds, containing
+    # at most one frozen atom.
+    elif not internals_with_frozen:
         use_atoms[freeze_atoms] = False
+    freeze_atom_set = set(freeze_atoms)
     atoms = [atom for mobile, atom in zip(use_atoms, atoms) if mobile]
     coords3d = coords3d[use_atoms]
+
     # Maps (different) indices of mobile atoms back to their original indices
     freeze_map = {
         sub_ind: org_ind for sub_ind, org_ind in enumerate(np.where(use_atoms)[0])
@@ -695,6 +700,8 @@ def keep_coords(prims, prim_cls):
         return_cdm=True,
         return_cbm=True,
     )
+    if internals_with_frozen:
+        bonds = [bond for bond in bonds if len(set(bond) & freeze_atom_set) <= 1]
     bonds = [tuple(bond) for bond in bonds]
     bonds = keep_coords(bonds, Stretch)
     bonds = [bond for bond in bonds if rm_for_frag.isdisjoint(set(bond))]
@@ -705,7 +712,7 @@ def keep_coords(prims, prim_cls):
     ]
     # Check for unbonded single atoms and create fragments for them.
     bonded_set = set(tuple(np.ravel(bonds)))
-    unbonded_set = set(range(len(atoms))) - bonded_set
+    unbonded_set = set(range(len(atoms))) - bonded_set - freeze_atom_set
     log(
         logger,
         f"Merging bonded atoms yielded {len(fragments)} fragment(s) and "