IMP  2.0.0
The Integrative Modeling Platform
fit_fft.py
1 #!/usr/bin/env python
2 
3 __doc__ = "Fit 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 multiproc_exception = None
14 try:
15  from multiprocessing import Pool
16  # Detect whether we are running Windows Python via Wine. Wine does not
17  # currently support some named pipe functions which the multiprocessing
18  # module needs: http://bugs.winehq.org/show_bug.cgi?id=17273
19  if sys.platform == 'win32' and 'WINELOADERNOEXEC' in os.environ:
20  multiproc_exception = "Wine does not currently support multiprocessing"
21 except ImportError, detail:
22  multiproc_exception = str(detail)
23 
24 class Fitter(object):
25  def __init__(self, em_map, spacing, resolution, origin, density_threshold,pdb, fits_fn, angle,num_fits,angles_per_voxel,ref_pdb=''):
26  self.em_map = em_map
27  self.spacing = spacing
28  self.resolution = resolution
29  self.threshold=density_threshold
30  self.originx = origin[0]
31  self.originy = origin[1]
32  self.originz = origin[2]
33  self.pdb = pdb
34  self.fits_fn = fits_fn
35  self.angle = angle
36  self.num_fits=num_fits
37  self.angles_per_voxel=angles_per_voxel
38  self.ref_pdb=ref_pdb
39  def run(self):
40  print "resolution is:",self.resolution
41  dmap = IMP.em.read_map(self.em_map)
42  dmap.get_header().set_resolution(self.resolution)
43  dmap.update_voxel_size(self.spacing)
44  dmap.set_origin(IMP.algebra.Vector3D(self.originx,
45  self.originy,
46  self.originz))
47  dmap.set_was_used(True)
48  dmap.get_header().show()
49  mdl = IMP.Model()
50  mol2fit = IMP.atom.read_pdb(self.pdb, mdl)
51  mh_xyz=IMP.core.XYZs(IMP.core.get_leaves(mol2fit))
52  rb = IMP.atom.create_rigid_body(mol2fit)
53  ff = IMP.multifit.FFTFitting()
54  ff.set_was_used(True)
55  fits = ff.do_global_fitting(dmap, self.threshold,mol2fit,
56  self.angle / 180.0 * math.pi,
57  self.num_fits, self.spacing, 0.5,
58  True, self.angles_per_voxel)
59  fits.set_was_used(True)
60  final_fits = fits.best_fits_
61  if self.ref_pdb!='':
62  ref_mh=IMP.atom.read_pdb(self.ref_pdb,mdl)
63  ref_mh_xyz=IMP.core.XYZs(IMP.core.get_leaves(ref_mh))
64  cur_low=[1e4,0]
65  for i, fit in enumerate(final_fits):
66  fit.set_index(i)
67  if self.ref_pdb!='':
68  trans=fit.get_fit_transformation()
69  IMP.atom.transform(mol2fit,trans)
70  rmsd=IMP.atom.get_rmsd(mh_xyz,ref_mh_xyz)
71  if rmsd<cur_low[0]:
72  cur_low[0]=rmsd
73  cur_low[1]=i
74  fit.set_rmsd_to_reference(rmsd)
75  IMP.atom.transform(mol2fit,trans.get_inverse())
76  if self.ref_pdb!='':
77  print 'from all fits, lowest rmsd to ref:',cur_low
78  IMP.multifit.write_fitting_solutions(self.fits_fn, final_fits)
79 
80 def do_work(f):
81  f.run()
82 
83 def parse_args():
84  usage = """%prog [options] <assembly input>
85 
86 Fit subunits into a density map with FFT."""
87  parser = OptionParser(usage)
88  parser.add_option("-c", "--cpu", dest="cpus", type="int", default=1,
89  help="number of cpus to use (default 1)")
90  parser.add_option("-a", "--angle", dest="angle", type="float",
91  default=30,
92  help="angle delta (degrees) for FFT rotational "
93  "search (default 30)")
94 
95  parser.add_option("-n", "--num", dest="num", type="int",
96  default=100,
97  help="Number of fits to report"
98  "(default 100)")
99 
100  parser.add_option("-v", "--angle_voxel", dest="angle_voxel", type="int",
101  default=10,
102  help="Number of angles to keep per voxel"
103  "(default 10)")
104 
105  #parser.add_option("-n", "--num", dest="num", type="int",
106  # default=100,
107  # help="Number of fits to report"
108  # "(default 100)")
109 
110  options, args = parser.parse_args()
111  if len(args) != 1:
112  parser.error("incorrect number of arguments")
113  return options,args
114 
115 def run(asmb_fn, options):
116  if multiproc_exception is None and options.cpus > 1:
117  work_units = []
118  asmb_input=IMP.multifit.read_settings(asmb_fn)
119  asmb_input.set_was_used(True)
120  em_map=asmb_input.get_assembly_header().get_dens_fn()
121  resolution=asmb_input.get_assembly_header().get_resolution()
122  spacing=asmb_input.get_assembly_header().get_spacing()
123  origin=asmb_input.get_assembly_header().get_origin()
124  for i in range(asmb_input.get_number_of_component_headers()):
125  fits_fn=asmb_input.get_component_header(i).get_transformations_fn()
126  pdb_fn=asmb_input.get_component_header(i).get_filename()
127  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)
128  if multiproc_exception is None and options.cpus > 1:
129  work_units.append(f)
130  else:
131  if options.cpus > 1:
132  options.cpus = 1
133  print >> sys.stderr, """
134 The Python 'multiprocessing' module (available in Python 2.6 and later) is
135 needed to run on multiple CPUs, and could not be found
136 (Python error: '%s').
137 Running on a single processor.""" % multiproc_exception
138  f.run()
139  if multiproc_exception is None and options.cpus > 1:
140  # No point in spawning more processes than components
141  nproc = min(options.cpus, asmb_input.get_number_of_component_headers())
142  p = Pool(processes=nproc)
143  out = list(p.imap_unordered(do_work, work_units))
144 
145 def main():
146  options,args = parse_args()
147  asmb_input = args[0]
148  run(asmb_input, options)
149 
150 if __name__=="__main__":
151  main()