tooba_f.py 31.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
import pickle5 as pkl
Stelios Karozis's avatar
Stelios Karozis committed
7
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
import joblib
Stelios Karozis's avatar
Stelios Karozis committed
12

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

Stelios Karozis's avatar
Stelios Karozis committed
19 20 21 22 23 24 25 26 27 28 29 30
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
31
def fitPlaneLTSQ(XYZ):
32
    #Source: https://gist.github.com/RustingSword/e22a11e1d391f2ab1f2c
Stelios Karozis's avatar
Stelios Karozis committed
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 63
    (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
64
def center_gravity(a):
65
    #Todo somatidia diaforetikis mazas -> weighted me maza ????
Stelios Karozis's avatar
Stelios Karozis committed
66 67 68 69
    m=len(a)
    cg = np.sum(a)/m
    return cg

Stelios Karozis's avatar
Stelios Karozis committed
70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100
class StreamFile(object):
    def __init__(self, f):
        self.f = f
    def __getattr__(self, item):
        return getattr(self.f, item)
    def read(self, n):
        # print("reading total_bytes=%s" % n, flush=True)
        if n >= (1 << 31):
            buffer = bytearray(n)
            idx = 0
            while idx < n:
                batch_size = min(n - idx, 1 << 31 - 1)
                # print("reading bytes [%s,%s)..." % (idx, idx + batch_size), end="", flush=True)
                buffer[idx:idx + batch_size] = self.f.read(batch_size)
                # print("done.", flush=True)
                idx += batch_size
            return buffer
        return self.f.read(n)
    def write(self, buffer):
        n = len(buffer)
        print("writing total_bytes=%s..." % n, flush=True)
        idx = 0
        while idx < n:
            batch_size = min(n - idx, 1 << 31 - 1)
            print("writing bytes [%s, %s)... " % (idx, idx + batch_size), end="", flush=True)
            self.f.write(buffer[idx:idx + batch_size])
            print("done.", flush=True)
            idx += batch_size



Stelios Karozis's avatar
Stelios Karozis committed
101
def topickle(fl, sv_name):
102 103
    print(' ')
    print('Save to pickle |################################| 1/1')
Stelios Karozis's avatar
Stelios Karozis committed
104 105
    with open(sv_name+'.pkl', 'wb') as handle:
        pkl.dump(fl, handle, protocol=pkl.HIGHEST_PROTOCOL)
Stelios Karozis's avatar
Stelios Karozis committed
106
        #joblib.dump(fl, handle)
Stelios Karozis's avatar
Stelios Karozis committed
107 108 109
def frompickle(fl):
    with open(fl, 'rb') as handle:
        b = pkl.load(handle)
Stelios Karozis's avatar
Stelios Karozis committed
110
        #b = joblib.load(handle)
Stelios Karozis's avatar
Stelios Karozis committed
111 112 113
    return b


Stelios Karozis's avatar
Stelios Karozis committed
114
def tojson(fl, sv_name):
115 116
    print(' ')
    print('Save to json |################################| 1/1')
Stelios Karozis's avatar
Stelios Karozis committed
117 118
    with open(sv_name+'.json', 'w') as file:
        file.write(json.dumps(str(fl)))
Stelios Karozis's avatar
Stelios Karozis committed
119 120 121 122 123 124 125 126
        #file.write(json.dumps(fl))
def fromjson(fl):
    print(' ')
    print('Load to json |################################| 1/1')
    with open(fl, 'r') as file:
        data = file.read()
        b = json.dumps(data)
    return b
Stelios Karozis's avatar
Stelios Karozis committed
127

128
def plot_surf(data, normal, c, save_name):
Stelios Karozis's avatar
Stelios Karozis committed
129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153
    #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')
154 155
    #plt.savefig(save_name+'.png')
    plt.show()
Stelios Karozis's avatar
Stelios Karozis committed
156 157 158
    plt.clf()
    plt.cla()
    plt.close()
159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179

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
180

Stelios Karozis's avatar
Stelios Karozis committed
181
def count_frames(trajfile='traj.trr', pr_time=False):
182 183 184 185 186
    """
    Count total frames of .trr file

    parameters:     trajfile = [.trr]
    """
Stelios Karozis's avatar
Stelios Karozis committed
187
    time_index={}
Stelios Karozis's avatar
Stelios Karozis committed
188 189
    cnt_fr=0
    with open(trajfile, 'rb') as inputfile:
Stelios Karozis's avatar
Stelios Karozis committed
190
        for i in range(1000000):
Stelios Karozis's avatar
Stelios Karozis committed
191 192 193
            try: 
                header = read_trr_header(inputfile)
                skip_trr_data(inputfile, header)
Stelios Karozis's avatar
Stelios Karozis committed
194 195 196 197
                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
198 199 200
                cnt_fr=cnt_fr+1
            except EOFError:
                pass
Stelios Karozis's avatar
Stelios Karozis committed
201 202 203 204
    if pr_time==True:
        return cnt_fr,time_index
    else:
        return cnt_fr
Stelios Karozis's avatar
Stelios Karozis committed
205

Stelios Karozis's avatar
Stelios Karozis committed
206
def fr_export(trajfile='traj.trr', num_frames=1):
207 208 209 210 211 212 213
    """
    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
214 215 216
    print(' ')
    print('Read gmx files progress:')
    print('==============================')
Stelios Karozis's avatar
Stelios Karozis committed
217
    cnt_fr=count_frames(trajfile)
218
    data_all={}
Stelios Karozis's avatar
Stelios Karozis committed
219
    with open(trajfile, 'rb') as inputfile:
Stelios Karozis's avatar
Stelios Karozis committed
220
        bar = Bar('Skipping frames from .trr', max=cnt_fr-num_frames)
Stelios Karozis's avatar
Stelios Karozis committed
221
        for i in range(cnt_fr-num_frames):
Stelios Karozis's avatar
Stelios Karozis committed
222 223 224
            header = read_trr_header(inputfile)
            #print('Step: {step}, time: {time}'.format(**header))
            skip_trr_data(inputfile, header)
Stelios Karozis's avatar
Stelios Karozis committed
225 226 227
            bar.next()
        bar.finish()
        bar = Bar('Read frames from .trr', max=num_frames)
228 229
        for i in range(cnt_fr-num_frames,cnt_fr):
            header = read_trr_header(inputfile)
230
            #print('Step: {step}, time: {time}'.format(**header))
231 232 233
            data = read_trr_data(inputfile, header)
            step='{step}'.format(**header)
            data_all[step]=data
Stelios Karozis's avatar
Stelios Karozis committed
234 235
            bar.next()
        bar.finish()
236
        return data_all
Stelios Karozis's avatar
Stelios Karozis committed
237

Stelios Karozis's avatar
Stelios Karozis committed
238
def read_gro(gro):
239 240 241 242 243 244 245 246 247 248 249 250 251
    """
    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
252 253
    cnt=0
    data_num=0
Stelios Karozis's avatar
Stelios Karozis committed
254 255 256 257 258
    res_num = [] 
    res_type = []  
    atom_type = [] 
    atom_num  = []
    rest_dt = []
259
    cnt_atoms=0
Stelios Karozis's avatar
Stelios Karozis committed
260
    num_lines = sum(1 for line in open(gro))
Stelios Karozis's avatar
Stelios Karozis committed
261
    with open(gro, 'r') as F:
Stelios Karozis's avatar
Stelios Karozis committed
262
        bar = Bar('Read .gro', max=num_lines)
Stelios Karozis's avatar
Stelios Karozis committed
263 264
        for line in F:
            cnt=cnt+1
265 266
            #print(cnt)
            if cnt>2 and cnt<=data_num+2:
Stelios Karozis's avatar
Stelios Karozis committed
267 268
                res_num.append(line[:5])  
                res_type.append(line[5:10])  
269 270 271 272 273 274
                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
275
                rest_dt.append(line[20:])
Stelios Karozis's avatar
Stelios Karozis committed
276 277 278 279
            elif cnt==1:
                system=line[:10]
            elif cnt==2:
                data_num=int(line[:7])
280
            elif cnt>data_num+2:
Stelios Karozis's avatar
Stelios Karozis committed
281
                box_size=line[:50]
Stelios Karozis's avatar
Stelios Karozis committed
282 283
            bar.next()
        bar.finish()
Stelios Karozis's avatar
Stelios Karozis committed
284
        #print(system,data_num,box_size)
Stelios Karozis's avatar
Stelios Karozis committed
285
    return system,data_num,box_size,res_num,res_type,atom_type,atom_num,rest_dt
286

287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349
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

350 351
def domain_decomposition(data,dx,dy,dz):
    """
352 353 354 355 356 357 358 359 360 361 362
    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
363 364 365
    print(' ')
    print('Domain decomposition progress:')
    print('==============================')
366
    for step in data.keys():
Stelios Karozis's avatar
Stelios Karozis committed
367
        print('Step: '+step)
368 369 370 371 372 373 374 375
        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
376

377
        xx=[]
Stelios Karozis's avatar
Stelios Karozis committed
378
        bar = Bar('X-axis', max=box_x)
Stelios Karozis's avatar
Stelios Karozis committed
379 380
#        for i in range(0,xs+1,dx):
        for i in range(0,box_x):
381
            xx.append(i)
Stelios Karozis's avatar
Stelios Karozis committed
382
            bar.next()
383
        box_p[step]['x']=xx
Stelios Karozis's avatar
Stelios Karozis committed
384
        bar.finish()
385 386

        yy=[]
Stelios Karozis's avatar
Stelios Karozis committed
387
        bar = Bar('Y-axis', max=box_y)
Stelios Karozis's avatar
Stelios Karozis committed
388 389
#        for i in range(0,ys+1,dy):
        for i in range(0,box_y):
390
            yy.append(i)
Stelios Karozis's avatar
Stelios Karozis committed
391
            bar.next()
392
        box_p[step]['y']=yy
Stelios Karozis's avatar
Stelios Karozis committed
393
        bar.finish()
394 395

        zz=[]
Stelios Karozis's avatar
Stelios Karozis committed
396
        bar = Bar('Z-axis', max=box_z)
Stelios Karozis's avatar
Stelios Karozis committed
397 398
#        for i in range(0,zs+1,dz):
        for i in range(0,box_z):
399
            zz.append(i)
Stelios Karozis's avatar
Stelios Karozis committed
400
            bar.next()
401
        box_p[step]['z']=zz
Stelios Karozis's avatar
Stelios Karozis committed
402
        bar.finish()
403 404 405


        xyz=[]
Stelios Karozis's avatar
Stelios Karozis committed
406
        bar = Bar('XYZ-axis', max=box_x*box_y*box_z)
Stelios Karozis's avatar
Stelios Karozis committed
407 408 409 410 411 412
#        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):
413
                xyz.append([ix,iy,iz])
Stelios Karozis's avatar
Stelios Karozis committed
414 415
                bar.next()
        bar.finish()
416 417 418 419
        box_p[step]['xyz']=xyz

    return box_p

Stelios Karozis's avatar
Stelios Karozis committed
420
def atomid_data(res_num, res_type, atom_type, atom_num, group={}):
421 422 423 424
    """
    Finds the index in list that 'def read_gro' returns,
    and correspond to the atom types in group list

425
    parameters:     res_num=[]
Stelios Karozis's avatar
Stelios Karozis committed
426
                    res_type=[]
427 428 429
                    atom_type=[]
                    atom_num=[]
                    group={}
430

431
    output:         dictionary {resid:{res_type:{atom_type:[atom_num]}}}
432
    """
Stelios Karozis's avatar
Stelios Karozis committed
433 434 435 436 437
    print(' ')
    print('Create index  progress:')
    print('======================')
    #bar = Bar('Create index step', max=len(group.keys()))
    res_ndx={}
438
    for res,atom in group.items():
Stelios Karozis's avatar
Stelios Karozis committed
439
        res_ndx[res]=[]
440
        for element in atom:
Stelios Karozis's avatar
Stelios Karozis committed
441
            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
442 443 444 445 446 447 448 449
        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]

450
    ndx={}
Stelios Karozis's avatar
Stelios Karozis committed
451 452 453 454 455 456 457 458 459 460 461 462 463 464 465
    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()
466
    return ndx
467

468
def atom2grid(data, box_p, ndx):
469
    """
470
    Assign atom number that corresponds to ndx (see 'def atomid_data')
471 472 473 474 475
    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'}
476
                    ndx = {dictionary input from 'def atomid_data'}
477

478 479 480
    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]}}}}
