Merge pull request #168 from OpenCOMPES/default_jitter_cols

add default jitter axes to default config, and add tests for that.

sed/config/default.yaml

@@ -31,6 +31,10 @@ dataframe:
   tof_binning: 1
   # binning factor used for the adc coordinate (2^(adc_binning-1))
   adc_binning: 1
+  # list of columns to apply jitter to.
+  jitter_cols: ["@x_column", "@y_column", "@tof_column"]
+  # Jitter amplitude or list of jitter amplitudes. Should equal half the digitial step size of each jitter_column
+  jitter_amps: [0.5, 0.5, 0.5]
 
 energy:
   # Number of bins to use for energy calibration traces

sed/core/dfops.py

@@ -30,6 +30,8 @@ def apply_jitter(
             with added jitter. Defaults to None.
         amps (Union[float, Sequence[float]], optional): Amplitude scalings for the
             jittering noise. If one number is given, the same is used for all axes.
+            For normal noise, the added noise will have sdev [-amp, +amp], for
+            uniform noise it will cover the interval [-amp, +amp].
             Defaults to 0.5.
         jitter_type (str, optional): the type of jitter to add. 'uniform' or 'normal'
             distributed noise. Defaults to "uniform".

sed/core/processor.py

@@ -1180,23 +1180,38 @@ def calibrate_delay_axis(
             else:
                 print(self._dataframe)
 
-    def add_jitter(self, cols: Sequence[str] = None):
+    def add_jitter(
+        self,
+        cols: List[str] = None,
+        amps: Union[float, Sequence[float]] = None,
+        **kwds,
+    ):
         """Add jitter to the selected dataframe columns.
 
         Args:
-            cols (Sequence[str], optional): The colums onto which to apply jitter.
+            cols (List[str], optional): The colums onto which to apply jitter.
                 Defaults to config["dataframe"]["jitter_cols"].
+            amps (Union[float, Sequence[float]], optional): Amplitude scalings for the
+                jittering noise. If one number is given, the same is used for all axes.
+                For uniform noise (default) it will cover the interval [-amp, +amp].
+                Defaults to config["dataframe"]["jitter_amps"].
+            **kwds: additional keyword arguments passed to apply_jitter
         """
         if cols is None:
-            cols = self._config["dataframe"].get(
-                "jitter_cols",
-                self._dataframe.columns,
-            )  # jitter all columns
+            cols = self._config["dataframe"]["jitter_cols"]
+        for loc, col in enumerate(cols):
+            if col.startswith("@"):
+                cols[loc] = self._config["dataframe"].get(col.strip("@"))
+
+        if amps is None:
+            amps = self._config["dataframe"]["jitter_amps"]
 
         self._dataframe = self._dataframe.map_partitions(
             apply_jitter,
             cols=cols,
             cols_jittered=cols,
+            amps=amps,
+            **kwds,
         )
         metadata = []
         for col in cols:

tests/test_processor.py

@@ -581,11 +581,15 @@ def test_add_jitter():
         system_config={},
     )
     res1 = processor.dataframe["X"].compute()
+    res1a = processor.dataframe["ADC"].compute()
     processor.add_jitter()
     res2 = processor.dataframe["X"].compute()
+    res2a = processor.dataframe["ADC"].compute()
     np.testing.assert_allclose(res1, np.round(res1))
     np.testing.assert_allclose(res1, np.round(res2))
     assert (res1 != res2).all()
+    # test that jittering is not applied on ADC column
+    np.testing.assert_allclose(res1a, res2a)
 
 
 def test_event_histogram():