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

…so, add **kwds arguments to add_jitter

sed/config/default.yaml

@@ -31,6 +31,8 @@ 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"]
 
 energy:
   # Number of bins to use for energy calibration traces

sed/core/processor.py

@@ -1180,23 +1180,25 @@ def calibrate_delay_axis(
             else:
                 print(self._dataframe)
 
-    def add_jitter(self, cols: Sequence[str] = None):
+    def add_jitter(self, cols: List[str] = 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"].
+            **kwds: 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("@"))
 
         self._dataframe = self._dataframe.map_partitions(
             apply_jitter,
             cols=cols,
             cols_jittered=cols,
+            **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():