tooba_f.py 25.8 KB
Newer Older
Stelios Karozis's avatar
Stelios Karozis committed
1
import os
Stelios Karozis's avatar
Stelios Karozis committed
2
import numpy as np
3 4
#from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
Stelios Karozis's avatar
Stelios Karozis committed
5
import scipy.optimize
Stelios Karozis's avatar
Stelios Karozis committed
6 7
import pickle as pkl
import json
Stelios Karozis's avatar
Stelios Karozis committed
8
import re
Stelios Karozis's avatar
Stelios Karozis committed
9
import pandas as pd
Stelios Karozis's avatar
Stelios Karozis committed
10
from progress.bar import Bar
Stelios Karozis's avatar
Stelios Karozis committed
11

Stelios Karozis's avatar
Stelios Karozis committed
12 13 14 15 16 17
from pytrr import (
    read_trr_header,
    read_trr_data,
    skip_trr_data,
    )

Stelios Karozis's avatar
Stelios Karozis committed
18 19 20 21 22 23 24 25 26 27 28 29
def search_pattern(fl,pattern):
    textfile = open(fl, 'r')
    filetext = textfile.read()
    textfile.close()
    matches = re.findall(pattern, filetext)
    
    if len(matches)==0:
        ind='not exist'
    else:
        ind='exist'
    return ind

Stelios Karozis's avatar
Stelios Karozis committed
30
def fitPlaneLTSQ(XYZ):
31
    #Source: https://gist.github.com/RustingSword/e22a11e1d391f2ab1f2c
Stelios Karozis's avatar
Stelios Karozis committed
32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62
    (rows, cols) = XYZ.shape
    G = np.ones((rows, 3))
    G[:, 0] = XYZ[:, 0]  #X
    G[:, 1] = XYZ[:, 1]  #Y
    Z = XYZ[:, 2]
    (a, b, c),resid,rank,s = np.linalg.lstsq(G, Z, rcond=None)
    normal = (a, b, -1)
    nn = np.linalg.norm(normal)
    normal = normal / nn
    return (c, normal)

def angle_between2D(p1, p2):
    #arctan2 is anticlockwise
    ang1 = np.arctan2(*p1[::-1])
    ang2 = np.arctan2(*p2[::-1])
    angle=np.rad2deg((ang1 - ang2) % (2 * np.pi)) #degree
    if angle<180:
        return angle
    else:
        return 360 - angle

def angle_between3D(v1, v2):
# v1 is your firsr vector
# v2 is your second vector
    angle = np.arccos(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))) #rad
    angle=180*angle/np.pi #degrees
    if angle<180:
        return angle
    else:
        return 360 - angle

Stelios Karozis's avatar
Stelios Karozis committed
63
def center_gravity(a):
64
    #Todo somatidia diaforetikis mazas -> weighted me maza ????
Stelios Karozis's avatar
Stelios Karozis committed
65 66 67 68 69
    m=len(a)
    cg = np.sum(a)/m
    return cg

def topickle(fl, sv_name):
70 71
    print(' ')
    print('Save to pickle |################################| 1/1')
Stelios Karozis's avatar
Stelios Karozis committed
72 73 74
    with open(sv_name+'.pkl', 'wb') as handle:
        pkl.dump(fl, handle, protocol=pkl.HIGHEST_PROTOCOL)

Stelios Karozis's avatar
Stelios Karozis committed
75 76 77 78 79 80
def frompickle(fl):
    with open(fl, 'rb') as handle:
        b = pkl.load(handle)
    return b


Stelios Karozis's avatar
Stelios Karozis committed
81
def tojson(fl, sv_name):
82 83
    print(' ')
    print('Save to json |################################| 1/1')
Stelios Karozis's avatar
Stelios Karozis committed
84 85 86
    with open(sv_name+'.json', 'w') as file:
        file.write(json.dumps(str(fl)))

