def test_registration(self):
        """Test the registration of 3 subjects from 1gyt.pdb at 0.5 SNR"""
        # Get model from PDB file
        smodel = IMP.Model()
        ssel = IMP.atom.ATOMPDBSelector()
        fn_model = self.get_input_file_name("1gyt.pdb")
        prot = IMP.atom.read_pdb(fn_model, smodel, ssel)
        particles = IMP.core.get_leaves(prot)
        # Read subject images
        srw = em2d.SpiderImageReaderWriter()
        selection_file = self.get_input_file_name("1gyt-subjects-0.5-SNR.sel")
        images_to_read_names = em2d.read_selection_file(selection_file)
        for i in range(0, len(images_to_read_names)):
            images_to_read_names[i] = self.get_input_file_name(
                images_to_read_names[i])
        subjects = em2d.read_images(images_to_read_names, srw)
        self.assertEqual(len(subjects), 3, "Problem reading subject images")

        # Generate 20 evenly distributed projections from the PDB file
        n_projections = 20
        proj_params = em2d.get_evenly_distributed_registration_results(
            n_projections)
        rows = 128
        cols = 128
        pixel_size = 1.5
        # for generating projections, use a very high resolution
        resolution = 1
        options = em2d.ProjectingOptions(pixel_size, resolution)
        projections = em2d.get_projections(particles, proj_params,
                                           rows, cols, options)
        self.assertEqual(len(projections), n_projections,
                         "Problem generating projections")
        # Prepare registration
        # IMP.set_log_level(IMP.VERBOSE)
        finder = em2d.ProjectionFinder()
        score_function = em2d.EM2DScore()

        params = em2d.Em2DRestraintParameters(pixel_size,
                                              resolution, n_projections)
        params.save_match_images = False
        params.coarse_registration_method = em2d.ALIGN2D_PREPROCESSING
        params.optimization_steps = 30
        params.simplex_initial_length = 0.1
        params.simplex_minimum_size = 0.01

        finder.setup(score_function, params)
        finder.set_model_particles(particles)
        finder.set_subjects(subjects)
        finder.set_projections(projections)
        finder.set_fast_mode(2)
        finder.get_complete_registration()
        # Recover the registration results:
        registration_parameters = finder.get_registration_results()
        fn_registration_results = "my_1gyt_registration.params"
        em2d.write_registration_results(fn_registration_results,
                                        registration_parameters)
        # Read the correct registration results:
        correct_parameters = em2d.read_registration_results(
            self.get_input_file_name("1gyt-subjects-0.5-SNR.params"))

        print("determined: ")
        for r in registration_parameters:
            print(r.get_rotation(), r.get_shift())
        print("correct: ")
        for r in correct_parameters:
            print(r.get_rotation(), r.get_shift())
        for i in range(0, len(registration_parameters)):
            # Generate the registered projection
            imgx = em2d.Image()
            imgx.set_size(rows, cols)
            em2d.get_projection(imgx, particles,
                                registration_parameters[i], options)
            ccc = em2d.get_cross_correlation_coefficient(
                subjects[i].get_data(),
                imgx.get_data())
            print(i, "ccc", ccc)
            snr = 0.5
            theoretical_ccc = (snr / (1. + snr)) ** .5
            self.assertAlmostEqual(ccc, theoretical_ccc, delta=0.02,
                                   msg="Error in registration of subject %d: ccc %8.3f "
                                   "theoretical_ccc %8.3f " % (i, ccc, theoretical_ccc))
        os.remove(fn_registration_results)
