Rebin vanadium if needed

Author: jl-wynen

src/ess/powder/correction.py

@@ -5,9 +5,10 @@
 import enum
 from typing import TypeVar
 
-import ess.reduce
 import sciline
 import scipp as sc
+
+import ess.reduce
 from ess.reduce.uncertainty import broadcast_uncertainties
 
 from ._util import event_or_outer_coord
@@ -151,9 +152,9 @@ def _normalize_by_vanadium(
     uncertainty_broadcast_mode: UncertaintyBroadcastMode,
 ) -> sc.DataArray:
     norm = (
-        vanadium.hist(dspacing=data.coords['dspacing'])
+        vanadium.hist(data.coords)
         if vanadium.is_binned
-        else vanadium.rebin(dspacing=data.coords['dspacing'])
+        else vanadium.rebin(data.coords)
     )
     norm = broadcast_uncertainties(
         norm, prototype=data, mode=uncertainty_broadcast_mode
@@ -177,8 +178,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
     ----------
@@ -196,6 +207,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)
@@ -207,9 +223,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
     ----------
@@ -227,6 +252,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

@@ -461,39 +461,57 @@ def test_normalize_by_monitor_integrated_assigns_mask_if_monitor_range_too_narro
         )
 
 
-def random_variable(
-    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(
-    rng: np.random.Generator,
-    dim: str,
-    n_events: int,
-    n_bins: int,
-    unit: str,
-    coord_unit: str,
-    with_variances: bool = False,
-) -> sc.DataArray:
-    return sc.DataArray(
-        random_variable(rng, 'event', n_events, unit, with_variances=with_variances),
-        coords={dim: random_variable(rng, 'event', n_events, coord_unit)},
-    ).bin({dim: n_bins})
-
-
-def make_sample_and_vanadium() -> tuple[sc.DataArray, sc.DataArray]:
-    rng = np.random.default_rng(seed=495)
-    sample = random_binned_data(rng, 'dspacing', 84, 35, 'counts', 'Å', True)
-    vanadium = random_binned_data(rng, 'dspacing', 146, 79, 'counts', 'Å', True)
-    return sample, vanadium
+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
 
-class TestNormalizeByVanadium:
     def test_1d_binned_vanadium(self) -> None:
-        sample, vanadium = make_sample_and_vanadium()
+        sample, vanadium = self.make_sample_and_vanadium_1d()
         normed = normalize_by_vanadium_dspacing(
             FocussedDataDspacing[SampleRun](sample),
             FocussedDataDspacing[VanadiumRun](vanadium),
@@ -507,7 +525,7 @@ def test_1d_binned_vanadium(self) -> None:
         sc.testing.assert_allclose(normed, expected)
 
     def test_1d_histogrammed_vanadium(self) -> None:
-        sample, vanadium = make_sample_and_vanadium()
+        sample, vanadium = self.make_sample_and_vanadium_1d()
         vanadium = vanadium.hist()
         normed = normalize_by_vanadium_dspacing(
             FocussedDataDspacing[SampleRun](sample),
@@ -521,8 +539,8 @@ def test_1d_histogrammed_vanadium(self) -> None:
         expected = sample / sc.values(norm)
         sc.testing.assert_allclose(normed, expected)
 
-    def test_binned_vanadium_binning_has_no_effect(self) -> None:
-        sample, vanadium = make_sample_and_vanadium()
+    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),
@@ -537,7 +555,7 @@ def test_binned_vanadium_binning_has_no_effect(self) -> None:
         sc.testing.assert_allclose(normed_a, normed_b)
 
     def test_1d_masks_zero_vanadium_bins(self) -> None:
-        sample, vanadium = make_sample_and_vanadium()
+        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),
@@ -551,4 +569,71 @@ def test_1d_masks_zero_vanadium_bins(self) -> None:
             normed.masks['zero_vanadium'], norm.data == sc.scalar(0.0, unit='counts')
         )
 
-# TODO vanadium and 2theta
+    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')
+        )