Merge pull request #124 from European-XFEL/scan-bin

Add support for binning data to the Scan component

docs/changelog.md

@@ -6,37 +6,45 @@
 # Changelog
 
 ## [Unreleased]
-Added:
 
-- `pulse_periods()`, `pulse_repetition_rates()` and `train_durations()` methods to obtain statistics about the pulses in all `PulsePattern`-based components
+!!! note
+    All of the changes here are deployed to our current environment, even though
+    a release hasn't been made for them yet. If you want to have these updates
+    in a personal environment you'll need to install the package from git.
+
+    ```bash title="Installation command"
+    pip install git+https://github.com/European-XFEL/EXtra.git
+    ```
 
-<!-- !!! note -->
-<!--     All of the changes here are deployed to our current environment, even though -->
-<!--     a release hasn't been made for them yet. If you want to have these updates -->
-<!--     in a personal environment you'll need to install the package from git. -->
+Added:
 
-<!--     ```bash title="Installation command" -->
-<!--     pip install git+https://github.com/European-XFEL/EXtra.git -->
-<!--     ``` -->
+- Implemented [Scan.bin_by_steps()][extra.components.Scan.bin_by_steps] and
+  [Scan.plot_bin_by_steps()][extra.components.Scan.plot_bin_by_steps] to help
+  with averaging data over scan steps (!124).
+- `pulse_periods()`, `pulse_repetition_rates()` and `train_durations()` methods
+  to obtain statistics about the pulses in all [pulse
+  pattern](components/pulse-patterns.md) components (!114).
 
 ## [2024.1]
 Added:
 
 - An [XGM][extra.components.XGM] component to access XGM devices (!53).
-- [PumpProbePulses][extra.components.PumpProbePulses] to combine X-ray and optical laser pulses in a single pattern (!24).
+- [PumpProbePulses][extra.components.PumpProbePulses] to combine X-ray and
+  optical laser pulses in a single pattern (!24).
 - The [Scan][extra.components.Scan] component to automatically detect steps
   within a motor scan (!4).
-- [DldPulses][extra.components.DldPulses] to access pulse information saved during delay line detector event reconstruction (!42).
-- The helper function [imshow2][extra.utils.imshow2] to provide good defaults when
-  plotting images (!38).
+- [DldPulses][extra.components.DldPulses] to access pulse information saved
+  during delay line detector event reconstruction (!42).
+- The helper function [imshow2][extra.utils.imshow2] to provide good defaults
+  when plotting images (!38).
 
 Changed:
 
 - The `get_` prefix was deprecated for some method names in the [pulse pattern
   components](components/pulse-patterns.md) (!106).
 - All methods in [XrayPulses][extra.components.XrayPulses] and
-  [OpticalLaserPulses][extra.components.OpticalLaserPulses] now support labelled results
-  and default to it (!40).
+  [OpticalLaserPulses][extra.components.OpticalLaserPulses] now support labelled
+  results and default to it (!40).
 - [Scantool][extra.components.Scantool]'s `__repr__()` functionality to print
   information was moved to [Scantool.info()][extra.components.Scantool.info]
   (!29).

src/extra/components/scan.py

@@ -147,7 +147,6 @@ def plot(self, ax=None):
 
         Example plot:
         ![](../images/scan-plot.png)
