refactor: share parts of the workflow between amor and estia

Author: jokasimr

src/ess/amor/conversions.py

@@ -6,15 +6,15 @@
 from ..reflectometry.conversions import reflectometry_q
 from ..reflectometry.types import (
     BeamDivergenceLimits,
+    CoordTransformationGraph,
     WavelengthBins,
     YIndexLimits,
     ZIndexLimits,
 )
 from .geometry import Detector
-from .types import CoordTransformationGraph
 
 
-def theta(wavelength, divergence_angle, L2, sample_rotation, detector_rotation):
+def theta(wavelength, pixel_divergence_angle, L2, sample_rotation, detector_rotation):
     '''
     Angle of reflection.
 
@@ -61,14 +61,15 @@ def theta(wavelength, divergence_angle, L2, sample_rotation, detector_rotation):
     '''
     c = sc.constants.g * sc.constants.m_n**2 / sc.constants.h**2
     out = (c * L2 * wavelength**2).to(unit='dimensionless') + sc.sin(
-        divergence_angle.to(unit='rad', copy=False) + detector_rotation.to(unit='rad')
+        pixel_divergence_angle.to(unit='rad', copy=False)
+        + detector_rotation.to(unit='rad')
     )
     out = sc.asin(out, out=out)
     out -= sample_rotation.to(unit='rad')
     return out
 
 