Example #2
0
    def test_registration(self):
        """Test the registration of 3 subjects from 1gyt.pdb at 0.5 SNR"""
        # Get model from PDB file
        smodel = IMP.Model()
        ssel = IMP.atom.ATOMPDBSelector()
        fn_model = self.get_input_file_name("1gyt.pdb")
        prot = IMP.atom.read_pdb(fn_model, smodel, ssel)
        particles = IMP.core.get_leaves(prot)
        # Read subject images
        srw = em2d.SpiderImageReaderWriter()
        selection_file = self.get_input_file_name("1gyt-subjects-0.5-SNR.sel")
        images_to_read_names = em2d.read_selection_file(selection_file)
        for i in range(0, len(images_to_read_names)):
            images_to_read_names[i] = self.get_input_file_name(
                images_to_read_names[i])
        subjects = em2d.read_images(images_to_read_names, srw)
        self.assertEqual(len(subjects), 3, "Problem reading subject images")

        # Generate 20 evenly distributed projections from the PDB file
        n_projections = 20
        proj_params = em2d.get_evenly_distributed_registration_results(
            n_projections)
        rows = 128
        cols = 128
        pixel_size = 1.5
        # for generating projections, use a very high resolution
        resolution = 8.5
        options = em2d.ProjectingOptions(pixel_size, resolution)
        projections = em2d.get_projections(particles, proj_params, rows, cols,
                                           options)
        self.assertEqual(len(projections), n_projections,
                         "Problem generating projections")
        # Prepare registration
        # IMP.set_log_level(IMP.VERBOSE)
        finder = em2d.ProjectionFinder()
        score_function = em2d.EM2DScore()

        params = em2d.Em2DRestraintParameters(pixel_size, resolution,
                                              n_projections)
        params.save_match_images = False
        params.coarse_registration_method = em2d.ALIGN2D_PREPROCESSING
        params.optimization_steps = 30
        params.simplex_initial_length = 0.1
        params.simplex_minimum_size = 0.01

        finder.setup(score_function, params)
        finder.set_model_particles(particles)
        finder.set_subjects(subjects)
        finder.set_projections(projections)
        finder.set_fast_mode(2)
        finder.get_complete_registration()
        # Recover the registration results:
        registration_parameters = finder.get_registration_results()
        fn_registration_results = "my_1gyt_registration.params"
        em2d.write_registration_results(fn_registration_results,
                                        registration_parameters)
        # Read the correct registration results:
        correct_parameters = em2d.read_registration_results(
            self.get_input_file_name("1gyt-subjects-0.5-SNR.params"))

        print("determined: ")
        for r in registration_parameters:
            print(r.get_rotation(), r.get_shift())
        print("correct: ")
        for r in correct_parameters:
            print(r.get_rotation(), r.get_shift())
        for i in range(0, len(registration_parameters)):
            # Generate the registered projection
            imgx = em2d.Image()
            imgx.set_size(rows, cols)
            em2d.get_projection(imgx, particles, registration_parameters[i],
                                options)
            ccc = em2d.get_cross_correlation_coefficient(
                subjects[i].get_data(), imgx.get_data())
            print(i, "ccc", ccc)
            snr = 0.5
            theoretical_ccc = (snr / (1. + snr))**.5
            self.assertAlmostEqual(
                ccc,
                theoretical_ccc,
                delta=0.02,
                msg="Error in registration of subject %d: ccc %8.3f "
                "theoretical_ccc %8.3f " % (i, ccc, theoretical_ccc))
        os.remove(fn_registration_results)
Example #3
0
def score_model(complete_fn_model,
                images_sel_file,
                pixel_size,
                n_projections=20,
                resolution=1,
                images_per_batch=250):
    """ Score a model

    Scores a model against the images in the selection file.
    Reads the images in batchs to avoid memory problems
    resolution and pixel_size are used for generating projections of the model
    The finder is an em2d.ProjectionFinder used for projection matching and optimizations
    """
    print "SCORING MODEL:",complete_fn_model
    cwd = os.getcwd()

    images_dir, nil  =  os.path.split(images_sel_file)
    images_names = em2d.read_selection_file(images_sel_file)
    n_images = len(images_names)
    if(n_images == 0):
        raise ValueError(" Scoring with a empty set of images")

    # TYPICAL OPTIMIZER PARAMETERS
    params = em2d.Em2DRestraintParameters(pixel_size, resolution)
    params.coarse_registration_method = em2d.ALIGN2D_PREPROCESSING
    params.optimization_steps = 4
    params.simplex_initial_length = 0.1
    params.simplex_minimum_size=0.02
    params.save_match_images = True
    score_function = em2d.EM2DScore()
    finder = em2d.ProjectionFinder()
    finder.setup(score_function, params)

    # Get the number of rows and cols from the 1st image
    srw = em2d.SpiderImageReaderWriter()
    test_imgs = em2d.read_images([os.path.join(images_dir, images_names[0])], srw)
    rows = test_imgs[0].get_header().get_number_of_columns()
    cols = test_imgs[0].get_header().get_number_of_rows()

    model = IMP.Model()
    ssel = atom.ATOMPDBSelector()
    prot =  atom.read_pdb(complete_fn_model, model, ssel)
    particles = IMP.core.get_leaves(prot)
    # generate projections
    proj_params = em2d.get_evenly_distributed_registration_results(n_projections)
    opts = em2d.ProjectingOptions(pixel_size, resolution)
    projections = em2d.get_projections(particles, proj_params, rows, cols, opts)

    finder.set_model_particles(particles)
    finder.set_projections(projections)
    optimized_solutions= 2
    finder.set_fast_mode(optimized_solutions)
    # read the images in blocks to avoid memory problems
    all_registration_results = []
    init_set = 0
    init_time = time.time()
    while(init_set < n_images):
        end_set = min( init_set + images_per_batch, n_images )
        if(images_dir != ""):
            os.chdir(images_dir)
        subjects = em2d.read_images(images_names[init_set:end_set], srw)
        # register
        finder.set_subjects(subjects)
        os.chdir(cwd)
        finder.get_complete_registration()
        # Recover the registration results:
        registration_results = finder.get_registration_results()
        for reg in registration_results:
            all_registration_results.append(reg)
        init_set += images_per_batch
    os.chdir(cwd)
    em2d.write_registration_results("registration.params", all_registration_results)
    print "score_model: time complete registration",time.time()-init_time
    print "coarse registration time",finder.get_coarse_registration_time()
    print "fine registration time",finder.get_fine_registration_time()
    return all_registration_results