IMP logo
IMP Reference Guide  2.21.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  import multiprocessing
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 def _get_context():
27  if hasattr(multiprocessing, 'get_context'):
28  # Use 'forkserver' rather than 'fork' start method if we can;
29  # 'fork' does not work well with multithreaded processes or CUDA
30  if 'forkserver' in multiprocessing.get_all_start_methods():
31  return multiprocessing.get_context('forkserver')
32  else:
33  return multiprocessing.get_context()
34  else:
35  # For Python < 3.4, just use the original module
36  return multiprocessing
37 
38 
39 class Fitter(object):
40 
41  def __init__(
42  self,
43  em_map,
44  spacing,
45  resolution,
46  origin,
47  density_threshold,
48  pdb,
49  fits_fn,
50  angle,
51  num_fits,
52  angles_per_voxel,
53  ref_pdb=''):
54  self.em_map = em_map
55  self.spacing = spacing
56  self.resolution = resolution
57  self.threshold = density_threshold
58  self.originx = origin[0]
59  self.originy = origin[1]
60  self.originz = origin[2]
61  self.pdb = pdb
62  self.fits_fn = fits_fn
63  self.angle = angle
64  self.num_fits = num_fits
65  self.angles_per_voxel = angles_per_voxel
66  self.ref_pdb = ref_pdb
67 
68  def run(self):
69  print("resolution is:", self.resolution)
70  dmap = IMP.em.read_map(self.em_map)
71  dmap.get_header().set_resolution(self.resolution)
72  dmap.update_voxel_size(self.spacing)
73  dmap.set_origin(IMP.algebra.Vector3D(self.originx,
74  self.originy,
75  self.originz))
76  dmap.set_was_used(True)
77  dmap.get_header().show()
78  mdl = IMP.Model()
79  mol2fit = IMP.atom.read_pdb(self.pdb, mdl)
80  mh_xyz = IMP.core.XYZs(IMP.core.get_leaves(mol2fit))
81  _ = IMP.atom.create_rigid_body(mol2fit)
83  ff.set_was_used(True)
84  fits = ff.do_global_fitting(dmap, self.threshold, mol2fit,
85  self.angle / 180.0 * math.pi,
86  self.num_fits, self.spacing, 0.5,
87  True, self.angles_per_voxel)
88  fits.set_was_used(True)
89  final_fits = fits.best_fits_
90  if self.ref_pdb != '':
91  ref_mh = IMP.atom.read_pdb(self.ref_pdb, mdl)
92  ref_mh_xyz = IMP.core.XYZs(IMP.core.get_leaves(ref_mh))
93  cur_low = [1e4, 0]
94  for i, fit in enumerate(final_fits):
95  fit.set_index(i)
96  if self.ref_pdb != '':
97  trans = fit.get_fit_transformation()
98  IMP.atom.transform(mol2fit, trans)
99  rmsd = IMP.atom.get_rmsd(mh_xyz, ref_mh_xyz)
100  if rmsd < cur_low[0]:
101  cur_low[0] = rmsd
102  cur_low[1] = i
103  fit.set_rmsd_to_reference(rmsd)
104  IMP.atom.transform(mol2fit, trans.get_inverse())
105  if self.ref_pdb != '':
106  print('from all fits, lowest rmsd to ref:', cur_low)
107  IMP.multifit.write_fitting_solutions(self.fits_fn, final_fits)
108 
109 
110 def do_work(f):
111  f.run()
112 
113 
114 def parse_args():
115  desc = """Fit subunits into a density map with FFT."""
116  p = ArgumentParser(description=desc)
117  p.add_argument("-c", "--cpu", dest="cpus", type=int, default=1,
118  help="number of cpus to use (default 1)")
119  p.add_argument("-a", "--angle", dest="angle", type=float, default=30,
120  help="angle delta (degrees) for FFT rotational "
121  "search (default 30)")
122 
123  p.add_argument("-n", "--num", dest="num", type=int,
124  default=100, help="Number of fits to report (default 100)")
125 
126  p.add_argument("-v", "--angle_voxel", dest="angle_voxel", type=int,
127  default=10,
128  help="Number of angles to keep per voxel (default 10)")
129 
130  p.add_argument("assembly_file", help="assembly file name")
131 
132  # p.add_argument("-n", "--num", dest="num", type="int",
133  # default=100,
134  # help="Number of fits to report"
135  # "(default 100)")
136 
137  return p.parse_args()
138 
139 
140 def run(asmb_fn, options):
141  if multiproc_exception is None and options.cpus > 1:
142  work_units = []
143  asmb_input = IMP.multifit.read_settings(asmb_fn)
144  asmb_input.set_was_used(True)
145  em_map = asmb_input.get_assembly_header().get_dens_fn()
146  resolution = asmb_input.get_assembly_header().get_resolution()
147  spacing = asmb_input.get_assembly_header().get_spacing()
148  origin = asmb_input.get_assembly_header().get_origin()
149  for i in range(asmb_input.get_number_of_component_headers()):
150  fits_fn = asmb_input.get_component_header(i).get_transformations_fn()
151  pdb_fn = asmb_input.get_component_header(i).get_filename()
152  f = Fitter(
153  em_map,
154  spacing,
155  resolution,
156  origin,
157  asmb_input.get_assembly_header().get_threshold(),
158  pdb_fn,
159  fits_fn,
160  options.angle,
161  options.num,
162  options.angle_voxel)
163  if multiproc_exception is None and options.cpus > 1:
164  work_units.append(f)
165  else:
166  if options.cpus > 1:
167  options.cpus = 1
168  print("""
169 The Python 'multiprocessing' module (available in Python 2.6 and later) is
170 needed to run on multiple CPUs, and could not be found
171 (Python error: '%s').
172 Running on a single processor.""" % multiproc_exception, file=sys.stderr)
173  f.run()
174  if multiproc_exception is None and options.cpus > 1:
175  # No point in spawning more processes than components
176  nproc = min(options.cpus, asmb_input.get_number_of_component_headers())
177  ctx = _get_context()
178  p = ctx.Pool(processes=nproc)
179  _ = list(p.imap_unordered(do_work, work_units))
180  p.close()
181 
182 
183 def main():
184  args = parse_args()
185  run(args.assembly_file, args)
186 
187 
188 if __name__ == "__main__":
189  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:86
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:408
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.