Merge pull request #904 from HEXRD/apply-fancing-indexing-after-dark

Apply fancy indexing after dark subtraction

Author: psavery

hexrd/core/imageseries/process.py

@@ -76,34 +76,37 @@ def _process_frame(self, key):
         # note: key refers to original imageseries
         oplist = self.oplist
 
+        # Separate the frame index from any fancy indexing (row/col slices).
+        # Fancy indexing must be deferred until after all operations are
+        # applied, otherwise operations like dark subtraction will receive
+        # a spatially-indexed subarray that doesn't match the full-frame
+        # operation data (e.g. full-size dark image).
+        if isinstance(key, int):
+            idx = key
+            rest = []
+        else:
+            idx = key[0]
+            rest = key[1:]
+
         # when rectangle is the first operation we can try to call the
         # optimized version. If the adapter provides one it should be
         # significantly faster if not it will fallback to the same
         # implementation that _rectangle provides.
         if oplist and oplist[0][0] == self.RECT:
-            region = oplist[0][1]
-            if isinstance(key, int):
-                idx = key
-                rest = []
-            else:
-                # Handle fancy indexing
-                idx = key[0]
-                rest = key[1:]
-
-            img = self._rectangle_optimized(idx, region)
-
-            if rest:
-                img = img[*rest]
-
+            img = self._rectangle_optimized(idx, oplist[0][1])
             # remove the first operation since we already used it
             oplist = oplist[1:]
         else:
-            img = self._imser[key]
+            img = self._imser[idx]
 
         for k, d in oplist:
             func = self._opdict[k]
             img = func(img, d)
 
+        # Apply fancy indexing after all operations
+        if rest:
+            img = img[*rest]
+
         return img
 
     def _subtract_dark(self, img, dark):

tests/core/imageseries/test_process.py

@@ -114,6 +114,35 @@ def test_op_rectangle(mock_series):
     assert res.shape == (2, 2)
 
 
+def test_dark_before_rectangle_with_fancy_indexing():
+    """Regression: fancy indexing with dark+rectangle must apply ops first.
+
+    This reproduces the actual bug from pull_spots, which accesses pixels
+    via omega_image_series[frame_idx, row_indices, col_indices]. When dark
+    is before rectangle in the oplist (the normal GUI ordering), fancy
+    indexing was being applied to the raw image BEFORE dark subtraction,
+    producing a small subarray that couldn't broadcast with the full dark.
+    """
+    data = np.arange(16, dtype=np.float32).reshape(4, 4) + 10
+    imser = MockImageSeries(shape=(4, 4), data={0: data})
+    dark = np.full((4, 4), 2.0)
+    roi = ((1, 3), (1, 3))
+    oplist = [('dark', dark), ('rectangle', roi)]
+    ps = ProcessedImageSeries(imser, oplist)
+
+    # Access with fancy indexing like pull_spots does
+    res = ps[0, 0, 0]
+    # Frame 0 data is all 10+offset, dark is 2, rectangle crops (1:3,1:3)
+    # After dark: data - 2, then rectangle crops rows 1-2, cols 1-2
+    # So pixel [0,0] of the result = data[1,1] - 2 = 15 - 2 = 13
+    expected_full = data[1:3, 1:3] - 2.0
+    assert res == expected_full[0, 0]
+
+    # Also test with slice indexing
+    res_slice = ps[(0, slice(0, 2))]
+    np.testing.assert_array_equal(res_slice, expected_full)
+
+
 def test_op_rectangle_optimized(mock_series):
     oplist = [('rectangle', ((0, 2), (0, 2)))]
     ps = ProcessedImageSeries(mock_series, oplist)