From d3f59b8eba5bfa480e07d3cdc913662230d730c7 Mon Sep 17 00:00:00 2001 From: James Beilsten-Edmands <30625594+jbeilstenedmands@users.noreply.github.com> Date: Fri, 1 Dec 2023 15:58:25 +0000 Subject: [PATCH] Output unindexed experiments in index, add expeditor to handle experiment list without xtals --- src/dials/command_line/index.py | 226 +++++++++++++----- src/dials/command_line/integrate.py | 7 + src/dials/command_line/refine.py | 7 + .../command_line/refine_bravais_settings.py | 9 + src/dials/util/multi_dataset_handling.py | 79 ++++++ 5 files changed, 266 insertions(+), 62 deletions(-) diff --git a/src/dials/command_line/index.py b/src/dials/command_line/index.py index ded6254d65e..4097420da0a 100644 --- a/src/dials/command_line/index.py +++ b/src/dials/command_line/index.py @@ -81,6 +81,12 @@ .type = path reflections = indexed.refl .type = path + retain_unindexed_experiment = False + .type = bool + .help = "If True, the input experiment models are extended with new crystal" + "models, thus the output contains a crystal-less experiment for each" + "imageset that matches the unindexed reflections. Setting this option" + "to False recovers the old behaviour of versions DIALS 3.17 and earlier." log = dials.index.log .type = str } @@ -103,6 +109,152 @@ working_phil = phil_scope.fetch(sources=[phil_overrides]) +from dials.util.multi_dataset_handling import generate_experiment_identifiers + + +def _index_single_imageset(experiments, reflections, params, log_text=None): + if log_text: + logger.info(log_text) + + unindexed = [i for i in experiments if not i.crystal] + # handle legacy cases where no unindexed experiments: + if not unindexed: # i.e. all have crystals + unindexed = [ + Experiment( + crystal=None, + beam=experiments[0].beam, + detector=experiments[0].detector, + goniometer=experiments[0].goniometer, + scan=experiments[0].scan, + imageset=experiments[0].imageset, + profile=experiments[0].profile, + scaling_model=experiments[0].scaling_model, + ) + ] + generate_experiment_identifiers(unindexed) + input_expts = experiments + known_crystal_models = [expt.crystal for expt in experiments] + else: + if any(experiments.crystals()): + input_expts = ExperimentList([i for i in experiments if i.crystal]) + + known_crystal_models = [expt.crystal for expt in input_expts] + else: + input_expts = experiments + known_crystal_models = None + + idxr = indexer.Indexer.from_parameters( + reflections, + input_expts, + known_crystal_models=known_crystal_models, + params=params, + ) + idxr.index() + if not params.output.retain_unindexed_experiment: + idx_refl = copy.deepcopy(idxr.refined_reflections) + idx_refl.extend(idxr.unindexed_reflections) + return idxr.refined_experiments, idx_refl + + # update the identifiers so that the unindexed has id 0, and the rest 1,2,3.. etc + idxr.unindexed_reflections["id"] = flex.int(idxr.unindexed_reflections.size(), 0) + idxr.unindexed_reflections.experiment_identifiers()[0] = unindexed[0].identifier + idxr.unindexed_reflections.clean_experiment_identifiers_map() + + idx_refl = idxr.refined_reflections + for id_ in sorted(set(idx_refl["id"]), reverse=True): + identifier = idx_refl.experiment_identifiers()[id_] + del idx_refl.experiment_identifiers()[id_] + idx_refl.experiment_identifiers()[id_ + 1] = identifier + idx_refl["id"] += 1 + + # the refiner copies the input beam, detector and gonio, we want to share these + unindexed[0].detector = idxr.refined_experiments[0].detector + unindexed[0].beam = idxr.refined_experiments[0].beam + if idxr.refined_experiments[ + 0 + ].goniometer: # might be using the stills indexer which deletes the gonio + unindexed[0].goniometer = idxr.refined_experiments[0].goniometer + + # Now join everything together for output + indexed_experiments = ExperimentList(unindexed) + indexed_experiments.extend(idxr.refined_experiments) + idx_refl.extend(idxr.unindexed_reflections) + idx_refl.assert_experiment_identifiers_are_consistent() + return indexed_experiments, idx_refl + + +def _index_joint_indexing(experiments, reflections, params): + + # first make an unindexed experiments for each imageset. + unindexed = [i for i in experiments if not i.crystal] + original_imagesets = experiments.imagesets() + # handle legacy cases where no unindexed experiments: + if not unindexed: # i.e. all have crystals + unindexed = [] + for iset in experiments.imagesets(): + i_expt = experiments.where(imageset=iset) + expt_to_copy = experiments[i_expt[0]] + unindexed.append( + Experiment( + crystal=None, + beam=expt_to_copy.beam, + detector=expt_to_copy.detector, + goniometer=expt_to_copy.goniometer, + scan=expt_to_copy.scan, + imageset=expt_to_copy.imageset, + profile=expt_to_copy.profile, + scaling_model=expt_to_copy.scaling_model, + ) + ) + generate_experiment_identifiers(unindexed) + input_expts = experiments + known_crystal_models = [expt.crystal for expt in experiments] + else: + if any(experiments.crystals()): + input_expts = ExperimentList([i for i in experiments if i.crystal]) + # note not just experiments.crystals(), as models may be shared. + known_crystal_models = [expt.crystal for expt in input_expts] + else: + input_expts = experiments + known_crystal_models = None + + for id_ in sorted(set(reflections["id"]), reverse=True): + del reflections.experiment_identifiers()[id_] + reflections["original_id"] = copy.deepcopy(reflections["imageset_id"]) + reflections["id"] = copy.deepcopy(reflections["imageset_id"]) + idxr = indexer.Indexer.from_parameters( + reflections, + input_expts, + known_crystal_models=known_crystal_models, + params=params, + ) + idxr.index() + if not params.output.retain_unindexed_experiment: + idx_refl = copy.deepcopy(idxr.refined_reflections) + idx_refl.extend(idxr.unindexed_reflections) + return idxr.refined_experiments, idx_refl + + # now want to split up so that the output is in imageset order + indexed_experiments = ExperimentList([]) + indexed_reflections = flex.reflection_table() + for i, iset in enumerate(original_imagesets): + # first sort out the unindexed + identifier = unindexed[i].identifier + sel = idxr.unindexed_reflections["original_id"] == i + unindexed_refl = idxr.unindexed_reflections.select(sel) + unindexed_refl["id"] = flex.int(unindexed_refl.size(), i) + del unindexed_refl["original_id"] + unindexed_refl.experiment_identifiers()[i] == identifier + indexed_reflections.extend(unindexed_refl) + # now get the indexed + i_expts = idxr.refined_experiments.where(imageset=iset) + identifiers = [idxr.refined_experiments[i].identifier for i in i_expts] + refls = idxr.refined_reflections.select_on_experiment_identifiers(identifiers) + + pass + # FIXME + return indexed_experiments, refls + def _index_experiments( experiments, @@ -111,8 +263,6 @@ def _index_experiments( known_crystal_models=None, log_text=None, ): - # general input of experiments: - # [no-xtal1, xtal1, xtal2, no-xtal2, xtal3,xtal4,..] if log_text: logger.info(log_text) idxr = indexer.Indexer.from_parameters( @@ -147,11 +297,6 @@ def index(experiments, reflections, params): combination of sequence and stills data. dials.algorithms.indexing.DialsIndexError: Indexing failed. """ - if experiments.crystals()[0] is not None: - # note not just experiments.crystals(), as models may be shared. - known_crystal_models = [expt.crystal for expt in experiments] - else: - known_crystal_models = None if len(reflections) == 0: raise ValueError("No reflection lists found in input") @@ -169,53 +314,18 @@ def index(experiments, reflections, params): if params.indexing.image_range: reflections = slice_reflections(reflections, params.indexing.image_range) - # if we only have one imageset, or joint_indexing - if len(experiments.imagesets()) == 1 or params.indexing.joint_indexing: - print("here") - # we want to retain an unindexed experiment per imageset. - unindexed = [i for i in experiments if not i.crystal] - # handle legacy cases where no unindexed experiments: - if not unindexed: # i.e. all have crstals - unindexed = [] - for iset in experiments.imagesets(): - i_expt = experiments.where(imageset=iset) - expt_to_copy = experiments[i_expt[0]] - unindexed.append( - Experiment( - identifier=expt_to_copy.identifier, - crystal=None, - beam=expt_to_copy.beam, - detector=expt_to_copy.detector, - goniometer=expt_to_copy.goniometer, - scan=expt_to_copy.scan, - imageset=expt_to_copy.imageset, - profile=expt_to_copy.profile, - scaling_model=expt_to_copy.scaling_model, - ) - ) - input_expts = experiments - known_crystal_models = [expt.crystal for expt in experiments] - else: - input_expts = ExperimentList([i for i in experiments if i.crystal]) - known_crystal_models = [expt.crystal for expt in input_expts] - - # reflections["id"] = copy.deepcopy(reflections["imageset_id"]) - indexed_experiments, indexed_reflections = _index_experiments( - input_expts, + if len(experiments.imagesets()) == 1: + indexed_experiments, indexed_reflections = _index_single_imageset( + experiments, + reflections, + copy.deepcopy(params), + ) + elif params.indexing.joint_indexing: + indexed_experiments, indexed_reflections = _index_joint_indexing( + experiments, reflections, copy.deepcopy(params), - known_crystal_models=known_crystal_models, ) - unindexed[0].beam = indexed_experiments[0].beam - unindexed[0].detector = indexed_experiments[0].detector - unindexed[0].goniometer = indexed_experiments[0].goniometer - - print(id(unindexed[0].beam)) - print(unindexed[0].identifier) - print(list(indexed_experiments.identifiers())) - print(id(indexed_experiments[0].beam)) - # unindexed.extend(indexed_experiments) - # indexed_experiments = unindexed else: indexed_experiments = ExperimentList() @@ -228,25 +338,16 @@ def index(experiments, reflections, params): refl = reflections.select(reflections["imageset_id"] == iset_id) i_expts = experiments.where(imageset=imgset) elist = ExperimentList([experiments[i] for i in i_expts]) - known_crystal_models_this = None - """if any(elist.crystals()): - known_crystal_models_this = e - else:""" - - if known_crystal_models: - known_crystal_models_this = [ - known_crystal_models[i] for i in i_expts - ] + assert len(elist.imagesets()) == 1 refl["imageset_id"] = flex.int( len(refl), 0 ) # _index_experiments functions requires ids numbered from 0 futures[ pool.submit( - _index_experiments, + _index_single_imageset, elist, refl, copy.deepcopy(params), - known_crystal_models=known_crystal_models_this, log_text=f"Indexing imageset id {iset_id} ({iset_id + 1}/{len(experiments.imagesets())})", ) ] = iset_id @@ -256,6 +357,7 @@ def index(experiments, reflections, params): try: iset_id = futures[future] idx_expts, idx_refl = future.result() + # the result will contain an unindexed experiment followed by up to N=max_lattices experiments with crystals except Exception as e: print(e) else: diff --git a/src/dials/command_line/integrate.py b/src/dials/command_line/integrate.py index 70e12abaf88..5e3c9430389 100644 --- a/src/dials/command_line/integrate.py +++ b/src/dials/command_line/integrate.py @@ -747,6 +747,10 @@ def run(args=None, phil=working_phil): if reference and "shoebox" not in reference: sys.exit("Error: shoebox data missing from reflection table") + from dials.util.multi_dataset_handling import Expeditor + + expeditor = Expeditor(experiments, reference) + experiments, reference = expeditor.filter_experiments_with_crystals() try: experiments, reflections, report = run_integration( params, experiments, reference @@ -761,6 +765,9 @@ def run(args=None, phil=working_phil): logger.info( "Saving %d reflections to %s", reflections.size(), params.output.reflections ) + experiments, reflections = expeditor.combine_experiments_for_output( + experiments, reflections + ) reflections.as_file(params.output.reflections) logger.info("Saving the experiments to %s", params.output.experiments) experiments.as_file(params.output.experiments) diff --git a/src/dials/command_line/refine.py b/src/dials/command_line/refine.py index cbf48ab795a..ad2b7493766 100644 --- a/src/dials/command_line/refine.py +++ b/src/dials/command_line/refine.py @@ -35,6 +35,7 @@ ) from dials.algorithms.refinement.corrgram import create_correlation_plots from dials.array_family import flex +from dials.util.multi_dataset_handling import Expeditor from dials.util.options import ArgumentParser, reflections_and_experiments_from_files from dials.util.version import dials_version @@ -608,12 +609,18 @@ def run(args=None, phil=working_phil): logger.info(diff_phil) # Run refinement + expeditor = Expeditor(experiments, reflections) + experiments, reflections = expeditor.filter_experiments_with_crystals() try: experiments, reflections, refiner, history = run_dials_refine( experiments, reflections, params ) except (DialsRefineConfigError, DialsRefineRuntimeError) as e: sys.exit(str(e)) + else: + experiments, reflections = expeditor.combine_experiments_for_output( + experiments, reflections + ) # For the usual case of refinement of one crystal, print that model for information crystals = experiments.crystals() diff --git a/src/dials/command_line/refine_bravais_settings.py b/src/dials/command_line/refine_bravais_settings.py index 7f613da21e9..6b9db1b2be7 100644 --- a/src/dials/command_line/refine_bravais_settings.py +++ b/src/dials/command_line/refine_bravais_settings.py @@ -49,6 +49,7 @@ ) from dials.array_family import flex from dials.util import log +from dials.util.multi_dataset_handling import Expeditor from dials.util.options import ArgumentParser, reflections_and_experiments_from_files from dials.util.version import dials_version @@ -200,6 +201,9 @@ def run(args=None): if len(experiments) == 0: parser.print_help() return + expeditor = Expeditor(experiments, reflections) + experiments, reflections = expeditor.filter_experiments_with_crystals() + if len(experiments.crystals()) > 1: if params.crystal_id is not None: experiments, reflections = select_datasets_on_crystal_id( @@ -237,6 +241,11 @@ def run(args=None): soln = int(subgroup.setting_number) bs_json = "%sbravais_setting_%i.expt" % (prefix, soln) logger.info("Saving solution %i as %s", soln, bs_json) + + expts = expeditor.generate_experiments_with_updated_crystal( + expts, params.crystal_id + ) + expts.as_file(os.path.join(params.output.directory, bs_json)) diff --git a/src/dials/util/multi_dataset_handling.py b/src/dials/util/multi_dataset_handling.py index f14823701eb..ccf3163534e 100644 --- a/src/dials/util/multi_dataset_handling.py +++ b/src/dials/util/multi_dataset_handling.py @@ -11,6 +11,7 @@ from orderedset import OrderedSet +from dxtbx.model import ExperimentList from dxtbx.util import ersatz_uuid4 from dials.array_family import flex @@ -41,6 +42,84 @@ ) +class Expeditor(object): + def __init__(self, experiments, reflection_table): + self.experiments = experiments + self.reflection_table = reflection_table + self.crystal_locs = [ + i for i, expt in enumerate(self.experiments) if expt.crystal + ] + self.input_has_crystalless_expts = bool( + len(self.crystal_locs) != len(self.experiments) + ) + self.crystalless_reflection_tables = [] + + def filter_experiments_with_crystals(self): + if not self.input_has_crystalless_expts: + return self.experiments, self.reflection_table + + expts_with_crystals = ExperimentList( + [expt for expt in self.experiments if expt.crystal] + ) + reflection_table = self.reflection_table.select_on_experiment_identifiers( + expts_with_crystals.identifiers() + ) + reflection_table.reset_ids() + + # record some quantities to help later when recombining + self.crystalless_reflection_tables = [ + self.reflection_table.select_on_experiment_identifiers([expt.identifier]) + if not expt.crystal + else None + for expt in self.experiments + ] + return expts_with_crystals, reflection_table + + def combine_experiments_for_output(self, experiments, reflection_table): + if not self.input_has_crystalless_expts: + return experiments, reflection_table + + if -1 in set(reflection_table["id"]): + # need to match it up to the right imageset + reflection_table = reflection_table.select(reflection_table["id"] >= 0) + + # if we have a model for a gonio, detector or beam which belongs to a scan, update + # all experiments that share that scan + tables = reflection_table.split_by_experiment_id() + for i, expt, table in zip(self.crystal_locs, experiments, tables): + other_expts_sharing_scan = self.experiments.where(scan=expt.scan) + for j in other_expts_sharing_scan: + self.experiments[j].beam = expt.beam + self.experiments[j].detector = expt.detector + self.experiments[j].goniometer = expt.goniometer + self.experiments[i] = expt + self.crystalless_reflection_tables[i] = table + reflections = flex.reflection_table.concat(self.crystalless_reflection_tables) + reflections.assert_experiment_identifiers_are_consistent(self.experiments) + return self.experiments, reflections + + def generate_experiments_with_updated_crystal(self, experiments, crystal_id): + # special function to handle the awkward case of dials.rbs + crystals = [ + expt.crystal for expt in self.experiments if expt.crystal is not None + ] + crystal = crystals[crystal_id] + expts_ids = self.experiments.where(crystal=crystal) + # make a copy + elist = copy.copy(self.experiments) + + for id_, new_expt in zip(expts_ids, experiments): + elist[id_].crystal = new_expt.crystal + # also copy across beam, detector and gonio for anything sharing those models + scan = elist[id_].scan + j_expt = elist.where(scan=scan) + for j in j_expt: + elist[j].beam = new_expt.beam + elist[j].goniometer = new_expt.goniometer + elist[j].detector = new_expt.detector + return elist + + def generate_experiment_identifiers(experiments, identifier_type="uuid"): """Generate unique identifiers for each experiment.""" if identifier_type == "uuid":