tooba_f.py 30.4 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
        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])
506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603
                        #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()]
            bar.next()
        bar.finish()
    return box_res, box_res_rev

def atom2group(data, box_p, ndx):
    """
    Assign atom number that corresponds to ndx (see 'def atomid_data')
    to sudomains created from 'def domain_decomposition' only for the first step.
    Then the same domain id is kept for the same group of molecules. 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'}
                    ndx = {dictionary input from 'def atomid_data'}

    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]}}}}
    """
    print(' ')
    print('Assing atom to group progress:')
    print('==============================')

    step_cnt=0
    domain_list={}
    box_res={}
    box_res_rev = {}
    for step in data.keys():
        step_cnt = step_cnt + 1
        #print('Step: '+step)
        bar = Bar('Step: '+step, max=len(ndx.keys()))
        box_res[step]={}
        box_res_rev[step] = {}
        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)]
                        if step_cnt == 1:
                            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
                            
                            domain_list[atomnum]=(cnt_x,cnt_y,cnt_z)
                            box_res[step][resid][res][atom][atomnum]=(cnt_x,cnt_y,cnt_z)
                        else:
                            box_res[step][resid][res][atom][atomnum]=domain_list[atomnum]

                    #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])
604 605 606 607 608 609 610 611
                        #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
612 613
            bar.next()
        bar.finish()
614 615 616 617 618 619 620 621 622
    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]}}
623
    
624 625 626
    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
627 628 629 630
    print(' ')
    print('Groupby subdomain progress:')
    print('==============================')

631 632 633
    coord_norm={}
    coord_vector={}
    for step in box.keys():
Stelios Karozis's avatar
Stelios Karozis committed
634

635 636
        coord_norm[step]={}
        coord_vector[step]={}
Stelios Karozis's avatar
Stelios Karozis committed
637 638 639 640 641 642 643 644
        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()))
645 646 647
        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
648 649 650 651
                    #tmp=[]
                    if (resid,subdomain) not in rs:
                        rs.append((resid,subdomain))
                        coord_vector[step][subdomain][resid]=[]
652 653
                    for atomnum in box[step][resid][res][subdomain]:
                        #tmp_atom=[]
Stelios Karozis's avatar
Stelios Karozis committed
654
                        #tmp.append(data[step]['x'][int(atomnum)-1])
655 656
                        #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
657 658 659
                        coord_vector[step][subdomain][resid].append(list(data[step]['x'][int(atomnum)-1]))
            bar.next()
        bar.finish()
660 661 662 663 664 665 666 667 668 669 670 671 672
    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
673 674 675 676
    print(' ')
    print('Finding surface progress:')
    print('==============================')

677 678 679
    surf={}
    for step in coord.keys():
        surf[step]={}
Stelios Karozis's avatar
Stelios Karozis committed
680
        bar = Bar('Step: '+step, max=len(coord[step].keys()))
681 682
        for subdomain in coord[step].keys():
            c,normal = fitPlaneLTSQ(np.array(coord[step][subdomain]))
683
            surf[step][subdomain]={'c':c, 'normal':normal}
684 685 686 687 688 689 690
            #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
691 692 693
            bar.next()
        bar.finish()
    return surf
694

Stelios Karozis's avatar
Stelios Karozis committed
695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715
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]={}
716 717 718
            xx=[]
            yy=[]
            zz=[]
Stelios Karozis's avatar
Stelios Karozis committed
719 720 721 722 723
            #surf_points=[]
            res_cg={}
            cnt=0
            for resid in coord_vector[step][subdomain].keys():
                cnt=cnt+1
724 725 726 727 728 729 730 731 732 733
                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
734 735 736 737 738
                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
739

Stelios Karozis's avatar
Stelios Karozis committed
740 741 742
            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
743
            
Stelios Karozis's avatar
Stelios Karozis committed
744
            if img==True:
745 746 747 748
                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
749 750
                try:
                    plot_surf(np.array(surf_points), normal, c, save)
751 752
                except Exception as e:
                    print(e)
Stelios Karozis's avatar
Stelios Karozis committed
753 754
            bar.next()
        bar.finish()
755 756 757 758 759 760 761
    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
762 763 764
    parameters:         coord = {step:{subdomain:{resid:[[X1,Y1,Z1] ... [Xn,Yn,Zn]]}}
        
    output:             dictionary = {step:{subdomain:{resid:[x, y, z]}}
765
    """
Stelios Karozis's avatar
Stelios Karozis committed
766
    print(' ')
Stelios Karozis's avatar
Stelios Karozis committed
767
    print('Finding vector progress:')
Stelios Karozis's avatar
Stelios Karozis committed
768
    print('==============================')
769 770 771
    vector={}
    for step in coord_vector.keys():
        vector[step]={}
Stelios Karozis's avatar
Stelios Karozis committed
772
        bar = Bar('Step: '+step, max=len(coord_vector[step].keys()))
773 774 775
        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
776 777 778 779
                grp=[]
                for ii in coord_vector[step][subdomain][resid]:
                    grp.append(ii)
                vv = points2vector(np.array(grp))
780
                vector[step][subdomain][resid]=vv
Stelios Karozis's avatar
Stelios Karozis committed
781 782
            bar.next()
        bar.finish()
Stelios Karozis's avatar
Stelios Karozis committed
783 784 785 786 787 788 789 790 791 792 793 794
    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}}
    """
795 796 797
    print(' ')
    print('Tilt progress:')
    print('==============================')
Stelios Karozis's avatar
Stelios Karozis committed
798 799 800
    angle={}
    for step in surf.keys():
        angle[step]={}
801
        bar = Bar('Step: '+step, max=len(surf[step].keys()))
Stelios Karozis's avatar
Stelios Karozis committed
802
        for sudomain in surf[step].keys():
Stelios Karozis's avatar
Stelios Karozis committed
803 804 805 806 807
            try:
                vector[step][sudomain].keys()
            except KeyError:
                #print(step, sudomain)
                continue
Stelios Karozis's avatar
Stelios Karozis committed
808
            angle[step][sudomain]={}
809
            tot=[]
Stelios Karozis's avatar
Stelios Karozis committed
810 811 812 813 814
            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)
815 816
                tot.append(angle_between3D(P1,P2))
            angle[step][sudomain]['avg/frame']=sum(tot)/len(tot)
817 818 819 820 821 822
            bar.next()
        bar.finish()
    return angle

def togmxndx(box_res, fld, sv_name):
    print(' ')
823
    print('Save to ndx progress:')
824 825
    print('==============================')
    cnt=0
826
    fl_save=fld+'/'
827
    if not os.path.exists(fl_save):
828 829 830 831
        os.makedirs(fl_save)
    else:
        print('WARNING: .ndx files exists. Nothing to do!')
        return
832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865
 
    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
866 867 868 869
    bar.finish()

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