MyScripts/make_matrix.py at master · ffallrain/MyScripts · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
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
63
64
65
#!/usr/bin/python
import sys,os
import numpy as np


def next_pose(infile):
    n = 0
    pose = dict()
    for line in open(infile):
        if "MODEL" in line:
            n = int(line.split()[1])
            pose = dict()
        elif "ENDMDL" or "END" in line:
            yield n,pose
        elif "ATOM" in line:
            name = line[12:16].strip()
            if 'H' not in name:
                x = float(line[30:38])
                y = float(line[38:46])
                z = float(line[46:54])
                pose[name] = (x,y,z)
            else:
                pass
        else:
            pass

def distance_2(a,b):
    return (a[0]-b[0])**2 + (a[1]-b[1])**2 + (a[2]-b[2])**2

def rmsd(pose1,pose2):
    keys = pose1.keys()
    n = len(keys)
    s = 0.0
    for key in keys:
        coord1 = pose1[key]
        coord2 = pose2[key]
        s += distance_2(coord1,coord2)
    return (s/n)**0.5

if len(sys.argv) > 1:
    infile = sys.argv[1]

    ## read pose
    poses = list()
    indexes = list()
    for index, pose in next_pose(infile):
        poses.append(pose)
        indexes.append(index)

    nframe = len(poses)

    ## make matrix
    matrix = np.ndarray( shape = (nframe,nframe) )
    for i in range(nframe):
        for j in range(nframe):
            matrix[i][j] = rmsd(poses[i],poses[j])

    ## output
    for i in range(nframe):
        for j in range(nframe):
            print "%8.5f "%matrix[i][j],
        print ""
else :
    print "Usage $0 input.pdb"