IMP logo
IMP Reference Guide  2.10.0
The Integrative Modeling Platform
fit_fft.py
1 #!/usr/bin/env python
2 
3 from __future__ import print_function
4 import math
5 import IMP.multifit
6 import IMP.atom
7 import IMP.em
8 from IMP import ArgumentParser
9 import os
10 import sys
11 
12 __doc__ = "Fit subunits into a density map with FFT."
13 
14 multiproc_exception = None
15 try:
16  from multiprocessing import Pool
17  # Detect whether we are running Windows Python via Wine. Wine does not
18  # currently support some named pipe functions which the multiprocessing
19  # module needs: http://bugs.winehq.org/show_bug.cgi?id=17273
20  if sys.platform == 'win32' and 'WINELOADERNOEXEC' in os.environ:
21  multiproc_exception = "Wine does not currently support multiprocessing"
22 except ImportError as detail:
23  multiproc_exception = str(detail)
24 
25 
26 class Fitter(object):
27 
28  def __init__(
29  self,
30  em_map,
31  spacing,
32  resolution,
33  origin,
34  density_threshold,
35  pdb,
36  fits_fn,
37  angle,
38  num_fits,
39  angles_per_voxel,
40  ref_pdb=''):
41  self.em_map = em_map
42  self.spacing = spacing
43  self.resolution = resolution
44  self.threshold = density_threshold
45  self.originx = origin[0]
46  self.originy = origin[1]
47  self.originz = origin[2]
48  self.pdb = pdb
49  self.fits_fn = fits_fn
50  self.angle = angle
51  self.num_fits = num_fits
52  self.angles_per_voxel = angles_per_voxel
53  self.ref_pdb = ref_pdb
54 
55  def run(self):
56  print("resolution is:", self.resolution)
57  dmap = IMP.em.read_map(self.em_map)
58  dmap.get_header().set_resolution(self.resolution)
59  dmap.update_voxel_size(self.spacing)
60  dmap.set_origin(IMP.algebra.Vector3D(self.originx,
61  self.originy,
62  self.originz))
63  dmap.set_was_used(True)
64  dmap.get_header().show()
65  mdl = IMP.Model()
66  mol2fit = IMP.atom.read_pdb(self.pdb, mdl)
67  mh_xyz = IMP.core.XYZs(IMP.core.get_leaves(mol2fit))
68  rb = IMP.atom.create_rigid_body(mol2fit)
70  ff.set_was_used(True)
71  fits = ff.do_global_fitting(dmap, self.threshold, mol2fit,
72  self.angle / 180.0 * math.pi,
73  self.num_fits, self.spacing, 0.5,
74  True, self.angles_per_voxel)
75  fits.set_was_used(True)
76  final_fits = fits.best_fits_
77  if self.ref_pdb != '':
78  ref_mh = IMP.atom.read_pdb(self.ref_pdb, mdl)
79  ref_mh_xyz = IMP.core.XYZs(IMP.core.get_leaves(ref_mh))
80  cur_low = [1e4, 0]
81  for i, fit in enumerate(final_fits):
82  fit.set_index(i)
83  if self.ref_pdb != '':
84  trans = fit.get_fit_transformation()
85  IMP.atom.transform(mol2fit, trans)
86  rmsd = IMP.atom.get_rmsd(mh_xyz, ref_mh_xyz)
87  if rmsd < cur_low[0]:
88  cur_low[0] = rmsd
89  cur_low[1] = i
90  fit.set_rmsd_to_reference(rmsd)
91  IMP.atom.transform(mol2fit, trans.get_inverse())
92  if self.ref_pdb != '':
93  print('from all fits, lowest rmsd to ref:', cur_low)
94  IMP.multifit.write_fitting_solutions(self.fits_fn, final_fits)
95 
96 
97 def do_work(f):
98  f.run()
99 
100 
101 def parse_args():
102  desc = """Fit subunits into a density map with FFT."""
103  p = ArgumentParser(description=desc)
104  p.add_argument("-c", "--cpu", dest="cpus", type=int, default=1,
105  help="number of cpus to use (default 1)")
106  p.add_argument("-a", "--angle", dest="angle", type=float, default=30,
107  help="angle delta (degrees) for FFT rotational "
108  "search (default 30)")
109 
110  p.add_argument("-n", "--num", dest="num", type=int,
111  default=100, help="Number of fits to report (default 100)")
112 
113  p.add_argument("-v", "--angle_voxel", dest="angle_voxel", type=int,
114  default=10,
115  help="Number of angles to keep per voxel (default 10)")
116 
117  p.add_argument("assembly_file", help="assembly file name")
118 
119  # p.add_argument("-n", "--num", dest="num", type="int",
120  # default=100,
121  # help="Number of fits to report"
122  # "(default 100)")
123 
124  return p.parse_args()
125 
126 
127 def run(asmb_fn, options):
128  if multiproc_exception is None and options.cpus > 1:
129  work_units = []
130  asmb_input = IMP.multifit.read_settings(asmb_fn)
131  asmb_input.set_was_used(True)
132  em_map = asmb_input.get_assembly_header().get_dens_fn()
133  resolution = asmb_input.get_assembly_header().get_resolution()
134  spacing = asmb_input.get_assembly_header().get_spacing()
135  origin = asmb_input.get_assembly_header().get_origin()
136  for i in range(asmb_input.get_number_of_component_headers()):
137  fits_fn = asmb_input.get_component_header(i).get_transformations_fn()
138  pdb_fn = asmb_input.get_component_header(i).get_filename()
139  f = Fitter(
140  em_map,
141  spacing,
142  resolution,
143  origin,
144  asmb_input.get_assembly_header().get_threshold(),
145  pdb_fn,
146  fits_fn,
147  options.angle,
148  options.num,
149  options.angle_voxel)
150  if multiproc_exception is None and options.cpus > 1:
151  work_units.append(f)
152  else:
153  if options.cpus > 1:
154  options.cpus = 1
155  print("""
156 The Python 'multiprocessing' module (available in Python 2.6 and later) is
157 needed to run on multiple CPUs, and could not be found
158 (Python error: '%s').
159 Running on a single processor.""" % multiproc_exception, file=sys.stderr)
160  f.run()
161  if multiproc_exception is None and options.cpus > 1:
162  # No point in spawning more processes than components
163  nproc = min(options.cpus, asmb_input.get_number_of_component_headers())
164  p = Pool(processes=nproc)
165  out = list(p.imap_unordered(do_work, work_units))
166 
167 
168 def main():
169  args = parse_args()
170  run(args.assembly_file, args)
171 
172 if __name__ == "__main__":
173  main()
Fit a molecule inside its density by local or global FFT.
SettingsData * read_settings(const char *filename)
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
double get_rmsd(const Selection &s0, const Selection &s1)
void transform(Hierarchy h, const algebra::Transformation3D &tr)
Transform a hierarchy. This is aware of rigid bodies.
Fitting atomic structures into a cryo-electron microscopy density map.
Basic utilities for handling cryo-electron microscopy 3D density maps.
void write_fitting_solutions(const char *fitting_fn, const FittingSolutionRecords &fit_sols, int num_sols=-1)
Write fitting solutions to a file.
std::ostream & show(Hierarchy h, std::ostream &out=std::cout)
Print the hierarchy using a given decorator to display each node.
VectorD< 3 > Vector3D
Definition: VectorD.h:395
IMP::core::RigidBody create_rigid_body(Hierarchy h)
double get_resolution(Model *m, ParticleIndex pi)
Estimate the resolution of the hierarchy as used by Representation.
Functionality for loading, creating, manipulating and scoring atomic structures.