From e5e72d36d29ee6b9e7622716a2a34e7e7adf075e Mon Sep 17 00:00:00 2001
From: Cameron Abrams <cfa22@drexel.edu>
Date: Wed, 20 Dec 2023 11:20:32 -0500
Subject: [PATCH] new stereoisomer generator

---
 HTPolyNet/coordinates.py |  15 ++++-
 HTPolyNet/matrix4.py     | 110 +++++++++++++++++++++++++++++++++
 HTPolyNet/molecule.py    |  70 +--------------------
 HTPolyNet/topocoord.py   | 127 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 253 insertions(+), 69 deletions(-)
 create mode 100644 HTPolyNet/matrix4.py

diff --git a/HTPolyNet/coordinates.py b/HTPolyNet/coordinates.py
index 4cbafba0..d52b7ff7 100644
--- a/HTPolyNet/coordinates.py
+++ b/HTPolyNet/coordinates.py
@@ -18,6 +18,7 @@
 from HTPolyNet.linkcell import Linkcell
 from HTPolyNet.ring import Ring,Segment
 from HTPolyNet.dataframetools import *
+from HTPolyNet.matrix4 import Matrix4
 
 logger=logging.getLogger(__name__)
 
@@ -675,6 +676,18 @@ def minimum_distance(self,other,self_excludes=[],other_excludes=[]):
                 if D<minD:
                     minD=D
         return minD
+    
+    def homog_trans(self,M:Matrix4,indices=[]):
+        """dfrotate applies homogeneous transformation matrix M [4x4] to coordinates
+
+        :param M: homogeneous transformation matrix
+        :type M: Matrix4
+        """
+        df=self.A
+        for i,srow in df.iterrows():
+            if len(indices)==0 or (srow['globalIdx'] in indices):
+                ri=np.array(list(srow[['posX','posY','posZ']].values))
+                df.loc[i,'posX':'posZ']=M.transform(ri)
 
     def rotate(self,R):
         """rotate Rotates all coordinate vectors by rotation matrix R
@@ -752,7 +765,7 @@ def get_R(self,idx):
         :rtype: numpy.ndarray(3,float)
         """
         df=self.A