-def angle_of_divergence(theta, sample_rotation, angle_to_center_of_beam):
+def divergence_angle(theta, sample_rotation, angle_to_center_of_beam):
     """
     Difference between the incident angle and the center of the incident beam.
     Useful for filtering parts of the beam that have too high divergence.
@@ -84,7 +85,7 @@ def angle_of_divergence(theta, sample_rotation, angle_to_center_of_beam):
 
 
 def wavelength(
-    event_time_offset, divergence_angle, L1, L2, chopper_phase, chopper_frequency
+    event_time_offset, pixel_divergence_angle, L1, L2, chopper_phase, chopper_frequency
 ):
     "Converts event_time_offset to wavelength using the chopper settings."
     out = event_time_offset.to(unit="ns", dtype="float64", copy=True)
@@ -108,37 +109,22 @@ def wavelength(
     )
     # Correction for path length through guides being different
     # depending on incident angle.
-    out -= (divergence_angle.to(unit="rad") / (np.pi * sc.units.rad)) * tau
+    out -= (pixel_divergence_angle.to(unit="rad") / (np.pi * sc.units.rad)) * tau
     out *= (sc.constants.h / sc.constants.m_n) / (L1 + L2)
     return out.to(unit='angstrom', copy=False)
 
 
 def coordinate_transformation_graph() -> CoordTransformationGraph:
     return {
-        "divergence_angle": "pixel_divergence_angle",
         "wavelength": wavelength,
         "theta": theta,
-        "angle_of_divergence": angle_of_divergence,
+        "divergence_angle": divergence_angle,
         "Q": reflectometry_q,
         "L1": lambda chopper_distance: sc.abs(chopper_distance),
         "L2": lambda distance_in_detector: distance_in_detector + Detector.distance,
     }
 
 
-def add_coords(
-    da: sc.DataArray,
-    graph: dict,
-) -> sc.DataArray:
-    "Adds scattering coordinates to the raw detector data."
-    return da.transform_coords(
-        ("wavelength", "theta", "angle_of_divergence", "Q", "L1", "L2"),
-        graph,
-        rename_dims=False,
-        keep_intermediate=False,
-        keep_aliases=False,
-    )
-
-
 def _not_between(v, a, b):
     return (v < a) | (v > b)
 
@@ -161,9 +147,9 @@ def add_masks(
     )
     da = da.bins.assign_masks(
         divergence_too_large=_not_between(
-            da.bins.coords["angle_of_divergence"],
-            bdlim[0].to(unit=da.bins.coords["angle_of_divergence"].bins.unit),
-            bdlim[1].to(unit=da.bins.coords["angle_of_divergence"].bins.unit),
+            da.bins.coords["divergence_angle"],
+            bdlim[0].to(unit=da.bins.coords["divergence_angle"].bins.unit),
+            bdlim[1].to(unit=da.bins.coords["divergence_angle"].bins.unit),
         ),
         wavelength=_not_between(
             da.bins.coords['wavelength'],

src/ess/amor/normalization.py

@@ -8,6 +8,7 @@
     supermirror_reflectivity,
 )
 from ..reflectometry.types import (
+    CoordTransformationGraph,
     DetectorSpatialResolution,
     ReducedReference,
     ReducibleData,
@@ -22,7 +23,6 @@
     sample_size_resolution,
     wavelength_resolution,
 )
-from .types import CoordTransformationGraph
 
 
 def mask_events_where_supermirror_does_not_cover(

src/ess/amor/types.py

@@ -9,8 +9,6 @@
 AngularResolution = NewType("AngularResolution", sc.Variable)
 SampleSizeResolution = NewType("SampleSizeResolution", sc.Variable)
 
-CoordTransformationGraph = NewType("CoordTransformationGraph", dict)
-
 
 class ChopperFrequency(sciline.Scope[RunType, sc.Variable], sc.Variable):
     """Frequency of the choppers used in the run."""

src/ess/amor/workflow.py

@@ -1,10 +1,12 @@
 from ..reflectometry.conversions import (
+    add_coords,
     add_proton_current_coord,
     add_proton_current_mask,
 )
 from ..reflectometry.corrections import correct_by_footprint, correct_by_proton_current
 from ..reflectometry.types import (
     BeamDivergenceLimits,
+    CoordTransformationGraph,
     ProtonCurrent,
     RawDetectorData,
     ReducibleData,
@@ -13,8 +15,7 @@
     YIndexLimits,
     ZIndexLimits,
 )
-from .conversions import add_coords, add_masks
-from .types import CoordTransformationGraph
+from .conversions import add_masks
 
 
 def add_coords_masks_and_apply_corrections(

src/ess/estia/conversions.py

@@ -5,11 +5,11 @@
 from ..reflectometry.conversions import reflectometry_q
 from ..reflectometry.types import (
     BeamDivergenceLimits,
+    CoordTransformationGraph,
     WavelengthBins,
     YIndexLimits,
     ZIndexLimits,
 )
-from .types import CoordTransformationGraph
 
 
 def theta(
@@ -65,20 +65,6 @@ def coordinate_transformation_graph() -> CoordTransformationGraph:
     }
 
 
-def add_coords(
-    da: sc.DataArray,
-    graph: dict,
-) -> sc.DataArray:
-    "Adds scattering coordinates to the raw detector data."
-    return da.transform_coords(
-        ("wavelength", "theta", "divergence_angle", "Q", "L1", "L2"),
-        graph,
-        rename_dims=False,
-        keep_intermediate=False,
-        keep_aliases=False,
-    )
-
-
 def _not_between(v, a, b):
     return (v < a) | (v > b)
 

src/ess/estia/normalization.py

@@ -8,6 +8,7 @@
     supermirror_reflectivity,
 )
 from ..reflectometry.types import (
+    CoordTransformationGraph,
     DetectorSpatialResolution,
     ReducedReference,
     ReducibleData,
@@ -21,7 +22,6 @@
     q_resolution,
     sample_size_resolution,
 )
-from .types import CoordTransformationGraph
 
 
 def mask_events_where_supermirror_does_not_cover(

src/ess/estia/workflow.py

@@ -1,10 +1,12 @@
 from ..reflectometry.conversions import (
+    add_coords,
     add_proton_current_coord,
     add_proton_current_mask,
 )
 from ..reflectometry.corrections import correct_by_proton_current
 from ..reflectometry.types import (
     BeamDivergenceLimits,
+    CoordTransformationGraph,
     ProtonCurrent,
     RawDetectorData,
     ReducibleData,
@@ -13,9 +15,8 @@
     YIndexLimits,
     ZIndexLimits,
 )
-from .conversions import add_coords, add_masks
+from .conversions import add_masks
 from .corrections import correct_by_footprint
-from .types import CoordTransformationGraph
 
 
 def add_coords_masks_and_apply_corrections(

src/ess/reflectometry/conversions.py

@@ -4,7 +4,10 @@
 from scipp.constants import pi
 from scippneutron._utils import elem_dtype
 
-from .types import ProtonCurrent, RunType
+from .types import (
+    ProtonCurrent,
+    RunType,
+)
 
 
 def reflectometry_q(wavelength: sc.Variable, theta: sc.Variable) -> sc.Variable:
@@ -61,4 +64,18 @@ def add_proton_current_mask(da: sc.DataArray) -> sc.DataArray:
     return da
 
 
+def add_coords(
+    da: sc.DataArray,
+    graph: dict,
+) -> sc.DataArray:
+    "Adds scattering coordinates to the raw detector data."
+    return da.transform_coords(
+        ("wavelength", "theta", "divergence_angle", "Q", "L1", "L2"),
+        graph,
+        rename_dims=False,
+        keep_intermediate=False,
+        keep_aliases=False,
+    )
+
+
 providers = ()

src/ess/reflectometry/types.py

@@ -7,6 +7,8 @@
 SampleRun = NewType("SampleRun", str)
 RunType = TypeVar("RunType", ReferenceRun, SampleRun)
 
+CoordTransformationGraph = NewType("CoordTransformationGraph", dict)
+
 
 class NeXusDetectorName(sciline.Scope[RunType, str], str):
     """Name of the detector in the nexus file containing the events of the RunType"""