diff --git a/sed/config/default.yaml b/sed/config/default.yaml index 85066a04..2972c345 100644 --- a/sed/config/default.yaml +++ b/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 diff --git a/sed/core/dfops.py b/sed/core/dfops.py index ecef954f..b7dbd21e 100644 --- a/sed/core/dfops.py +++ b/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". diff --git a/sed/core/processor.py b/sed/core/processor.py index 9069215c..93a0413f 100644 --- a/sed/core/processor.py +++ b/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: diff --git a/tests/test_processor.py b/tests/test_processor.py index fba83d77..270d1614 100644 --- a/tests/test_processor.py +++ b/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():