-
         """
         if ax is None:
             import matplotlib.pyplot as plt
@@ -163,6 +162,107 @@ def plot(self, ax=None):
 
         return ax
 
+    def bin_by_steps(self, data, uncertainty_method="std"):
+        """Average train-resolved data within each scan step.
+
+        This will return a [DataArray][xarray.DataArray] containing the averaged
+        value of `data` over each step in the scan, along with `position`,
+        `counts`, and `uncertainty` coordinates containing the
+        position/counts/uncertainty for each step. The `.name` property will be
+        set to `data.name`, and the following attributes will be set:
+
+        - `motor`: `scan.name`
+        - `uncertainty_method`: the `uncertainty_method` that was passed,
+          indicating either the standard deviation or standard error.
+
+        Args:
+            data (xarray.DataArray): A train-resolved (i.e. with a `trainId`
+                coordinate) array to bin.
+            uncertainty_method (str): Can be set to `std` to use the standard
+                deviation (i.e. the 1-sigma region) or `stderr` to use the
+                [standard error](https://en.wikipedia.org/wiki/Standard_error).
+        """
+        import xarray as xr
+        if not isinstance(data, xr.DataArray) or "trainId" not in data.coords:
+            raise TypeError("Input must be a DataArray with a `trainId` coordinate")
+        elif uncertainty_method != "std" and uncertainty_method != "stderr":
+            raise ValueError(f"`uncertainty_method` must be either 'std' or 'stderr', not: {uncertainty_method}")
+
+        common_tids = [np.intersect1d(tids, data.trainId)
+                       for tids in self.positions_train_ids]
+
+        counts = np.array([len(tids) for tids in common_tids])
+        signal = np.array([data.sel(trainId=tids).mean().item()
+                           for tids in common_tids])
+        uncertainty = np.array([data.sel(trainId=tids).std().item()
+                          for tids in common_tids])
+        if uncertainty_method == "stderr":
+            uncertainty /= np.sqrt(counts)
+
+        return xr.DataArray(signal, dims=("position",),
+                            coords={"position": ("position", self.positions),
+                                    "uncertainty": ("position", uncertainty),
+                                    "counts": ("position", counts)},
+                            name=data.name,
+                            attrs={
+                                "motor": self.name,
+                                "uncertainty_method": uncertainty_method
+                            })
+
+    def plot_bin_by_steps(self, data, uncertainty_method="std",
+                          title=None, xlabel=None, ylabel=None,
+                          ax=None, figsize=(9, 5)):
+        """Plot step-averaged data.
+
+        This calls [Scan.bin_by_steps()][extra.components.Scan.bin_by_steps] and
+        plots the result. Note that while it's possible to explicitly specify
+        the title/xlabel/ylabel, it's recommended to set `data.name` and
+        `scan.name` and let the plot settings be inferred automatically.
+
+        Example plot with `data.name == "ROI intensity"` and `scan.name ==
+        "Theta"`:
+        ![](../images/scan-plot-bin-by-steps.png)
+
+        Args:
+            data (xarray.DataArray): A train-resolved array to pass to
+                [Scan.bin_by_steps()][extra.components.Scan.bin_by_steps].
+            uncertainty_method (str): Same as in
+                [Scan.bin_by_steps()][extra.components.Scan.bin_by_steps].
+            title (str): The title of the plot.
+            xlabel (str): The xlabel of the plot.
+            ylabel (str): The ylabel of the plot.
+            ax (matplotlib.axes.Axes): The axes to plot on. A figure will be
+                created if this is not set.
+            figsize (tuple): The figure size. Only used if `ax=None`.
+
+        """
+        if ax is None:
+            import matplotlib.pyplot as plt
+            _, ax = plt.subplots(figsize=figsize)
+
+        binned_data = self.bin_by_steps(data)
+
+        uncertainty_label = "standard deviation" if uncertainty_method == "std" else "standard error"
+        ax.plot(binned_data.position, binned_data, "-o", markersize=4, label=f"Uncertainty: {uncertainty_label}")
+        ax.fill_between(binned_data.position,
+                        binned_data - binned_data.uncertainty,
+                        binned_data + binned_data.uncertainty,
+                        alpha=0.5)
+        ax.grid()
+        ax.set_xlabel(self.name)
+
+        if binned_data.name is not None:
+            ax.set_ylabel(binned_data.name)
+            yaxis = binned_data.name
+        else:
+            ax.set_ylabel("Signal [arb. u.]")
+            yaxis = "Signal"
+
+        ax.legend()
+        ax.set_title(f"{yaxis} vs {self.name}")
+
+        return ax
+
     def _plot_resolution_data(self):
         """Plot the data points that used to guess the resolution.
 

tests/test_components_scan.py

@@ -68,9 +68,32 @@ def test_scan(mock_spb_aux_run):
     # This should not throw, and should not detect any steps
     assert len(Scan(motor_slice).steps) == 0
 
+    # Test scan binning with some fake data
+    s, _ = Scan._mkscan(20)
+    data = xr.DataArray(np.random.rand(len(s._input_pos.trainId)),
+                        name="fake-data",
+                        dims=("trainId",),
+                        coords=dict(trainId=s._input_pos.trainId))
+
+    binned_data = s.bin_by_steps(data, uncertainty_method="stderr")
+    assert len(binned_data) == 20
+    assert set(binned_data.coords.keys()) == {"position", "uncertainty", "counts"}
+    assert binned_data.name == data.name
+    assert binned_data.attrs["motor"] == s.name
+
+    # Only 'std' and 'stderr' should be accepted as uncertainty methods
+    with pytest.raises(ValueError):
+        s.bin_by_steps(data, uncertainty_method="foo")
+
+    # Remove some trains from `data` to make sure they aren't used while binning
+    data_sel = data.sel(trainId=data.trainId[:-5])
+    # This should not throw an exception
+    s.bin_by_steps(data_sel)
+
     # Smoke tests
     s.plot()
     s._plot_resolution_data()
     repr(s)
     s.format()
     s.info()
+    s.plot_bin_by_steps(data)