Merge pull request #152 from OpenCOMPES/flash_energy_calibration

Flash energy calibration

sed/config/default.yaml

@@ -11,6 +11,8 @@ dataframe:
   tof_column: "t"
   # dataframe column containing analog-to-digital data
   adc_column: "ADC"
+  # dataframe column containing bias voltage data
+  bias_column: "sampleBias"
   # dataframe column containing corrected x coordinates
   corrected_x_column: "Xm"
   # dataframe column containing corrected y coordinates

sed/core/processor.py

@@ -838,7 +838,8 @@ def apply_energy_correction(
     # 1. Load and normalize data
     def load_bias_series(
         self,
-        data_files: List[str],
+        binned_data: Union[xr.DataArray, Tuple[np.ndarray, np.ndarray, np.ndarray]] = None,
+        data_files: List[str] = None,
         axes: List[str] = None,
         bins: List = None,
         ranges: Sequence[Tuple[float, float]] = None,
@@ -849,10 +850,14 @@ def load_bias_series(
         order: int = None,
     ):
         """1. step of the energy calibration workflow: Load and bin data from
-        single-event files.
+        single-event files, or load binned bias/TOF traces.
 
         Args:
-            data_files (List[str]): list of file paths to bin
+            binned_data (Union[xr.DataArray, Tuple[np.ndarray, np.ndarray, np.ndarray]], optional):
+                Binned data If provided as DataArray, Needs to contain dimensions
+                config["dataframe"]["tof_column"] and config["dataframe"]["bias_column"]. If
+                provided as tuple, needs to contain elements tof, biases, traces.
+            data_files (List[str], optional): list of file paths to bin
             axes (List[str], optional): bin axes.
                 Defaults to config["dataframe"]["tof_column"].
             bins (List, optional): number of bins.
@@ -872,14 +877,44 @@ def load_bias_series(
                 (see ``scipy.signal.savgol_filter()``).
                 Defaults to config["energy"]["normalize_order"].
         """
-        self.ec.bin_data(
-            data_files=cast(List[str], self.cpy(data_files)),
-            axes=axes,
-            bins=bins,
-            ranges=ranges,
-            biases=biases,
-            bias_key=bias_key,
-        )
+        if binned_data is not None:
+            if isinstance(binned_data, xr.DataArray):
+                if (
+                    self._config["dataframe"]["tof_column"] not in binned_data.dims
+                    or self._config["dataframe"]["bias_column"] not in binned_data.dims
+                ):
+                    raise ValueError(
+                        "If binned_data is provided as an xarray, it needs to contain dimensions "
+                        f"'{self._config['dataframe']['tof_column']}' and "
+                        f"'{self._config['dataframe']['bias_column']}'!.",
+                    )
+                tof = binned_data.coords[self._config["dataframe"]["tof_column"]].values
+                biases = binned_data.coords[self._config["dataframe"]["bias_column"]].values
+                traces = binned_data.values[:, :]
+            else:
+                try:
+                    (tof, biases, traces) = binned_data
+                except ValueError as exc:
+                    raise ValueError(
+                        "If binned_data is provided as tuple, it needs to contain "
+                        "(tof, biases, traces)!",
+                    ) from exc
+            self.ec.load_data(biases=biases, traces=traces, tof=tof)
+
+        elif data_files is not None:
+
+            self.ec.bin_data(
+                data_files=cast(List[str], self.cpy(data_files)),
+                axes=axes,
+                bins=bins,
+                ranges=ranges,
+                biases=biases,
+                bias_key=bias_key,
+            )
+
+        else:
+            raise ValueError("Either binned_data or data_files needs to be provided!")
+
         if (normalize is not None and normalize is True) or (
             normalize is None and self._config["energy"]["normalize"]
         ):
@@ -1158,7 +1193,7 @@ def calibrate_delay_axis(
                         datafile = self._files[0]
                     except IndexError:
                         print(
-                            "No datafile available, specify eihter",
+                            "No datafile available, specify either",
                             " 'datafile' or 'delay_range'",
                         )
                         raise

tests/test_processor.py

@@ -15,6 +15,7 @@
 import dask.dataframe as ddf
 import numpy as np
 import pytest
+import xarray as xr
 
 from sed import SedProcessor
 from sed.core.config import parse_config
@@ -471,6 +472,8 @@ def test_energy_calibration_workflow(energy_scale: str, calibration_method: str)
         user_config={},
         system_config={},
     )
+    with pytest.raises(ValueError):
+        processor.load_bias_series()
     with pytest.raises(ValueError):
         processor.load_bias_series(data_files=glob.glob(df_folder + "../mpes/*.h5"), normalize=True)
     processor.load_bias_series(
@@ -479,8 +482,22 @@ def test_energy_calibration_workflow(energy_scale: str, calibration_method: str)
         bias_key="@KTOF:Lens:Sample:V",
     )
     assert len(processor.ec.biases) == 2
-    # load test data into class
-    processor.ec.load_data(biases=biases, traces=traces, tof=tof)
+    # load data as tuple
+    with pytest.raises(ValueError):
+        processor.load_bias_series(binned_data=(tof, traces))  # type: ignore
+    processor.load_bias_series(binned_data=(tof, biases, traces))
+    assert processor.ec.traces.shape == traces.shape
+    assert len(processor.ec.biases) == processor.ec.traces.shape[0]
+    assert len(processor.ec.tof) == processor.ec.traces.shape[1]
+    # load data as xarray
+    with pytest.raises(ValueError):
+        bias_series = xr.DataArray(data=traces, coords={"biases": biases, "tof": tof})
+        processor.load_bias_series(binned_data=bias_series)
+    bias_series = xr.DataArray(data=traces, coords={"sampleBias": biases, "t": tof})
+    processor.load_bias_series(binned_data=bias_series)
+    assert processor.ec.traces.shape == traces.shape
+    assert len(processor.ec.biases) == processor.ec.traces.shape[0]
+    assert len(processor.ec.tof) == processor.ec.traces.shape[1]
     processor.ec.normalize()
     ref_id = 5
     rng = (66100, 67000)

tutorial/Flash energy calibration.ipynb

@@ -0,0 +1,206 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "39b2e62a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "\n",
+    "from sed import SedProcessor\n",
+    "import sed\n",
+    "import numpy as np\n",
+    "\n",
+    "# %matplotlib inline\n",
+    "%matplotlib widget\n",
+    "import matplotlib.pyplot as plt"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4d78d236",
+   "metadata": {},
+   "source": [
+    "# Try to calibrate energy"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a62f084f",
+   "metadata": {},
+   "source": [
+    "## Spin-integrated branch, E_TOF=10eV\n",
+    "single scan, move sample bias manually every 2000 pulses."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7dabbe92",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sp = SedProcessor(runs=[44638], config=\"config_flash_energy_calib.yaml\", system_config={})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "248a41a7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sp.add_jitter()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2b867e40",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "axes = ['sampleBias', 'dldTime']\n",
+    "bins = [6, 500]\n",
+    "ranges = [[0,6], [40000, 55000]]\n",
+    "res = sp.compute(bins=bins, axes=axes, ranges=ranges)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "62081458",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sp.load_bias_series(binned_data=res)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "424af94e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ranges=(44500, 46000)\n",
+    "ref_id=3\n",
+    "sp.find_bias_peaks(ranges=ranges, ref_id=ref_id)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "034eff42",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ref_id=3\n",
+    "ref_energy=-.3\n",
+    "sp.calibrate_energy_axis(ref_id=ref_id, ref_energy=ref_energy, method=\"lstsq\", order=3)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bbbfe992",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ref_id=3\n",
+    "ref_energy=-.3\n",
+    "sp.calibrate_energy_axis(ref_id=ref_id, ref_energy=ref_energy, method=\"lmfit\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e14d6cef",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sp.append_energy_axis(preview=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "59c83544",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "axes = ['sampleBias', 'energy']\n",
+    "bins = [6, 1000]\n",
+    "ranges = [[0,6], [-5, 5]]\n",
+    "res = sp.compute(bins=bins, axes=axes, ranges=ranges)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "addba4cb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.figure()\n",
+    "res[3,:].plot()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1676ec57",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "axes = ['sampleBias', 'energy', 'dldPosX']\n",
+    "bins = [6, 100, 480]\n",
+    "ranges = [[0,6], [-2, 1], [420,900]]\n",
+    "res = sp.compute(bins=bins, axes=axes, ranges=ranges)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ad199c40",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.figure()\n",
+    "res[3, :, :].plot()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3a4ae88c",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".pyenv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}