IMP  2.0.0
The Integrative Modeling Platform
refine_fft.py
1 #!/usr/bin/env python
2 
3 __doc__ = "Refine fitting subunits into a density map with FFT."
4 
5 import math
6 import IMP.multifit
7 import IMP.atom
8 import IMP.em
9 from optparse import OptionParser
10 import os
11 import sys
12 
13 class Fitter(object):
14  def __init__(self, em_map, spacing, resolution, origin, density_threshold,pdb, fits_fn, angle,num_fits,angles_per_voxel,max_trans,max_angle,ref_pdb=''):
15  self.em_map = em_map
16  self.spacing = spacing
17  self.resolution = resolution
18  self.threshold=density_threshold
19  self.originx = origin[0]
20  self.originy = origin[1]
21  self.originz = origin[2]
22  self.pdb = pdb
23  self.fits_fn = fits_fn
24  self.angle = angle
25  self.num_fits=num_fits
26  self.angles_per_voxel=angles_per_voxel
27  self.max_trans=max_trans
28  self.max_angle=max_angle
29  self.ref_pdb=ref_pdb
30 
31  #TODO - update function
32  def run_local_fitting(self,mol2fit,rb,initial_transformation):
33  print "resolution is:",self.resolution
34  dmap = IMP.em.read_map(self.em_map)
35  dmap.get_header().set_resolution(self.resolution)
36  dmap.update_voxel_size(self.spacing)
37  dmap.set_origin(IMP.algebra.Vector3D(self.originx,
38  self.originy,
39  self.originz))
40  dmap.set_was_used(True)
41  dmap.get_header().show()
42  mh_xyz=IMP.core.XYZs(IMP.core.get_leaves(mol2fit))
43  ff = IMP.multifit.FFTFitting()
44  ff.set_was_used(True)
45  #####
46  do_cluster_fits=True
47  max_clustering_translation=3
48  max_clustering_rotation=5
49  num_fits_to_report=100
50  #####
51  fits = ff.do_local_fitting(dmap, self.threshold,mol2fit,
52  self.angle / 180.0 * math.pi,
53  self.max_angle / 180.0 * math.pi, self.max_trans, num_fits_to_report,
54  do_cluster_fits, self.angles_per_voxel,
55  max_clustering_translation,max_clustering_rotation)
56  fits.set_was_used(True)
57  final_fits = fits.best_fits_
58  if self.ref_pdb!='':
59  ref_mh=IMP.atom.read_pdb(self.ref_pdb,mdl)
60  ref_mh_xyz=IMP.core.XYZs(IMP.core.get_leaves(ref_mh))
61  cur_low=[1e4,0]
62  for i, fit in enumerate(final_fits):
63  fit.set_index(i)
64  if self.ref_pdb!='':
65  trans=fit.get_fit_transformation()
66  IMP.atom.transform(mol2fit,trans)
67  rmsd=IMP.atom.get_rmsd(mh_xyz,ref_mh_xyz)
68  if rmsd<cur_low[0]:
69  cur_low[0]=rmsd
70  cur_low[1]=i
71  fit.set_rmsd_to_reference(rmsd)
72  IMP.atom.transform(mol2fit,trans.get_inverse())
73  fit.set_fit_transformation(trans*initial_transformation)
74  if self.ref_pdb!='':
75  print 'from all fits, lowest rmsd to ref:',cur_low
76  IMP.multifit.write_fitting_solutions(self.fits_fn, final_fits)
77 
78 
79 
80 def do_work(f):
81  f.run()
82 
83 def parse_args():
84  usage = """%prog [options] <assembly input> <refined assembly input> <proteomics.input> <mapping.input> <combinations file> <combination index>
85 
86 Fit subunits locally around a combination solution with FFT."""
87  parser = OptionParser(usage)
88  parser.add_option("-a", "--angle", dest="angle", type="float",
89  default=5,
90  help="angle delta (degrees) for FFT rotational "
91  "search (default 5)")
92 
93  parser.add_option("-n", "--num", dest="num", type="int",
94  default=100,
95  help="Number of fits to report"
96  "(default 100)")
97 
98  parser.add_option("-v", "--angle_voxel", dest="angle_voxel", type="int",
99  default=10,
100  help="Number of angles to keep per voxel"
101  "(default 10)")
102 
103  parser.add_option("-t", "--max_trans", dest="max_trans", type="float",
104  default=10.,
105  help="maximum translational search in A"
106  "(default 10)")
107 
108  parser.add_option("-m", "--max_angle", dest="max_angle", type="float",
109  default=30.,
110  help="maximum angular search in degrees"
111  "(default 50)")
112 
113  options, args = parser.parse_args()
114  if len(args) != 6:
115  parser.error("incorrect number of arguments")
116  return options,args
117 
118 def run(asmb_fn, asmb_refined_fn, proteomics_fn,mapping_fn,combs_fn,comb_ind,options):
119  #get rmsd for subunits
120  mdl1=IMP.Model()
121  mdl2=IMP.Model()
122  combs=IMP.multifit.read_paths(combs_fn)
123  asmb_input=IMP.multifit.read_settings(asmb_fn)
124  asmb_input.set_was_used(True)
125  asmb_refined_input=IMP.multifit.read_settings(asmb_refined_fn)
126  asmb_refined_input.set_was_used(True)
127  prot_data=IMP.multifit.read_proteomics_data(proteomics_fn)
128  mapping_data=IMP.multifit.read_protein_anchors_mapping(prot_data,
129  mapping_fn)
130  ensmb=IMP.multifit.load_ensemble(asmb_input,mdl1,mapping_data)
131  ensmb.set_was_used(True)
132  mhs=ensmb.get_molecules()
133 
134  ensmb_ref=IMP.multifit.load_ensemble(asmb_input,mdl2,mapping_data)
135  ensmb_ref.set_was_used(True)
136  mhs_ref=ensmb_ref.get_molecules()
137 
138  ensmb.load_combination(combs[comb_ind])
139 
140  em_map=asmb_input.get_assembly_header().get_dens_fn()
141  resolution=asmb_input.get_assembly_header().get_resolution()
142  spacing=asmb_input.get_assembly_header().get_spacing()
143  origin=asmb_input.get_assembly_header().get_origin()
144 
145  rbs_ref=ensmb_ref.get_rigid_bodies()
146  rbs=ensmb.get_rigid_bodies()
147 
148  for i,mh in enumerate(mhs):
149  fits_fn=asmb_refined_input.get_component_header(i).get_transformations_fn()
150 
151  #todo - get the initial transformation
152  rb_ref=rbs_ref[i]
153  rb=rbs[i]
154 
155  initial_transformation=IMP.algebra.get_transformation_from_first_to_second(
156  rb_ref.get_reference_frame(),
157  rb.get_reference_frame())
158 
159  pdb_fn=asmb_input.get_component_header(i).get_filename()
160 
161  f = Fitter(em_map, spacing, resolution, origin, asmb_input.get_assembly_header().get_threshold(),pdb_fn, fits_fn, options.angle,options.num,options.angle_voxel,
162  options.max_trans,options.max_angle)
163  f.run_local_fitting(mh,rb,initial_transformation)
164 
165 def main():
166  options,args = parse_args()
167  asmb_input = args[0]
168  asmb_refined_input = args[1]
169  proteomics_fn=args[2]
170  mapping_fn=args[3]
171  combinations_fn = args[4]
172  combination_ind = int(args[5])
173  run(asmb_input, asmb_refined_input, proteomics_fn,mapping_fn,combinations_fn, combination_ind, options)
174 
175 if __name__=="__main__":
176  main()