Decomposition+Assign atoms to subdomains

CHANGELOG

@@ -6,13 +6,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
-## [0.0.2] - 2020-05-04
+## [0.0.2] - 2020-05-05
 ### Added
-- debug trr frames export
-- initial code fro domain decomposition function
+- code for domain decomposition function
+- code for identify desired atoms
+- code for assigning desired atoms to subdomains
 
 ### Changed
-- None
+- debug trr frames export
 
 ### Removed
 - None

main.py

@@ -4,7 +4,15 @@ import tooba_f as tbf
 #P2=(1,1,0)
 #print(tbf.angle_between3D(P1,P2))
 
+#Read .trr file
+data_all=tbf.fr_export(trajfile='traj.trr',num_frames=1)
+#Read .gro file
+_,data_num,_,res_num,res_type,atom_type,atom_num,_ = tbf.read_gro('initial.gro')
+#Create subdomains coordinates
+box_p=tbf.domain_decomposition(data=data_all,dx=2,dy=2,dz=2)
+#Find atom type index in lists created above
+group_ndx=tbf.resid_data(atom_type, group=['C1'])
+#Assign desired atoms (from above function) to subdomains
+box_res=tbf.res2grid(data_all,atom_num,box_p, group_ndx)
+
 
-data_all=tbf.fr_export(trajfile='traj.trr',num_frames=10)
-#system,data_num,box_size,res_num,res_type,atom_type,atom_num,rest_dt = tbf.read_gro('initial.gro')
-tbf.domain_decomposition(data=data_all,dx=-1,dy=-1,dz=-1)
\ No newline at end of file

tooba_f.py

 import numpy as np
 import scipy.optimize
 
-from mpl_toolkits.mplot3d import Axes3D
+#from mpl_toolkits.mplot3d import Axes3D
 import matplotlib.pyplot as plt
 
 from pytrr import (
@@ -106,11 +106,8 @@ def fr_export(trajfile='traj.trr', num_frames=1):
             skip_trr_data(inputfile, header)
         for i in range(cnt_fr-num_frames,cnt_fr):
             header = read_trr_header(inputfile)
-            print('Step: {step}, time: {time}'.format(**header))
+            #print('Step: {step}, time: {time}'.format(**header))
             data = read_trr_data(inputfile, header)
-            #print(data.keys())
-            #print(data['box'])
-            #print(data['x'][0])
             step='{step}'.format(**header)
             data_all[step]=data
         return data_all
@@ -139,8 +136,8 @@ def read_gro(gro):
     with open(gro, 'r') as F:
         for line in F:
             cnt=cnt+1
-            print(cnt)
-            if cnt>2:
+            #print(cnt)
+            if cnt>2 and cnt<=data_num+2:
                 res_num.append(line[:5])  
                 res_type.append(line[5:10])  
                 atom_type.append(line[10:15]) 
@@ -150,28 +147,120 @@ def read_gro(gro):
                 system=line[:10]
             elif cnt==2:
                 data_num=int(line[:7])
-            if cnt>data_num:
+            elif cnt>data_num+2:
                 box_size=line[:50]
         #print(system,data_num,box_size)
         return system,data_num,box_size,res_num,res_type,atom_type,atom_num,rest_dt
 
 def domain_decomposition(data,dx,dy,dz):
     """
-    Identify residues that rely inside a subdomain 
-    of a rectangular box
+    Identify subdomain for each frame of 
+    of .trr frame
+
+    parameters:     data = {dictionary input from 'def fr_export'}
+                    dx = desired size X of subdomain {units of input}
+                    dy = desired size Y of subdomain {units of input}
+                    dz = desired size Z of subdomain {units of input}
+    
+    output:         dictionary of x,y,z grid points per step (frame)
+    """
+    box_p={}
+    for step in data.keys():
+        box_p[step]={}
+        xs=int(round(data[step]['box'][0][0]+0.1)) #to round always to bigger number
+        ys=int(round(data[step]['box'][1][1]+0.1))
+        zs=int(round(data[step]['box'][2][2]+0.1))
+        box_x=int(xs/dx)
+        box_y=int(ys/dy)
+        box_z=int(zs/dz)
+
+        xx=[]
+        for i in range(0,xs+1,box_x):
+            xx.append(i)
+        box_p[step]['x']=xx
+
+        yy=[]
+        for i in range(0,ys+1,box_y):
+            yy.append(i)
+        box_p[step]['y']=yy
+
+        zz=[]
+        for i in range(0,zs+1,box_z):
+            zz.append(i)
+        box_p[step]['z']=zz
+
+
+        xyz=[]
+        for ix in range(0,xs+1,box_x):
+          for iy in range(0,ys+1,box_y):
+            for iz in range(0,zs+1,box_z):
+                xyz.append([ix,iy,iz])
+        box_p[step]['xyz']=xyz
+
+    return box_p
+
+def resid_data(atom_type, group=[]):
+    """
+    Finds the index in list that 'def read_gro' returns,
+    and correspond to the atom types in group list
+
+    parameters:     atom_type=[]
+                    group=[]
+
+    output:         dictionary
+    """
+    ndx={}
+    for element in group:
+        ndx[element]=[i for i, e in enumerate(atom_type) if e.strip() == element.strip()]
+        #print(element,ndx)
+    return ndx
 
+def res2grid(data, atom_num, box_p, ndx):
     """
-    print('xs ys zs x y z')
+    Assign residue that corresponds to ndx (see 'def resid_data')
+    to sudomains created from 'def domain_decomposition'. The output
+    is a the box location (non zero based) in xyz-space for each atom number.
+
+    parameters:     data = {dictionary input from 'def fr_export'}
+                    atom_num = [list from 'def read_gro']
+                    box_p = {dictionary input from 'def domain_decomposition'}
+                    ndx = {dictionary input from 'def resid_data'}
+
+    output:         dictionary
+    """
+
+    box_res={}
     for step in data.keys():
-        #print(data[step]['box'])
-        #data[step]['box'][row][element]
-        xs=data[step]['box'][0][0]
-        ys=data[step]['box'][1][1]
-        zs=data[step]['box'][2][2]
-        #data[step]['x'][atom_num][x(0),y(1),z(3)]
-        x=data[step]['x'][0][0]
-        y=data[step]['x'][0][1]
-        z=data[step]['x'][0][2]
-        print(xs,ys,zs, x, y, z)
-    #Todo: identify points inside box that belong to same residue
-    #return
\ No newline at end of file
+        box_res[step]={}
+        for key in ndx.keys():
+            for i in ndx[key]:
+                #data[step]['x'][atom_num][x(0),y(1),z(2)]
+
+                xx=data[step]['x'][i][0]
+                cnt_x=-1
+                prev=-1
+                for ix in box_p[step]['x']:
+                    cnt_x=cnt_x+1
+                    if xx<ix and xx>prev:
+                        prev=ix
+                        break
+
+                yy=data[step]['x'][i][1]
+                cnt_y=-1
+                prev=-1
+                for iy in box_p[step]['y']:
+                    cnt_y=cnt_y+1
+                    if yy<iy and yy>prev:
+                        prev=iy
+                        break
+
+                zz=data[step]['x'][i][2]
+                cnt_z=-1
+                prev=-1
+                for iz in box_p[step]['z']:
+                    cnt_z=cnt_z+1
+                    if zz<iz and zz>prev:
+                        prev=iz
+                        break
+                box_res[step][atom_num[i]]=[cnt_x,cnt_y,cnt_z]
+    return box_res
\ No newline at end of file