def setup_logging_from_params(params):
"""params: PHIL params, see simtbx.diffBragg.hopper phil string"""
if params.logging.logfiles:
if COMM.rank == 0:
utils.safe_makedirs(params.outdir)
COMM.barrier()
main_level = LEVELS[params.logging.logfiles_level]
main_logger = _make_logger("diffBragg.main",
os.path.join(params.outdir, HOST+"-"+params.logging.logname),
level=main_level,
overwrite=params.logging.overwrite,
formatter=logging.Formatter(DETAILED_FORMAT))
_make_logger("diffBragg.profile",
os.path.join(params.outdir, HOST+"-"+params.profile_name),
level=logging.INFO,
overwrite=params.logging.overwrite,
formatter=logging.Formatter(SIMPLE_FORMAT))
# for convenience we add a console logger for rank0 so we can optionally see output
# even when logging to files
if COMM.rank == 0:
console = logging.StreamHandler()
console.setFormatter(logging.Formatter(SIMPLE_FORMAT))
console.setLevel(LEVELS[params.logging.rank0_level])
main_logger.addHandler(console)
else:
if COMM.rank == 0:
level = LEVELS[params.logging.rank0_level]
else:
level = LEVELS[params.logging.other_ranks_level]
_make_logger("diffBragg.main", level=level, formatter=logging.Formatter(SIMPLE_FORMAT))
_make_logger("diffBragg.profile", level=level, formatter=logging.Formatter(SIMPLE_FORMAT))
def __init__(self):
from dials.util.options import OptionParser
self.params = self.parser = None
if COMM.rank == 0:
self.parser = OptionParser(
usage="", # stage 1 (per-shot) diffBragg refinement",
sort_options=True,
phil=phil_scope,
read_experiments=True,
read_reflections=True,
check_format=False,
epilog="PyCuties")
self.params, _ = self.parser.parse_args(show_diff_phil=True)
assert self.params.outdir is not None
self.params = COMM.bcast(self.params)
if COMM.rank == 0:
if not os.path.exists(self.params.outdir):
utils.safe_makedirs(self.params.outdir)
COMM.barrier()
if self.params.logging.logname is None:
self.params.logging.logname = "main_stage1.log"
if self.params.profile_name is None:
self.params.profile_name = "prof_stage1.log"
from simtbx.diffBragg import mpi_logger
mpi_logger.setup_logging_from_params(self.params)
def __init__(self):
from dials.util.options import ArgumentParser
self.params = None
if COMM.rank == 0:
self.parser = ArgumentParser(
usage="", # stage 1 (per-shot) diffBragg refinement",
sort_options=True,
phil=phil_scope,
read_experiments=True,
read_reflections=True,
check_format=False,
epilog="PyCuties")
self.params, _ = self.parser.parse_args(show_diff_phil=True)
assert self.params.outdir is not None
utils.safe_makedirs(self.params.outdir)
ts = time.strftime("%Y%m%d-%H%M%S")
diff_phil_outname = os.path.join(self.params.outdir,
"diff_phil_run_at_%s.txt" % ts)
with open(diff_phil_outname, "w") as o:
o.write("command line:\n%s\n" % (" ".join(sys.argv)))
o.write("workding directory: \n%s\n" % os.getcwd())
o.write("diff phil:\n")
o.write(self.parser.diff_phil.as_str())
self.params = COMM.bcast(self.params)
if self.params.logging.logname is None:
self.params.logging.logname = "main_stage1.log"
if self.params.profile_name is None:
self.params.profile_name = "prof_stage1.log"
from simtbx.diffBragg import mpi_logger
mpi_logger.setup_logging_from_params(self.params)
def run(self):
MAIN_LOGGER = logging.getLogger("diffBragg.main")
assert os.path.exists(self.params.exp_ref_spec_file)
input_lines = None
best_models = None
if COMM.rank == 0:
input_lines = open(self.params.exp_ref_spec_file, "r").readlines()
if self.params.skip is not None:
input_lines = input_lines[self.params.skip:]
if self.params.first_n is not None:
input_lines = input_lines[:self.params.first_n]
if self.params.sanity_test_input:
hopper_utils.sanity_test_input_lines(input_lines)
if self.params.best_pickle is not None:
logging.info("reading pickle %s" % self.params.best_pickle)
best_models = pandas.read_pickle(self.params.best_pickle)
if self.params.dump_gathers:
if self.params.gathers_dir is None:
raise ValueError("Need to provide a file dir path in order to dump_gathers")
utils.safe_makedirs(self.params.gathers_dir)
input_lines = COMM.bcast(input_lines)
best_models = COMM.bcast(best_models)
if self.params.ignore_existing:
exp_names_already =None
if COMM.rank==0:
exp_names_already = {os.path.basename(f) for f in glob.glob("%s/expers/rank*/*.expt" % self.params.outdir)}
exp_names_already = COMM.bcast(exp_names_already)
exp_gatheredRef_spec = [] # optional list of expt, refls, spectra
for i_exp, line in enumerate(input_lines):
if i_exp == self.params.max_process:
break
if i_exp % COMM.size != COMM.rank:
continue
logging.info("COMM.rank %d on shot %d / %d" % (COMM.rank, i_exp + 1, len(input_lines)))
line_fields = line.strip().split()
assert len(line_fields) in [2, 3]
if len(line_fields) == 2:
exp, ref = line_fields
spec = None
else:
exp, ref, spec = line_fields
if self.params.ignore_existing:
basename = os.path.splitext(os.path.basename(exp))[0]
opt_exp = "%s_%s_%d.expt" % (self.params.tag, basename, i_exp)
if opt_exp in exp_names_already:
continue
best = None
if best_models is not None:
best = best_models.query("exp_name=='%s'" % exp)
if len(best) == 0:
best = best_models.query("opt_exp_name=='%s'" % exp)
if len(best) != 1:
raise ValueError("Should be 1 entry for exp %s in best pickle %s" % (exp, self.params.best_pickle))
self.params.simulator.spectrum.filename = spec
MAIN_LOGGER.info("Modeling %s" % exp)
Modeler = hopper_utils.DataModeler(self.params)
if self.params.load_data_from_refls:
gathered = Modeler.GatherFromReflectionTable(exp, ref)
else:
gathered = Modeler.GatherFromExperiment(exp, ref)
if not gathered:
logging.warning("No refls in %s; CONTINUE; COMM.rank=%d" % (ref, COMM.rank))
continue
if self.params.dump_gathers:
output_name = os.path.splitext(os.path.basename(exp))[0]
output_name += "_withData.refl"
output_name = os.path.join(self.params.gathers_dir, output_name)
Modeler.dump_gathered_to_refl(output_name, do_xyobs_sanity_check=True) # NOTE do this is modelin strong spots only
if self.params.test_gathered_file:
all_data = Modeler.all_data.copy()
all_roi_id = Modeler.roi_id.copy()
all_bg = Modeler.all_background.copy()
all_trusted = Modeler.all_trusted.copy()
all_pids = np.array(Modeler.pids)
all_rois = np.array(Modeler.rois)
new_Modeler = hopper_utils.DataModeler(self.params)
assert new_Modeler.GatherFromReflectionTable(exp, output_name)
assert np.allclose(new_Modeler.all_data, all_data)
assert np.allclose(new_Modeler.all_background, all_bg)
assert np.allclose(new_Modeler.rois, all_rois)
assert np.allclose(new_Modeler.pids, all_pids)
assert np.allclose(new_Modeler.all_trusted, all_trusted)
assert np.allclose(new_Modeler.roi_id, all_roi_id)
exp_gatheredRef_spec.append((exp, os.path.abspath(output_name), spec))
if self.params.only_dump_gathers:
continue
if self.params.refiner.reference_geom is not None:
detector = ExperimentListFactory.from_json_file(self.params.refiner.reference_geom, check_format=False)[0].detector
Modeler.E.detector = detector
Modeler.SimulatorFromExperiment(best)
Modeler.SIM.D.store_ave_wavelength_image = True
if self.params.refiner.verbose is not None and COMM.rank==0:
Modeler.SIM.D.verbose = self.params.refiner.verbose
if self.params.profile:
Modeler.SIM.record_timings = True
if self.params.use_float32:
Modeler.all_data = Modeler.all_data.astype(np.float32)
Modeler.all_background = Modeler.all_background.astype(np.float32)
if self.params.refiner.randomize_devices:
dev = np.random.choice(self.params.refiner.num_devices)
logging.info("Rank %d will use randomly chosen device %d on host %s" % (COMM.rank, dev, socket.gethostname()))
else:
dev = COMM.rank % self.params.refiner.num_devices
logging.info("Rank %d will use device %d on host %s" % (COMM.rank, dev, socket.gethostname()))
Modeler.SIM.D.device_Id = dev
nparam = len(Modeler.SIM.P)
x0 = [1] * nparam
x = Modeler.Minimize(x0)
if self.params.profile:
Modeler.SIM.D.show_timings(COMM.rank) #, out)
save_up(Modeler, x, exp, i_exp, ref)
if self.params.dump_gathers and self.params.gathered_output_file is not None:
exp_gatheredRef_spec = COMM.reduce(exp_gatheredRef_spec)
if COMM.rank == 0:
o = open(self.params.gathered_output_file, "w")
for e, r, s in exp_gatheredRef_spec:
if s is not None:
o.write("%s %s %s\n" % (e,r,s))
else:
o.write("%s %s\n" % (e,r))
o.close()