Merge remote-tracking branch 'origin/estia' into polcal

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,8 +8,8 @@
     supermirror_reflectivity,
 )
 from ..reflectometry.types import (
+    CoordTransformationGraph,
     DetectorSpatialResolution,
-    QBins,
     ReducedReference,
     ReducibleData,
     Reference,
@@ -23,7 +23,6 @@
     sample_size_resolution,
     wavelength_resolution,
 )
-from .types import CoordTransformationGraph
 
 
 def mask_events_where_supermirror_does_not_cover(
@@ -64,14 +63,12 @@ def mask_events_where_supermirror_does_not_cover(
         m=mvalue,
         alpha=alpha,
     )
-    sam.bins.masks["supermirror_does_not_cover"] = sc.isnan(R)
-    return sam
+    return sam.bins.assign_masks(supermirror_does_not_cover=sc.isnan(R))
 
 
 def evaluate_reference_at_sample_coords(
     reference: ReducedReference,
     sample: ReducibleData[SampleRun],
-    qbins: QBins,
     detector_spatial_resolution: DetectorSpatialResolution[SampleRun],
     graph: CoordTransformationGraph,
 ) -> Reference:
@@ -103,6 +100,8 @@ def evaluate_reference_at_sample_coords(
             "Q_resolution": q_resolution,
         },
         rename_dims=False,
+        keep_intermediate=False,
+        keep_aliases=False,
     )
     return sc.values(ref)
 

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/__init__.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2023 Scipp contributors (https://github.com/scipp)
+import importlib.metadata
+
+import sciline
+import scipp as sc
+
+from ..reflectometry import providers as reflectometry_providers
+from ..reflectometry import supermirror
+from ..reflectometry.types import (
+    BeamDivergenceLimits,
+    BeamSize,
+    DetectorSpatialResolution,
+    NeXusDetectorName,
+    RunType,
+    SamplePosition,
+)
+from . import conversions, load, normalization, resolution, workflow
+from .types import (
+    AngularResolution,
+    SampleSizeResolution,
+    WavelengthResolution,
+)
+
+try:
+    __version__ = importlib.metadata.version(__package__ or __name__)
+except importlib.metadata.PackageNotFoundError:
+    __version__ = "0.0.0"
+
+
+providers = (
+    *reflectometry_providers,
+    # *load.providers,
+    *conversions.providers,
+    *workflow.providers,
+    *normalization.providers,
+)
+"""
+List of providers for setting up a Sciline pipeline.
+
+This provides a default Estia workflow including providers for loadings files.
+"""
+
+
+def default_parameters() -> dict:
+    return {
+        supermirror.MValue: sc.scalar(5, unit=sc.units.dimensionless),
+        supermirror.CriticalEdge: 0.022 * sc.Unit("1/angstrom"),
+        supermirror.Alpha: sc.scalar(0.25 / 0.088, unit=sc.units.angstrom),
+        BeamSize[RunType]: 2.0 * sc.units.mm,
+        DetectorSpatialResolution[RunType]: 0.0025 * sc.units.m,
+        SamplePosition[RunType]: sc.vector([0, 0, 0], unit="m"),
+        NeXusDetectorName[RunType]: "detector",
+        BeamDivergenceLimits: (
+            sc.scalar(-0.75, unit='deg'),
+            sc.scalar(0.75, unit='deg'),
+        ),
+    }
+
+
+def EstiaWorkflow() -> sciline.Pipeline:
+    """
+    Workflow with default parameters for the Estia instrument.
+    """
+    return sciline.Pipeline(providers=providers, params=default_parameters())
+
+
+__all__ = [
+    "AngularResolution",
+    "EstiaWorkflow",
+    "SampleSizeResolution",
+    "WavelengthResolution",
+    "conversions",
+    "default_parameters",
+    "load",
+    "providers",
+    "resolution",
+    "supermirror",
+]

src/ess/estia/conversions.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2023 Scipp contributors (https://github.com/scipp)
+import scipp as sc
+
+from ..reflectometry.conversions import reflectometry_q
+from ..reflectometry.types import (
+    BeamDivergenceLimits,
+    CoordTransformationGraph,
+    WavelengthBins,
+    YIndexLimits,
+    ZIndexLimits,
+)
+
+
+def theta(
+    divergence_angle,
+    sample_rotation,
+):
+    '''
+    Angle of reflection.
+
+    Computes the angle between the scattering direction of
+    the neutron and the sample surface.
+    '''
+    return divergence_angle + sample_rotation.to(unit=divergence_angle.unit)
+
+
+def divergence_angle(
+    position, sample_position, detector_rotation, incident_angle_of_center_of_beam
+):
+    """
+    Angle between the scattering direction and
+    the ray from the sample to the center of the detector.
+    """
+    p = position - sample_position.to(unit=position.unit)
+    return (
+        sc.atan2(y=p.fields.x, x=p.fields.z)
+        - detector_rotation.to(unit='rad')
+        - incident_angle_of_center_of_beam.to(unit='rad')
+    )
+
+
+def wavelength(
+    event_time_offset,
+    # Other inputs
+):
+    "Converts event_time_offset to wavelength"
+    # Use frame unwrapping from scippneutron
+    pass
+
+
+def coordinate_transformation_graph() -> CoordTransformationGraph:
+    return {
+        "wavelength": wavelength,
+        "theta": theta,
+        "divergence_angle": divergence_angle,
+        "Q": reflectometry_q,
+        "L1": lambda source_position, sample_position: sc.norm(
+            sample_position - source_position
+        ),  # + extra correction for guides?
+        "L2": lambda position, sample_position: sc.norm(position - sample_position),
+        "incident_angle_of_center_of_beam": lambda: sc.scalar(1.7, unit='deg').to(
+            unit='rad'
+        ),
+    }
+
+
+def _not_between(v, a, b):
+    return (v < a) | (v > b)
+
+
+def add_masks(
+    da: sc.DataArray,
+    ylim: YIndexLimits,
+    zlims: ZIndexLimits,
+    bdlim: BeamDivergenceLimits,
+    wbins: WavelengthBins,
+) -> sc.DataArray:
+    """
+    Masks the data by ranges in the detector
+    coordinates ``z`` and ``y``, and by the divergence of the beam,
+    and by wavelength.
+    """
+    da = da.assign_masks(
+        stripe_range=_not_between(da.coords["stripe"], *ylim),
+        z_range=_not_between(da.coords["z_index"], *zlims),
+        divergence_too_large=_not_between(
+            da.coords["divergence_angle"],
+            bdlim[0].to(unit=da.coords["divergence_angle"].unit),
+            bdlim[1].to(unit=da.coords["divergence_angle"].unit),
+        ),
+    )
+    da = da.bins.assign_masks(
+        wavelength=_not_between(
+            da.bins.coords['wavelength'],
+            wbins[0],
+            wbins[-1],
+        ),
+    )
+    return da
+
+
+providers = (coordinate_transformation_graph,)

src/ess/estia/corrections.py

@@ -0,0 +1,6 @@
+import scipp as sc
+
+
+def correct_by_footprint(da: sc.DataArray) -> sc.DataArray:
+    "Corrects the data by the size of the footprint on the sample."
+    return da / sc.sin(da.coords['theta'])