From 9b2d4d23f953d344e0ca3f7a7c04a9a8f3a1e460 Mon Sep 17 00:00:00 2001
From: James Beilsten-Edmands
 <30625594+jbeilstenedmands@users.noreply.github.com>
Date: Mon, 23 Oct 2023 15:01:10 +0100
Subject: [PATCH] Propagate unindexed as separate experiment

---
 src/dials/algorithms/refinement/target.py |  2 +-
 src/dials/command_line/index.py           | 15 ++++++++++-
 src/dials/command_line/refine.py          | 33 ++++++++++++++++++++---
 3 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/src/dials/algorithms/refinement/target.py b/src/dials/algorithms/refinement/target.py
index 1c8ab1f28f..40f7e18b17 100644
--- a/src/dials/algorithms/refinement/target.py
+++ b/src/dials/algorithms/refinement/target.py
@@ -525,7 +525,7 @@ def rmsds_for_experiment(self, exp):
         """calculate unweighted RMSDs for the selected experiment."""
 
         self.update_matches()
-        sel = self._matches.get_selection_for_experiment_identifier(exp)
+        sel = self._matches.get_selection_for_experiment_identifier(exp.identifier)
         n = sel.count(True)
         if n == 0:
             return None
diff --git a/src/dials/command_line/index.py b/src/dials/command_line/index.py
index f462746c8a..780e14baa4 100644
--- a/src/dials/command_line/index.py
+++ b/src/dials/command_line/index.py
@@ -109,6 +109,7 @@ def _index_experiments(
 ):
     if log_text:
         logger.info(log_text)
+    reflections["original_id"] = copy.deepcopy(reflections["id"])
     idxr = indexer.Indexer.from_parameters(
         reflections,
         experiments,
@@ -117,8 +118,20 @@ def _index_experiments(
     )
     idxr.index()
     idx_refl = copy.deepcopy(idxr.refined_reflections)
+    n_unindex = len(experiments)
+    idxr.unindexed_reflections["id"] = idxr.unindexed_reflections["original_id"]
+    for i, expt in enumerate(experiments):
+        idxr.unindexed_reflections.experiment_identifiers()[i] = expt.identifier
+    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_ + n_unindex] = identifier
+
+    idx_refl["id"] += n_unindex
     idx_refl.extend(idxr.unindexed_reflections)
-    return idxr.refined_experiments, idx_refl
+    del idx_refl["original_id"]
+    experiments.extend(idxr.refined_experiments)
+    return experiments, idx_refl
 
 
 def index(experiments, reflections, params):
diff --git a/src/dials/command_line/refine.py b/src/dials/command_line/refine.py
index 4db987be05..2d9c3efadc 100644
--- a/src/dials/command_line/refine.py
+++ b/src/dials/command_line/refine.py
@@ -508,7 +508,12 @@ def run_dials_refine(experiments, reflections, params):
         table_rows = {}
         for (el, refl, refiner, _), ids in zip(refinement_results, disjoint_sets):
             header, rows = refiner.calc_exp_rmsd_table()
+            for new_id, orig_id in enumerate(ids):
+                experiments[orig_id] = el[new_id]
+                rows[new_id][0] = str(orig_id)
+                table_rows[orig_id] = rows[new_id]
             table_headers.append(header)
+
         reflections = flex.reflection_table.concat([i[1] for i in refinement_results])
 
         experiments = ExperimentList([experiments[i] for i in range(len(experiments))])
@@ -603,13 +608,35 @@ def run(args=None, phil=working_phil):
         logger.info("The following parameters have been modified:\n")
         logger.info(diff_phil)
 
+    if not all(experiments.identifiers()):
+        from dials.util.multi_dataset_handling import generate_experiment_identifiers
+
+        generate_experiment_identifiers(experiments)
+        for i, expt in enumerate(experiments):
+            reflections.experiment_identifiers()[i] = expt.identifier
+
+    experiments_to_refine = ExperimentList([e for e in experiments if e.crystal])
+    experiments = ExperimentList([e for e in experiments if not e.crystal])
+    reflections_to_refine = reflections.select_on_experiment_identifiers(
+        experiments_to_refine.identifiers()
+    )
+    reflections_to_refine.reset_ids()
+    reflections = reflections.select_on_experiment_identifiers(
+        experiments.identifiers()
+    )
+
     # Run refinement
     try:
-        experiments, reflections, refiner, history = run_dials_refine(
-            experiments, reflections, params
-        )
+        (
+            experiments_to_refine,
+            reflections_to_refine,
+            refiner,
+            history,
+        ) = run_dials_refine(experiments_to_refine, reflections_to_refine, params)
     except (DialsRefineConfigError, DialsRefineRuntimeError) as e:
         sys.exit(str(e))
+    experiments.extend(experiments_to_refine)
+    reflections = flex.reflection_table.concat([reflections, reflections_to_refine])
 
     # For the usual case of refinement of one crystal, print that model for information
     crystals = experiments.crystals()