added new detection of symmetric atoms to molecule

cctk/__init__.py

@@ -11,3 +11,4 @@
 from .mol2_file import MOL2File
 from .mae_file import MAEFile
 from .pdb_file import PDBFile
+

cctk/group.py

@@ -22,17 +22,21 @@ class Group(cctk.Molecule):
         _map_from_truncated(dict): a dictionary mapping atom numbers of the group without ``attach_to`` to the atom numbers of the normal group
     """
 
-    def __init__(self, attach_to, **kwargs):
+    def __init__(self, attach_to, isomorphic=None, **kwargs):
         super().__init__(**kwargs)
         self.add_attachment_point(attach_to)
         self._map_from_truncated = None
 
+        if isomorphic is not None:
+            assert isinstance(isomorphic, list), "group.isomorphic must be list of lists!"
+        self.isomorphic = isomorphic
+
     @classmethod
-    def new_from_molecule(cls, molecule, attach_to):
+    def new_from_molecule(cls, molecule, attach_to, **kwargs):
         """
         Convenient method to convert ``molecule`` to ``group`` directly.
         """
-        group = Group(attach_to, atomic_numbers=molecule.atomic_numbers, geometry=molecule.geometry, bonds=molecule.bonds.edges())
+        group = Group(attach_to, atomic_numbers=molecule.atomic_numbers, geometry=molecule.geometry, bonds=molecule.bonds.edges(), **kwargs)
         return group
 
     def add_attachment_point(self, attach_to):

cctk/load_groups.py

@@ -56,10 +56,10 @@
 ]
 
 isomorphic = [
-    [3, 4, 5],
-    None,
-    None,
+    [[3, 4, 5]],
     None,
+    [[4, 8], [9, 10, 11, 5, 6, 7]],
+    [[3, 7, 11], [4, 5, 6, 8, 9, 10, 12, 13, 14]],
     None,
     None,
     None,
@@ -81,10 +81,12 @@
 
 def load_group(name):
     filename = None
+    iso = None
 
     for row in names:
         if name in row:
             filename = filenames[names.index(row)]
+            iso = isomorphic[names.index(row)]
             break
 
     assert filename is not None, f"can't find name {name}!"
@@ -94,5 +96,15 @@ def load_group(name):
         mol.assign_connectivity()
 
         #### every molecule is set so you need to attach to atom 2
-        new_group = Group.new_from_molecule(attach_to=2, molecule=mol)
+        new_group = Group.new_from_molecule(attach_to=2, molecule=mol, isomorphic=iso)
         return new_group
+
+def group_iterator(symmetric_only=False):
+    """
+    Returns a generator over all *cctk*-predefined groups.
+    """
+    for row, iso in zip(names, isomorphic):
+        if symmetric_only:
+            if iso is None:
+                continue
+        yield load_group(row[0])

cctk/molecule.py

@@ -17,7 +17,6 @@
     get_isotopic_distribution,
     compute_chirality,
 )
-
 import cctk.topology as top
 
 class Molecule:
@@ -1566,12 +1565,42 @@ def new_from_smiles(cls, smiles):
             raise ValueError(f"something went wrong auto-generating molecule {smiles}:\n{e}")
 
     def fragment(self):
+        """
+        Returns list of ``cctk.Molecule`` objects based on the bond-connected components of ``self``.
+        """
         fragments = list()
         indices = self.get_components()
         for idx in indices:
             mol = cctk.Molecule(self.atomic_numbers[idx], self.geometry[idx]).assign_connectivity()
             fragments.append(mol)
-
         return fragments
 
+    def get_symmetric_atoms(self):
+        """
+        Returns lists of symmetric atoms, as defined in ``cctk.load_group``.
+
+        Useful for NMR spectroscopy, etc.
+        """
+        from cctk.load_groups import group_iterator
+        symmetric_sets = []
+        for group in group_iterator(symmetric_only=True):
+            # this gives us a list of dictionaries mapping from self.atomic_numbers to group numbers
+            matches = top.find_group(self, group)
+
+            for m in matches:
+                i = {v: k for k,v in m.items()}
+                for n in group.isomorphic:
+                    symmetric_sets.append([i[idx] for idx in n])
+
+        #### some groups overlap (e.g. methyl and t-butyl), so now we collapse the overlapping sets
+        for i, s1 in enumerate(symmetric_sets):
+            for j, s2 in enumerate(symmetric_sets[i+1:]):
+                if set(s1).intersection(set(s2)):
+                    symmetric_sets[i + j + 1] = list(set(s1).union(s2))
+                    symmetric_sets[i] = None # can't delete yet - messes up indexing
+
+        #### now we delete
+        symmetric_sets = list(filter(None, symmetric_sets))
+        return symmetric_sets
+
 

test/test_group.py

@@ -77,6 +77,12 @@ def test_subgroup_isomorphism(self):
         molecule = cctk.Molecule.new_from_name("camphor").assign_connectivity()
 
         subgroups = top.find_group(molecule, group)
+        self.assertEqual(len(subgroups), 3)
+
+        molecule = cctk.Molecule.new_from_name("tert-butanol").assign_connectivity()
+        sets = molecule.get_symmetric_atoms()
+        self.assertEqual(len(sets[0]), 3)
+        self.assertEqual(len(sets[1]), 9)
 
 if __name__ == '__main__':
     unittest.main()