481
    """
Stelios Karozis's avatar
Stelios Karozis committed
482 483 484
    print(' ')
    print('Assing atom to grid progress:')
    print('==============================')
485 486

    box_res={}
487
    box_res_rev = {}
488
    for step in data.keys():
Stelios Karozis's avatar
Stelios Karozis committed
489 490
        #print('Step: '+step)
        bar = Bar('Step: '+step, max=len(ndx.keys()))
491
        box_res[step]={}
492
        box_res_rev[step] = {}
493 494 495 496 497 498 499 500 501 502 503 504 505 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
        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])
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 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644
                        #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])
645 646 647 648 649 650 651 652
                        #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
653 654
            bar.next()
        bar.finish()
655 656 657 658 659 660 661 662 663
    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]}}
664
    
665 666 667
    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
668 669 670 671
    print(' ')
    print('Groupby subdomain progress:')
    print('==============================')

672 673 674
    coord_norm={}
    coord_vector={}
    for step in box.keys():
Stelios Karozis's avatar
Stelios Karozis committed
675

676 677
        coord_norm[step]={}
        coord_vector[step]={}
Stelios Karozis's avatar
Stelios Karozis committed
678 679 680 681 682 683 684 685
        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()))
686 687 688
        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
689 690 691 692
                    #tmp=[]
                    if (resid,subdomain) not in rs:
                        rs.append((resid,subdomain))
                        coord_vector[step][subdomain][resid]=[]
693 694
                    for atomnum in box[step][resid][res][subdomain]:
                        #tmp_atom=[]
Stelios Karozis's avatar
Stelios Karozis committed
695
                        #tmp.append(data[step]['x'][int(atomnum)-1])
696 697
                        #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
698 699 700
                        coord_vector[step][subdomain][resid].append(list(data[step]['x'][int(atomnum)-1]))
            bar.next()
        bar.finish()
701 702 703 704 705 706 707 708 709 710 711 712 713
    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
714 715 716 717
    print(' ')
    print('Finding surface progress:')
    print('==============================')

718 719 720
    surf={}
    for step in coord.keys():
        surf[step]={}
Stelios Karozis's avatar
Stelios Karozis committed
721
        bar = Bar('Step: '+step, max=len(coord[step].keys()))
722 723
        for subdomain in coord[step].keys():
            c,normal = fitPlaneLTSQ(np.array(coord[step][subdomain]))
724
            surf[step][subdomain]={'c':c, 'normal':normal}
725 726 727 728 729 730 731
            #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
732 733 734
            bar.next()
        bar.finish()
    return surf
735

Stelios Karozis's avatar
Stelios Karozis committed
736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756
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]={}
757 758 759
            xx=[]
            yy=[]
            zz=[]
Stelios Karozis's avatar
Stelios Karozis committed
760 761 762 763 764
            #surf_points=[]
            res_cg={}
            cnt=0
            for resid in coord_vector[step][subdomain].keys():
                cnt=cnt+1
765 766 767 768 769 770 771 772 773 774
                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
775 776 777 778 779
                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
780

Stelios Karozis's avatar
Stelios Karozis committed
781 782 783
            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
784
            
Stelios Karozis's avatar
Stelios Karozis committed
785
            if img==True:
786 787 788 789
                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
790 791
                try:
                    plot_surf(np.array(surf_points), normal, c, save)
792 793
                except Exception as e:
                    print(e)
Stelios Karozis's avatar
Stelios Karozis committed
794 795
            bar.next()
        bar.finish()
796 797 798 799 800 801 802
    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
803 804 805
    parameters:         coord = {step:{subdomain:{resid:[[X1,Y1,Z1] ... [Xn,Yn,Zn]]}}
        
    output:             dictionary = {step:{subdomain:{resid:[x, y, z]}}
806
    """