87
def plot_surf(data, normal, c, save_name):
Stelios Karozis's avatar
Stelios Karozis committed
88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112
    #Plot surface
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    # plot fitted plane
    maxx = np.max(data[:,0])
    maxy = np.max(data[:,1])
    minx = np.min(data[:,0])
    miny = np.min(data[:,1])

    point = np.array([0.0, 0.0, c])
    d = -point.dot(normal)

    # plot original points
    ax.scatter(data[:, 0], data[:, 1], data[:, 2])

    # compute needed points for plane plotting
    xx, yy = np.meshgrid([minx, maxx], [miny, maxy])
    z = (-normal[0]*xx - normal[1]*yy - d)*1. / normal[2]

    # plot plane
    ax.plot_surface(xx, yy, z, alpha=0.2)

    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
113 114
    #plt.savefig(save_name+'.png')
    plt.show()
Stelios Karozis's avatar
Stelios Karozis committed
115 116 117
    plt.clf()
    plt.cla()
    plt.close()
118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138

def points2vector(data):
    """
    Defines the line whose direction vector is the eigenvector 
    of the covariance matrix corresponding to the largest eigenvalue, 
    that passes through the mean of the data. 

    parameters:     data = array[[x y z]]

    output:         vector = [X,Y,Z]
    """
    # Calculate the mean of the points, i.e. the 'center' of the cloud
    datamean = data.mean(axis=0)

    # Do an SVD on the mean-centered data.
    uu, dd, vv = np.linalg.svd(data - datamean)

    # Now vv[0] contains the first principal component, i.e. the direction
    # vector of the 'best fit' line in the least squares sense.

    return vv[0]
Stelios Karozis's avatar
Stelios Karozis committed
139

Stelios Karozis's avatar
Stelios Karozis committed
140
def count_frames(trajfile='traj.trr', pr_time=False):
141 142 143 144 145
    """
    Count total frames of .trr file

    parameters:     trajfile = [.trr]
    """
Stelios Karozis's avatar
Stelios Karozis committed
146
    time_index={}
Stelios Karozis's avatar
Stelios Karozis committed
147 148
    cnt_fr=0
    with open(trajfile, 'rb') as inputfile:
Stelios Karozis's avatar
Stelios Karozis committed
149
        for i in range(1000000):
Stelios Karozis's avatar
Stelios Karozis committed
150 151 152
            try: 
                header = read_trr_header(inputfile)
                skip_trr_data(inputfile, header)
Stelios Karozis's avatar
Stelios Karozis committed
153 154 155 156
                if pr_time==True:
                    #print('Step: {step}, time: {time}'.format(**header))
                    #print('Step: {step}, time: {time}'.format(**header))
                    time_index[cnt_fr]='{time}'.format(**header)
Stelios Karozis's avatar
Stelios Karozis committed
157 158 159
                cnt_fr=cnt_fr+1
            except EOFError:
                pass
Stelios Karozis's avatar
Stelios Karozis committed
160 161 162 163
    if pr_time==True:
        return cnt_fr,time_index
    else:
        return cnt_fr
Stelios Karozis's avatar
Stelios Karozis committed
164

Stelios Karozis's avatar
Stelios Karozis committed
165
def fr_export(trajfile='traj.trr', num_frames=1):
166 167 168 169 170 171 172
    """
    Export frames from gromacs .trr file. 

    parameters:     trajfile = [.trr]
                    num_frames = [number of frames to keep
                                  counting from the end of file]
    """
Stelios Karozis's avatar
Stelios Karozis committed
173 174 175
    print(' ')
    print('Read gmx files progress:')
    print('==============================')
Stelios Karozis's avatar
Stelios Karozis committed
176
    cnt_fr=count_frames(trajfile)
177
    data_all={}
Stelios Karozis's avatar
Stelios Karozis committed
178
    with open(trajfile, 'rb') as inputfile:
Stelios Karozis's avatar
Stelios Karozis committed
179
        bar = Bar('Skipping frames from .trr', max=cnt_fr-num_frames)
Stelios Karozis's avatar
Stelios Karozis committed
180
        for i in range(cnt_fr-num_frames):
Stelios Karozis's avatar
Stelios Karozis committed
181 182 183
            header = read_trr_header(inputfile)
            #print('Step: {step}, time: {time}'.format(**header))
            skip_trr_data(inputfile, header)
Stelios Karozis's avatar
Stelios Karozis committed
184 185 186
            bar.next()
        bar.finish()
        bar = Bar('Read frames from .trr', max=num_frames)
