Merge branch 'master' of github.com:ekwan/cctk

cctk/__init__.py

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

cctk/array.py

@@ -16,7 +16,16 @@ def __new__(cls, obj, **kwargs):
     def __getitem__(self, index):
         index = copy.deepcopy(index)
         if isinstance(index, slice):
-            new_index = slice(index.start - 1, index.stop - 1, index.step)
+            start, stop = None, None
+            if index.start is None:
+                start = 0
+            else:
+                start = index.start - 1
+            if index.stop is None:
+                stop = -1
+            else:
+                stop = index.stop - 1
+            new_index = slice(start, stop, index.step)
             return super().__getitem__(new_index)
         elif isinstance(index, int):
             if index > 0:

cctk/group.py

@@ -1,5 +1,6 @@
 import copy
 import numpy as np
+import networkx as nx
 
 from abc import abstractmethod
 
@@ -16,18 +17,26 @@ class Group(cctk.Molecule):
     Attributes:
         attach_to (int): atom number to replace with larger fragment. must have only one bond! (e.g. H in F3C-H)
         adjacent (int): atom number that will be bonded to new molecule. (e.g. C in F3C-H)
+        isomorphic (list of lists): list of lists of atoms that should be considered symmetry equivalent.
+            For instance, the three methyl protons can be considered symmetry equivalent, so ``methane.isomorphic = [[3, 4, 5]]``.
+        _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):
@@ -245,3 +254,26 @@ def remove_group_from_molecule(molecule, atom1, atom2, return_mapping=False):
             return new_mol, group, molecule_to_molecule, molecule_to_group
         else:
             return new_mol, group
+
+    def map_from_truncated(self):
+        """
+        Returns a dictionary mapping atomic numbers without ``attach_to`` to atomic_numbers with ``attach_to``.
+        """
+        if self._map_from_truncated is not None:
+            return self._map_from_truncated
+
+        assert self.bonds.number_of_edges() > 0, "need a bond graph to perform this operation -- try calling self.assign_connectivity()!"
+        g = copy.deepcopy(self)
+        g._add_atomic_numbers_to_nodes()
+        tg = copy.deepcopy(g)
+        tg.remove_atom(g.attach_to)
+
+        nm = nx.algorithms.isomorphism.categorical_node_match("atomic_number", 0)
+        match = nx.algorithms.isomorphism.GraphMatcher(g.bonds, tg.bonds, node_match=nm)
+
+        for sg in match.subgraph_isomorphisms_iter():
+            if self.attach_to in sg.keys():
+                continue
+            sg = {v: k for k, v in sg.items()} # invert
+            self._map_from_truncated = sg
+            return sg

cctk/load_groups.py

@@ -55,13 +55,38 @@
     ["formyl", "CHO",],
 ]
 
+isomorphic = [
+    [[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,
+    None,
+    None,
+    None,
+    None,
+    None,
+    None,
+    None,
+    None,
+    None,
+    None,
+    None,
+    None,
+    None,
+    None,
+]
 
 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}!"
@@ -71,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])