IMP  2.0.1
The Integrative Modeling Platform
emagefit_score.py
1 #!/usr/bin/env python
2 
3 import IMP
4 import IMP.core as core
5 import IMP.atom as atom
6 import IMP.em2d as em2d
7 import os
8 import sys
9 import time
10 import logging
11 log = logging.getLogger("score_model")
12 
13 def score_model(complete_fn_model,
14  images_sel_file,
15  pixel_size,
16  n_projections=20,
17  resolution=1,
18  images_per_batch=250):
19  """ Score a model
20 
21  Scores a model against the images in the selection file.
22  Reads the images in batchs to avoid memory problems
23  resolution and pixel_size are used for generating projections of the model
24  The finder is an em2d.ProjectionFinder used for projection matching and optimizations
25  """
26  print "SCORING MODEL:",complete_fn_model
27  cwd = os.getcwd()
28 
29  images_dir, nil = os.path.split(images_sel_file)
30  images_names = em2d.read_selection_file(images_sel_file)
31  n_images = len(images_names)
32  if(n_images == 0):
33  raise ValueError(" Scoring with a empty set of images")
34 
35  # TYPICAL OPTIMIZER PARAMETERS
36  params = em2d.Em2DRestraintParameters(pixel_size, resolution)
37  params.coarse_registration_method = em2d.ALIGN2D_PREPROCESSING
38  params.optimization_steps = 4
39  params.simplex_initial_length = 0.1
40  params.simplex_minimum_size=0.02
41  params.save_match_images = True
42  score_function = em2d.EM2DScore()
43  finder = em2d.ProjectionFinder()
44  finder.setup(score_function, params)
45 
46  # Get the number of rows and cols from the 1st image
47  srw = em2d.SpiderImageReaderWriter()
48  test_imgs = em2d.read_images([os.path.join(images_dir, images_names[0])], srw)
49  rows = test_imgs[0].get_header().get_number_of_columns()
50  cols = test_imgs[0].get_header().get_number_of_rows()
51 
52  model = IMP.Model()
53  ssel = atom.ATOMPDBSelector()
54  prot = atom.read_pdb(complete_fn_model, model, ssel)
55  particles = IMP.core.get_leaves(prot)
56  # generate projections
57  proj_params = em2d.get_evenly_distributed_registration_results(n_projections)
58  opts = em2d.ProjectingOptions(pixel_size, resolution)
59  projections = em2d.get_projections(particles, proj_params, rows, cols, opts)
60 
61  finder.set_model_particles(particles)
62  finder.set_projections(projections)
63  optimized_solutions= 2
64  finder.set_fast_mode(optimized_solutions)
65  # read the images in blocks to avoid memory problems
66  all_registration_results = []
67  init_set = 0
68  init_time = time.time()
69  while(init_set < n_images):
70  end_set = min( init_set + images_per_batch, n_images )
71  if(images_dir != ""):
72  os.chdir(images_dir)
73  subjects = em2d.read_images(images_names[init_set:end_set], srw)
74  # register
75  finder.set_subjects(subjects)
76  os.chdir(cwd)
77  finder.get_complete_registration()
78  # Recover the registration results:
79  registration_results = finder.get_registration_results()
80  for reg in registration_results:
81  all_registration_results.append(reg)
82  init_set += images_per_batch
83  os.chdir(cwd)
84  em2d.write_registration_results("registration.params", all_registration_results)
85  print "score_model: time complete registration",time.time()-init_time
86  print "coarse registration time",finder.get_coarse_registration_time()
87  print "fine registration time",finder.get_fine_registration_time()
88  return all_registration_results
89 
90 def parse_args():
91  parser = IMP.OptionParser(usage="""%prog model selfile pxsize nproj res nimg
92 
93 model: PDB file of the model to score
94 selfile: Selection file containing the images used for scoring
95 pxsize: Pixel size of the images in Angstrom/pixel
96 nproj: Number of projections of the model used for start registering
97 res: Resolution to generate the projections
98 nimg: Images per batch used when scoring a lot of images (to avoid memory
99  problems)
100 """,
101  description="EMageFit scoring",
102  imp_module=em2d)
103  opts, args = parser.parse_args()
104  if len(args) != 6:
105  parser.error("wrong number of arguments")
106  return args
107 
108 
109 if __name__ == "__main__":
110  args = parse_args()
111  complete_fn_model = args[0]
112  images_sel_file = args[1]
113  pixel_size = float(args[2])
114  n_projections= int(args[3])
115  resolution= float(args[4])
116  images_per_batch= int(args[5])
117 
118  all_regs = score_model(complete_fn_model,
119  images_sel_file,
120  pixel_size,
121  n_projections,
122  resolution,
123  images_per_batch)
124  for i, reg in enumerate(all_regs):
125  print "Score for image %s: %f" % (i, reg.get_score())
126  print "Global score ", em2d.get_global_score(all_regs)