Stelios Karozis's avatar
Stelios Karozis committed
807
    print(' ')
Stelios Karozis's avatar
Stelios Karozis committed
808
    print('Finding vector progress:')
Stelios Karozis's avatar
Stelios Karozis committed
809
    print('==============================')
810 811 812
    vector={}
    for step in coord_vector.keys():
        vector[step]={}
Stelios Karozis's avatar
Stelios Karozis committed
813
        bar = Bar('Step: '+step, max=len(coord_vector[step].keys()))
814 815 816
        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
817 818 819 820
                grp=[]
                for ii in coord_vector[step][subdomain][resid]:
                    grp.append(ii)
                vv = points2vector(np.array(grp))
821
                vector[step][subdomain][resid]=vv
Stelios Karozis's avatar
Stelios Karozis committed
822 823
            bar.next()
        bar.finish()
Stelios Karozis's avatar
Stelios Karozis committed
824 825 826 827 828 829 830 831 832 833 834 835
    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}}
    """
836 837 838
    print(' ')
    print('Tilt progress:')
    print('==============================')
Stelios Karozis's avatar
Stelios Karozis committed
839 840 841
    angle={}
    for step in surf.keys():
        angle[step]={}
842
        bar = Bar('Step: '+step, max=len(surf[step].keys()))
Stelios Karozis's avatar
Stelios Karozis committed
843
        for sudomain in surf[step].keys():
Stelios Karozis's avatar
Stelios Karozis committed
844 845 846 847 848
            try:
                vector[step][sudomain].keys()
            except KeyError:
                #print(step, sudomain)
                continue
Stelios Karozis's avatar
Stelios Karozis committed
849
            angle[step][sudomain]={}
850
            tot=[]
Stelios Karozis's avatar
Stelios Karozis committed
851 852 853 854 855
            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)
856 857
                tot.append(angle_between3D(P1,P2))
            angle[step][sudomain]['avg/frame']=sum(tot)/len(tot)
858 859 860 861 862 863
            bar.next()
        bar.finish()
    return angle

def togmxndx(box_res, fld, sv_name):
    print(' ')
864
    print('Save to ndx progress:')
865 866
    print('==============================')
    cnt=0
867
    fl_save=fld+'/'
868
    if not os.path.exists(fl_save):
869 870 871 872
        os.makedirs(fl_save)
    else:
        print('WARNING: .ndx files exists. Nothing to do!')
        return
873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906
 
    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
907 908 909
    bar.finish()

def dict2pd(d, col=[]):
Stelios Karozis's avatar
Stelios Karozis committed
910
    return pd.DataFrame.from_dict(d, orient='index', columns=col)