IMP logo
IMP Reference Guide  2.13.0
The Integrative Modeling Platform
rmsd_calculation.py
1 from __future__ import print_function
2 import pyRMSD.RMSDCalculator
3 from pyRMSD.matrixHandler import MatrixHandler
4 from pyRMSD.condensedMatrix import CondensedMatrix
5 
6 import numpy as np
7 import sys, os, glob
8 
9 import IMP
10 import IMP.atom
11 import IMP.rmf
12 import RMF
13 
14 def get_pdbs_coordinates(path, idfile_A, idfile_B):
15 
16  pts = []
17  conform = []
18  num = 0
19  masses = []
20  radii = []
21 
22  models_name = []
23 
24  with open(idfile_A, 'w+') as f1:
25  for str_file in sorted(glob.glob("%s/sample_A/*.pdb" % path),
26  key=lambda x:int(x.split('/')[-1].split('.')[0])):
27  print(str_file, num, file=f1)
28  models_name.append(str_file)
29 
30  m = IMP.Model()
31  mh = IMP.atom.read_pdb(str_file, m,
33  mps = IMP.core.get_leaves(mh)
34  pts = [IMP.core.XYZ(p).get_coordinates() for p in mps]
35  if num == 0:
36  masses = [IMP.atom.Mass(p).get_mass() for p in mps]
37  radii = [IMP.core.XYZR(p).get_radius() for p in mps]
38  conform.append(pts)
39  pts = []
40  num = num + 1
41 
42  with open(idfile_B, 'w+') as f2:
43  for str_file in sorted(glob.glob("%s/sample_B/*.pdb" % path),
44  key=lambda x:int(x.split('/')[-1].split('.')[0])):
45  print(str_file, num, file=f2)
46  models_name.append(str_file)
47 
48  m = IMP.Model()
49  mh = IMP.atom.read_pdb(str_file, m,
51  mps = IMP.core.get_leaves(mh)
52  pts = [IMP.core.XYZ(p).get_coordinates() for p in mps]
53  conform.append(pts)
54  pts = []
55  num = num + 1
56 
57  return np.array(conform), masses, radii, models_name
58 
59 def get_rmfs_coordinates(path, idfile_A, idfile_B, subunit_name):
60 
61  conform = []
62  num = 0
63  masses = []
64  radii = []
65  ps_names = []
66 
67  f1=open(idfile_A, 'w+')
68  f2=open(idfile_B, 'w+')
69 
70  models_name = []
71 
72  for sample_name,sample_id_file in zip(['A','B'],[f1,f2]):
73 
74  for str_file in sorted(glob.glob("%s/sample_%s/*.rmf3" % (path,sample_name)),key=lambda x:int(x.split('/')[-1].split('.')[0])):
75  print(str_file, num, file=sample_id_file)
76  models_name.append(str_file)
77 
78  m = IMP.Model()
79  inf = RMF.open_rmf_file_read_only(str_file)
80  h = IMP.rmf.create_hierarchies(inf, m)[0]
81  IMP.rmf.load_frame(inf, 0)
82 
83  pts = []
84 
85  if subunit_name:
86  s0 = IMP.atom.Selection(h, resolution=1,molecule=subunit_name)
87  else:
88  s0 = IMP.atom.Selection(h, resolution=1)
89 
90  for leaf in s0.get_selected_particles():
91 
92  p=IMP.core.XYZR(leaf)
93  pts.append([p.get_coordinates()[i] for i in range(3)])
94 
95  if num == 0 and sample_name=='A':
96  masses.append(IMP.atom.Mass(leaf).get_mass())
97  radii.append(p.get_radius())
99  # traverse up the Hierarchy to get copy number of the molecule that the bead belongs to
101 
102  if IMP.atom.Fragment.get_is_setup(leaf): #TODO not tested on non-fragment systems
103  residues_in_bead = IMP.atom.Fragment(leaf).get_residue_indexes()
104 
105  ps_names.append(mol_name+"_"+str(min(residues_in_bead))+"_"+str(max(residues_in_bead))+"_"+str(copy_number))
106 
107  else:
108  residue_in_bead = str(IMP.atom.Residue(leaf).get_index())
109 
110  ps_names.append(mol_name+"_"+residue_in_bead+"_"+residue_in_bead+"_"+str(copy_number))
111 
112  conform.append(pts)
113  pts = []
114  num = num + 1
115  f1.close()
116  f2.close()
117 
118  return ps_names, masses, radii, np.array(conform), models_name
119 
120 
121 def get_rmfs_coordinates_one_rmf(path, rmf_A, rmf_B, subunit_name):
122 
123  # Open RMFs and get total number of models
124  rmf_fh = RMF.open_rmf_file_read_only(rmf_A)
125  n_models = [rmf_fh.get_number_of_frames()]
126  rmf_fh = RMF.open_rmf_file_read_only(rmf_B)
127  n_models.append(rmf_fh.get_number_of_frames())
128 
129  masses = []
130  radii = []
131  ps_names = []
132 
133  models_name = []
134 
135  # Build hierarchy from the RMF file
136  m = IMP.Model()
137  h = IMP.rmf.create_hierarchies(rmf_fh, m)[0]
138  IMP.rmf.load_frame(rmf_fh, 0)
139  m.update()
140 
141  ######
142  # Initialize output array
143 
144  pts = 0 # number of particles in each model
145 
146  # Get selection
147  if subunit_name:
148  s0 = IMP.atom.Selection(h, resolution=1, molecule=subunit_name)
149  else:
150  s0 = IMP.atom.Selection(h, resolution=1)
151  # Count particles
152  for leaf in s0.get_selected_particles():
153  p=IMP.core.XYZR(leaf)
154  pts+=1
155  # Initialize array
156  conform = np.empty([n_models[0]+n_models[1], pts, 3])
157 
158  mod_id = 0 # index for each model in conform.
159  for rmf_file in [rmf_A, rmf_B]:
160 
161  models_name.append(rmf_file)
162 
163  rmf_fh = RMF.open_rmf_file_read_only(rmf_file)
164  h = IMP.rmf.create_hierarchies(rmf_fh, m)[0]
165 
166  print("Opening RMF file:", rmf_file, "with", rmf_fh.get_number_of_frames(), "frames")
167  for f in range(rmf_fh.get_number_of_frames()):
168  if f%100==0:
169  #pass
170  print(" -- Opening frame", f, "of", rmf_fh.get_number_of_frames())
171  IMP.rmf.load_frame(rmf_fh, f)
172 
173  m.update()
174  if subunit_name:
175  s0 = IMP.atom.Selection(h, resolution=1, molecule=subunit_name)
176  else:
177  s0 = IMP.atom.Selection(h, resolution=1)
178  particles = s0.get_selected_particles()
179  # Copy particle coordinates
180  for i in range(len(particles)):
181  leaf=particles[i]
182  p=IMP.core.XYZR(leaf)
183  pxyz = p.get_coordinates()
184  conform[mod_id][i][0] = pxyz[0]
185  conform[mod_id][i][1] = pxyz[1]
186  conform[mod_id][i][2] = pxyz[2]
187 
188  # Just for the first model, update the masses and radii and log the particle name in ps_names
189  if mod_id == 0 and rmf_file==rmf_A:
190  masses.append(IMP.atom.Mass(leaf).get_mass())
191  radii.append(p.get_radius())
193  # traverse up the Hierarchy to get copy number of the molecule that the bead belongs to
195  if IMP.atom.Fragment.get_is_setup(leaf): #TODO not tested on non-fragment systems
196  residues_in_bead = IMP.atom.Fragment(leaf).get_residue_indexes()
197  ps_names.append(mol_name+"_"+str(min(residues_in_bead))+"_"+str(max(residues_in_bead))+"_"+str(copy_number))
198  else:
199  residue_in_bead = str(IMP.atom.Residue(leaf).get_index())
200  ps_names.append(mol_name+"_"+residue_in_bead+"_"+residue_in_bead+"_"+str(copy_number))
201  mod_id+=1
202 
203  return ps_names, masses, radii, conform, models_name, n_models
204 
205 def get_rmsds_matrix(conforms, mode, sup, cores):
206  print("Mode:",mode,"Superposition:",sup,"Number of cores:",cores)
207 
208  if(mode=="cpu_serial" and not sup):
209  calculator = pyRMSD.RMSDCalculator.RMSDCalculator("NOSUP_OMP_CALCULATOR", conforms)
210 
211  elif(mode=="cpu_omp" and not sup):
212  calculator = pyRMSD.RMSDCalculator.RMSDCalculator("NOSUP_OMP_CALCULATOR", conforms)
213  calculator.setNumberOfOpenMPThreads(cores)
214 
215  elif(mode=="cpu_omp" and sup):
216  calculator = pyRMSD.RMSDCalculator.RMSDCalculator("QCP_OMP_CALCULATOR", conforms)
217  calculator.setNumberOfOpenMPThreads(cores)
218 
219  elif(mode=="cuda" and sup):
220  calculator = pyRMSD.RMSDCalculator.RMSDCalculator("QCP_CUDA_MEM_CALCULATOR", conforms)
221 
222  else:
223  print("Wrong values to pyRMSD ! Please Fix")
224  exit()
225 
226  rmsd = calculator.pairwiseRMSDMatrix()
227  rmsd_matrix=CondensedMatrix(rmsd)
228  inner_data = rmsd_matrix.get_data()
229  np.save("Distances_Matrix.data", inner_data)
230 
231  return inner_data
Select non water and non hydrogen atoms.
Definition: pdb.h:245
Add mass to a particle.
Definition: Mass.h:23
A decorator to associate a particle with a part of a protein/DNA/RNA.
Definition: Fragment.h:20
atom::Hierarchies create_hierarchies(RMF::FileConstHandle fh, Model *m)
double get_mass(ResidueType c)
Get the mass from the residue type.
GenericHierarchies get_leaves(Hierarchy mhd)
Get all the leaves of the bit of hierarchy.
void read_pdb(TextInput input, int model, Hierarchy h)
Class for storing model, its restraints, constraints, and particles.
Definition: Model.h:72
static bool get_is_setup(Model *m, ParticleIndex pi)
Definition: Fragment.h:46
The standard decorator for manipulating molecular structures.
Ints get_index(const ParticlesTemp &particles, const Subset &subset, const Subsets &excluded)
void load_frame(RMF::FileConstHandle file, RMF::FrameID frame)
Load the given RMF frame into the state of the linked objects.
A decorator for a particle with x,y,z coordinates.
Definition: XYZ.h:30
std::string get_molecule_name(Hierarchy h)
A decorator for a residue.
Definition: Residue.h:135
int get_copy_index(Hierarchy h)
Walk up the hierarchy to find the current copy index.
Functionality for loading, creating, manipulating and scoring atomic structures.
Select hierarchy particles identified by the biological name.
Definition: Selection.h:66
Support for the RMF file format for storing hierarchical molecular data and markup.
A decorator for a particle with x,y,z coordinates and a radius.
Definition: XYZR.h:27