Merge pull request #218 from scipp/robust-vana-norm

Rebin vanadium to data

Author: jl-wynen

pyproject.toml

@@ -37,7 +37,7 @@ dependencies = [
   "plopp>=25.07.0",
   "pythreejs>=2.4.1",
   "sciline>=25.04.1",
-  "scipp>=25.05.1",
+  "scipp>=25.11.0",
   "scippneutron>=25.02.0",
   "scippnexus>=23.12.0",
   "tof>=25.11.1",

src/ess/powder/correction.py

@@ -3,6 +3,7 @@
 """Correction algorithms for powder diffraction."""
 
 import enum
+from collections.abc import Mapping, Sequence
 from typing import TypeVar
 
 import sciline
@@ -31,6 +32,8 @@
     WavelengthMonitor,
 )
 
+_T = TypeVar("_T")
+
 
 def normalize_by_monitor_histogram(
     detector: CorrectedDspacing[RunType],
@@ -146,19 +149,41 @@ def _mask_out_of_monitor_range_data(
     return detector, False
 
 
+def _filter_keys(mapping: Mapping[str, _T], keys: Sequence[str]) -> dict[str, _T]:
+    return {key: value for key, value in mapping.items() if key in keys}
+
+
+def _histogram_vanadium(vanadium: sc.DataArray, sample: sc.DataArray) -> sc.DataArray:
+    target_coords = _filter_keys(sample.coords, ("dspacing", "two_theta"))
+    if "two_theta" in target_coords and "two_theta" not in vanadium.bins.coords:
+        # Need to provide a two_theta event coord so we can histogram.
+        if sc.identical(vanadium.coords["two_theta"], target_coords["two_theta"]):
+            # Skip the expensive event coord if possible:
+            return vanadium.hist(dspacing=sample.coords["dspacing"])
+        return vanadium.bins.assign_coords(
+            two_theta=sc.bins_like(vanadium, vanadium.coords["two_theta"])
+        ).hist(target_coords)
+    return vanadium.hist(target_coords)
+
+
 def _normalize_by_vanadium(
     data: sc.DataArray,
     vanadium: sc.DataArray,
     uncertainty_broadcast_mode: UncertaintyBroadcastMode,
 ) -> sc.DataArray:
-    norm = vanadium.hist() if vanadium.bins is not None else vanadium
+    norm = (
+        _histogram_vanadium(vanadium, sample=data)
+        if vanadium.is_binned
+        else vanadium.rebin(_filter_keys(data.coords, ("dspacing", "two_theta")))
+    )
     norm = broadcast_uncertainties(
         norm, prototype=data, mode=uncertainty_broadcast_mode
     )
-    # Converting to unit 'one' because the division might produce a unit
-    # with a large scale if the proton charges in data and vanadium were
-    # measured with different units.
-    normed = (data / norm).to(unit="one", copy=False)
+    # Convert the unit such that the end-result has unit 'one' because the division
+    # might otherwise produce a unit with a large scale if the proton charges in data
+    # and vanadium were measured with different units.
+    norm = norm.to(unit=data.unit, copy=False)
+    normed = data / norm
     mask = norm.data == sc.scalar(0.0, unit=norm.unit)
     if mask.any():
         normed.masks['zero_vanadium'] = mask
@@ -173,8 +198,18 @@ def normalize_by_vanadium_dspacing(
     vanadium: FocussedDataDspacing[VanadiumRun],
     uncertainty_broadcast_mode: UncertaintyBroadcastMode,
 ) -> IntensityDspacing[_RunTypeNoVanadium]:
-    """
-    Normalize sample data by a vanadium measurement and return intensity vs. d-spacing.
+    """Normalize sample data binned in d-spacing by a vanadium measurement.
+
+    If the vanadium data is binned, it gets :func:`histogrammed <scipp.hist>` to the
+    same bins as ``data``. If it is not binned, it gets :func:`rebinned <scipp.rebin>`
+    to the same coordinates as ``data``. Then, the result is computed as
+
+    .. code-block:: python
+
+        data / vanadium
+
+    And any bins where vanadium is zero are masked out
+    with a mask called "zero_vanadium".
 
     Parameters
     ----------
@@ -192,6 +227,11 @@ def normalize_by_vanadium_dspacing(
         ``data / vanadium``.
         May contain a mask "zero_vanadium" which is ``True``
         for bins where vanadium is zero.
+
+    See Also
+    --------
+    normalize_by_vanadium_dspacing_and_two_theta:
+        Normalization for 2d data binned in d-spacing and :math`2\\theta`.
     """
     return IntensityDspacing(
         _normalize_by_vanadium(data, vanadium, uncertainty_broadcast_mode)
@@ -203,9 +243,18 @@ def normalize_by_vanadium_dspacing_and_two_theta(
     vanadium: FocussedDataDspacingTwoTheta[VanadiumRun],
     uncertainty_broadcast_mode: UncertaintyBroadcastMode,
 ) -> IntensityDspacingTwoTheta[_RunTypeNoVanadium]:
-    """
-    Normalize sample data by a vanadium measurement and return intensity vs.
-    (d-spacing, 2theta).
+    """Normalize sample data binned in (d-spacing, 2theta) by a vanadium measurement.
+
+    If the vanadium data is binned, it gets :func:`histogrammed <scipp.hist>` to the
+    same bins as ``data``. If it is not binned, it gets :func:`rebinned <scipp.rebin>`
+    to the same coordinates as ``data``. Then, the result is computed as
+
+    .. code-block:: python
+
+        data / vanadium
+
+    And any bins where vanadium is zero are masked out
+    with a mask called "zero_vanadium".
 
     Parameters
     ----------
@@ -223,6 +272,11 @@ def normalize_by_vanadium_dspacing_and_two_theta(
         ``data / vanadium``.
         May contain a mask "zero_vanadium" which is ``True``
         for bins where vanadium is zero.
+
+    See Also
+    --------
+    normalize_by_vanadium_dspacing:
+        Normalization for 1d data binned in d-spacing.
     """
     return IntensityDspacingTwoTheta(
         _normalize_by_vanadium(data, vanadium, uncertainty_broadcast_mode)

tests/powder/correction_test.py

@@ -10,13 +10,18 @@
     merge_calibration,
     normalize_by_monitor_histogram,
     normalize_by_monitor_integrated,
+    normalize_by_vanadium_dspacing,
+    normalize_by_vanadium_dspacing_and_two_theta,
 )
 from ess.powder.types import (
     CaveMonitor,
     CorrectedDspacing,
+    FocussedDataDspacing,
+    FocussedDataDspacingTwoTheta,
     NormalizedDspacing,
     SampleRun,
     UncertaintyBroadcastMode,
+    VanadiumRun,
     WavelengthMonitor,
 )
 
@@ -454,3 +459,181 @@ def test_normalize_by_monitor_integrated_assigns_mask_if_monitor_range_too_narro
             monitor=WavelengthMonitor[SampleRun, CaveMonitor](monitor),
             uncertainty_broadcast_mode=UncertaintyBroadcastMode.fail,
         )
+
+
+class TestNormalizeByVanadium:
+    def random_variable(
+        self,
+        rng: np.random.Generator,
+        dim: str,
+        n: int,
+        unit: str,
+        with_variances: bool = False,
+    ) -> sc.Variable:
+        values = rng.uniform(0.1, 2.0, n)
+        variances = values * rng.uniform(0.1, 0.5, n) if with_variances else None
+        return sc.array(dims=[dim], values=values, variances=variances, unit=unit)
+
+    def random_binned_data(
+        self,
+        rng: np.random.Generator,
+        n_events: int,
+        unit: str,
+        with_variances: bool = False,
+        *coords: tuple[str, int, str],
+    ) -> sc.DataArray:
+        return sc.DataArray(
+            self.random_variable(
+                rng, 'event', n_events, unit, with_variances=with_variances
+            ),
+            coords={
+                dim: self.random_variable(rng, 'event', n_events, coord_unit)
+                for (dim, _, coord_unit) in coords
+            },
+        ).bin({dim: n_bins for (dim, n_bins, _) in coords})
+
+    def make_sample_and_vanadium_1d(self) -> tuple[sc.DataArray, sc.DataArray]:
+        rng = np.random.default_rng(seed=495)
+        sample = self.random_binned_data(rng, 84, 'count', True, ('dspacing', 35, 'Å'))
+        vanadium = self.random_binned_data(
+            rng, 146, 'count', True, ('dspacing', 79, 'Å')
+        )
+        return sample, vanadium
+
+    def make_sample_and_vanadium_2d(self) -> tuple[sc.DataArray, sc.DataArray]:
+        rng = np.random.default_rng(seed=3193)
+        sample = self.random_binned_data(
+            rng, 138, 'count', True, ('dspacing', 35, 'Å'), ('two_theta', 13, 'rad')
+        )
+        vanadium = self.random_binned_data(
+            rng, 170, 'count', True, ('dspacing', 79, 'Å'), ('two_theta', 14, 'rad')
+        )
+        return sample, vanadium
+
+    def test_1d_binned_vanadium(self) -> None:
+        sample, vanadium = self.make_sample_and_vanadium_1d()
+        normed = normalize_by_vanadium_dspacing(
+            FocussedDataDspacing[SampleRun](sample),
+            FocussedDataDspacing[VanadiumRun](vanadium),
+            UncertaintyBroadcastMode.drop,
+        )
+        # we test masks separately
+        normed = normed.drop_masks(list(normed.masks.keys()))
+
+        norm = vanadium.hist(dspacing=sample.coords['dspacing'])
+        expected = sample / sc.values(norm)
+        sc.testing.assert_allclose(normed, expected)
+
+    def test_1d_histogrammed_vanadium(self) -> None:
+        sample, vanadium = self.make_sample_and_vanadium_1d()
+        vanadium = vanadium.hist()
+        normed = normalize_by_vanadium_dspacing(
+            FocussedDataDspacing[SampleRun](sample),
+            FocussedDataDspacing[VanadiumRun](vanadium),
+            UncertaintyBroadcastMode.drop,
+        )
+        # we test masks separately
+        normed = normed.drop_masks(list(normed.masks.keys()))
+
+        norm = vanadium.rebin(dspacing=sample.coords['dspacing'])
+        expected = sample / sc.values(norm)
+        sc.testing.assert_allclose(normed, expected)
+
+    def test_1d_binned_vanadium_binning_has_no_effect(self) -> None:
+        sample, vanadium = self.make_sample_and_vanadium_1d()
+        vana_binned_like_sample = vanadium.bin(dspacing=sample.coords['dspacing'])
+        normed_a = normalize_by_vanadium_dspacing(
+            FocussedDataDspacing[SampleRun](sample),
+            FocussedDataDspacing[VanadiumRun](vanadium),
+            UncertaintyBroadcastMode.drop,
+        )
+        normed_b = normalize_by_vanadium_dspacing(
+            FocussedDataDspacing[SampleRun](sample),
+            FocussedDataDspacing[VanadiumRun](vana_binned_like_sample),
+            UncertaintyBroadcastMode.drop,
+        )
+        sc.testing.assert_allclose(normed_a, normed_b)
+
+    def test_1d_masks_zero_vanadium_bins(self) -> None:
+        sample, vanadium = self.make_sample_and_vanadium_1d()
+        vanadium['dspacing', 5] = sc.scalar(0.0, variance=0.0, unit='counts')
+        normed = normalize_by_vanadium_dspacing(
+            FocussedDataDspacing[SampleRun](sample),
+            FocussedDataDspacing[VanadiumRun](vanadium),
+            UncertaintyBroadcastMode.drop,
+        )
+
+        norm = vanadium.hist(dspacing=sample.coords['dspacing'])
+
+        sc.testing.assert_allclose(
+            normed.masks['zero_vanadium'], norm.data == sc.scalar(0.0, unit='counts')
+        )
+
+    def test_2d_binned_vanadium(self) -> None:
+        sample, vanadium = self.make_sample_and_vanadium_2d()
+        normed = normalize_by_vanadium_dspacing_and_two_theta(
+            FocussedDataDspacingTwoTheta[SampleRun](sample),
+            FocussedDataDspacingTwoTheta[VanadiumRun](vanadium),
+            UncertaintyBroadcastMode.drop,
+        )
+        # we test masks separately
+        normed = normed.drop_masks(list(normed.masks.keys()))
+
+        norm = vanadium.hist(
+            dspacing=sample.coords['dspacing'], two_theta=sample.coords['two_theta']
+        )
+        expected = sample / sc.values(norm)
+        sc.testing.assert_allclose(normed, expected)
+
+    def test_2d_histogrammed_vanadium(self) -> None:
+        sample, vanadium = self.make_sample_and_vanadium_2d()
+        vanadium = vanadium.hist()
+        normed = normalize_by_vanadium_dspacing_and_two_theta(
+            FocussedDataDspacingTwoTheta[SampleRun](sample),
+            FocussedDataDspacingTwoTheta[VanadiumRun](vanadium),
+            UncertaintyBroadcastMode.drop,
+        )
+        # we test masks separately
+        normed = normed.drop_masks(list(normed.masks.keys()))
+
+        norm = vanadium.rebin(
+            dspacing=sample.coords['dspacing'], two_theta=sample.coords['two_theta']
+        )
+        expected = sample / sc.values(norm)
+        sc.testing.assert_allclose(normed, expected)
+
+    def test_2d_binned_vanadium_binning_has_no_effect(self) -> None:
+        sample, vanadium = self.make_sample_and_vanadium_2d()
+        vana_binned_like_sample = vanadium.bin(
+            dspacing=sample.coords['dspacing'], two_theta=sample.coords['two_theta']
+        )
+        normed_a = normalize_by_vanadium_dspacing_and_two_theta(
+            FocussedDataDspacingTwoTheta[SampleRun](sample),
+            FocussedDataDspacingTwoTheta[VanadiumRun](vanadium),
+            UncertaintyBroadcastMode.drop,
+        )
+        normed_b = normalize_by_vanadium_dspacing_and_two_theta(
+            FocussedDataDspacingTwoTheta[SampleRun](sample),
+            FocussedDataDspacingTwoTheta[VanadiumRun](vana_binned_like_sample),
+            UncertaintyBroadcastMode.drop,
+        )
+        sc.testing.assert_allclose(normed_a, normed_b)
+
+    def test_2d_masks_zero_vanadium_bins(self) -> None:
+        sample, vanadium = self.make_sample_and_vanadium_2d()
+        vanadium['dspacing', 5]['two_theta', 7] = sc.scalar(
+            0.0, variance=0.0, unit='counts'
+        )
+        normed = normalize_by_vanadium_dspacing_and_two_theta(
+            FocussedDataDspacingTwoTheta[SampleRun](sample),
+            FocussedDataDspacingTwoTheta[VanadiumRun](vanadium),
+            UncertaintyBroadcastMode.drop,
+        )
+
+        norm = vanadium.hist(
+            dspacing=sample.coords['dspacing'], two_theta=sample.coords['two_theta']
+        )
+
+        sc.testing.assert_allclose(
+            normed.masks['zero_vanadium'], norm.data == sc.scalar(0.0, unit='counts')
+        )