From 63e8028e3cebe1037eb23d8d574a811023e1f540 Mon Sep 17 00:00:00 2001
From: davidmcdonagh <davidnmcdonagh@gmail.com>
Date: Tue, 19 Mar 2024 12:15:08 +0000
Subject: [PATCH] Enabled find_max_cell to work with tof data.

---
 src/dials/algorithms/indexing/indexer.py      | 20 +++++++++++++++++++
 src/dials/algorithms/indexing/max_cell.py     |  2 ++
 .../algorithms/indexing/nearest_neighbor.py   | 18 ++++++++++-------
 3 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/src/dials/algorithms/indexing/indexer.py b/src/dials/algorithms/indexing/indexer.py
index 3b3c1c4bb0..30463afc6e 100644
--- a/src/dials/algorithms/indexing/indexer.py
+++ b/src/dials/algorithms/indexing/indexer.py
@@ -941,6 +941,25 @@ def find_max_cell(self):
                 self.params.max_cell = params.multiplier * max(uc_params[:3])
                 logger.info("Using max_cell: %.1f Angstrom", self.params.max_cell)
             else:
+
+                convert_reflections_z_to_deg = True
+                all_tof_experiments = False
+                for expt in self.experiments:
+                    if expt.scan is not None and expt.scan.has_property(
+                        "time_of_flight"
+                    ):
+                        all_tof_experiments = True
+                    elif all_tof_experiments:
+                        raise ValueError(
+                            "Cannot find max cell for ToF and non-ToF experiments at the same time"
+                        )
+
+                if all_tof_experiments:
+                    if params.step_size < 100:
+                        logger.info("Setting default ToF step size to 500 usec")
+                        params.step_size = 500
+                        convert_reflections_z_to_deg = False
+
                 self.params.max_cell = find_max_cell(
                     self.reflections,
                     max_cell_multiplier=params.multiplier,
@@ -952,6 +971,7 @@ def find_max_cell(self):
                     filter_ice=params.filter_ice,
                     filter_overlaps=params.filter_overlaps,
                     overlaps_border=params.overlaps_border,
+                    convert_reflections_z_to_deg=convert_reflections_z_to_deg,
                 ).max_cell
                 logger.info("Found max_cell: %.1f Angstrom", self.params.max_cell)
 
diff --git a/src/dials/algorithms/indexing/max_cell.py b/src/dials/algorithms/indexing/max_cell.py
index 7c45cc6d5d..dccb542abf 100644
--- a/src/dials/algorithms/indexing/max_cell.py
+++ b/src/dials/algorithms/indexing/max_cell.py
@@ -22,6 +22,7 @@ def find_max_cell(
     filter_ice=True,
     filter_overlaps=True,
     overlaps_border=0,
+    convert_reflections_z_to_deg=True,
 ):
     logger.debug("Finding suitable max_cell based on %i reflections", len(reflections))
     # Exclude potential ice-ring spots from nearest neighbour analysis if needed
@@ -63,6 +64,7 @@ def find_max_cell(
             percentile=nearest_neighbor_percentile,
             histogram_binning=histogram_binning,
             nn_per_bin=nn_per_bin,
+            convert_reflections_z_to_deg=convert_reflections_z_to_deg,
         )
     except AssertionError as e:
         raise DialsIndexError("Failure in nearest neighbour analysis:\n" + str(e))
diff --git a/src/dials/algorithms/indexing/nearest_neighbor.py b/src/dials/algorithms/indexing/nearest_neighbor.py
index 4d5de82d88..2e5c0af2e1 100644
--- a/src/dials/algorithms/indexing/nearest_neighbor.py
+++ b/src/dials/algorithms/indexing/nearest_neighbor.py
@@ -13,6 +13,7 @@ def __init__(
         percentile=None,
         histogram_binning="linear",
         nn_per_bin=5,
+        convert_reflections_z_to_deg=True,
     ):
         self.tolerance = tolerance  # Margin of error for max unit cell estimate
         from scitbx.array_family import flex
@@ -28,7 +29,10 @@ def __init__(
         else:
             entering_flags = flex.bool(reflections.size(), True)
         rs_vectors = reflections["rlp"]
-        phi_deg = reflections["xyzobs.mm.value"].parts()[2] * (180 / math.pi)
+
+        z = reflections["xyzobs.mm.value"].parts()[2]
+        if convert_reflections_z_to_deg:
+            z = z * (180 / math.pi)
 
         d_spacings = flex.double()
         # nearest neighbor analysis
@@ -38,16 +42,16 @@ def __init__(
             sel_imageset = reflections["imageset_id"] == imageset_id
             if sel_imageset.count(True) == 0:
                 continue
-            phi_min = flex.min(phi_deg.select(sel_imageset))
-            phi_max = flex.max(phi_deg.select(sel_imageset))
-            d_phi = phi_max - phi_min
-            n_steps = max(int(math.ceil(d_phi / step_size)), 1)
+            z_min = flex.min(z.select(sel_imageset))
+            z_max = flex.max(z.select(sel_imageset))
+            d_z = z_max - z_min
+            n_steps = max(int(math.ceil(d_z / step_size)), 1)
 
             for n in range(n_steps):
                 sel_step = (
                     sel_imageset
-                    & (phi_deg >= (phi_min + n * step_size))
-                    & (phi_deg < (phi_min + (n + 1) * step_size))
+                    & (z >= (z_min + n * step_size))
+                    & (z < (z_min + (n + 1) * step_size))
                 )
 
                 for entering in (True, False):