コード例 #1
0
ファイル: pixel_refinement.py プロジェクト: monarin/LS49
    def run(self):

        self.parse_input()

        N_total = 100000  # self.params.N_total # number of items to simulate, nominally 100000
        logical_rank = self.mpi_helper.rank
        logical_size = self.mpi_helper.size
        if self.mpi_helper.rank == 0 and self.mpi_helper.size == 1:  # special case of testing it
            try:
                logical_rank = self.params.tester.rank
                logical_size = self.params.tester.size
            except Exception:
                pass
        N_stride = int(math.ceil(
            N_total / logical_size))  # total number of tasks per rank
        print("hello from rank %d of %d with stride %d" %
              (logical_rank, logical_size, N_stride))

        #from scitbx.lbfgs.tst_mpi_split_evaluator import mpi_split_evaluator_run
        #from scitbx.lbfgs.tst_mpi_split_evaluator import run_mpi as simple_tester
        #simple_tester()

        if self.params.starting_model.algorithm == "to_file":
            if self.mpi_helper.rank == 0:
                HKL_lookup, static_fcalcs = get_static_fcalcs_with_HKL_lookup()
                from LS49.work2_for_aca_lsq.remake_range_intensities_with_complex \
                   import get_intensity_structure
                model_intensities = get_intensity_structure(
                    static_fcalcs,
                    FE1_model=Fe_oxidized_model,
                    FE2_model=Fe_reduced_model)
                with (open(self.params.starting_model.filename, "wb")) as out:
                    pickle.dump(HKL_lookup, out, pickle.HIGHEST_PROTOCOL)
                    pickle.dump(static_fcalcs, out, pickle.HIGHEST_PROTOCOL)
                    pickle.dump(model_intensities, out,
                                pickle.HIGHEST_PROTOCOL)
            return
        else:
            if self.mpi_helper.rank == 0:
                with (open(self.params.starting_model.filename, "rb")) as inp:
                    print("the starting model (used for channel weighting) is",
                          self.params.starting_model.filename)
                    HKL_lookup = pickle.load(inp)
                    static_fcalcs = pickle.load(inp)
                    model_intensities = pickle.load(inp)
                    from LS49.spectra.generate_spectra import spectra_simulation
                    from LS49.sim.step5_pad import microcrystal

                shuffA = list(range(N_total))
                import random
                random.shuffle(shuffA)

                transmitted_info = dict(
                    HKL_lookup=HKL_lookup,
                    static_fcalcs=static_fcalcs,
                    model_intensities=model_intensities,
                    spectra_simulation=spectra_simulation(),
                    crystal=microcrystal(Deff_A=4000,
                                         length_um=4.,
                                         beam_diameter_um=1.0),
                    shuffA=shuffA)
            else:
                transmitted_info = None
        print("before braodcast with ", self.mpi_helper.rank,
              self.mpi_helper.size)
        transmitted_info = self.mpi_helper.comm.bcast(transmitted_info, root=0)
        self.mpi_helper.comm.barrier()
        print("after barrier")

        # -----------------------------------------------------------------------
        if self.mpi_helper.rank == 0:
            print("Finding initial G and abc factors")
        per_rank_items = []
        per_rank_keys = []
        per_rank_G = []
        min_spots = 3
        N_input = 0
        import os, omptbx  # cori workaround, which does not get OMP_NUM_THREADS from environment
        workaround_nt = int(os.environ.get("OMP_NUM_THREADS", 1))
        omptbx.omp_set_num_threads(workaround_nt)
        for item, key in get_items(logical_rank, N_total, N_stride,
                                   transmitted_info["shuffA"],
                                   self.params.cohort):
            N_input += 1
            if len(item) >= min_spots:
                try:
                    FOI = fit_one_image_multispot(
                        key=key,
                        list_of_images=item,
                        HKL_lookup=transmitted_info["HKL_lookup"],
                        model_intensities=transmitted_info[
                            "model_intensities"],
                        spectra=transmitted_info["spectra_simulation"],
                        crystal=transmitted_info["crystal"])
                except RuntimeError as e:
                    # no recovery from LBFGS error, skip event
                    continue
                metric_P1, metric_C2 = FOI.DRM.get_current_angular_offsets(
                    FOI.x[-4:-1])
                print(
                    """LLG Image %06d on %d Bragg spots NLL    channels F = %9.1f angular offsets in P1 and C2 (degrees): %8.5f %8.5f"""
                    %
                    (key, len(item), FOI.compute_functional_and_gradients()[0],
                     metric_P1, metric_C2), FOI.x[-4:-1])

                # reporting out results to new abc_coverage pickles.  Use the old one "item" as a template:
                #    item := [<LS49.work2_for_aca_lsq.abc_background.fit_roi_multichannel>,... one for each spot]
                # each fit_roi has the following attributes and we modify them as follows:
                #        'a', the abcG parameters of the original one-spot fit, unused
                #        'asu_idx_C2_setting', unmodified
                #        'image_no', same as key, unmodified
                #        'n', number of parameters for one-spot fit, unused
                #        'orig_idx_C2_setting', unmodified
                #        'sb_data', original shoebox data [integers as floats], passed along unmodified
                #        'simtbx_P1_miller', unmodified
                #        'simtbx_intensity_7122', unmodified
                #        'x', the abcG parameters of the original one-spot fit, unused
                # modify these:
                #        'bkgrd_a', the spot abc parameters output here, to be passed on to global data fit
                for ispot in range(len(item)):
                    item[ispot].bkgrd_a = FOI.x[3 * ispot:3 * (ispot + 1)]
                    #        'channels', need to save the new data that was set during update_roi_model_pixels_with_current_rotation()
                    #                    but only for the Amat, not the derivatives.
                    item[ispot].channels = FOI.new_calc2_dict_last_round[
                        ispot]["channels"]
                    #        'roi', get new region of interest summation that was set during update_roi_model_pixels_with_current_rotation()
                    item[ispot].roi = FOI.roi_model_pixels[ispot]

                # put the newly refined background model back into the item
                per_rank_keys.append(key)
                per_rank_G.append(FOI.a[-1])

                print(
                    "pickling modified abc_coverage file for key %d in rank %d"
                    % (key, logical_rank), )
                with open(abc_glob_pixel_ref % (key), "wb") as F:
                    pickle.dump(item, F, pickle.HIGHEST_PROTOCOL)

        print("rank %d has %d refined images" %
              (logical_rank, len(per_rank_keys)))

        N_ranks = self.mpi_helper.comm.reduce(1, self.mpi_helper.MPI.SUM, 0)
        N_refined_images = self.mpi_helper.comm.reduce(len(per_rank_keys),
                                                       self.mpi_helper.MPI.SUM,
                                                       0)
        N_input_images = self.mpi_helper.comm.reduce(N_input,
                                                     self.mpi_helper.MPI.SUM,
                                                     0)
        self.mpi_helper.comm.barrier()
        if self.mpi_helper.rank == 0:
            print("final report %d ranks, %d input images, %d refined models" %
                  (N_ranks, N_input_images, N_refined_images))
            print("Finished finding initial G and abc factors")

        self.mpi_helper.comm.barrier()