187 188
        for i in range(cnt_fr-num_frames,cnt_fr):
            header = read_trr_header(inputfile)
189
            #print('Step: {step}, time: {time}'.format(**header))
190 191 192
            data = read_trr_data(inputfile, header)
            step='{step}'.format(**header)
            data_all[step]=data
Stelios Karozis's avatar
Stelios Karozis committed
193 194
            bar.next()
        bar.finish()
195
        return data_all
Stelios Karozis's avatar
Stelios Karozis committed
196

Stelios Karozis's avatar
Stelios Karozis committed
197
def read_gro(gro):
198 199 200 201 202 203 204 205 206 207 208 209 210
    """
    Read .gro file and exports 
    1. system name
    2. data number
    3. box size
    4. residue number
    5. residue type
    6. atom type
    7. atom number
    8. free format data (x,y,z,v,u,w)

    parameters:     gro = [.gro]
    """
Stelios Karozis's avatar
Stelios Karozis committed
211 212
    cnt=0
    data_num=0
Stelios Karozis's avatar
Stelios Karozis committed
213 214 215 216 217
    res_num = [] 
    res_type = []  
    atom_type = [] 
    atom_num  = []
    rest_dt = []
218
    cnt_atoms=0
Stelios Karozis's avatar
Stelios Karozis committed
219
    num_lines = sum(1 for line in open(gro))
Stelios Karozis's avatar
Stelios Karozis committed
220
    with open(gro, 'r') as F:
Stelios Karozis's avatar
Stelios Karozis committed
221
        bar = Bar('Read .gro', max=num_lines)
Stelios Karozis's avatar
Stelios Karozis committed
222 223
        for line in F:
            cnt=cnt+1
224 225
            #print(cnt)
            if cnt>2 and cnt<=data_num+2:
Stelios Karozis's avatar
Stelios Karozis committed
226 227
                res_num.append(line[:5])  
                res_type.append(line[5:10])  
228 229 230 231 232 233
                atom_type.append(line[10:15])
                #Compensate for large .gro files
                #ToDo check if data index is the same in the function that follows
                #atom_num.append(line[15:20])
                cnt_atoms=cnt_atoms+1
                atom_num.append(str(cnt_atoms))
Stelios Karozis's avatar
Stelios Karozis committed
234
                rest_dt.append(line[20:])
Stelios Karozis's avatar
Stelios Karozis committed
235 236 237 238
            elif cnt==1:
                system=line[:10]
            elif cnt==2:
                data_num=int(line[:7])
239
            elif cnt>data_num+2:
Stelios Karozis's avatar
Stelios Karozis committed
240
                box_size=line[:50]
Stelios Karozis's avatar
Stelios Karozis committed
241 242
            bar.next()
        bar.finish()
Stelios Karozis's avatar
Stelios Karozis committed
243
        #print(system,data_num,box_size)
Stelios Karozis's avatar
Stelios Karozis committed
244
    return system,data_num,box_size,res_num,res_type,atom_type,atom_num,rest_dt
245

246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308
def read_itp(itp):
    """
    Read mass and atomtype from .itp file.

    parameters:     itp = [.itp]

    output:         dict {atomtype{mass}}

    ATTENTION:  the format of itp is free but the current function
                has some restrictions
                1. sections starting points doesn't have spaces [ex. [atoms] ]
                2. the sections have to be seperated by an empty line
    """

    ll=0
    #mass=[]
    #atomtype=[]
    weights={}
    num_lines = sum(1 for line in open(itp))
    with open(itp, 'r') as F:
        bar = Bar('Read .itp', max=num_lines)
        for line in F:

            if line.strip()=='[atoms]':
                ll=1
                #print(line.strip())
                continue
            elif line.strip()=='[atomtypes]':
                ll=2
                #print(line.strip())
                continue
            elif len(line.strip()) == 0:
                #print(len(line.strip()))
                ll=0
                continue
            elif line.strip()[0] == ';':
                continue

            if ll==1: #molecule structure itp
                word=line.split()
                try:
                    if float(word[7])==0.0:
                        print('WARNING: Atomtype mass of '+word[4]+' is 0.0')
                except:
                    print('Skip line: '+word[7])
                    continue
                #print(word[4],word[7])
                weights[word[4]]=word[7]
                #mass.append(word[1]) 
                #atomtype.append(word[0])
            if ll==2: #molecule forcefield itp
                word=line.split()
                try:
                    if float(word[1])==0.0:
                        print('WARNING: Atomtype mass of '+word[0]+' is 0.0')
                except:
                    print('Skip line: '+word[1])
                    continue
                weights[word[0]]=word[1]
            bar.next()
        bar.finish()
    return weights