-        return get_row_attribute(df,['posX','posY','posZ'],{'globalIdx':idx})
+        return np.array(get_row_attribute(df,['posX','posY','posZ'],{'globalIdx':idx}))
     
     def get_atom_attribute(self,name,attributes):
         """get_atom_attribute return values of attributes listed in name from atoms specified by attribute:value pairs in attributes
diff --git a/HTPolyNet/matrix4.py b/HTPolyNet/matrix4.py
new file mode 100644
index 00000000..31640422
--- /dev/null
+++ b/HTPolyNet/matrix4.py
@@ -0,0 +1,110 @@
+import numpy as np
+
+class Matrix4:
+    """Class for handling 4x4 homogeneous transformation matrices
+    """
+    def __init__(self,*args):
+        """
+        Three argument cases:
+
+        1. Two arguments: first is assumed to be a 3x3 matrix and second a 3-element array
+        2. One argument: argument's size/shape is used to determine whether it is a 3x3 matrix
+        or a 3-element array;
+        3. No arguments: the identity transformation matrix is returned
+        
+        Given rotation matrix R and translation vector T, the homogeneous transformation 
+        matrix H is defined as
+
+            ┌                                   ┐
+            │  R[0][0]  R[0][1]  R[0][2]  T[0]  │
+            │                                   │
+            │  R[1][0]  R[1][1]  R[1][2]  T[1]  │
+        H = │                                   │
+            │  R[2][0]  R[2][1]  R[2][2]  T[2]  │
+            │                                   │
+            │  0.0      0.0      0.0      1.0   │
+            └                                   ┘
+
+        """
+        if len(args)==0:
+            self.m=np.identity(4)
+        elif len(args)==2:
+            self.m=np.identity(4)
+            for i in range(3):
+                for j in range(3):
+                    self.m[i,j]=args[0][i][j]
+                self.m[3,i]=0.0
+                self.m[i,3]=args[1][i]
+        elif len(args)==1:
+            tok=args[0]
+            if tok.shape==(3,3):
+                self.__init__(tok,np.zeros(3))
+            elif len(tok.shape)==1 and tok.shape[0]==3:
+                self.__init__(np.identity(3),tok)
+        else:
+            raise Exception(f'cannot parse arguments to Matrix4.__init__: {args}')
+
+    def rot(self,degrees,axis):
+        """Apply a rotation around a Cartesian axis by a certain number of degrees;
+           right-hand rule applies
+        """
+        assert axis in 'xyz'
+        ca=np.cos(degrees*np.pi/180.0)
+        sa=np.sin(degrees*np.pi/180.0)
+        m=np.identity(3)
+        if axis=='x':
+            m[1][1]=ca
+            m[1][2]=sa
+            m[2][2]=ca
+            m[2][1]=-sa
+        elif axis=='y':
+            m[0][0]=ca
+            m[0][2]=-sa
+            m[2][2]=ca
+            m[2][0]=sa
+        elif axis=='z':
+            m[0][0]=ca
+            m[0][1]=sa
+            m[1][1]=ca
+            m[1][0]=-sa
+        v=np.zeros(3)
+        self.m=np.matmul(Matrix4(m,v).m,self.m)
+        return self
+
+    def translate(self,*args):
+        if len(args)==3:
+            tvec=np.array(args)
+        elif len(args)==1:
+            tvec=np.array(args[0])
+        self.m=np.matmul(Matrix4(tvec).m,self.m)
+        return self
+
+    def transvec(self,x,y,z):
+        theta=np.arctan2(y,x)*180.0/np.pi
+        length=np.sqrt(y*y+x*x)
+        phi=np.arctan2(z,length)*180.0/np.pi
+        self.rot(theta,'z')
+        self.rot(-phi,'y')
+        return self
+
+    def transinvec(self,x,y,z):
+        theta=np.arctan2(y,x)*180.0/np.pi
+        length=np.sqrt(y*y+x*x)
+        phi=np.arctan2(z,length)*180.0/np.pi
+        self.rot(phi,'y')
+        self.rot(-theta,'z')
+        return self
+
+    def rotate_axis(self,degrees,axis):
+        assert axis.shape==(3,)
+        self.transvec(axis[0],axis[1],axis[2])
+        self.rot(degrees,'x')
+        self.transinvec(axis[0],axis[1],axis[2])
+        return self
+
+    def transform(self,point3D):
+        point4D=np.array([point3D[0],point3D[1],point3D[2],1.0])
+        return np.dot(self.m,point4D)[:3]
+
+    def __str__(self):
+        return np.array2string(self.m)
diff --git a/HTPolyNet/molecule.py b/HTPolyNet/molecule.py
index 4c36c193..148d96dc 100644
--- a/HTPolyNet/molecule.py
+++ b/HTPolyNet/molecule.py
@@ -19,6 +19,7 @@
 from HTPolyNet.topocoord import TopoCoord
 from HTPolyNet.bondtemplate import BondTemplate,BondTemplateList,ReactionBond,ReactionBondList
 from HTPolyNet.coordinates import dfrotate, GRX_ATTRIBUTES
+from HTPolyNet.matrix4 import Matrix4
 from HTPolyNet.ambertools import GAFFParameterize
 from HTPolyNet.gromacs import mdp_modify,gro_from_trr
 from HTPolyNet.command import Command
@@ -949,74 +950,7 @@ def flip_stereocenters(self,idxlist):
         :param idxlist: global indices of chiral atoms
         :type idxlist: list
         """
