-
Notifications
You must be signed in to change notification settings - Fork 0
/
2d_rmsd.py
53 lines (42 loc) · 1.37 KB
/
2d_rmsd.py
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
import mdtraj as md
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import time
start = time.time()
# Load your trajectory file (.dcd) and topology file (.pdb or .gro)
traj = md.load('strided_trajectory_mdtraj.dcd', top='1stframe.pdb')
# Compute the backbone atoms selection
backbone_atoms = traj.top.select('backbone')
# Align the frames to a reference structure (optional but recommended)
ref_frame = traj[0]
traj.superpose(reference=ref_frame, atom_indices=backbone_atoms)
print('superpose done')
# Get the number of frames
num_frames = traj.n_frames
# Initialize an empty RMSD matrix
rmsd_matrix = np.zeros((num_frames, num_frames))
print("dummy matrix created")
# Calculate RMSD values for all pairs of frames
k = 1
for i in range(num_frames):
for j in range(i, num_frames):
rmsd_value = md.rmsd(traj[i], traj[j], atom_indices=backbone_atoms)
rmsd_matrix[i, j] = rmsd_value
rmsd_matrix[j, i] = rmsd_value
k += 1
print(k)
np.savetxt('2d_rmsd_matrix.dat', rmsd_matrix)
# Create a heatmap
plt.figure(figsize=(15, 10))
sns.heatmap(rmsd_matrix, cmap='viridis', cbar=True, square=True)
# Add labels and title
plt.xlabel('Frame Index')
plt.ylabel('Frame Index')
plt.title('RMSD Heatmap')
# Save the heatmap as a .png image
plt.savefig('rmsd_heatmap.png')
end = time.time()
print(end - start)
# Show the heatmap
#plt.show()