309 310
def domain_decomposition(data,dx,dy,dz):
    """
311 312 313 314 315 316 317 318 319 320 321
    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={}
Stelios Karozis's avatar
Stelios Karozis committed
322 323 324
    print(' ')
    print('Domain decomposition progress:')
    print('==============================')
325
    for step in data.keys():
Stelios Karozis's avatar
Stelios Karozis committed
326
        print('Step: '+step)
327 328 329 330 331 332 333 334
        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)

Stelios Karozis's avatar
Stelios Karozis committed
335

336
        xx=[]
Stelios Karozis's avatar
Stelios Karozis committed
337
        bar = Bar('X-axis', max=box_x)
Stelios Karozis's avatar
Stelios Karozis committed
338 339
#        for i in range(0,xs+1,dx):
        for i in range(0,box_x):
340
            xx.append(i)
Stelios Karozis's avatar
Stelios Karozis committed
341
            bar.next()
342
        box_p[step]['x']=xx
Stelios Karozis's avatar
Stelios Karozis committed
343
        bar.finish()
344 345

        yy=[]
Stelios Karozis's avatar
Stelios Karozis committed
346
        bar = Bar('Y-axis', max=box_y)
Stelios Karozis's avatar
Stelios Karozis committed
347 348
#        for i in range(0,ys+1,dy):
        for i in range(0,box_y):
349
            yy.append(i)
Stelios Karozis's avatar
Stelios Karozis committed
350
            bar.next()
351
        box_p[step]['y']=yy
Stelios Karozis's avatar
Stelios Karozis committed
352
        bar.finish()
353 354

        zz=[]
Stelios Karozis's avatar
Stelios Karozis committed
355
        bar = Bar('Z-axis', max=box_z)
Stelios Karozis's avatar
Stelios Karozis committed
356 357
#        for i in range(0,zs+1,dz):
        for i in range(0,box_z):
358
            zz.append(i)
Stelios Karozis's avatar
Stelios Karozis committed
359
            bar.next()
360
        box_p[step]['z']=zz
Stelios Karozis's avatar
Stelios Karozis committed
361
        bar.finish()
362 363 364


        xyz=[]
Stelios Karozis's avatar
Stelios Karozis committed
365
        bar = Bar('XYZ-axis', max=box_x*box_y*box_z)
Stelios Karozis's avatar
Stelios Karozis committed
366 367 368 369 370 371
#        for ix in range(0,xs+1,dx):
        for ix in range(0,box_x):
#          for iy in range(0,ys+1,dy):
          for iy in range(0,box_y):
#            for iz in range(0,zs+1,dz):
            for iz in range(0,box_z):
372
                xyz.append([ix,iy,iz])
Stelios Karozis's avatar
Stelios Karozis committed
373 374
                bar.next()
        bar.finish()
375 376 377 378
        box_p[step]['xyz']=xyz

    return box_p

Stelios Karozis's avatar
Stelios Karozis committed
379
def atomid_data(res_num, res_type, atom_type, atom_num, group={}):
380 381 382 383
    """
    Finds the index in list that 'def read_gro' returns,
    and correspond to the atom types in group list

384
    parameters:     res_num=[]
Stelios Karozis's avatar
Stelios Karozis committed
385
                    res_type=[]
386 387 388
                    atom_type=[]
                    atom_num=[]
                    group={}
389

390
    output:         dictionary {resid:{res_type:{atom_type:[atom_num]}}}
