feat: add lorentz_correction to powder workflow

src/ess/powder/__init__.py

@@ -15,6 +15,7 @@
     masking,
     smoothing,
     transform,
+    workflow,
 )
 from .correction import RunNormalization
 from .masking import with_pixel_mask_filenames
@@ -33,6 +34,7 @@
     *filtering.providers,
     *grouping.providers,
     *masking.providers,
+    *workflow.providers,
 )
 
 """Sciline providers for powder diffraction."""
@@ -50,4 +52,5 @@
     "smoothing",
     "transform",
     "with_pixel_mask_filenames",
+    "workflow",
 ]

src/ess/powder/calibration.py

@@ -9,7 +9,7 @@
 import scipp as sc
 import scipp.constants
 
-from .types import DspacingDetector, SampleRun
+from .types import CorrectedDetector, SampleRun
 
 
 class OutputCalibrationData(Mapping[int, sc.Variable]):
@@ -94,7 +94,7 @@ def to_cif_format(self) -> sc.DataArray:
 
 
 def assemble_output_calibration(
-    data: DspacingDetector[SampleRun],
+    data: CorrectedDetector[SampleRun],
 ) -> OutputCalibrationData:
     """Construct output calibration data from average pixel positions."""
     # Use nanmean because pixels without events have position=NaN.

src/ess/powder/conversion.py