コード例 #2
0
ファイル: macrocycle_refinery.py プロジェクト: monarin/LS49
    def run(self):

        self.parse_input()

        N_total = self.params.N_total  # number of items to simulate, nominally 100000
        logical_rank = self.mpi_helper.rank
        logical_size = self.mpi_helper.size
        if self.mpi_helper.rank == 0 and self.mpi_helper.size == 1:  # special case of testing it
            try:
                logical_rank = self.params.tester.rank
                logical_size = self.params.tester.size
            except Exception:
                pass
        N_stride = int(math.ceil(
            N_total / logical_size))  # total number of tasks per rank
        print("hello from rank %d of %d with stride %d" %
              (logical_rank, logical_size, N_stride))

        from scitbx.lbfgs.tst_mpi_split_evaluator import mpi_split_evaluator_run

        assert self.params.starting_model.algorithm!="to_file", \
               "run new_global__fdp_refinery.py first, to generate starting model"
        if self.mpi_helper.rank == 0:
            with (open(self.params.starting_model.filename, "rb")) as inp:
                HKL_lookup = pickle.load(inp)
                static_fcalcs = pickle.load(inp)
                model_intensities = pickle.load(inp)

            transmitted_info = dict(HKL_lookup=HKL_lookup,
                                    static_fcalcs=static_fcalcs,
                                    model_intensities=model_intensities)
        else:
            transmitted_info = None
        transmitted_info = self.mpi_helper.comm.bcast(transmitted_info, root=0)
        self.mpi_helper.comm.barrier()
        # macrocycle 1 ---------------------------------------------------------
        # generate model_intensities table based on initial conditions
        if logical_rank == 0 or self.mpi_helper.size == 1:
            FE1 = local_data.get(self.params.starting_model.preset.FE1)
            FE2 = local_data.get(self.params.starting_model.preset.FE2)

            from LS49.work2_for_aca_lsq.remake_range_intensities_with_complex \
               import get_intensity_structure
            new_model_intensities = get_intensity_structure(
                transmitted_info["static_fcalcs"],
                FE1_model=FE1,
                FE2_model=FE2)
            broadcast_info = new_model_intensities
        else:
            broadcast_info = None
        current_model_intensities = self.mpi_helper.comm.bcast(broadcast_info,
                                                               root=0)
        self.mpi_helper.comm.barrier()

        # -----------------------------------------------------------------------
        if self.mpi_helper.rank == 0:
            print("Finding initial G and abc factors")
        per_rank_items = []
        per_rank_keys = []
        per_rank_G = []
        min_spots = 3
        N_input = 0
        for item, key in get_items(logical_rank, N_total, N_stride,
                                   self.params.cohort):
            N_input += 1
            if len(item) >= min_spots:
                try:
                    FOI = fit_one_image_multispot(
                        list_of_images=item,
                        HKL_lookup=transmitted_info["HKL_lookup"],
                        model_intensities=current_model_intensities)
                except Exception as e:
                    print("FAILing fit_roi_multichannel on", e)
                    continue
                print(
                    """LLG Image %06d on %d Bragg spots NLL    channels F = %9.1f"""
                    % (key, len(item),
                       FOI.compute_functional_and_gradients()[0]))
                # put the newly refined background model back into the item
                per_rank_items.append(item)
                per_rank_keys.append(key)
                for ihkl in range(FOI.n_spots):
                    per_rank_items[-1][ihkl].bkgrd_a = flex.double([
                        FOI.a[3 * ihkl + 0], FOI.a[3 * ihkl + 1],
                        FOI.a[3 * ihkl + 2]
                    ])
                per_rank_G.append(FOI.a[-1])

        print("rank %d has %d refined images" %
              (logical_rank, len(per_rank_items)))

        N_ranks = self.mpi_helper.comm.reduce(1, self.mpi_helper.MPI.SUM, 0)
        N_refined_images = self.mpi_helper.comm.reduce(len(per_rank_items),
                                                       self.mpi_helper.MPI.SUM,
                                                       0)
        N_input_images = self.mpi_helper.comm.reduce(N_input,
                                                     self.mpi_helper.MPI.SUM,
                                                     0)
        self.mpi_helper.comm.barrier()
        if self.mpi_helper.rank == 0:
            print("final report %d ranks, %d input images, %d refined models" %
                  (N_ranks, N_input_images, N_refined_images))
            print("Finished finding initial G and abc factors")
            print("Initiating the full minimization")
        # -----------------------------------------------------------------------

        W = rank_0_fit_all_f(
            self.params,
            FE1_model=local_data.get(self.params.starting_model.preset.FE1),
            FE2_model=local_data.get(self.params.starting_model.preset.FE2))
        W.reinitialize(logical_rank,
                       self.mpi_helper.size,
                       per_rank_items,
                       per_rank_keys,
                       per_rank_G,
                       transmitted_info["HKL_lookup"],
                       transmitted_info["static_fcalcs"],
                       current_model_intensities,
                       force_recompute=True)
        W.set_macrocycle(1)
        minimizer = mpi_split_evaluator_run(
            target_evaluator=W,
            termination_params=scitbx.lbfgs.termination_parameters(
                traditional_convergence_test=True,
                traditional_convergence_test_eps=1.e-2,
                max_calls=self.params.LLG_evaluator.max_calls))
        if logical_rank == 0:
            print("Minimizer ended at iteration", W.iteration)
        self.mpi_helper.comm.barrier()
        # 2nd macrocycle---------------------------------------------------------
        # generate model_intensities table based on initial conditions
        FE1 = george_sherrell_star(fp=W.x[0:100], fdp=W.x[100:200])
        FE2 = george_sherrell_star(fp=W.x[200:300], fdp=W.x[300:400])
        if logical_rank == 0 or self.mpi_helper.size == 1:

            from LS49.work2_for_aca_lsq.remake_range_intensities_with_complex \
               import get_intensity_structure
            new_model_intensities = get_intensity_structure(
                transmitted_info["static_fcalcs"],
                FE1_model=FE1,
                FE2_model=FE2)
            broadcast_info = new_model_intensities
        else:
            broadcast_info = None
        current_model_intensities = self.mpi_helper.comm.bcast(broadcast_info,
                                                               root=0)
        self.mpi_helper.comm.barrier()

        # -----------------------------------------------------------------------
        if self.mpi_helper.rank == 0:
            print("Finding initial G and abc factors")
        per_rank_items = []
        per_rank_keys = []
        per_rank_G = []
        min_spots = 3
        N_input = 0
        for item, key in get_items(logical_rank, N_total, N_stride,
                                   self.params.cohort):
            N_input += 1
            if len(item) >= min_spots:
                try:
                    FOI = fit_one_image_multispot(
                        list_of_images=item,
                        HKL_lookup=transmitted_info["HKL_lookup"],
                        model_intensities=current_model_intensities)
                except Exception as e:
                    print("FAILing fit_roi_multichannel on", e)
                    continue
                print(
                    """LLG Image %06d on %d Bragg spots NLL    channels F = %9.1f"""
                    % (key, len(item),
                       FOI.compute_functional_and_gradients()[0]))
                # put the newly refined background model back into the item
                per_rank_items.append(item)
                per_rank_keys.append(key)
                for ihkl in range(FOI.n_spots):
                    per_rank_items[-1][ihkl].bkgrd_a = flex.double([
                        FOI.a[3 * ihkl + 0], FOI.a[3 * ihkl + 1],
                        FOI.a[3 * ihkl + 2]
                    ])
                per_rank_G.append(FOI.a[-1])

        print("rank %d has %d refined images" %
              (logical_rank, len(per_rank_items)))

        N_ranks = self.mpi_helper.comm.reduce(1, self.mpi_helper.MPI.SUM, 0)
        N_refined_images = self.mpi_helper.comm.reduce(len(per_rank_items),
                                                       self.mpi_helper.MPI.SUM,
                                                       0)
        N_input_images = self.mpi_helper.comm.reduce(N_input,
                                                     self.mpi_helper.MPI.SUM,
                                                     0)
        self.mpi_helper.comm.barrier()
        if self.mpi_helper.rank == 0:
            print("final report %d ranks, %d input images, %d refined models" %
                  (N_ranks, N_input_images, N_refined_images))
            print("Finished finding initial G and abc factors")
            print("Initiating the full minimization")
        # -----------------------------------------------------------------------
        W_previous = W
        W = rank_0_fit_all_f(self.params, FE1_model=FE1, FE2_model=FE2)
        W.reinitialize(logical_rank,
                       self.mpi_helper.size,
                       per_rank_items,
                       per_rank_keys,
                       per_rank_G,
                       transmitted_info["HKL_lookup"],
                       transmitted_info["static_fcalcs"],
                       current_model_intensities,
                       force_recompute=True)
        W.set_macrocycle(2, W_previous.starting_params_FE1,
                         W_previous.starting_params_FE2)
        minimizer = mpi_split_evaluator_run(
            target_evaluator=W,
            termination_params=scitbx.lbfgs.termination_parameters(
                traditional_convergence_test=True,
                traditional_convergence_test_eps=1.e-2,
                max_calls=self.params.LLG_evaluator.max_calls))
        if logical_rank == 0:
            print("Minimizer ended at iteration", W.iteration)
        self.mpi_helper.comm.barrier()
        # 3rd macrocycle-----------------------------------------------------------
        # generate model_intensities table based on initial conditions
        FE1 = george_sherrell_star(fp=W.x[0:100], fdp=W.x[100:200])
        FE2 = george_sherrell_star(fp=W.x[200:300], fdp=W.x[300:400])
        if logical_rank == 0 or self.mpi_helper.size == 1:

            from LS49.work2_for_aca_lsq.remake_range_intensities_with_complex \
               import get_intensity_structure
            new_model_intensities = get_intensity_structure(
                transmitted_info["static_fcalcs"],
                FE1_model=FE1,
                FE2_model=FE2)
            broadcast_info = new_model_intensities
        else:
            broadcast_info = None
        current_model_intensities = self.mpi_helper.comm.bcast(broadcast_info,
                                                               root=0)
        self.mpi_helper.comm.barrier()

        # -----------------------------------------------------------------------
        if self.mpi_helper.rank == 0:
            print("Finding initial G and abc factors")
        per_rank_items = []
        per_rank_keys = []
        per_rank_G = []
        min_spots = 3
        N_input = 0
        for item, key in get_items(logical_rank, N_total, N_stride,
                                   self.params.cohort):
            N_input += 1
            if len(item) >= min_spots:
                try:
                    FOI = fit_one_image_multispot(
                        list_of_images=item,
                        HKL_lookup=transmitted_info["HKL_lookup"],
                        model_intensities=current_model_intensities)
                except Exception as e:
                    print("FAILing fit_roi_multichannel on", e)
                    continue
                print(
                    """LLG Image %06d on %d Bragg spots NLL    channels F = %9.1f"""
                    % (key, len(item),
                       FOI.compute_functional_and_gradients()[0]))
                # put the newly refined background model back into the item
                per_rank_items.append(item)
                per_rank_keys.append(key)
                for ihkl in range(FOI.n_spots):
                    per_rank_items[-1][ihkl].bkgrd_a = flex.double([
                        FOI.a[3 * ihkl + 0], FOI.a[3 * ihkl + 1],
                        FOI.a[3 * ihkl + 2]
                    ])
                per_rank_G.append(FOI.a[-1])

        print("rank %d has %d refined images" %
              (logical_rank, len(per_rank_items)))

        N_ranks = self.mpi_helper.comm.reduce(1, self.mpi_helper.MPI.SUM, 0)
        N_refined_images = self.mpi_helper.comm.reduce(len(per_rank_items),
                                                       self.mpi_helper.MPI.SUM,
                                                       0)
        N_input_images = self.mpi_helper.comm.reduce(N_input,
                                                     self.mpi_helper.MPI.SUM,
                                                     0)
        self.mpi_helper.comm.barrier()
        if self.mpi_helper.rank == 0:
            print("final report %d ranks, %d input images, %d refined models" %
                  (N_ranks, N_input_images, N_refined_images))
            print("Finished finding initial G and abc factors")
            print("Initiating the full minimization")
        # -----------------------------------------------------------------------
        W_previous = W
        W = rank_0_fit_all_f(self.params, FE1_model=FE1, FE2_model=FE2)
        W.reinitialize(logical_rank,
                       self.mpi_helper.size,
                       per_rank_items,
                       per_rank_keys,
                       per_rank_G,
                       transmitted_info["HKL_lookup"],
                       transmitted_info["static_fcalcs"],
                       current_model_intensities,
                       force_recompute=True)
        W.set_macrocycle(3, W_previous.starting_params_FE1,
                         W_previous.starting_params_FE2)
        minimizer = mpi_split_evaluator_run(
            target_evaluator=W,
            termination_params=scitbx.lbfgs.termination_parameters(
                traditional_convergence_test=True,
                traditional_convergence_test_eps=1.e-2,
                max_calls=self.params.LLG_evaluator.max_calls))
        if logical_rank == 0:
            print("Minimizer ended at iteration", W.iteration)
        self.mpi_helper.comm.barrier()