391
    """
Stelios Karozis's avatar
Stelios Karozis committed
392 393 394 395 396
    print(' ')
    print('Create index  progress:')
    print('======================')
    #bar = Bar('Create index step', max=len(group.keys()))
    res_ndx={}
397
    for res,atom in group.items():
Stelios Karozis's avatar
Stelios Karozis committed
398
        res_ndx[res]=[]
399
        for element in atom:
Stelios Karozis's avatar
Stelios Karozis committed
400
            tmp=[res_num[i].strip() for i, e in enumerate(atom_type) if e.strip() == element.strip() and res.strip()==res_type[i].strip()]
Stelios Karozis's avatar
Stelios Karozis committed
401 402 403 404 405 406 407 408
        res_ndx[res].append(tmp)
        #bar.next()
    #bar.finish()

    res_ndx_flat={}
    for res in res_ndx.keys():
        res_ndx_flat[res] = [y for x in res_ndx[res] for y in x]

409
    ndx={}
Stelios Karozis's avatar
Stelios Karozis committed
410 411 412 413 414 415 416 417 418 419 420 421 422 423 424
    for rs in res_ndx_flat.keys():
        bar = Bar(rs, max=len(res_ndx_flat[rs]))
        for resid in res_ndx_flat[rs]:
            ndx[resid.strip()]={}
            for res,atom in group.items():
                ndx[resid][res]={}
                for element in atom:
                    tmp=[atom_num[i].strip()  for i, e in enumerate(atom_type) if e.strip() == element.strip() and resid.strip()==res_num[i].strip() and res.strip()==res_type[i].strip()]
                    if not tmp:
                        ndx[resid].pop(res,None)
                        pass
                    else:
                        ndx[resid][res][element]=tmp
            bar.next()
        bar.finish()
425
    return ndx
426

427
def atom2grid(data, box_p, ndx):
428
    """
429
    Assign atom number that corresponds to ndx (see 'def atomid_data')
430 431 432 433 434
    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'}
                    box_p = {dictionary input from 'def domain_decomposition'}
435
                    ndx = {dictionary input from 'def atomid_data'}
436

437 438 439
    output:         dictionaries 
                        1. box_res = {step:{resid:{res_type:{atom_type:{atom_num:(subX,subYsubZ)}}}}}
                        2. box_res_rev = {step:{resid:{res:{(subX,subYsubZ):[atom_num]}}}}
440
    """
Stelios Karozis's avatar
Stelios Karozis committed
441 442 443
    print(' ')
    print('Assing atom to grid progress:')
    print('==============================')
444 445

    box_res={}
446
    box_res_rev = {}
447
    for step in data.keys():
Stelios Karozis's avatar
Stelios Karozis committed
448 449
        #print('Step: '+step)
        bar = Bar('Step: '+step, max=len(ndx.keys()))
450
        box_res[step]={}
451
        box_res_rev[step] = {}
452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513
        for resid in ndx.keys():
            box_res[step][resid]={}
            box_res_rev[step][resid] = {}
            for res in ndx[resid].keys():
                box_res[step][resid][res]={}
                box_res_rev[step][resid][res]={}
                for atom in ndx[resid][res].keys():
                    box_res[step][resid][res][atom]={}
                    box_res_rev[step][resid][res][atom] = {}
                    for atomnum in ndx[resid][res][atom]:
                        #data[step]['x'][atom_num-1][x(0),y(1),z(2)]

                        xx=data[step]['x'][int(atomnum)-1][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'][int(atomnum)-1][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'][int(atomnum)-1][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][resid][res][atom][atomnum]=(cnt_x,cnt_y,cnt_z)

                    #Make subdomain position the key and group the residues
                    for key, value in sorted(box_res[step][resid][res][atom].items()):
                        box_res_rev[step][resid][res].setdefault(value, []).append(key.strip())
                    #Remove  atom_type as key of dictionary
                    box_res_rev[step][resid][res].pop(atom,None)
                    #If atoms of residue lie in more than one subdomain, put the all in the first one
                    if len(box_res_rev[step][resid][res]) > 1:
                        tmp=[]
                        flattened_list=[]
                        kk=[]
                        for sb in box_res_rev[step][resid][res].keys():
                            kk.append(sb)
                            tmp.append(box_res_rev[step][resid][res][sb])
                        #Remove keys
                        for k in kk: 
                            box_res_rev[step][resid][res].pop(k,None)
                        #flatten the lists
                        flattened_list = [y for x in tmp for y in x]
                        #Keep first key and results
                        box_res_rev[step][resid][res][kk[0]]=flattened_list
                    #box_res_rev[step][resid][res]=[i  for i, e in enumerate(box_res_rev[step][resid][res]) if e.strip() != .strip()]
Stelios Karozis's avatar
Stelios Karozis committed
514 515
            bar.next()
        bar.finish()
516 517 518 519 520 521 522 523 524
    return box_res, box_res_rev

def sub_coord(box, data, res_num=[]):
    """
    Use the box_res_rev from 'def atom2grid' and data from 'def fr_export'
    and groups all the XYZ coordinates per subdomain

    parameters:         box = {step:{resid:{res:{(subX,subYsubZ):[atom_num]}}}}
                        data = {step:{obj:[x,y,z]}}