@@ -14,7 +14,6 @@
 from .types import (
     CalibrationData,
     CorrectedDetector,
-    DspacingDetector,
     ElasticCoordTransformGraph,
     EmptyCanSubtractedIntensityTof,
     EmptyCanSubtractedIofDspacing,
@@ -150,7 +149,7 @@ def to_dspacing_with_calibration(
         graph["_tag_positions_consumed"] = lambda: sc.scalar(0)
     out = out.transform_coords("dspacing", graph=graph, keep_intermediate=False)
     out.coords.pop("_tag_positions_consumed", None)
-    return DspacingDetector[RunType](out)
+    return CorrectedDetector[RunType](out)
 
 
 def powder_coordinate_transformation_graph(
@@ -203,10 +202,12 @@ def _restore_tof_if_in_wavelength(data: sc.DataArray) -> sc.DataArray:
 
 
 def add_scattering_coordinates_from_positions(
-    data: TofDetector[RunType], graph: ElasticCoordTransformGraph[RunType]
+    data: TofDetector[RunType],
+    graph: ElasticCoordTransformGraph[RunType],
+    calibration: CalibrationData,
 ) -> WavelengthDetector[RunType]:
     """
-    Add ``wavelength`` and ``two_theta`` coordinates to the data.
+    Add ``wavelength``, ``two_theta`` and ``dspacing`` coordinates to the data.
     The input ``data`` must have a ``tof`` coordinate, as well as the necessary
     positions of the beamline components (source, sample, detectors) to compute
     the scattering coordinates.
@@ -223,30 +224,31 @@ def add_scattering_coordinates_from_positions(
         graph=graph,
         keep_intermediate=False,
     )
+    out = convert_to_dspacing(out, graph, calibration)
     return WavelengthDetector[RunType](out)
 
 
 def convert_to_dspacing(
-    data: CorrectedDetector[RunType],
+    data: sc.DataArray,
     graph: ElasticCoordTransformGraph[RunType],
     calibration: CalibrationData,
-) -> DspacingDetector[RunType]:
+) -> sc.DataArray:
     if calibration is None:
         out = data.transform_coords(["dspacing"], graph=graph, keep_intermediate=False)
     else:
         out = to_dspacing_with_calibration(data, calibration=calibration)
         for key in ("wavelength", "two_theta"):
             if key in out.coords.keys():
                 out.coords.set_aligned(key, False)
-    out.bins.coords.pop("tof", None)
-    # See scipp/essreduce#249
-    out.bins.coords.pop("event_time_offset", None)
-    return DspacingDetector[RunType](out)
+    return out
 
 
 def _convert_reduced_to_tof_impl(
     data: sc.DataArray, calibration: OutputCalibrationData
 ) -> sc.DataArray:
+    if data.bins is not None:
+        if 'tof' in data.bins.coords:
+            data.bins.coords.pop('tof')
     return data.transform_coords(
         tof=calibration.d_to_tof_transformer(), keep_inputs=False
     )
@@ -281,7 +283,6 @@ def convert_monitor_to_wavelength(
 providers = (
     powder_coordinate_transformation_graph,
     add_scattering_coordinates_from_positions,
-    convert_to_dspacing,
     convert_reduced_to_tof,
     convert_reduced_to_empty_can_subtracted_tof,
     convert_monitor_to_wavelength,

src/ess/powder/correction.py

@@ -332,30 +332,12 @@ def apply_lorentz_correction(da: sc.DataArray) -> sc.DataArray:
         ``da`` multiplied by :math:`L`.
         Has the same dtype as ``da``.
     """
-    # The implementation is optimized under the assumption that two_theta
-    # is small and dspacing and the data are large.
-    out = _shallow_copy(da)
     dspacing = event_or_outer_coord(da, "dspacing")
     two_theta = event_or_outer_coord(da, "two_theta")
-    theta = 0.5 * two_theta
-
-    d4 = dspacing.broadcast(sizes=out.sizes) ** 4
-    if out.bins is None:
-        out.data = d4.to(dtype=out.dtype, copy=False)
-        out_data = out.data
-    else:
-        out.bins.data = d4.to(dtype=out.bins.dtype, copy=False)
-        out_data = out.bins.data
-    out_data *= sc.sin(theta, out=theta)
-    out_data *= da.data if da.bins is None else da.bins.data
-    return out
-
-
-def _shallow_copy(da: sc.DataArray) -> sc.DataArray:
-    # See https://github.com/scipp/scipp/issues/2773
-    out = da.copy(deep=False)
-    if da.bins is not None:
-        out.data = sc.bins(**da.bins.constituents)
+    sin_theta = sc.sin(0.5 * two_theta)
+    d4 = dspacing**4
+    out = da * sin_theta.to(dtype=da.bins.dtype if da.bins else da.dtype, copy=False)
+    out *= d4.to(dtype=da.bins.dtype if da.bins else da.dtype, copy=False)
     return out
 
 

src/ess/powder/grouping.py

@@ -6,9 +6,9 @@
 import scipp as sc
 
 from .types import (
+    CorrectedDetector,
     CorrectedDspacing,
     DspacingBins,
-    DspacingDetector,
     FocussedDataDspacing,
     FocussedDataDspacingTwoTheta,
     KeepEvents,
@@ -27,7 +27,7 @@ def _reconstruct_wavelength(
 
 
 def focus_data_dspacing_and_two_theta(
-    data: DspacingDetector[RunType],
+    data: CorrectedDetector[RunType],
     dspacing_bins: DspacingBins,
     keep_events: KeepEvents[RunType],
 ) -> CorrectedDspacing[RunType]:

src/ess/powder/masking.py

@@ -11,7 +11,6 @@
 import scipp as sc
 
 from .types import (
-    CorrectedDetector,
     MaskedDetectorIDs,
     PixelMaskFilename,
     RunType,
@@ -39,7 +38,7 @@ def apply_masks(
     tof_mask_func: TofMask,
     wavelength_mask_func: WavelengthMask,
     two_theta_mask_func: TwoThetaMask,
-) -> CorrectedDetector[RunType]:
+) -> sc.DataArray:
     """ """
     out = data.copy(deep=False)
     if len(masked_pixel_ids) > 0:
@@ -68,10 +67,10 @@ def apply_masks(
                 )
                 out.masks[dim] = mask(coord)
 
-    return CorrectedDetector[RunType](out)
+    return out
 
 
-providers = (read_pixel_masks, apply_masks)
+providers = (read_pixel_masks,)
 
 
 def _merge(*dicts: dict) -> dict:

src/ess/powder/types.py

@@ -99,10 +99,6 @@ class WavelengthDetector(sciline.Scope[RunType, sc.DataArray], sc.DataArray):
     d-spacing."""
 
 
-class DspacingDetector(sciline.Scope[RunType, sc.DataArray], sc.DataArray):
-    """Data converted to d-spacing."""
-
-
 DspacingHistogram = NewType("DspacingHistogram", sc.DataArray)
 """Histogrammed intensity vs d-spacing."""
 

src/ess/powder/workflow.py

@@ -0,0 +1,32 @@
+from .correction import apply_lorentz_correction
+from .masking import apply_masks
+from .types import (
+    CorrectedDetector,
+    MaskedDetectorIDs,
+    RunType,
+    TofMask,
+    TwoThetaMask,
+    WavelengthDetector,
+    WavelengthMask,
+)
+
+
+def add_masks_and_corrections(
+    da: WavelengthDetector[RunType],
+    masked_pixel_ids: MaskedDetectorIDs,
+    tof_mask_func: TofMask,
+    wavelength_mask_func: WavelengthMask,
+    two_theta_mask_func: TwoThetaMask,
+) -> CorrectedDetector[RunType]:
+    masked = apply_masks(
+        data=da,
+        masked_pixel_ids=masked_pixel_ids,
+        tof_mask_func=tof_mask_func,
+        wavelength_mask_func=wavelength_mask_func,
+        two_theta_mask_func=two_theta_mask_func,
+    )
+    out = apply_lorentz_correction(masked)
+    return CorrectedDetector[RunType](out)
+
+
+providers = (add_masks_and_corrections,)

tests/dream/geant4_reduction_test.py

@@ -26,7 +26,6 @@
     CIFAuthors,
     CorrectedDetector,
     DspacingBins,
-    DspacingDetector,
     EmptyCanRun,
     EmptyCanSubtractedIofDspacing,
     Filename,
@@ -222,8 +221,8 @@ def test_workflow_is_deterministic(workflow):
 
 def test_pipeline_can_compute_intermediate_results(workflow):
     workflow = powder.with_pixel_mask_filenames(workflow, [])
-    results = workflow.compute((DspacingDetector[SampleRun], NeXusDetectorName))
-    result = results[DspacingDetector[SampleRun]]
+    results = workflow.compute((CorrectedDetector[SampleRun], NeXusDetectorName))
+    result = results[CorrectedDetector[SampleRun]]
 
     detector_name = results[NeXusDetectorName]
     expected_dims = {'segment', 'wire', 'counter', 'strip', 'module'}
@@ -349,6 +348,6 @@ class MyWorkflow: ...
 
 def test_dream_workflow_parameters_returns_filtered_params():
     wf = DreamGeant4ProtonChargeWorkflow()
-    parameters = reduce_workflow.get_parameters(wf, (DspacingDetector[SampleRun],))
+    parameters = reduce_workflow.get_parameters(wf, (CorrectedDetector[SampleRun],))
     assert Filename[SampleRun] in parameters
     assert Filename[EmptyCanRun] not in parameters

tests/powder/conversion_test.py

@@ -159,7 +159,7 @@ def test_add_scattering_coordinates_from_positions():
         **scn.conversion.graph.tof.elastic('tof'),
     }
 
-    result = add_scattering_coordinates_from_positions(tof, graph)
+    result = add_scattering_coordinates_from_positions(tof, graph, calibration=None)
 
     assert 'wavelength' in result.coords
     assert 'two_theta' in result.coords

tests/snspowder/powgen/powgen_reduction_test.py

@@ -13,7 +13,6 @@
     CorrectedDetector,
     DetectorBankSizes,
     DspacingBins,
-    DspacingDetector,
     Filename,
     GravityVector,
     IntensityDspacing,
@@ -114,7 +113,7 @@ def test_workflow_is_deterministic(providers, params):
 def test_pipeline_can_compute_intermediate_results(providers, params):
     pipeline = sciline.Pipeline(providers, params=params)
     pipeline = powder.with_pixel_mask_filenames(pipeline, [])
-    result = pipeline.compute(DspacingDetector[SampleRun])
+    result = pipeline.compute(CorrectedDetector[SampleRun])
     assert set(result.dims) == {'bank', 'column', 'row'}