Optimize spectra merge and normalization

src/esssans/conversions.py

@@ -204,7 +204,7 @@ def monitor_to_wavelength(
     monitor: CalibratedMonitor[RunType, MonitorType], graph: MonitorCoordTransformGraph
 ) -> WavelengthMonitor[RunType, MonitorType]:
     return WavelengthMonitor[RunType, MonitorType](
-        monitor.transform_coords('wavelength', graph=graph)
+        monitor.transform_coords('wavelength', graph=graph, keep_inputs=False)
     )
 
 
@@ -229,7 +229,7 @@ def detector_to_wavelength(
     graph: ElasticCoordTransformGraph,
 ) -> CleanWavelength[ScatteringRunType, Numerator]:
     return CleanWavelength[ScatteringRunType, Numerator](
-        detector.transform_coords('wavelength', graph=graph)
+        detector.transform_coords('wavelength', graph=graph, keep_inputs=False)
     )
 
 

src/esssans/i_of_q.py

@@ -171,7 +171,14 @@ def merge_spectra(
 
     if data.bins is not None:
         q_all_pixels = data.bins.concat(dims_to_reduce)
-        out = q_all_pixels.bin(**edges)
+        # q_all_pixels may just have a single bin now, which currently yields
+        # inferior performance when binning (no/bad multi-threading?).
+        # We operate on the content buffer for better multi-threaded performance.
+        if q_all_pixels.ndim == 0:
+            content = q_all_pixels.bins.constituents['data']
+            out = content.bin(**edges).assign_coords(q_all_pixels.coords)
+        else:
+            out = q_all_pixels.bin(**edges)
     else:
         # We want to flatten data to make histogramming cheaper (avoiding allocation of
         # large output before summing). We strip unnecessary content since it makes
@@ -209,6 +216,7 @@ def merge_spectra(
             out = (
                 flat.flatten(to=str(uuid.uuid4()))
                 .group(*[flat.coords[dim] for dim in flat.dims if dim != helper_dim])
+                .drop_coords(dims_to_keep or ())
                 .hist(**edges)
             )
     return CleanSummedQ[ScatteringRunType, IofQPart](out.squeeze())

src/esssans/normalization.py

@@ -379,6 +379,8 @@ def normalize(
         numerator = numerator.bin(wavelength=wavelength_bounds)
 
     def _reduce(da: sc.DataArray) -> sc.DataArray:
+        if da.sizes[wav] == 1:  # Can avoid costly event-data da.bins.concat
+            return da.squeeze(wav)
         return da.sum(wav) if da.bins is None else da.bins.concat(wav)
 
     num_parts = []
@@ -408,7 +410,10 @@ def _reduce(da: sc.DataArray) -> sc.DataArray:
                 )
     elif numerator.bins is not None:
         numerator = numerator.hist()
-    return IofQ[ScatteringRunType](numerator / denominator)
+    numerator /= denominator.drop_coords(
+        [name for name in denominator.coords if name not in denominator.dims]
+    )
+    return IofQ[ScatteringRunType](numerator)
 
 
 providers = (

src/esssans/uncertainty.py

@@ -67,27 +67,12 @@ def broadcast_to_events_with_upper_bound_variances(
     """
     if da.variances is None:
         return da
-    constituents = events.bins.constituents
-
-    if 'Q' in constituents['data'].coords:
-        Q = constituents['data'].coords['Q']
-        constituents['data'] = sc.DataArray(
-            sc.ones(sizes=Q.sizes, dtype='float32'), coords={'Q': Q}
-        )
-        edges = {'Q': da.coords['Q']}
-    else:
-        Qx = constituents['data'].coords['Qx']
-        Qy = constituents['data'].coords['Qy']
-        constituents['data'] = sc.DataArray(
-            sc.ones(sizes=Qx.sizes, dtype='float32'),
-            coords={'Qx': Qx, 'Qy': Qy},
-        )
-        edges = {'Qy': da.coords['Qy'], 'Qx': da.coords['Qx']}
-    # Combine all bins of the events that correspond to the same bin in the input data
-    to_concat = set(events.dims) - set(da.dims)
-    binned = sc.DataArray(sc.bins(**constituents).bins.concat(to_concat))
-    counts = binned.hist(**edges)
-    da = da.copy()
-    da.variances *= counts.values
+    if da.sizes != events.sizes:
+        # This is a safety check, but we should never get here.
+        raise ValueError(f"Sizes {da.sizes} do not match event sizes {events.sizes}")
+    # Given how this function is used currently (in the context of normalization
+    # with matching binning in numerator and denominator, not using scipp.lookup),
+    # we can simply count the events in the existing binning.
+    da.variances *= events.bins.size().values
     da.data = sc.bins_like(events, da.data)
     return da

tests/loki/iofq_test.py

@@ -81,13 +81,17 @@ def test_pipeline_can_compute_IofQ_in_event_mode(uncertainties, target, qxy: boo
     assert sc.allclose(
         sc.values(reference.data),
         sc.values(result.hist().data),
-        rtol=sc.scalar(1e-11),
-        atol=sc.scalar(1e-11),
+        # Could be 1e-11, but currently the workflow defaults to float32 data, as
+        # returned by ScippNexus.
+        rtol=sc.scalar(1e-7),
+        atol=sc.scalar(1e-7),
     )
     if uncertainties == UncertaintyBroadcastMode.drop:
-        tol = sc.scalar(1e-14)
+        # Could both be 1e-14 if using float64
+        tol = sc.scalar(1e-7 if qxy else 1e-10)
     else:
-        tol = sc.scalar(1e-8 if qxy else 1e-9)
+        # Could be 1e-8 if using float64
+        tol = sc.scalar(1e-7 if qxy else 1e-9)
     assert sc.allclose(
         sc.variances(reference).data,
         sc.variances(result.hist()).data,