525
    
526 527 528
    output:             norm = {step:{subdomain:[[X1,Y1,Z1] ... [Xn,Yn,Zn]]}}
                        vector = {step:{subdomain:resid:{[[X1,Y1,Z1] ... [Xn,Yn,Zn]]}}}
    """
Stelios Karozis's avatar
Stelios Karozis committed
529 530 531 532
    print(' ')
    print('Groupby subdomain progress:')
    print('==============================')

533 534 535
    coord_norm={}
    coord_vector={}
    for step in box.keys():
Stelios Karozis's avatar
Stelios Karozis committed
536

537 538
        coord_norm[step]={}
        coord_vector[step]={}
Stelios Karozis's avatar
Stelios Karozis committed
539 540 541 542 543 544 545 546
        for resid in box[step].keys():
                for res in box[step][resid].keys():
                    for subdomain in box[step][resid][res].keys():
                        coord_norm[step][subdomain]=[]
                        coord_vector[step][subdomain]={}

        rs=[]
        bar = Bar('Step: '+step, max=len(box[step].keys()))
547 548 549
        for resid in box[step].keys():
            for res in box[step][resid].keys():
                for subdomain in box[step][resid][res].keys():
Stelios Karozis's avatar
Stelios Karozis committed
550 551 552 553
                    #tmp=[]
                    if (resid,subdomain) not in rs:
                        rs.append((resid,subdomain))
                        coord_vector[step][subdomain][resid]=[]
554 555
                    for atomnum in box[step][resid][res][subdomain]:
                        #tmp_atom=[]
Stelios Karozis's avatar
Stelios Karozis committed
556
                        #tmp.append(data[step]['x'][int(atomnum)-1])
557 558
                        #tmp_atom.append(list(data[step]['x'][int(atomnum)-1])) #not append to previous
                        coord_norm[step][subdomain].append(list(data[step]['x'][int(atomnum)-1]))
Stelios Karozis's avatar
Stelios Karozis committed
559 560 561
                        coord_vector[step][subdomain][resid].append(list(data[step]['x'][int(atomnum)-1]))
            bar.next()
        bar.finish()
562 563 564 565 566 567 568 569 570 571 572 573 574
    return coord_norm, coord_vector

def coord2norm(coord,img=True):
    """
    Use the coord from 'def sub_coord' and replaces XYZ coordinates
    with c, normal of surface that fits the points. Also it can plot
    the surfaces for each subdomain

    parameters:         coord = {step:{subdomain:[[X1,Y1,Z1] ... [Xn,Yn,Zn]]}}
                        img = True or False
    
    output:             dictionary = {step:{subdomain:{c:int, normal:array[]}}}
    """
Stelios Karozis's avatar
Stelios Karozis committed
575 576 577 578
    print(' ')
    print('Finding surface progress:')
    print('==============================')

579 580 581
    surf={}
    for step in coord.keys():
        surf[step]={}
Stelios Karozis's avatar
Stelios Karozis committed
582
        bar = Bar('Step: '+step, max=len(coord[step].keys()))
583 584
        for subdomain in coord[step].keys():
            c,normal = fitPlaneLTSQ(np.array(coord[step][subdomain]))
585
            surf[step][subdomain]={'c':c, 'normal':normal}
586 587 588 589 590 591 592
            #Change save variable if you want to save .png elsewere
            save='png/'+str(subdomain)
            if img==True:
                try:
                    plot_surf(np.array(coord[step][subdomain]), normal, c, save)
                except:
                    print("ERROR: Folder png/ doesn't exist in root")
Stelios Karozis's avatar
Stelios Karozis committed
593 594 595
            bar.next()
        bar.finish()
    return surf
596

Stelios Karozis's avatar
Stelios Karozis committed
597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617
def coord2norm2cg(coord_vector,img=True):
    """
    Use the coord_vector from 'def sub_coord' and replaces XYZ coordinates
    with c, normal of surface that fits the points and calculates center of gravity
    for each resid. Also it can plot the surfaces for each subdomain

    parameters:         coord = {step:{subdomain:{resid:[[X1,Y1,Z1] ... [Xn,Yn,Zn]]}}
                        img = True or False
    
    output:             dictionary = {step:{subdomain:{c:int, normal:array[], resid:[cgx, cgy, cgz]}}}
    """
    print(' ')
    print('Finding surface and gravity point progress:')
    print('==========================================')

    surf={}
    for step in coord_vector.keys():
        surf[step]={}
        bar = Bar('Step: '+step, max=len(coord_vector[step].keys()))
        for subdomain in coord_vector[step].keys():
            surf[step][subdomain]={}
618 619 620
            xx=[]
            yy=[]
            zz=[]
Stelios Karozis's avatar
Stelios Karozis committed
621 622 623 624 625
            #surf_points=[]
            res_cg={}
            cnt=0
            for resid in coord_vector[step][subdomain].keys():
                cnt=cnt+1
626 627 628 629 630 631 632 633 634 635
                for ii in coord_vector[step][subdomain][resid]:
                    #print(ii)
                    #print(ii[0],ii[1],ii[2])
                    #print(coord_vector[step][subdomain][resid][0])
                    xx.append(ii[0])
                    yy.append(ii[1])
                    zz.append(ii[2])
                cgx = center_gravity(xx)
                cgy = center_gravity(yy)
                cgz = center_gravity(zz)
Stelios Karozis's avatar
Stelios Karozis committed
636 637 638 639 640
                res_cg[resid]=[cgx,cgy,cgz]
                if cnt==1:
                    surf_points=coord_vector[step][subdomain][resid]
                else:
                    surf_points=surf_points+coord_vector[step][subdomain][resid]
Stelios Karozis's avatar
Stelios Karozis committed
641

Stelios Karozis's avatar
Stelios Karozis committed
642 643 644
            c,normal = fitPlaneLTSQ(np.array(surf_points))
            surf[step][subdomain]={'c':c, 'normal':normal}
            surf[step][subdomain].update(res_cg)
Stelios Karozis's avatar
Stelios Karozis committed
645
            
Stelios Karozis's avatar
Stelios Karozis committed
646
            if img==True:
647 648 649 650
                fl_save='png/'
                save=fl_save+str(subdomain)
                if not os.path.exists(fl_save):
                    os.makedirs(fl_save) 
Stelios Karozis's avatar
Stelios Karozis committed
651 652
                try:
                    plot_surf(np.array(surf_points), normal, c, save)
653 654
                except Exception as e:
                    print(e)
Stelios Karozis's avatar
Stelios Karozis committed
655 656
            bar.next()
        bar.finish()
657 658 659 660 661 662 663
    return surf

def coord2vector(coord_vector):
    """
    Use the coord_vector from 'def sub_coord' and replaces XYZ coordinates
    with vector of best fit line

Stelios Karozis's avatar
Stelios Karozis committed
664 665 666
    parameters:         coord = {step:{subdomain:{resid:[[X1,Y1,Z1] ... [Xn,Yn,Zn]]}}
        
    output:             dictionary = {step:{subdomain:{resid:[x, y, z]}}
667
    """
Stelios Karozis's avatar
Stelios Karozis committed
668
    print(' ')
Stelios Karozis's avatar
Stelios Karozis committed
669
    print('Finding vector progress:')
Stelios Karozis's avatar
Stelios Karozis committed
670
    print('==============================')
671 672 673
    vector={}
    for step in coord_vector.keys():
        vector[step]={}
Stelios Karozis's avatar
Stelios Karozis committed
674
        bar = Bar('Step: '+step, max=len(coord_vector[step].keys()))
675 676 677
        for subdomain in coord_vector[step].keys():
            vector[step][subdomain]={}
            for resid in coord_vector[step][subdomain].keys():
Stelios Karozis's avatar
Stelios Karozis committed
678 679 680 681
                grp=[]
                for ii in coord_vector[step][subdomain][resid]:
                    grp.append(ii)
                vv = points2vector(np.array(grp))
682
                vector[step][subdomain][resid]=vv
Stelios Karozis's avatar
Stelios Karozis committed
683 684
            bar.next()
        bar.finish()
Stelios Karozis's avatar
Stelios Karozis committed
685 686 687 688 689 690 691 692 693 694 695 696
    return vector

def SurfVector_angle(surf,vector):
    """
    Use the surf and vector from 'def coord2norm' and 'def coord2vector'
    and calculate the angle between them for every step, sudomain and resid

    parameters:         surf = {step:{subdomain:{c:int, normal:array[]}}}
                        vector = {step:{subdomain:{resid:[x, y, z]}}
    
    output:             angle = {step:{subdomain:{resid:angle}}
    """
697 698 699
    print(' ')
    print('Tilt progress:')
    print('==============================')
Stelios Karozis's avatar
Stelios Karozis committed
700 701 702
    angle={}
    for step in surf.keys():
        angle[step]={}
703
        bar = Bar('Step: '+step, max=len(surf[step].keys()))
Stelios Karozis's avatar
Stelios Karozis committed
704
        for sudomain in surf[step].keys():
Stelios Karozis's avatar
Stelios Karozis committed
705 706 707 708 709
            try:
                vector[step][sudomain].keys()
            except KeyError:
                #print(step, sudomain)
                continue
Stelios Karozis's avatar
Stelios Karozis committed
710
            angle[step][sudomain]={}
711
            tot=[]
Stelios Karozis's avatar
Stelios Karozis committed
712 713 714 715 716
            for resid in vector[step][sudomain].keys():
                P1=tuple(surf[step][sudomain]['normal'])
                P2=tuple(vector[step][sudomain][resid])
                #print(tbf.angle_between3D(P1,P2))
                angle[step][sudomain][resid]=angle_between3D(P1,P2)
717 718
                tot.append(angle_between3D(P1,P2))
            angle[step][sudomain]['avg/frame']=sum(tot)/len(tot)
719 720 721 722 723 724
            bar.next()
        bar.finish()
    return angle

def togmxndx(box_res, fld, sv_name):
    print(' ')
725
    print('Save to ndx progress:')
726 727
    print('==============================')
    cnt=0
728
    fl_save=fld+'/'
729
    if not os.path.exists(fl_save):
730 731 732 733
        os.makedirs(fl_save)
    else:
        print('WARNING: .ndx files exists. Nothing to do!')
        return
734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767
 
    ndx={}
    for step in box_res.keys():
        if cnt==1:
            break
        for resid in box_res[step].keys():
            for res in box_res[step][resid].keys():
                for subdomain in box_res[step][resid][res].keys():
                    ndx[subdomain]={}

        rs=[]
        bar = Bar('Create format: ', max=len(box_res[step].keys()))
        for resid in box_res[step].keys():
            for res in box_res[step][resid].keys():
                for subdomain in box_res[step][resid][res]:
                    if (res,subdomain) not in rs:
                        rs.append((res,subdomain))
                        ndx[subdomain][res]=[]
                    for atomnum in box_res[step][resid][res][subdomain]:    
                        ndx[subdomain][res].append(atomnum)
            bar.next()
        bar.finish()                 
        cnt=cnt+1

    bar = Bar('Save file: ', max=len(ndx.keys()))
    for subdomain in ndx.keys():
        save=fl_save+sv_name+'_'+str(subdomain)+'.ndx'
        with open(save, "a") as myfile:
            for res in ndx[subdomain].keys():
                myfile.write("["+res+"]\n")
                for atomnum in ndx[subdomain][res]:
                    myfile.write(atomnum+'\n')

        bar.next()
Stelios Karozis's avatar
Stelios Karozis committed
768 769 770 771
    bar.finish()

def dict2pd(d, col=[]):
    return pd.DataFrame.from_dict(d, orient='index', columns=col)