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()
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()