-        TC=self.TopoCoord
-        A=TC.Coordinates.A
-        g=self.TopoCoord.Topology.bondlist.graph()
-        for idx in idxlist:
-            ligand_idx=self.TopoCoord.Topology.bondlist.partners_of(idx)
-            if len(ligand_idx)!=4:
-                logger.debug(f'Atom {idx} cannot be a stereocenter; it only has {len(ligand_idx)} ligands.')
-            else:
-                logger.debug(f'Flipping stereochemistry at atom {idx}')
-                # do stuff
-                cg=g.copy()
-                cg.remove_node(idx)
-                cc=[cg.subgraph(c).copy() for c in nx.connected_components(cg)]
-                ncc=len(cc)
-                associates={k:[] for k in range(ncc)}
-                for i in ligand_idx:
-                    for icc in range(ncc):
-                        tcc=cc[icc]
-                        if i in tcc:
-                            associates[icc].append(i)
-                logger.debug(f'Atom {idx} grouped associate ligands {associates}')
-
-                O=TC.get_R(idx)
-                TC.translate(-1*O)
-                r0=TC.get_R(idx)
-                # define the rotation axis
-                # define the rotating group
-                # rotate that group
-                TC.translate(O)
-        # determine the two lightest ligands
-        # branches={}
-        # for n in ligand_idx:
-        #     bl=deepcopy(self.TopoCoord.Topology.bondlist)
-        #     branches[n]=bl.half_as_list([idx,n],3)
-        # ligand_idx.sort(key=lambda x: len(branches[x]))
-        # a=ligand_idx[0]
-        # aset=list(set(list(chain.from_iterable(branches[a]))))
-        # b=ligand_idx[1]
-        # bset=list(set(list(chain.from_iterable(branches[b]))))
-        # # translate origin to location of stereocenter
-        # O=TC.get_R(idx)
-        # TC.translate(-1*O)
-        # rO=TC.get_R(idx)
-        # ra=TC.get_R(a)
-        # rb=TC.get_R(b)
-        # adf=A[A['globalIdx'].isin(aset)]
-        # bdf=A[A['globalIdx'].isin(bset)]
-        # # rotate the two branches to swap them
-        # rOa=rO-ra
-        # rOa*=1.0/np.linalg.norm(rOa)
-        # rOb=rO-rb
-        # rOb*=1.0/np.linalg.norm(rOb)
-        # cp=np.cross(rOa,rOb)
-        # c=np.dot(rOa,rOb)
-        # v=np.array([[0,-cp[2],cp[1]],[cp[2],0,-cp[0]],[-cp[1],cp[0],0]])
-        # v2=np.dot(v,v)
-        # I=np.array([[1.,0.,0.],[0.,1.,0.],[0.,0.,1.]])
-        # R=I+v+v2/(1.+c)
-        # dfrotate(adf,R)
-        # A.loc[A['globalIdx'].isin(aset),['posX','posY','posZ']]=adf[['posX','posY','posZ']]
-        # cp=np.cross(rOb,rOa)
-        # v=np.array([[0,-cp[2],cp[1]],[cp[2],0,-cp[0]],[-cp[1],cp[0],0]])
-        # v2=np.dot(v,v)
-        # R=I+v+v2/(1.+c)
-        # dfrotate(bdf,R)
-        # A.loc[A['globalIdx'].isin(bset),['posX','posY','posZ']]=bdf[['posX','posY','posZ']]
-        # # translate back to original coordinate frame
-        # TC.translate(O)
+        self.TopoCoord.flip_stereocenters(idxlist)
 
     def rotate_bond(self,a,b,deg):
         """rotate_bond rotates all atoms in molecule on b-side of a-b bond by deg degrees
diff --git a/HTPolyNet/topocoord.py b/HTPolyNet/topocoord.py
index 1b32bef6..f39df05c 100644
--- a/HTPolyNet/topocoord.py
+++ b/HTPolyNet/topocoord.py
@@ -14,9 +14,11 @@
 import os
 import shutil
 from copy import deepcopy
+import networkx as nx
 from HTPolyNet.coordinates import Coordinates, GRX_ATTRIBUTES, GRX_GLOBALLY_UNIQUE, GRX_UNSET_DEFAULTS
 from HTPolyNet.topology import Topology
 from HTPolyNet.bondtemplate import BondTemplate,ReactionBond
+from HTPolyNet.matrix4 import Matrix4
 from HTPolyNet.gromacs import grompp_and_mdrun,mdp_get, mdp_modify, gmx_energy_trace
 import HTPolyNet.projectfilesystem as pfs
 
@@ -1835,6 +1837,131 @@ def check_your_topology(self):
         print(pdf.sort_values(by='dz').head(3).to_string())
         self.write_top('checked.top')
 
+    def flip_stereocenters(self,idxlist):
+        """flip_stereocenters flips stereochemistry of atoms in idxlist
+
+        :param idxlist: global indices of chiral atoms
+        :type idxlist: list
+        """
+        A=self.Coordinates.A
+        g=self.Topology.bondlist.graph()
+        for idx in idxlist:
+            ligand_idx=self.Topology.bondlist.partners_of(idx)
+            O=self.get_R(idx)
+            name=self.get_gro_attribute_by_attributes('atomName',{'globalIdx':idx})
+            if len(ligand_idx)!=4:
+                logger.debug(f'Atom {idx} cannot be a stereocenter; it only has {len(ligand_idx)} ligands.')
+            else:
+                logger.debug(f'Flipping stereochemistry at atom {idx}')
+                # do stuff
+                cg=g.copy()
+                # remove the stereocenter and see what unconnected components are generated
+                cg.remove_node(idx)
+                cc=[cg.subgraph(c).copy() for c in nx.connected_components(cg)]
+                ncc=len(cc)
+                logger.debug(f'At stereocenter {idx} there are {ncc} unconnected ligands')
+                # for each ligand-group, assign one or more stereocenter ligand atoms
+                associates={k:[] for k in range(ncc)}
+                for i in ligand_idx:
+                    for icc in range(ncc):
+                        tcc=cc[icc]
+                        if i in tcc:
+                            associates[icc].append(i)
+                sizes=[]
+                for icc in range(ncc):
+                    sizes.append(len(cc[icc]))
+                sizes=np.array(sizes)
+                sarg=sizes.argsort()
+                logger.debug(f'Atom {idx} grouped associate ligand subgraphs {associates}')
+                logger.debug(f'Subgroups sorted by size:')
+                for sa in sarg:
+                    logger.debug(f'  {sa}: {" ".join([self.get_gro_attribute_by_attributes("atomName",{"globalIdx":x}) for x in cc[sa]])}')
+
+                # The flip is a 180-degree rotation of "half" of the atoms connected to the stereocenter
+                # about an axis defined by (a) the location of the stereocenter and (b) the midpoint on
+                # the vector connecting two specially-selected ligand atoms.
+                #
+                # How the two ligand atoms are selected:
+                # 1. If there are four unconnected ligand groups, then the two groups with the fewest atoms
+                #    are selected.
+                # 2. If there are three unconnected groups, then one of them has two ligand atoms.
+                #    These two atoms are selected if that group is smaller than the union of
+                #    the other two groups.  If that group is not smaller than the union of the
+                #    other two groups, then the ligand atoms of those other two groups are selected.
+                # 3. If there are two unconnected groups, AND both have two ligand atoms, the group with
+                #    the smaller number of atoms has its ligand atoms selected.  If one group has three and
+                #    the other has one, we cannot flip this stereocenter using this method.  However,
+                #    we could perform a reflection of all atoms in this group through the plane defined by
+                #    (a) the stereocenter, (b) one of three ligand atoms, and (c) the midpoint between the
+                #    the other two.  However, this would also flip all stereocenters in that group.
+                #    So for now, we throw an error.
+                # 4. If there is only one unconnected group, there is no way we can do this.
+                if ncc==4:
+                    l1idx=associates[sarg[0]][0]
+                    l2idx=associates[sarg[1]][0]
+                    l1name=self.get_gro_attribute_by_attributes('atomName',{'globalIdx':l1idx})
+                    l2name=self.get_gro_attribute_by_attributes('atomName',{'globalIdx':l1idx})
+                    logger.debug(f'ncc {ncc}: Atoms {l1name} and {l2name} are the selected ligands of stereocenter {name}')
+                    p1=self.get_R(l1idx)
+                    p2=self.get_R(l2idx)
+                    mp=0.5*(p1+p2)
+                    ax=O-mp
+                    g1idx=list(cc[sarg[0]])
+                    g2idx=list(cc[sarg[1]])
+                    gidx=g1idx+g2idx # set of indices of rotating group
+                    M=Matrix4()
+                    M.translate(-O).rotate_axis(180.0,ax).translate(O)
+                    self.Coordinates.homog_trans(M,indices=gidx)
+                elif ncc==3:
+                    # largest group must have the two connected atoms?  not necessarily.
+                    the_idx_of_the_double=[x for x in sarg if len(associates[x])==2][0]
+                    other_idx=[x for x in sarg if len(associates[x])==1]
+                    assert len(other_idx)==2
+                    double_size=len(cc[the_idx_of_the_double])
+                    joint_size=sum([len(cc[x]) for x in other_idx])
+                    logger.debug(f'ncc {ncc}: Size of UG with two ligands: {double_size}; side of union of UGs with one ligand each: {joint_size}')
+                    if joint_size<double_size:
+                        l1idx=associates[other_idx[0]][0]
+                        l2idx=associates[other_idx[1]][0]
+                        gidx=list(cc[other_idx[0]])+list(cc[other_idx[1]])
+                    else:
+                        l1idx=associates[the_idx_of_the_double][0]
+                        l1idx=associates[the_idx_of_the_double][1]
+                        gidx=list(cc[the_idx_of_the_double])
+                    l1name=self.get_gro_attribute_by_attributes('atomName',{'globalIdx':l1idx})
+                    l2name=self.get_gro_attribute_by_attributes('atomName',{'globalIdx':l2idx})
+                    logger.debug(f'ncc {ncc}: Atoms {l1name} and {l2name} are the selected ligands of stereocenter {name}')
+                    p1=self.get_R(l1idx)
+                    p2=self.get_R(l2idx)
+                    mp=0.5*(p1+p2)
+                    ax=O-mp
+                    M=Matrix4()
+                    M.translate(-O).rotate_axis(180.0,ax).translate(O)
+                    self.Coordinates.homog_trans(M,indices=gidx)
+                elif ncc==2:
+                    if all([len(x)==2 for x in associates.values()]):
+                        l1idx=associates[sarg[0]][0]
+                        l2idx=associates[sarg[0]][1]
+                        l1name=self.get_gro_attribute_by_attributes('atomName',{'globalIdx':l1idx})
+                        l2name=self.get_gro_attribute_by_attributes('atomName',{'globalIdx':l2idx})
+                        logger.debug(f'ncc {ncc}: Atoms {l1name} and {l2name} are the selected ligands of stereocenter {name}')
+                        p1=self.get_R(l1idx)
+                        p2=self.get_R(l2idx)
+                        mp=0.5*(p1+p2)
+                        ax=mp-O
+                        logger.debug(f'p1 {p1} p2 {p2} mp {mp} ax {ax} O {O}')
+                        gidx=list(cc[sarg[0]])
+                        logger.debug(f'The rotating group is {" ".join([self.get_gro_attribute_by_attributes("atomName",{"globalIdx":x}) for x in gidx])}')
+                        M=Matrix4()
+                        M.translate(-O).rotate_axis(180.0,ax).translate(O)
+                        self.Coordinates.homog_trans(M,indices=gidx)
+                    else:
+                        logger.debug(f'Arrangment of unconnected ligand groups at stereocenter {idx} does not permit rotation.')
+                    pass
+                elif ncc==4:
+                    logger.debug(f'Cannot flip at stereocenter {idx}.')
+
+
 def find_template(BT:BondTemplate,moldict):
     """find_template searches the dictionary of available molecules to identify a bond template that matches the passed-in template, returning the corresponding template molecule and reaction-bond