From 5847398c81c1eef477926a7e505ef853bb85a8d3 Mon Sep 17 00:00:00 2001
From: sjoerd-bouma <sjoerdbouma@xs4all.nl>
Date: Thu, 11 Aug 2022 16:01:58 +0200
Subject: [PATCH 01/19] added block offset fitter for RNO-G

---
 .../modules/RNO_G/channelBlockOffsetFitter.py | 240 ++++++++++++++++++
 1 file changed, 240 insertions(+)
 create mode 100644 NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py

diff --git a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
new file mode 100644
index 000000000..67412e62f
--- /dev/null
+++ b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
@@ -0,0 +1,240 @@
+from NuRadioReco.utilities import units
+from NuRadioReco.framework.base_trace import BaseTrace
+# from iminuit import Minuit
+import numpy as np
+import scipy
+
+class channelBlockOffsets:
+
+    def __init__(self, block_size=128, max_frequency=51*units.MHz):
+        """
+        Initialize block offset fitter.
+
+        Parameters:
+        -----------
+        block_size: int (default: 128)
+            The size (in samples) of the blocks
+        max_frequency: float (default: 51 MHz)
+            The maximum frequency to include in the out-of-band
+            block offset fit
+
+        """
+        self.sampling_rate = None
+        self.block_size = block_size # the size (in samples) of the blocks
+        self._offset_guess = dict()
+        self._offset_fit = dict()
+        self._offset_inject = dict()
+        self._max_frequency = max_frequency
+
+    def add_offsets(self, event, station, offsets=1*units.mV, channel_ids=None):
+        """
+        Add (simulated or reconstructed) block offsets to an event.
+
+        Parameters
+        ----------
+        event: `NuRadioReco.framework.event.Event` | None
+        station: `NuRadioReco.framework.station.Station`
+            The station to add block offsets to
+        offsets: float | array | dict
+            offsets to add to the event. Default: 1 mV
+
+            - if a float, add gaussian-distributed of amplitude `offsets`
+              to all channels specified;
+            - if an array, the length should be the same as the number
+              of blocks in a single trace, and the entries will be
+              interpreted as the amplitudes of the offsets;
+            - if a dict, the keys should be the channel ids, and each
+              value should contain either a float or an array to add to
+              each channel as specified above.
+
+        channel_ids: list | None
+            either a list of channel ids to apply the offsets to, or
+            None to apply the offsets to all channels in the station
+            (default: None).
+
+        """
+
+        if channel_ids is None:
+            channel_ids = station.get_channel_ids()
+        for channel_id in channel_ids:
+            channel = station.get_channel(channel_id)
+            if isinstance(offsets, dict):
+                add_offsets = offsets[channel_id]
+            elif len(np.atleast_1d(offsets)) == 1:
+                add_offsets = np.random.normal(
+                    0, offsets, (channel.get_number_of_samples() // self.block_size)
+                )
+            else:
+                add_offsets = offsets
+            if channel_id in self._offset_inject.keys():
+                self._offset_inject[channel_id] += add_offsets
+            else:
+                self._offset_inject[channel_id] = add_offsets
+
+            channel.set_trace(
+                channel.get_trace() + np.repeat(add_offsets, self.block_size),
+                channel.get_sampling_rate()
+            )
+
+    def remove_offsets(self, event, station, offsets='fit', channel_ids=None):
+        """
+        Remove block offsets from an event
+
+        Fits and removes the block offsets from an event.
+
+        Parameters
+        ----------
+        event: `NuRadioReco.framework.event.Event` | None
+        station: `NuRadioReco.framework.station.Station`
+            The station to remove the block offsets from
+        offsets: str
+            How to remove the offsets. Options are:
+
+            - 'fit': fit the offsets out of band
+            - 'guess': similar to 'fit', but just take
+              a first guess at the offsets from the out-of-band region
+              without actually performing the fit.
+            - 'injected': if offsets were injected using the `add_offsets`
+              method, this removes those offsets. Otherwise, this does nothing.
+
+            Default: 'fit'
+        channel_ids: list | None
+            List of channel ids to remove offsets from. If None (default),
+            remove offsets from all channels in `station`
+
+        """
+        if offsets=='fit':
+            if not len(self._offset_fit):
+                self.fit_offsets(event, station, channel_ids)
+            offsets = self._offset_fit
+        elif offsets=='guess':
+            offsets = self._offset_guess
+        elif offsets=='injected':
+            if not len(self._offset_inject):
+                offsets = np.zeros(16) #TODO - ensure this works for different trace lengths
+            else:
+                offsets = self._offset_inject
+
+        if isinstance(offsets, dict):
+            remove_offsets = {key: -offsets[key] for key in offsets.keys()}
+        else:
+            remove_offsets = -offsets
+        self.add_offsets(event, station, remove_offsets, channel_ids)
+
+    def fit_offsets(self, event, station, channel_ids=None):
+        """
+        Fit the block offsets using an out-of-band fit
+
+        This function fits the block offsets present in a given
+        event / station using an out-of-band fit in frequency space.
+
+        Parameters
+        ----------
+        event: `NuRadioReco.framework.event.Event` | None
+        station: `NuRadioReco.framework.station.Station`
+            The station to fit the block offsets to
+        channel_ids: list | None
+            List of channel ids to fit block offsets for. If None (default),
+            fit offsets for all channels in `station`
+
+        """
+        block_size = self.block_size
+        if channel_ids is None:
+            channel_ids = station.get_channel_ids()
+        for channel_id in channel_ids:
+            channel = station.get_channel(channel_id)
+            trace = channel.get_trace()
+
+            # we use the bandpass-filtered trace to get a first estimate of
+            # the block offsets, by simply averaging over each block.
+            filtered_trace = channel.get_filtered_trace([0, self._max_frequency], 'rectangular')
+            self.sampling_rate = channel.get_sampling_rate()
+            n_blocks = len(trace) // block_size
+
+            # obtain guesses for block offsets
+            a_guess = np.array([
+                np.mean(filtered_trace[i*block_size:(i+1)*block_size])
+                for i in range(n_blocks)
+            ])
+            self._offset_guess[channel_id] = a_guess
+            # we can get rid of one parameter through a global shift
+            a_guess = a_guess[:-1] - a_guess[-1]
+            frequencies = channel.get_frequencies()
+            spectrum = channel.get_frequency_spectrum()
+
+            # we perform the fit out-of-band, in order to avoid
+            # distorting any actual signal
+            mask = (frequencies > 0) & (frequencies < self._max_frequency)
+
+            frequencies_oob = frequencies[mask]
+            spectrum_oob = spectrum[mask]
+
+            ns = self.block_size
+            dt = 1/self.sampling_rate
+
+            # most of the terms in the fit depend only on the frequencies,
+            # sampling rate and number of blocks. We therefore calculate these
+            # only once, outside the fit function.
+            pre_factor_exponent = np.array([
+                -2.j * np.pi * frequencies_oob * dt * ((j+.5) * ns - .5)
+                 for j in range(len(a_guess))
+            ])
+            self._const_fft_term = (
+                 1 / self.sampling_rate * np.sqrt(2) # NuRadio FFT normalization
+                * np.exp(pre_factor_exponent)
+                * np.sin(np.pi*frequencies_oob*ns*dt)[None]
+                / np.sin(np.pi*frequencies_oob*dt)[None]
+            )
+            self._spectrum = spectrum_oob
+            res = scipy.optimize.fmin(
+                self._pedestal_fit, a_guess, disp=0)
+            ### Minuit seems a lot quicker than scipy - not sure why?
+            # m = Minuit(self._pedestal_fit_minuit, a_guess * nufft_conversion_factor)
+            # m.errordef = 1
+            # m.errors = 0.01 * np.ones_like(a_guess)
+            # m.migrad(ncall=20000)
+            # res = m.values
+            block_offsets = np.zeros(len(res) + 1)
+            block_offsets[:-1] = res
+            # the fit is not sensitive to an overall shift,
+            # so we include the zero-meaning here
+            block_offsets += np.mean(trace) - np.mean(block_offsets)
+
+            self._offset_fit[channel_id] = block_offsets
+
+    def get_offsets(self, channel_id, offset_type='fit'):
+        """
+        Return the block offsets for a given channel.
+
+        Parameters
+        ----------
+        channel_id: int
+            channel id that specifies the channel to return block offsets for
+        offset_type: str
+            Options:
+
+            - 'fit': return the fitted block offsets
+            - 'guess': return the first guess for the block offsets
+            - 'injected': return the block offsets that were injected
+              using the `add_offsets` method.
+
+        Returns
+        -------
+        trace: `BaseTrace`
+            A `BaseTrace` object with the same length as the channel trace,
+            containing only the block offsets.
+
+        """
+        trace = BaseTrace()
+        if offset_type == 'fit':
+            trace.set_trace(np.repeat(self._offset_fit[channel_id], self.block_size), self.sampling_rate)
+        elif offset_type == 'guess':
+            trace.set_trace(np.repeat(self._offset_guess[channel_id], self.block_size), self.sampling_rate)
+        elif offset_type == 'injected':
+            trace.set_trace(np.repeat(self._offset_inject[channel_id], self.block_size), self.sampling_rate)
+        return trace
+
+    def _pedestal_fit(self, a):
+        fit = np.sum(a[:, None] * self._const_fft_term, axis=0)
+        chi2 = np.sum(np.abs(fit-self._spectrum)**2)
+        return chi2
\ No newline at end of file

From 27babda38d9ccfa66e5e1a0d8f2ba8b83065d9b9 Mon Sep 17 00:00:00 2001
From: sjoerd-bouma <sjoerdbouma@xs4all.nl>
Date: Thu, 11 Aug 2022 16:36:35 +0200
Subject: [PATCH 02/19] fix docstrings

---
 .../modules/RNO_G/channelBlockOffsetFitter.py | 35 ++++++++++---------
 1 file changed, 19 insertions(+), 16 deletions(-)

diff --git a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
index 67412e62f..a4444b535 100644
--- a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
+++ b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
@@ -8,10 +8,13 @@ class channelBlockOffsets:
 
     def __init__(self, block_size=128, max_frequency=51*units.MHz):
         """
-        Initialize block offset fitter.
+        Add or remove block offsets to channel traces
 
-        Parameters:
-        -----------
+        This module adds, fits or removes 'block offsets' by fitting
+        them in a specified out-of-band region in frequency space.
+
+        Parameters
+        ----------
         block_size: int (default: 128)
             The size (in samples) of the blocks
         max_frequency: float (default: 51 MHz)
@@ -32,13 +35,13 @@ def add_offsets(self, event, station, offsets=1*units.mV, channel_ids=None):
 
         Parameters
         ----------
-        event: `NuRadioReco.framework.event.Event` | None
-        station: `NuRadioReco.framework.station.Station`
+        event: Event object | None
+        station: Station
             The station to add block offsets to
         offsets: float | array | dict
             offsets to add to the event. Default: 1 mV
 
-            - if a float, add gaussian-distributed of amplitude `offsets`
+            - if a float, add gaussian-distributed of amplitude ``offsets``
               to all channels specified;
             - if an array, the length should be the same as the number
               of blocks in a single trace, and the entries will be
@@ -84,8 +87,8 @@ def remove_offsets(self, event, station, offsets='fit', channel_ids=None):
 
         Parameters
         ----------
-        event: `NuRadioReco.framework.event.Event` | None
-        station: `NuRadioReco.framework.station.Station`
+        event: NuRadioReco.framework.event.Event | None
+        station: NuRadioReco.framework.station.Station
             The station to remove the block offsets from
         offsets: str
             How to remove the offsets. Options are:
@@ -94,13 +97,13 @@ def remove_offsets(self, event, station, offsets='fit', channel_ids=None):
             - 'guess': similar to 'fit', but just take
               a first guess at the offsets from the out-of-band region
               without actually performing the fit.
-            - 'injected': if offsets were injected using the `add_offsets`
+            - 'injected': if offsets were injected using the ``add_offsets``
               method, this removes those offsets. Otherwise, this does nothing.
 
             Default: 'fit'
         channel_ids: list | None
             List of channel ids to remove offsets from. If None (default),
-            remove offsets from all channels in `station`
+            remove offsets from all channels in ``station``
 
         """
         if offsets=='fit':
@@ -130,12 +133,12 @@ def fit_offsets(self, event, station, channel_ids=None):
 
         Parameters
         ----------
-        event: `NuRadioReco.framework.event.Event` | None
-        station: `NuRadioReco.framework.station.Station`
+        event: NuRadioReco.framework.event.Event | None
+        station: NuRadioReco.framework.station.Station
             The station to fit the block offsets to
         channel_ids: list | None
             List of channel ids to fit block offsets for. If None (default),
-            fit offsets for all channels in `station`
+            fit offsets for all channels in ``station``
 
         """
         block_size = self.block_size
@@ -216,12 +219,12 @@ def get_offsets(self, channel_id, offset_type='fit'):
             - 'fit': return the fitted block offsets
             - 'guess': return the first guess for the block offsets
             - 'injected': return the block offsets that were injected
-              using the `add_offsets` method.
+              using the ``add_offsets`` method.
 
         Returns
         -------
-        trace: `BaseTrace`
-            A `BaseTrace` object with the same length as the channel trace,
+        trace: BaseTrace
+            A :class:`NuRadioReco.framework.base_trace.BaseTrace` object with the same length as the channel trace,
             containing only the block offsets.
 
         """

From 37cd0702908040fc19f5fcf3f0ddf20a664374b0 Mon Sep 17 00:00:00 2001
From: sjoerd-bouma <sjoerdbouma@xs4all.nl>
Date: Sat, 22 Oct 2022 13:41:29 +0200
Subject: [PATCH 03/19] changed block offset fit function to standalone

---
 .../modules/RNO_G/channelBlockOffsetFitter.py | 178 ++++++++++++------
 1 file changed, 118 insertions(+), 60 deletions(-)

diff --git a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
index a4444b535..c873f9d96 100644
--- a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
+++ b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
@@ -1,4 +1,4 @@
-from NuRadioReco.utilities import units
+from NuRadioReco.utilities import units, fft
 from NuRadioReco.framework.base_trace import BaseTrace
 # from iminuit import Minuit
 import numpy as np
@@ -24,7 +24,6 @@ def __init__(self, block_size=128, max_frequency=51*units.MHz):
         """
         self.sampling_rate = None
         self.block_size = block_size # the size (in samples) of the blocks
-        self._offset_guess = dict()
         self._offset_fit = dict()
         self._offset_inject = dict()
         self._max_frequency = max_frequency
@@ -148,61 +147,10 @@ def fit_offsets(self, event, station, channel_ids=None):
             channel = station.get_channel(channel_id)
             trace = channel.get_trace()
 
-            # we use the bandpass-filtered trace to get a first estimate of
-            # the block offsets, by simply averaging over each block.
-            filtered_trace = channel.get_filtered_trace([0, self._max_frequency], 'rectangular')
-            self.sampling_rate = channel.get_sampling_rate()
-            n_blocks = len(trace) // block_size
-
-            # obtain guesses for block offsets
-            a_guess = np.array([
-                np.mean(filtered_trace[i*block_size:(i+1)*block_size])
-                for i in range(n_blocks)
-            ])
-            self._offset_guess[channel_id] = a_guess
-            # we can get rid of one parameter through a global shift
-            a_guess = a_guess[:-1] - a_guess[-1]
-            frequencies = channel.get_frequencies()
-            spectrum = channel.get_frequency_spectrum()
-
-            # we perform the fit out-of-band, in order to avoid
-            # distorting any actual signal
-            mask = (frequencies > 0) & (frequencies < self._max_frequency)
-
-            frequencies_oob = frequencies[mask]
-            spectrum_oob = spectrum[mask]
-
-            ns = self.block_size
-            dt = 1/self.sampling_rate
-
-            # most of the terms in the fit depend only on the frequencies,
-            # sampling rate and number of blocks. We therefore calculate these
-            # only once, outside the fit function.
-            pre_factor_exponent = np.array([
-                -2.j * np.pi * frequencies_oob * dt * ((j+.5) * ns - .5)
-                 for j in range(len(a_guess))
-            ])
-            self._const_fft_term = (
-                 1 / self.sampling_rate * np.sqrt(2) # NuRadio FFT normalization
-                * np.exp(pre_factor_exponent)
-                * np.sin(np.pi*frequencies_oob*ns*dt)[None]
-                / np.sin(np.pi*frequencies_oob*dt)[None]
+            block_offsets = fit_block_offsets(
+                trace, block_size,
+                channel.get_sampling_rate(), self._max_frequency
             )
-            self._spectrum = spectrum_oob
-            res = scipy.optimize.fmin(
-                self._pedestal_fit, a_guess, disp=0)
-            ### Minuit seems a lot quicker than scipy - not sure why?
-            # m = Minuit(self._pedestal_fit_minuit, a_guess * nufft_conversion_factor)
-            # m.errordef = 1
-            # m.errors = 0.01 * np.ones_like(a_guess)
-            # m.migrad(ncall=20000)
-            # res = m.values
-            block_offsets = np.zeros(len(res) + 1)
-            block_offsets[:-1] = res
-            # the fit is not sensitive to an overall shift,
-            # so we include the zero-meaning here
-            block_offsets += np.mean(trace) - np.mean(block_offsets)
-
             self._offset_fit[channel_id] = block_offsets
 
     def get_offsets(self, channel_id, offset_type='fit'):
@@ -217,7 +165,6 @@ def get_offsets(self, channel_id, offset_type='fit'):
             Options:
 
             - 'fit': return the fitted block offsets
-            - 'guess': return the first guess for the block offsets
             - 'injected': return the block offsets that were injected
               using the ``add_offsets`` method.
 
@@ -231,8 +178,6 @@ def get_offsets(self, channel_id, offset_type='fit'):
         trace = BaseTrace()
         if offset_type == 'fit':
             trace.set_trace(np.repeat(self._offset_fit[channel_id], self.block_size), self.sampling_rate)
-        elif offset_type == 'guess':
-            trace.set_trace(np.repeat(self._offset_guess[channel_id], self.block_size), self.sampling_rate)
         elif offset_type == 'injected':
             trace.set_trace(np.repeat(self._offset_inject[channel_id], self.block_size), self.sampling_rate)
         return trace
@@ -240,4 +185,117 @@ def get_offsets(self, channel_id, offset_type='fit'):
     def _pedestal_fit(self, a):
         fit = np.sum(a[:, None] * self._const_fft_term, axis=0)
         chi2 = np.sum(np.abs(fit-self._spectrum)**2)
-        return chi2
\ No newline at end of file
+        return chi2
+
+def fit_block_offsets(
+        trace, block_size=128, sampling_rate=3.2*units.GHz,
+        max_frequency=50*units.MHz, return_trace = False,
+        xtol=1e-6, maxiter=100000):
+    """
+    Fit 'block' offsets for a voltage trace
+
+    Fit block offsets ('rect'-shaped offsets from a baseline)
+    using a fit to the out-of-band spectrum of a voltage trace.
+
+    Parameters
+    ----------
+    trace: numpy Array
+        the voltage trace
+    block_size: int (default: 128)
+        the number of samples in one block
+    sampling_rate: float (default: 3.2 GHz)
+        the sampling rate of the trace
+    max_frequency: float (default: 50 MHz)
+        the fit to the block offsets is performed
+        in the frequency domain, in the band up to
+        max_frequency
+    return_trace: bool (default: False)
+        if True, return the tuple (offsets, output_trace)
+        where the output_trace is the input trace with
+        fitted block offsets removed
+
+    Returns
+    -------
+    block_offsets: numpy array
+        The fitted block offsets.
+    output_trace: numpy array or None
+        The input trace with the fitted block offsets removed.
+        Returned only if return_trace=True
+
+    Other Parameters
+    ----------------
+    xtol: float (default: 1e-6)
+        tolerance parameter passed on to scipy.optimize.fmin
+    maxiter: int (default: 100000)
+        maximum number of iterations for scipy.optimize.fmin
+
+    """
+    dt = 1. / sampling_rate
+    spectrum = fft.time2freq(trace, sampling_rate)
+    frequencies = np.fft.rfftfreq(len(trace), dt)
+    n_blocks = len(trace) // block_size
+
+    mask = (frequencies > 0) & (frequencies < max_frequency) #  a simple rectangular filter
+    frequencies_oob = frequencies[mask]
+    spectrum_oob = spectrum[mask]
+
+    # we use the bandpass-filtered trace to get a first estimate of
+    # the block offsets, by simply averaging over each block.
+    filtered_trace_fft = np.copy(spectrum)
+    filtered_trace_fft[~mask] = 0
+    filtered_trace = fft.freq2time(filtered_trace_fft, dt)
+
+    # obtain guesses for block offsets
+    a_guess = np.array([
+        np.mean(filtered_trace[i*block_size:(i+1)*block_size])
+        for i in range(n_blocks)
+    ])
+    # self._offset_guess[channel_id] = a_guess
+    # we can get rid of one parameter through a global shift
+    a_guess = a_guess[:-1] - a_guess[-1]
+
+    # we perform the fit out-of-band, in order to avoid
+    # distorting any actual signal
+
+    # most of the terms in the fit depend only on the frequencies,
+    # sampling rate and number of blocks. We therefore calculate these
+    # only once, outside the fit function.
+    pre_factor_exponent = np.array([
+        -2.j * np.pi * frequencies_oob * dt * ((j+.5) * block_size - .5)
+            for j in range(len(a_guess))
+    ])
+    const_fft_term = (
+            1 / sampling_rate * np.sqrt(2) # NuRadio FFT normalization
+        * np.exp(pre_factor_exponent)
+        * np.sin(np.pi*frequencies_oob*block_size*dt)[None]
+        / np.sin(np.pi*frequencies_oob*dt)[None]
+    )
+
+    def pedestal_fit(self, a):
+        fit = np.sum(a[:, None] * const_fft_term, axis=0)
+        chi2 = np.sum(np.abs(fit-spectrum_oob)**2)
+        return chi2
+
+    # self._spectrum = spectrum_oob
+    res = scipy.optimize.fmin(
+        pedestal_fit, a_guess, disp=0,
+        xtol=xtol, maxiter=maxiter)
+    ### maybe TODO - include option to use Minuit, which seems a lot quicker?
+    # m = Minuit(self._pedestal_fit_minuit, a_guess * nufft_conversion_factor)
+    # m.errordef = 1
+    # m.errors = 0.01 * np.ones_like(a_guess)
+    # m.migrad(ncall=20000)
+    # res = m.values
+
+    block_offsets = np.zeros(len(res) + 1)
+    block_offsets[:-1] = res
+
+    # the fit is not sensitive to an overall shift,
+    # so we include the zero-meaning here
+    block_offsets += np.mean(trace) - np.mean(block_offsets)
+
+    if return_trace:
+        output_trace = trace - np.repeat(block_offsets, block_size)
+        return block_offsets, output_trace
+
+    return block_offsets
\ No newline at end of file

From 9bc2e3f8518e4b6b4d48616340ad606117912f31 Mon Sep 17 00:00:00 2001
From: sjoerd-bouma <sjoerdbouma@xs4all.nl>
Date: Mon, 24 Oct 2022 13:26:46 +0200
Subject: [PATCH 04/19] fixed bug (dt instead of sampling rate)

---
 NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
index c873f9d96..eeb34e87d 100644
--- a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
+++ b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
@@ -93,9 +93,6 @@ def remove_offsets(self, event, station, offsets='fit', channel_ids=None):
             How to remove the offsets. Options are:
 
             - 'fit': fit the offsets out of band
-            - 'guess': similar to 'fit', but just take
-              a first guess at the offsets from the out-of-band region
-              without actually performing the fit.
             - 'injected': if offsets were injected using the ``add_offsets``
               method, this removes those offsets. Otherwise, this does nothing.
 
@@ -109,8 +106,6 @@ def remove_offsets(self, event, station, offsets='fit', channel_ids=None):
             if not len(self._offset_fit):
                 self.fit_offsets(event, station, channel_ids)
             offsets = self._offset_fit
-        elif offsets=='guess':
-            offsets = self._offset_guess
         elif offsets=='injected':
             if not len(self._offset_inject):
                 offsets = np.zeros(16) #TODO - ensure this works for different trace lengths
@@ -243,7 +238,7 @@ def fit_block_offsets(
     # the block offsets, by simply averaging over each block.
     filtered_trace_fft = np.copy(spectrum)
     filtered_trace_fft[~mask] = 0
-    filtered_trace = fft.freq2time(filtered_trace_fft, dt)
+    filtered_trace = fft.freq2time(filtered_trace_fft, sampling_rate)
 
     # obtain guesses for block offsets
     a_guess = np.array([
@@ -271,7 +266,7 @@ def fit_block_offsets(
         / np.sin(np.pi*frequencies_oob*dt)[None]
     )
 
-    def pedestal_fit(self, a):
+    def pedestal_fit(a):
         fit = np.sum(a[:, None] * const_fft_term, axis=0)
         chi2 = np.sum(np.abs(fit-spectrum_oob)**2)
         return chi2

From c0057268906e6dff69009327a06dfa8ab4ca7cb2 Mon Sep 17 00:00:00 2001
From: Sjoerd Bouma <sjoerdbouma@xs4all.nl>
Date: Wed, 19 Apr 2023 13:18:50 +0200
Subject: [PATCH 05/19] change to smarter minimization

---
 .../modules/RNO_G/channelBlockOffsetFitter.py | 104 +++++++++---------
 1 file changed, 54 insertions(+), 50 deletions(-)

diff --git a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
index eeb34e87d..4467fd0d8 100644
--- a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
+++ b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
@@ -184,8 +184,8 @@ def _pedestal_fit(self, a):
 
 def fit_block_offsets(
         trace, block_size=128, sampling_rate=3.2*units.GHz,
-        max_frequency=50*units.MHz, return_trace = False,
-        xtol=1e-6, maxiter=100000):
+        max_frequency=50*units.MHz, mode='fit', return_trace = False,
+        tol=1e-6,):
     """
     Fit 'block' offsets for a voltage trace
 
@@ -204,6 +204,10 @@ def fit_block_offsets(
         the fit to the block offsets is performed
         in the frequency domain, in the band up to
         max_frequency
+    mode: 'fit' | 'approximate'
+        Whether to fit the block offsets (default)
+        or just use the first guess from the out-of-band
+        component (faster)
     return_trace: bool (default: False)
         if True, return the tuple (offsets, output_trace)
         where the output_trace is the input trace with
@@ -219,11 +223,8 @@ def fit_block_offsets(
 
     Other Parameters
     ----------------
-    xtol: float (default: 1e-6)
-        tolerance parameter passed on to scipy.optimize.fmin
-    maxiter: int (default: 100000)
-        maximum number of iterations for scipy.optimize.fmin
-
+    tol: float (default: 1e-6)
+        tolerance parameter passed on to scipy.optimize.minimize
     """
     dt = 1. / sampling_rate
     spectrum = fft.time2freq(trace, sampling_rate)
@@ -245,49 +246,52 @@ def fit_block_offsets(
         np.mean(filtered_trace[i*block_size:(i+1)*block_size])
         for i in range(n_blocks)
     ])
-    # self._offset_guess[channel_id] = a_guess
-    # we can get rid of one parameter through a global shift
-    a_guess = a_guess[:-1] - a_guess[-1]
-
-    # we perform the fit out-of-band, in order to avoid
-    # distorting any actual signal
-
-    # most of the terms in the fit depend only on the frequencies,
-    # sampling rate and number of blocks. We therefore calculate these
-    # only once, outside the fit function.
-    pre_factor_exponent = np.array([
-        -2.j * np.pi * frequencies_oob * dt * ((j+.5) * block_size - .5)
-            for j in range(len(a_guess))
-    ])
-    const_fft_term = (
-            1 / sampling_rate * np.sqrt(2) # NuRadio FFT normalization
-        * np.exp(pre_factor_exponent)
-        * np.sin(np.pi*frequencies_oob*block_size*dt)[None]
-        / np.sin(np.pi*frequencies_oob*dt)[None]
-    )
-
-    def pedestal_fit(a):
-        fit = np.sum(a[:, None] * const_fft_term, axis=0)
-        chi2 = np.sum(np.abs(fit-spectrum_oob)**2)
-        return chi2
-
-    # self._spectrum = spectrum_oob
-    res = scipy.optimize.fmin(
-        pedestal_fit, a_guess, disp=0,
-        xtol=xtol, maxiter=maxiter)
-    ### maybe TODO - include option to use Minuit, which seems a lot quicker?
-    # m = Minuit(self._pedestal_fit_minuit, a_guess * nufft_conversion_factor)
-    # m.errordef = 1
-    # m.errors = 0.01 * np.ones_like(a_guess)
-    # m.migrad(ncall=20000)
-    # res = m.values
-
-    block_offsets = np.zeros(len(res) + 1)
-    block_offsets[:-1] = res
-
-    # the fit is not sensitive to an overall shift,
-    # so we include the zero-meaning here
-    block_offsets += np.mean(trace) - np.mean(block_offsets)
+    if mode == 'approximate':
+        block_offsets = a_guess
+    elif mode == 'fit':
+        # self._offset_guess[channel_id] = a_guess
+        # we can get rid of one parameter through a global shift
+        a_guess = a_guess[:-1] - a_guess[-1]
+
+        # we perform the fit out-of-band, in order to avoid
+        # distorting any actual signal
+
+        # most of the terms in the fit depend only on the frequencies,
+        # sampling rate and number of blocks. We therefore calculate these
+        # only once, outside the fit function.
+        pre_factor_exponent = np.array([
+            -2.j * np.pi * frequencies_oob * dt * ((j+.5) * block_size - .5)
+                for j in range(len(a_guess))
+        ])
+        const_fft_term = (
+                1 / sampling_rate * np.sqrt(2) # NuRadio FFT normalization
+            * np.exp(pre_factor_exponent)
+            * np.sin(np.pi*frequencies_oob*block_size*dt)[None]
+            / np.sin(np.pi*frequencies_oob*dt)[None]
+        )
+
+        def pedestal_fit(a):
+            fit = np.sum(a[:, None] * const_fft_term, axis=0)
+            chi2 = np.sum(np.abs(fit-spectrum_oob)**2)
+            return chi2
+
+        # self._spectrum = spectrum_oob
+        res = scipy.optimize.minimize(pedestal_fit, a_guess, tol=tol).x
+        ### maybe TODO - include option to use Minuit, which seems a lot quicker?
+        # m = Minuit(self._pedestal_fit_minuit, a_guess * nufft_conversion_factor)
+        # m.errordef = 1
+        # m.errors = 0.01 * np.ones_like(a_guess)
+        # m.migrad(ncall=20000)
+        # res = m.values
+
+        block_offsets = np.zeros(len(res) + 1)
+        block_offsets[:-1] = res
+
+        # the fit is not sensitive to an overall shift,
+        # so we include the zero-meaning here
+        block_offsets += np.mean(trace) - np.mean(block_offsets)
+    else:
+        raise ValueError(f'Invalid value for mode={mode}. Accepted values are {{"fit", "approximate"}}')
 
     if return_trace:
         output_trace = trace - np.repeat(block_offsets, block_size)

From 4f423f7a0bd63c76288e3aa7c5cdae4db3dfb115 Mon Sep 17 00:00:00 2001
From: Sjoerd Bouma <sjoerdbouma@xs4all.nl>
Date: Mon, 10 Jul 2023 16:36:47 +0200
Subject: [PATCH 06/19] store fitted block offsets as a channelParameter,
 simplify interface

---
 NuRadioReco/framework/parameters.py           |   1 +
 .../modules/RNO_G/channelBlockOffsetFitter.py | 146 ++++++------------
 2 files changed, 52 insertions(+), 95 deletions(-)

diff --git a/NuRadioReco/framework/parameters.py b/NuRadioReco/framework/parameters.py
index 575d7b402..b234ec112 100644
--- a/NuRadioReco/framework/parameters.py
+++ b/NuRadioReco/framework/parameters.py
@@ -57,6 +57,7 @@ class channelParameters(Enum):
     signal_receiving_zenith = 15    #: the zenith angle of direction at which the radio signal arrived at the antenna
     signal_ray_type = 16        #: type of the ray propagation path of the signal received by this channel. Options are direct, reflected and refracted
     signal_receiving_azimuth = 17   #: the azimuth angle of direction at which the radio signal arrived at the antenna
+    block_offsets = 18 #: 'block' or pedestal offsets. See `NuRadioReco.modules.RNO_G.channelBlockOffsetFitter`
 
 
 class electricFieldParameters(Enum):
diff --git a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
index 4467fd0d8..da0cd3e37 100644
--- a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
+++ b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
@@ -1,5 +1,17 @@
+"""
+Module to remove 'block offsets' from RNO-G voltage traces.
+
+The function ``fit_block_offsets`` can be used standalone to perform an out-of-band 
+fit to the block offsets. Alternatively, the ``channelBlockOffsets`` class contains convenience 
+``add_offsets`` (to add block offsets in simulation) and ``remove_offsets`` methods that can be run 
+directly on a NuRadioMC/imported ``Event``. The added/removed block offsets are stored per channel
+in the `NuRadioReco.framework.parameters.channelParameters.block_offsets` parameter.
+
+"""
+
 from NuRadioReco.utilities import units, fft
 from NuRadioReco.framework.base_trace import BaseTrace
+from NuRadioReco.framework.parameters import channelParameters
 # from iminuit import Minuit
 import numpy as np
 import scipy
@@ -32,6 +44,9 @@ def add_offsets(self, event, station, offsets=1*units.mV, channel_ids=None):
         """
         Add (simulated or reconstructed) block offsets to an event.
 
+        Added block offsets for each channel are stored in the 
+        ``channelParameters.block_offsets`` parameter.
+
         Parameters
         ----------
         event: Event object | None
@@ -68,119 +83,67 @@ def add_offsets(self, event, station, offsets=1*units.mV, channel_ids=None):
                 )
             else:
                 add_offsets = offsets
-            if channel_id in self._offset_inject.keys():
-                self._offset_inject[channel_id] += add_offsets
+
+            # save the added offsets as a channelParameter
+            if channel.has_parameter(channelParameters.block_offsets):
+                block_offsets_old = channel.get_parameter(channelParameters.block_offsets)
+                channel.set_parameter(channelParameters.block_offsets, block_offsets_old + offsets)
             else:
-                self._offset_inject[channel_id] = add_offsets
+                channel.set_parameter(channelParameters.block_offsets, offsets)
 
             channel.set_trace(
                 channel.get_trace() + np.repeat(add_offsets, self.block_size),
                 channel.get_sampling_rate()
             )
 
-    def remove_offsets(self, event, station, offsets='fit', channel_ids=None):
+    def remove_offsets(self, event, station, mode='fit', channel_ids=None):
         """
         Remove block offsets from an event
 
-        Fits and removes the block offsets from an event.
+        Fits and removes the block offsets from an event. The removed
+        offsets are stored in the ``channelParameters.block_offsets``
+        parameter.
 
         Parameters
         ----------
         event: NuRadioReco.framework.event.Event | None
         station: NuRadioReco.framework.station.Station
             The station to remove the block offsets from
-        offsets: str
-            How to remove the offsets. Options are:
-
-            - 'fit': fit the offsets out of band
-            - 'injected': if offsets were injected using the ``add_offsets``
-              method, this removes those offsets. Otherwise, this does nothing.
+        mode: 'fit' | 'approximate' | 'stored'
+            
+            - 'fit' (default): fit the block offsets with a minimizer
+            - 'approximate' : use the first guess from the out-of-band component,
+              without any fitting (slightly faster)
+            - 'stored': use the block offsets already stored in the 
+              ``channelParameters.block_offsets`` parameter. Will raise an error 
+              if this parameter is not present.
 
-            Default: 'fit'
         channel_ids: list | None
             List of channel ids to remove offsets from. If None (default),
             remove offsets from all channels in ``station``
 
         """
-        if offsets=='fit':
-            if not len(self._offset_fit):
-                self.fit_offsets(event, station, channel_ids)
-            offsets = self._offset_fit
-        elif offsets=='injected':
-            if not len(self._offset_inject):
-                offsets = np.zeros(16) #TODO - ensure this works for different trace lengths
-            else:
-                offsets = self._offset_inject
-
-        if isinstance(offsets, dict):
-            remove_offsets = {key: -offsets[key] for key in offsets.keys()}
-        else:
-            remove_offsets = -offsets
-        self.add_offsets(event, station, remove_offsets, channel_ids)
-
-    def fit_offsets(self, event, station, channel_ids=None):
-        """
-        Fit the block offsets using an out-of-band fit
-
-        This function fits the block offsets present in a given
-        event / station using an out-of-band fit in frequency space.
-
-        Parameters
-        ----------
-        event: NuRadioReco.framework.event.Event | None
-        station: NuRadioReco.framework.station.Station
-            The station to fit the block offsets to
-        channel_ids: list | None
-            List of channel ids to fit block offsets for. If None (default),
-            fit offsets for all channels in ``station``
-
-        """
-        block_size = self.block_size
-        if channel_ids is None:
+        if channel_ids  is None:
             channel_ids = station.get_channel_ids()
-        for channel_id in channel_ids:
-            channel = station.get_channel(channel_id)
-            trace = channel.get_trace()
-
-            block_offsets = fit_block_offsets(
-                trace, block_size,
-                channel.get_sampling_rate(), self._max_frequency
-            )
-            self._offset_fit[channel_id] = block_offsets
 
-    def get_offsets(self, channel_id, offset_type='fit'):
-        """
-        Return the block offsets for a given channel.
-
-        Parameters
-        ----------
-        channel_id: int
-            channel id that specifies the channel to return block offsets for
-        offset_type: str
-            Options:
-
-            - 'fit': return the fitted block offsets
-            - 'injected': return the block offsets that were injected
-              using the ``add_offsets`` method.
-
-        Returns
-        -------
-        trace: BaseTrace
-            A :class:`NuRadioReco.framework.base_trace.BaseTrace` object with the same length as the channel trace,
-            containing only the block offsets.
+        offsets = {}
+        if mode == 'stored': # remove offsets stored in channelParameters.block_offsets
+            offsets = {
+                channel_id: -station.get_channel(channel_id).get_parameter(channelParameters.block_offsets)
+                for channel_id in channel_ids}
+        else: # fit & remove offsets
+            for channel_id in channel_ids:
+                channel = station.get_channel(channel_id)
+                trace = channel.get_trace()
+
+                block_offsets = fit_block_offsets(
+                    trace, self.block_size,
+                    channel.get_sampling_rate(), self._max_frequency
+                )
+                offsets[channel_id] = -block_offsets
+        
+        self.add_offsets(event, station, offsets, channel_ids)
 
-        """
-        trace = BaseTrace()
-        if offset_type == 'fit':
-            trace.set_trace(np.repeat(self._offset_fit[channel_id], self.block_size), self.sampling_rate)
-        elif offset_type == 'injected':
-            trace.set_trace(np.repeat(self._offset_inject[channel_id], self.block_size), self.sampling_rate)
-        return trace
-
-    def _pedestal_fit(self, a):
-        fit = np.sum(a[:, None] * self._const_fft_term, axis=0)
-        chi2 = np.sum(np.abs(fit-self._spectrum)**2)
-        return chi2
 
 def fit_block_offsets(
         trace, block_size=128, sampling_rate=3.2*units.GHz,
@@ -275,14 +238,7 @@ def pedestal_fit(a):
             chi2 = np.sum(np.abs(fit-spectrum_oob)**2)
             return chi2
 
-        # self._spectrum = spectrum_oob
         res = scipy.optimize.minimize(pedestal_fit, a_guess, tol=tol).x
-        ### maybe TODO - include option to use Minuit, which seems a lot quicker?
-        # m = Minuit(self._pedestal_fit_minuit, a_guess * nufft_conversion_factor)
-        # m.errordef = 1
-        # m.errors = 0.01 * np.ones_like(a_guess)
-        # m.migrad(ncall=20000)
-        # res = m.values
 
         block_offsets = np.zeros(len(res) + 1)
         block_offsets[:-1] = res

From 256f21f23c3fc34bfe474dc816b50b2e934b1bc7 Mon Sep 17 00:00:00 2001
From: Sjoerd Bouma <sjoerdbouma@xs4all.nl>
Date: Mon, 10 Jul 2023 16:37:26 +0200
Subject: [PATCH 07/19] update changelog

---
 changelog.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/changelog.txt b/changelog.txt
index 8d6735957..1004f2da1 100644
--- a/changelog.txt
+++ b/changelog.txt
@@ -10,6 +10,7 @@ new features:
 - add positions array functionality to simple ice model in average and gradient functions
 - analytic ray tracing solutions are now sorted consistently from lowest to highest ray
 - added ability to generate high-low-triggered noise on a narrow band but return full-band waveforms
+- added 'block offset' removal/simulation module for RNO-G
 bugfixes:
 - fixed/improved C++ raytracer not finding solutions for some near-horizontal or near-shadowzone vertices
 - fixed wrong number in Feldman-Cousins upper limit

From 57f5a44c9efc5e57f1d8fd0b28859a089c34ed5e Mon Sep 17 00:00:00 2001
From: Sjoerd Bouma <sjoerdbouma@xs4all.nl>
Date: Wed, 30 Aug 2023 12:11:58 +0200
Subject: [PATCH 08/19] forgot to pass through optional kwarg

---
 NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
index da0cd3e37..130aed3df 100644
--- a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
+++ b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
@@ -138,7 +138,8 @@ def remove_offsets(self, event, station, mode='fit', channel_ids=None):
 
                 block_offsets = fit_block_offsets(
                     trace, self.block_size,
-                    channel.get_sampling_rate(), self._max_frequency
+                    channel.get_sampling_rate(), self._max_frequency,
+                    mode=mode
                 )
                 offsets[channel_id] = -block_offsets
         

From 162f0140817d228f0814c7e954be41ddebd745a9 Mon Sep 17 00:00:00 2001
From: Sjoerd Bouma <sjoerdbouma@xs4all.nl>
Date: Tue, 12 Sep 2023 16:15:11 +0200
Subject: [PATCH 09/19] remove commented import

---
 NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
index 130aed3df..9b62a7aeb 100644
--- a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
+++ b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
@@ -12,7 +12,6 @@
 from NuRadioReco.utilities import units, fft
 from NuRadioReco.framework.base_trace import BaseTrace
 from NuRadioReco.framework.parameters import channelParameters
-# from iminuit import Minuit
 import numpy as np
 import scipy
 

From 3d58f050089818db5f400a2ddd6abd9e71f15b2e Mon Sep 17 00:00:00 2001
From: Sjoerd Bouma <sjoerdbouma@xs4all.nl>
Date: Mon, 25 Sep 2023 13:04:17 +0200
Subject: [PATCH 10/19] treat offsets given as dict correctly

---
 NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
index 9b62a7aeb..7dce7929a 100644
--- a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
+++ b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
@@ -76,12 +76,12 @@ def add_offsets(self, event, station, offsets=1*units.mV, channel_ids=None):
             channel = station.get_channel(channel_id)
             if isinstance(offsets, dict):
                 add_offsets = offsets[channel_id]
-            elif len(np.atleast_1d(offsets)) == 1:
-                add_offsets = np.random.normal(
-                    0, offsets, (channel.get_number_of_samples() // self.block_size)
-                )
             else:
                 add_offsets = offsets
+            if len(np.atleast_1d(add_offsets)) == 1:
+                add_offsets = np.random.normal(
+                    0, add_offsets, (channel.get_number_of_samples() // self.block_size)
+                )
 
             # save the added offsets as a channelParameter
             if channel.has_parameter(channelParameters.block_offsets):

From 98f51e4f7d3a53cf1e8ad31f0109b14015b02ecc Mon Sep 17 00:00:00 2001
From: Sjoerd Bouma <sjoerdbouma@xs4all.nl>
Date: Mon, 25 Sep 2023 13:07:50 +0200
Subject: [PATCH 11/19] save the microseconds

---
 NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
index 7dce7929a..e686e3b75 100644
--- a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
+++ b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
@@ -205,10 +205,7 @@ def fit_block_offsets(
     filtered_trace = fft.freq2time(filtered_trace_fft, sampling_rate)
 
     # obtain guesses for block offsets
-    a_guess = np.array([
-        np.mean(filtered_trace[i*block_size:(i+1)*block_size])
-        for i in range(n_blocks)
-    ])
+    a_guess = np.mean(np.split(filtered_trace, n_blocks), axis=1)
     if mode == 'approximate':
         block_offsets = a_guess
     elif mode == 'fit':

From 8a77811310c3a5a05d66d0f972aa44693ce436f5 Mon Sep 17 00:00:00 2001
From: Sjoerd Bouma <sjoerdbouma@xs4all.nl>
Date: Mon, 25 Sep 2023 14:11:11 +0200
Subject: [PATCH 12/19] add different block offsets fit options to mattak
 reader and store the offsets

---
 .../modules/io/RNO_G/readRNOGDataMattak.py    | 53 ++++++++++++++-----
 1 file changed, 40 insertions(+), 13 deletions(-)

diff --git a/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py b/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
index db2835ecb..fbed78970 100644
--- a/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
+++ b/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
@@ -6,17 +6,19 @@
 import math
 
 from NuRadioReco.modules.base.module import register_run
+from NuRadioReco.modules.RNO_G.channelBlockOffsetFitter import channelBlockOffsets
 
 import NuRadioReco.framework.event
 import NuRadioReco.framework.station
 import NuRadioReco.framework.channel
 import NuRadioReco.framework.trigger
+import NuRadioReco.framework.parameters
 
 from NuRadioReco.utilities import units
 import mattak.Dataset
 
 
-def baseline_correction(wfs, n_bins=128, func=np.median):
+def baseline_correction(wfs, n_bins=128, func=np.median, return_offsets=False):
     """
     Simple baseline correction function. Determines baseline in discrete chuncks of "n_bins" with
     the function specified (i.e., mean or median).
@@ -33,11 +35,17 @@ def baseline_correction(wfs, n_bins=128, func=np.median):
     func: np.mean or np.median
         Function to calculate pedestal
     
+    return_offsets: bool, default False
+        if True, additionally return the baseline offsets
+    
     Returns
     -------
     
     wfs_corrected: np.array(n_events, n_channels, n_samples)
         Baseline/pedestal corrected waveforms
+
+    baseline_values: np.array of shape (n_samples // n_bins, n_events, n_channels)
+        (Only if return_offsets==True) The baseline offsets
     """
      
     # Example: Get baselines in chunks of 128 bins
@@ -57,8 +65,12 @@ def baseline_correction(wfs, n_bins=128, func=np.median):
     # np.moveaxis -> (n_events, n_channels, 2048)
     baseline_traces = np.moveaxis(baseline_traces, 0, -1)
      
+    if return_offsets:
+        return wfs - baseline_traces, baseline_values
+    
     return wfs - baseline_traces
 
+blockoffsetfitter = channelBlockOffsets()
 
 def get_time_offset(trigger_type):
     """ 
@@ -175,7 +187,7 @@ def begin(self,
             read_calibrated_data=False,
             select_triggers=None,
             select_runs=False,
-            apply_baseline_correction=True,
+            apply_baseline_correction='median',
             convert_to_voltage=True,
             selectors=[],
             run_types=["physics"],
@@ -206,9 +218,16 @@ def begin(self,
         Other Parameters
         ----------------
         
-        apply_baseline_correction: bool
-            Only applies when non-calibrated data are read. If true, correct for DC offset.
-            (Default: True)
+        apply_baseline_correction: 'median' | 'approximate' | 'fit' | 'none'
+            Only applies when non-calibrated data are read. Removes the DC (baseline)
+            block offsets (pedestals).
+            Options are:
+
+            * 'median' (Default) - subtract the median of each block
+            * 'approximate' - apply a bandpass filter before calculating the median of each block
+            * 'fit' - do a full out-of-band fit to determine the block offsets; for more details,
+              see ``NuRadioReco.modules.RNO_G.channelBlockOffsetFitter``
+            * 'none' - do not apply a baseline correction
 
         convert_to_voltage: bool
             Only applies when non-calibrated data are read. If true, convert ADC to voltage.
@@ -245,6 +264,12 @@ def begin(self,
         t0 = time.time()
         
         self._read_calibrated_data = read_calibrated_data
+        baseline_correction_valid_options = ['median', 'approximate', 'fit', 'none']
+        if apply_baseline_correction.lower() not in baseline_correction_valid_options:
+            raise ValueError(
+                f"Value for apply_baseline_correction ({apply_baseline_correction}) not recognized. "
+                f"Valid options are {baseline_correction_valid_options}"
+            )
         self._apply_baseline_correction = apply_baseline_correction
         self._convert_to_voltage = convert_to_voltage
         
@@ -614,17 +639,17 @@ def _get_event(self, event_info, waveforms):
             if self._read_calibrated_data:    
                 channel.set_trace(wf * units.mV, sampling_rate * units.GHz)
             else:
-                # wf stores ADC counts
-                
-                if self._apply_baseline_correction:
-                    # correct baseline
-                    wf = baseline_correction(wf)
-                
+                # wf stores ADC counts                
                 if self._convert_to_voltage:
                     # convert adc to voltage
                     wf *= (self._adc_ref_voltage_range / (2 ** (self._adc_n_bits) - 1))
-                    
+    
+                if self._apply_baseline_correction == 'median':
+                    # correct baseline
+                    wf, offsets = baseline_correction(wf, return_offsets=True)
+
                 channel.set_trace(wf, sampling_rate * units.GHz)
+                channel.set_parameter(NuRadioReco.framework.parameters.channelParameters.block_offsets, offsets)
             
             time_offset = get_time_offset(event_info.triggerType)
             channel.set_trace_start_time(-time_offset)  # relative to event/trigger time
@@ -632,7 +657,9 @@ def _get_event(self, event_info, waveforms):
             station.add_channel(channel)
         
         evt.set_station(station)
-        
+        if self._apply_baseline_correction in ['fit', 'approximate'] and not self._read_calibrated_data:
+            blockoffsetfitter.remove_offsets(evt, station, mode=self._apply_baseline_correction)
+                
         return evt      
 
 

From 87a505d379e438ca0af39f31dc9fb576d3f3960e Mon Sep 17 00:00:00 2001
From: Sjoerd Bouma <sjoerdbouma@xs4all.nl>
Date: Mon, 25 Sep 2023 14:12:10 +0200
Subject: [PATCH 13/19] change *= to prevent error due to Type change from int
 to float

---
 NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py b/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
index fbed78970..247bdb0b0 100644
--- a/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
+++ b/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
@@ -642,7 +642,7 @@ def _get_event(self, event_info, waveforms):
                 # wf stores ADC counts                
                 if self._convert_to_voltage:
                     # convert adc to voltage
-                    wf *= (self._adc_ref_voltage_range / (2 ** (self._adc_n_bits) - 1))
+                    wf = wf * (self._adc_ref_voltage_range / (2 ** (self._adc_n_bits) - 1))
     
                 if self._apply_baseline_correction == 'median':
                     # correct baseline

From 23843233f3f3ef89391639e8f85592ddc0575b21 Mon Sep 17 00:00:00 2001
From: Sjoerd Bouma <sjoerdbouma@xs4all.nl>
Date: Mon, 25 Sep 2023 14:19:45 +0200
Subject: [PATCH 14/19] store offsets only if they are actually computed

---
 NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py b/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
index 247bdb0b0..f31c1f80c 100644
--- a/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
+++ b/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
@@ -647,9 +647,9 @@ def _get_event(self, event_info, waveforms):
                 if self._apply_baseline_correction == 'median':
                     # correct baseline
                     wf, offsets = baseline_correction(wf, return_offsets=True)
+                    channel.set_parameter(NuRadioReco.framework.parameters.channelParameters.block_offsets, offsets)
 
                 channel.set_trace(wf, sampling_rate * units.GHz)
-                channel.set_parameter(NuRadioReco.framework.parameters.channelParameters.block_offsets, offsets)
             
             time_offset = get_time_offset(event_info.triggerType)
             channel.set_trace_start_time(-time_offset)  # relative to event/trigger time

From ba3b1b2d51db9266480f2e4cbbeed5071967f3a6 Mon Sep 17 00:00:00 2001
From: Sjoerd Bouma <sjoerdbouma@xs4all.nl>
Date: Tue, 26 Sep 2023 16:36:32 +0200
Subject: [PATCH 15/19] make blockoffsetfitter faster by setting maxiter=2 by
 default

---
 .../modules/RNO_G/channelBlockOffsetFitter.py | 20 ++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
index e686e3b75..b55247799 100644
--- a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
+++ b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
@@ -13,7 +13,7 @@
 from NuRadioReco.framework.base_trace import BaseTrace
 from NuRadioReco.framework.parameters import channelParameters
 import numpy as np
-import scipy
+import scipy.optimize
 
 class channelBlockOffsets:
 
@@ -95,7 +95,7 @@ def add_offsets(self, event, station, offsets=1*units.mV, channel_ids=None):
                 channel.get_sampling_rate()
             )
 
-    def remove_offsets(self, event, station, mode='fit', channel_ids=None):
+    def remove_offsets(self, event, station, mode='fit', channel_ids=None, maxiter=2):
         """
         Remove block offsets from an event
 
@@ -120,6 +120,11 @@ def remove_offsets(self, event, station, mode='fit', channel_ids=None):
         channel_ids: list | None
             List of channel ids to remove offsets from. If None (default),
             remove offsets from all channels in ``station``
+        maxiter: int, default 2
+            (Only if mode=='fit') The maximum number of fit iterations.
+            This can be increased to more accurately remove the block offsets
+            at the cost of performance. (The default value removes 'most' offsets
+            to about 1%)
 
         """
         if channel_ids  is None:
@@ -138,7 +143,7 @@ def remove_offsets(self, event, station, mode='fit', channel_ids=None):
                 block_offsets = fit_block_offsets(
                     trace, self.block_size,
                     channel.get_sampling_rate(), self._max_frequency,
-                    mode=mode
+                    mode=mode, maxiter=maxiter
                 )
                 offsets[channel_id] = -block_offsets
         
@@ -148,7 +153,7 @@ def remove_offsets(self, event, station, mode='fit', channel_ids=None):
 def fit_block_offsets(
         trace, block_size=128, sampling_rate=3.2*units.GHz,
         max_frequency=50*units.MHz, mode='fit', return_trace = False,
-        tol=1e-6,):
+        maxiter=2, tol=1e-6):
     """
     Fit 'block' offsets for a voltage trace
 
@@ -175,6 +180,11 @@ def fit_block_offsets(
         if True, return the tuple (offsets, output_trace)
         where the output_trace is the input trace with
         fitted block offsets removed
+    maxiter: int (default: 2)
+        (Only if mode=='fit') The maximum number of fit iterations.
+        This can be increased to more accurately remove the block offsets
+        at the cost of performance. (The default value removes 'most' offsets
+        to about 1%)
 
     Returns
     -------
@@ -235,7 +245,7 @@ def pedestal_fit(a):
             chi2 = np.sum(np.abs(fit-spectrum_oob)**2)
             return chi2
 
-        res = scipy.optimize.minimize(pedestal_fit, a_guess, tol=tol).x
+        res = scipy.optimize.minimize(pedestal_fit, a_guess, tol=tol, options=dict(maxiter=maxiter)).x
 
         block_offsets = np.zeros(len(res) + 1)
         block_offsets[:-1] = res

From 80ef76dfbd4bfc3e9e846968af1cf59781ae44f2 Mon Sep 17 00:00:00 2001
From: Sjoerd Bouma <sjoerdbouma@xs4all.nl>
Date: Tue, 26 Sep 2023 17:00:55 +0200
Subject: [PATCH 16/19] deprecate old baseline correction in readRNOGDataMattak

---
 .../modules/io/RNO_G/readRNOGDataMattak.py    | 25 ++++++++++++-------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py b/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
index f31c1f80c..5a349f8a8 100644
--- a/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
+++ b/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
@@ -20,8 +20,13 @@
 
 def baseline_correction(wfs, n_bins=128, func=np.median, return_offsets=False):
     """
-    Simple baseline correction function. Determines baseline in discrete chuncks of "n_bins" with
+    Simple baseline correction function. 
+    
+    Determines baseline in discrete chuncks of "n_bins" with
     the function specified (i.e., mean or median).
+
+    .. Warning:: This function has been deprecated, use :mod:`NuRadioReco.modules.RNO_G.channelBlockOffsetFitter`
+      instead
     
     Parameters
     ----------
@@ -47,7 +52,10 @@ def baseline_correction(wfs, n_bins=128, func=np.median, return_offsets=False):
     baseline_values: np.array of shape (n_samples // n_bins, n_events, n_channels)
         (Only if return_offsets==True) The baseline offsets
     """
-     
+    import warnings
+    warnings.warn(
+        'baseline_correction is deprecated, use NuRadioReco.modules.RNO_G.channelBlockOffsetFitter instead',
+        DeprecationWarning)
     # Example: Get baselines in chunks of 128 bins
     # wfs in (n_events, n_channels, 2048)
     # np.split -> (16, n_events, n_channels, 128) each waveform split in 16 chuncks
@@ -187,7 +195,7 @@ def begin(self,
             read_calibrated_data=False,
             select_triggers=None,
             select_runs=False,
-            apply_baseline_correction='median',
+            apply_baseline_correction='approximate',
             convert_to_voltage=True,
             selectors=[],
             run_types=["physics"],
@@ -218,16 +226,15 @@ def begin(self,
         Other Parameters
         ----------------
         
-        apply_baseline_correction: 'median' | 'approximate' | 'fit' | 'none'
+        apply_baseline_correction: 'approximate' | 'fit' | 'none'
             Only applies when non-calibrated data are read. Removes the DC (baseline)
             block offsets (pedestals).
             Options are:
 
-            * 'median' (Default) - subtract the median of each block
-            * 'approximate' - apply a bandpass filter before calculating the median of each block
+            * 'approximate' (default) - estimate block offsets by looking at the low-pass filtered trace
             * 'fit' - do a full out-of-band fit to determine the block offsets; for more details,
               see ``NuRadioReco.modules.RNO_G.channelBlockOffsetFitter``
-            * 'none' - do not apply a baseline correction
+            * 'none' - do not apply a baseline correction (faster)
 
         convert_to_voltage: bool
             Only applies when non-calibrated data are read. If true, convert ADC to voltage.
@@ -264,7 +271,7 @@ def begin(self,
         t0 = time.time()
         
         self._read_calibrated_data = read_calibrated_data
-        baseline_correction_valid_options = ['median', 'approximate', 'fit', 'none']
+        baseline_correction_valid_options = ['approximate', 'fit', 'none']
         if apply_baseline_correction.lower() not in baseline_correction_valid_options:
             raise ValueError(
                 f"Value for apply_baseline_correction ({apply_baseline_correction}) not recognized. "
@@ -644,7 +651,7 @@ def _get_event(self, event_info, waveforms):
                     # convert adc to voltage
                     wf = wf * (self._adc_ref_voltage_range / (2 ** (self._adc_n_bits) - 1))
     
-                if self._apply_baseline_correction == 'median':
+                if self._apply_baseline_correction == 'median': #TODO: deprecate / remove
                     # correct baseline
                     wf, offsets = baseline_correction(wf, return_offsets=True)
                     channel.set_parameter(NuRadioReco.framework.parameters.channelParameters.block_offsets, offsets)

From 291433e3fd71b69ef0e0c6052adddbf8ae19a920 Mon Sep 17 00:00:00 2001
From: Sjoerd Bouma <sjoerdbouma@xs4all.nl>
Date: Tue, 26 Sep 2023 17:08:55 +0200
Subject: [PATCH 17/19] remove inaccessible old baseline correction code

---
 NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py b/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
index 5a349f8a8..783ed512b 100644
--- a/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
+++ b/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
@@ -651,11 +651,6 @@ def _get_event(self, event_info, waveforms):
                     # convert adc to voltage
                     wf = wf * (self._adc_ref_voltage_range / (2 ** (self._adc_n_bits) - 1))
     
-                if self._apply_baseline_correction == 'median': #TODO: deprecate / remove
-                    # correct baseline
-                    wf, offsets = baseline_correction(wf, return_offsets=True)
-                    channel.set_parameter(NuRadioReco.framework.parameters.channelParameters.block_offsets, offsets)
-
                 channel.set_trace(wf, sampling_rate * units.GHz)
             
             time_offset = get_time_offset(event_info.triggerType)

From 30726a0ce6e920d7de58c656a8ad0a1286f0aa42 Mon Sep 17 00:00:00 2001
From: Sjoerd Bouma <sjoerdbouma@xs4all.nl>
Date: Fri, 29 Sep 2023 19:11:20 +0200
Subject: [PATCH 18/19] increase default number of iterations to 5

---
 NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py | 8 ++++----
 NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py    | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
index b55247799..e4711e094 100644
--- a/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
+++ b/NuRadioReco/modules/RNO_G/channelBlockOffsetFitter.py
@@ -95,7 +95,7 @@ def add_offsets(self, event, station, offsets=1*units.mV, channel_ids=None):
                 channel.get_sampling_rate()
             )
 
-    def remove_offsets(self, event, station, mode='fit', channel_ids=None, maxiter=2):
+    def remove_offsets(self, event, station, mode='fit', channel_ids=None, maxiter=5):
         """
         Remove block offsets from an event
 
@@ -120,7 +120,7 @@ def remove_offsets(self, event, station, mode='fit', channel_ids=None, maxiter=2
         channel_ids: list | None
             List of channel ids to remove offsets from. If None (default),
             remove offsets from all channels in ``station``
-        maxiter: int, default 2
+        maxiter: int, default 5
             (Only if mode=='fit') The maximum number of fit iterations.
             This can be increased to more accurately remove the block offsets
             at the cost of performance. (The default value removes 'most' offsets
@@ -153,7 +153,7 @@ def remove_offsets(self, event, station, mode='fit', channel_ids=None, maxiter=2
 def fit_block_offsets(
         trace, block_size=128, sampling_rate=3.2*units.GHz,
         max_frequency=50*units.MHz, mode='fit', return_trace = False,
-        maxiter=2, tol=1e-6):
+        maxiter=5, tol=1e-6):
     """
     Fit 'block' offsets for a voltage trace
 
@@ -180,7 +180,7 @@ def fit_block_offsets(
         if True, return the tuple (offsets, output_trace)
         where the output_trace is the input trace with
         fitted block offsets removed
-    maxiter: int (default: 2)
+    maxiter: int (default: 5)
         (Only if mode=='fit') The maximum number of fit iterations.
         This can be increased to more accurately remove the block offsets
         at the cost of performance. (The default value removes 'most' offsets
diff --git a/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py b/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
index 783ed512b..9b5783e8b 100644
--- a/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
+++ b/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
@@ -17,6 +17,7 @@
 from NuRadioReco.utilities import units
 import mattak.Dataset
 
+blockoffsetfitter = channelBlockOffsets()
 
 def baseline_correction(wfs, n_bins=128, func=np.median, return_offsets=False):
     """
@@ -78,7 +79,6 @@ def baseline_correction(wfs, n_bins=128, func=np.median, return_offsets=False):
     
     return wfs - baseline_traces
 
-blockoffsetfitter = channelBlockOffsets()
 
 def get_time_offset(trigger_type):
     """ 

From 69f6fda60b1638f9766a039ec2557a128c52aa8e Mon Sep 17 00:00:00 2001
From: Sjoerd Bouma <sjoerdbouma@xs4all.nl>
Date: Fri, 29 Sep 2023 19:13:22 +0200
Subject: [PATCH 19/19] make blockoffsetfitter a private attribute to reduce
 clutter in module

---
 NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py b/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
index 9b5783e8b..7ca5edd89 100644
--- a/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
+++ b/NuRadioReco/modules/io/RNO_G/readRNOGDataMattak.py
@@ -17,7 +17,6 @@
 from NuRadioReco.utilities import units
 import mattak.Dataset
 
-blockoffsetfitter = channelBlockOffsets()
 
 def baseline_correction(wfs, n_bins=128, func=np.median, return_offsets=False):
     """
@@ -169,6 +168,8 @@ def __init__(self, run_table_path=None, log_level=logging.INFO):
         """
         self.logger = logging.getLogger('NuRadioReco.readRNOGData')
         self.logger.setLevel(log_level)
+
+        self._blockoffsetfitter = channelBlockOffsets()
      
         # Initialize run table for run selection
         self.__run_table = None
@@ -233,7 +234,7 @@ def begin(self,
 
             * 'approximate' (default) - estimate block offsets by looking at the low-pass filtered trace
             * 'fit' - do a full out-of-band fit to determine the block offsets; for more details,
-              see ``NuRadioReco.modules.RNO_G.channelBlockOffsetFitter``
+              see :mod:`NuRadioReco.modules.RNO_G.channelBlockOffsetFitter`
             * 'none' - do not apply a baseline correction (faster)
 
         convert_to_voltage: bool
@@ -660,7 +661,7 @@ def _get_event(self, event_info, waveforms):
         
         evt.set_station(station)
         if self._apply_baseline_correction in ['fit', 'approximate'] and not self._read_calibrated_data:
-            blockoffsetfitter.remove_offsets(evt, station, mode=self._apply_baseline_correction)
+            self._blockoffsetfitter.remove_offsets(evt, station, mode=self._apply_baseline_correction)
                 
         return evt