Repository: cta-observatory/cta-lstchain

Issue:
Speeding-up the NSB addition at waveform level
The generation of random NSB p.e.'s to be added to waveforms in the R0 to DL1 stage of MC is currently done "on the fly" for every event & pixel. That takes a significant part of the total cpu for the case of nigh NSB - for 16 * dark we saw like a factor 6 slower processing.

One could just generate, at the beginning of the r0_to_dl1 execution, a database to be kept in memory of a few thousand waveforms containing different realizations of the additional NSB, and just pick one at random from there for every pixel. In the proton simulations we have ~800 events, so if we just simulate, say, 10*1855 waveforms and use them shuffled, that would be a factor 80 less effort used in the noise waveform production.  

Discussion:

Patch:
diff --git a/lstchain/data/lstchain_lhfit_config.json b/lstchain/data/lstchain_lhfit_config.json
index 05368110..c761c199 100644
--- a/lstchain/data/lstchain_lhfit_config.json
+++ b/lstchain/data/lstchain_lhfit_config.json
@@ -279,8 +279,9 @@
   },
   "waveform_nsb_tuning":{
     "nsb_tuning": false,
-    "nsb_tuning_ratio": 0.52,
-    "spe_location": null
+    "nsb_tuning_ratio": 0.5,
+    "spe_location": null,
+    "pre_computed_multiplicity": 10
   },
   "lh_fit_config": {
     "sigma_s": [
diff --git a/lstchain/data/lstchain_standard_config.json b/lstchain/data/lstchain_standard_config.json
index d63c9e0f..74312da6 100644
--- a/lstchain/data/lstchain_standard_config.json
+++ b/lstchain/data/lstchain_standard_config.json
@@ -274,8 +274,9 @@
   },
   "waveform_nsb_tuning":{
     "nsb_tuning": false,
-    "nsb_tuning_ratio": 0.52,
-    "spe_location": null
+    "nsb_tuning_ratio": 0.5,
+    "spe_location": null,
+    "pre_computed_multiplicity": 10
   },
   "write_interleaved_events":{
     "DataWriter": {
diff --git a/lstchain/image/modifier.py b/lstchain/image/modifier.py
index 1c0f28f8..9f523902 100644
--- a/lstchain/image/modifier.py
+++ b/lstchain/image/modifier.py
@@ -1,6 +1,5 @@
-
+import astropy.units as u
 import logging
-
 import numpy as np
 from ctapipe.calib.camera import CameraCalibrator
 from ctapipe.io import (
@@ -21,12 +20,11 @@ __all__ = [
     'calculate_noise_parameters',
     'random_psf_smearer',
     'set_numba_seed',
-    'tune_nsb_on_waveform',
+    'WaveformNsbTunner',
 ]
 
 log = logging.getLogger(__name__)
 
-
 # number of neighbors of completely surrounded pixels of hexagonal cameras:
 N_PIXEL_NEIGHBORS = 6
 SMEAR_PROBABILITIES = np.full(N_PIXEL_NEIGHBORS, 1 / N_PIXEL_NEIGHBORS)
@@ -191,22 +189,22 @@ def calculate_noise_parameters(simtel_filename, data_dl1_filename,
 
     # Real data DL1 tables:
     data_dl1_calibration = read_table(data_dl1_filename,
-                    '/dl1/event/telescope/monitoring/calibration')
-    data_dl1_pedestal =  read_table(data_dl1_filename,
-                    '/dl1/event/telescope/monitoring/pedestal')
-    data_dl1_flatfield =  read_table(data_dl1_filename,
-                    '/dl1/event/telescope/monitoring/flatfield')
-    data_dl1_parameters =  read_table(data_dl1_filename,
-                    '/dl1/event/telescope/parameters/LST_LSTCam')
+                                      '/dl1/event/telescope/monitoring/calibration')
+    data_dl1_pedestal = read_table(data_dl1_filename,
+                                   '/dl1/event/telescope/monitoring/pedestal')
+    data_dl1_flatfield = read_table(data_dl1_filename,
+                                    '/dl1/event/telescope/monitoring/flatfield')
+    data_dl1_parameters = read_table(data_dl1_filename,
+                                     '/dl1/event/telescope/parameters/LST_LSTCam')
     data_dl1_image = read_table(data_dl1_filename,
-                    '/dl1/event/telescope/image/LST_LSTCam')
+                                '/dl1/event/telescope/image/LST_LSTCam')
 
     unusable = data_dl1_calibration['unusable_pixels']
     # Locate pixels with HG declared unusable either in original calibration or
     # in interleaved events:
     bad_pixels = unusable[0][0]  # original calibration
     if unusable.shape[0] > 1:
-        for tf in unusable[1:][0]:   # calibrations with interleaveds
+        for tf in unusable[1:][0]:  # calibrations with interleaveds
             bad_pixels = np.logical_or(bad_pixels, tf)
     good_pixels = ~bad_pixels
 
@@ -223,7 +221,7 @@ def calculate_noise_parameters(simtel_filename, data_dl1_filename,
                       'monitoring table!')
         return None, None, None
 
-    if data_dl1_pedestal['charge_std'].shape[0] < 2 :
+    if data_dl1_pedestal['charge_std'].shape[0] < 2:
         logging.error('\nCould not find interleaved pedestal info in '
                       'monitoring table!')
         return None, None, None
@@ -238,7 +236,7 @@ def calculate_noise_parameters(simtel_filename, data_dl1_filename,
         dummy.append(x * data_HG_dc_to_pe[calibration_id[i],])
     dummy = np.array(dummy)
     # Average for all interleaved calibrations (in case there are more than one)
-    data_HG_FF_mean_pe = np.mean(dummy, axis=0) # one value per pixel
+    data_HG_FF_mean_pe = np.mean(dummy, axis=0)  # one value per pixel
 
     # Pixel-wise pedestal standard deviation (for an unbiased extractor),
     # in adc counts:
@@ -251,7 +249,7 @@ def calculate_noise_parameters(simtel_filename, data_dl1_filename,
         dummy.append(x * data_HG_dc_to_pe[calibration_id[i],])
     dummy = np.array(dummy)
     # Average for all interleaved calibrations (in case there are more than one)
-    data_HG_ped_std_pe = np.mean(dummy, axis=0) # one value per pixel
+    data_HG_ped_std_pe = np.mean(dummy, axis=0)  # one value per pixel
 
     # Identify noisy pixels, likely containing stars - we want to adjust MC to
     # the average diffuse NSB across the camera
@@ -296,9 +294,9 @@ def calculate_noise_parameters(simtel_filename, data_dl1_filename,
 
     # Find the peak of the pedestal biased charge distribution of real data.
     # Use an interpolated version of the histogram, for robustness:
-    func = interp1d(0.5*(dataq[1][1:]+dataq[1][:-1]), dataq[0],
+    func = interp1d(0.5 * (dataq[1][1:] + dataq[1][:-1]), dataq[0],
                     kind='quadratic', fill_value='extrapolate')
-    xx = np.linspace(qrange[0], qrange[1], 100*qbins)
+    xx = np.linspace(qrange[0], qrange[1], 100 * qbins)
     mode_data = xx[np.argmax(func(xx))]
 
     # Event reader for simtel file:
@@ -369,9 +367,9 @@ def calculate_noise_parameters(simtel_filename, data_dl1_filename,
                        density=True)
     # Find the peak of the pedestal biased charge distribution of MC. Use
     # an interpolated version of the histogram, for robustness:
-    func = interp1d(0.5*(mcq[1][1:]+mcq[1][:-1]), mcq[0],
+    func = interp1d(0.5 * (mcq[1][1:] + mcq[1][:-1]), mcq[0],
                     kind='quadratic', fill_value='extrapolate')
-    xx = np.linspace(qrange[0], qrange[1], 100*qbins)
+    xx = np.linspace(qrange[0], qrange[1], 100 * qbins)
     mode_mc = xx[np.argmax(func(xx))]
 
     mc_unbiased_std_ped_pe = np.std(mc_ped_charges)
@@ -384,7 +382,7 @@ def calculate_noise_parameters(simtel_filename, data_dl1_filename,
     # The noise is defined as the number of NSB photoelectrons, i.e. the extra
     # variance, rather than standard deviation, of the distribution
     extra_noise_in_bright_pixels = \
-        ((data_median_std_ped_pe**2 - mc_unbiased_std_ped_pe**2) *
+        ((data_median_std_ped_pe ** 2 - mc_unbiased_std_ped_pe ** 2) *
          shower_extractor_window_width / pedestal_extractor_window_width)
 
     # Just in case, makes sure we just add noise if the MC noise is smaller
@@ -400,60 +398,20 @@ def calculate_noise_parameters(simtel_filename, data_dl1_filename,
     # maximum distance (in pe) from peak, to avoid strong impact of outliers:
     maxq = 10
     # calculate widening of the noise bump:
-    added_noise = (np.sum(dq[dq<maxq]**2)/len(dq[dq<maxq]) -
-                   np.sum(dqmc[dqmc<maxq]**2)/len(dqmc[dqmc < maxq]))
+    added_noise = (np.sum(dq[dq < maxq] ** 2) / len(dq[dq < maxq]) -
+                   np.sum(dqmc[dqmc < maxq] ** 2) / len(dqmc[dqmc < maxq]))
     extra_noise_in_dim_pixels = max(0., added_noise)
 
     return extra_noise_in_dim_pixels, extra_bias_in_dim_pixels, \
-           extra_noise_in_bright_pixels
-
+        extra_noise_in_bright_pixels
 
-def tune_nsb_on_waveform(waveform, added_nsb_fraction, original_nsb,
-                         dt, pulse_templates, gain, charge_spe_cumulative_pdf):
-    """
-    Inject single photon pulses in existing R1 waveforms to increase NSB.
-    Parameters
-    ----------
-    waveform: charge (p.e. / ns) in each pixel and sampled time
-    added_nsb_fraction: fraction of the original NSB in simulation to be added
-    original_nsb: original NSB rate (astropy unit Hz)
-    dt: time between waveform samples (astropy unit s)
-    pulse_templates: `lstchain.data.NormalizedPulseTemplate` containing
-    the single p.e. pulse template used for the injection
-    gain: gain channel identifier for each pixel
-    charge_spe_cumulative_pdf: `scipy.interpolate.interp1d` Single p.e. gain
-    fluctuation cumulative pdf used to randomise the normalisation of
-    injected pulses
-    """
-    n = 25
-    n_pixels, n_samples = waveform.shape
-    duration = (n + n_samples) * dt  # TODO check needed time window, effect of edges
-    t = np.arange(-n, n_samples) * dt.value
-    mean_added_nsb = (added_nsb_fraction * original_nsb) * duration
-    rng = np.random.default_rng()
-    additional_nsb = rng.poisson(mean_added_nsb, n_pixels)
-    added_nsb_time = rng.uniform(-n * dt.value, -n * dt.value + duration.value, (n_pixels, max(additional_nsb)))
-    added_nsb_amp = charge_spe_cumulative_pdf(rng.uniform(size=(n_pixels, max(additional_nsb))))
-    baseline_correction = (added_nsb_fraction * original_nsb * dt).value
-    waveform -= baseline_correction
-    waveform += nsb_only_waveforms(
-        time=t[n:],
-        is_high_gain=gain,
-        additional_nsb=additional_nsb,
-        amplitude=added_nsb_amp,
-        t_0=added_nsb_time,
-        t0_template=pulse_templates.t0,
-        dt_template=pulse_templates.dt,
-        a_hg_template=pulse_templates.amplitude_HG,
-        a_lg_template=pulse_templates.amplitude_LG
-    )
 
 def calculate_required_additional_nsb(simtel_filename, data_dl1_filename, config=None):
-    # TODO check if good estimation
     # TODO reduce duplicated code with 'calculate_noise_parameters'
     """
     Calculates the additional NSB needed in the MC waveforms
     to match a real data DL1 file
+
     Parameters
     ----------
     simtel_filename: a simtel file containing showers, from the production
@@ -467,11 +425,13 @@ def calculate_required_additional_nsb(simtel_filename, data_dl1_filename, config
     agreement of data and simulations.
     config: configuration containing the calibration
     settings used for processing both the data and the MC files above
+
     Returns
     -------
     extra_nsb: Fraction of the additional NSB in data compared to MC.
     data_ped_variance: Pedestal variance from data
     mc_ped_variance: Pedestal variance from MC
+
     """
 
     log.setLevel(logging.INFO)
@@ -577,3 +537,165 @@ def calculate_required_additional_nsb(simtel_filename, data_dl1_filename, config
     extra_nsb = ((data_ped_variance - mc_ped_variance)
                  / mc_ped_variance)
     return extra_nsb, data_ped_variance, mc_ped_variance
+
+
+class WaveformNsbTunner:
+    """
+    Handles the injection of additional NSB pulses in waveforms.
+    """
+    def __init__(self, added_nsb_fraction, original_nsb, pulse_template, charge_spe_cumulative_pdf,
+                 pre_computed_multiplicity=10):
+        """
+        Parameters
+        ----------
+        added_nsb_fraction: `float`
+            Fraction of `original_nsb` to be added during NSB injection. Can be more than 1.
+        original_nsb: `dict`[`int`,`astropy.units.Quantity`]
+            NSB frequency used in the original waveforms per telescope
+        pulse_template: `dict`[int,`lstchain.data.NormalizedPulseTemplate`]
+            Single photo-electron pulse template per telescope
+        charge_spe_cumulative_pdf: `scipy.interpolate.interp1d`
+            Cumulative amplitude probability density function for single photo-electrons
+        pre_computed_multiplicity: `int`
+            Multiplicative factor on the number of pixels used to determine the number of pre-generated, nsb-only,
+            waveforms. Later used during event modification. Set to 0 to always compute the correction on the fly.
+
+        """
+        self.added_nsb_fraction = added_nsb_fraction
+        self.original_nsb = original_nsb
+        self.pulse_template = pulse_template
+        self.charge_spe_cumulative_pdf = charge_spe_cumulative_pdf
+        self.multiplicity = pre_computed_multiplicity
+
+        # Number of extra time sample added at the start of the time window used to inject NSB pulses
+        # It should be large enough to account for pulses created by NSB hits before the recorded time window
+        self.extra_samples = 25
+
+        self.nsb_waveforms = {}
+        self.nb_simulated = {}
+        self.rng = np.random.default_rng()
+        if self.multiplicity == 0:
+            self.tune_nsb_on_waveform = self._tune_nsb_on_waveform_direct
+        else:
+            self.tune_nsb_on_waveform = self._tune_nsb_on_waveform_precomputed
+
+    def initialise_waveforms(self, waveform, dt, tel_id):
+        """
+        Creates an array of nsb only waveforms later used to inject additional nsb in events.
+        Called once for each telescope id at first encounter in `self._tune_nsb_on_waveform_precomputed`.
+
+        Parameters
+        ----------
+        waveform: ndarray
+            Waveform used to know the number of pixels and time samples for a given telescope id
+        dt: `astropy.units.Quantity`
+            Time between samples
+        tel_id: `int`
+            Telescope identifier for which nsb waveforms are generated
+
+        """
+        log.info(f"Pre-generating nsb waveforms for nsb tuning and telescope id {tel_id}.")
+        n_pixels, n_samples = waveform.shape
+        baseline_correction = -(self.added_nsb_fraction * self.original_nsb[tel_id] * dt).to_value("")
+        nsb_waveforms = np.full((self.multiplicity * n_pixels, 2, n_samples), baseline_correction)
+        duration = (self.extra_samples + n_samples) * dt
+        t = np.arange(-self.extra_samples, n_samples) * dt.to_value(u.ns)
+        mean_added_nsb = (self.added_nsb_fraction * self.original_nsb[tel_id] * duration).to_value("")
+        additional_nsb = self.rng.poisson(mean_added_nsb, self.multiplicity * n_pixels)
+        added_nsb_time = self.rng.uniform(-self.extra_samples * dt.to_value(u.ns),
+                                          -self.extra_samples * dt.to_value(u.ns) + duration.to_value(u.ns),
+                                          (self.multiplicity * n_pixels, max(additional_nsb)))
+        added_nsb_amp = self.charge_spe_cumulative_pdf(
+            self.rng.uniform(size=(self.multiplicity * n_pixels, max(additional_nsb))))
+        nsb_waveforms[:, 0, :] += nsb_only_waveforms(
+            time=t[self.extra_samples:],
+            is_high_gain=np.zeros(self.multiplicity * n_pixels),
+            additional_nsb=additional_nsb,
+            amplitude=added_nsb_amp,
+            t_0=added_nsb_time,
+            t0_template=self.pulse_template[tel_id].t0,
+            dt_template=self.pulse_template[tel_id].dt,
+            a_hg_template=self.pulse_template[tel_id].amplitude_HG,
+            a_lg_template=self.pulse_template[tel_id].amplitude_LG
+        )
+        nsb_waveforms[:, 1, :] += nsb_only_waveforms(
+            time=t[self.extra_samples:],
+            is_high_gain=np.ones(self.multiplicity * n_pixels),
+            additional_nsb=additional_nsb,
+            amplitude=added_nsb_amp,
+            t_0=added_nsb_time,
+            t0_template=self.pulse_template[tel_id].t0,
+            dt_template=self.pulse_template[tel_id].dt,
+            a_hg_template=self.pulse_template[tel_id].amplitude_HG,
+            a_lg_template=self.pulse_template[tel_id].amplitude_LG
+        )
+        self.nsb_waveforms[tel_id] = nsb_waveforms
+        self.nb_simulated[tel_id] = self.multiplicity * n_pixels
+
+    def _tune_nsb_on_waveform_precomputed(self, waveform, tel_id, is_high_gain, subarray):
+        """
+        Inject single photon pulses in existing R1 waveforms to increase NSB using pre-computed pure nsb waveforms.
+
+        Parameters
+        ----------
+        waveform: ndarray
+            Charge (p.e. / ns) in each pixel and sampled time
+        tel_id: `int`
+            Telescope id associated to the waveform to tune
+        is_high_gain: ndarray of boolean
+            Gain channel used per pixel: True=hg, False=lg
+        subarray: `ctapipe.instrument.subarray.SubarrayDescription`
+
+        """
+        if tel_id not in self.nsb_waveforms:
+            readout = subarray.tel[tel_id].camera.readout
+            sampling_rate = readout.sampling_rate
+            dt = (1.0 / sampling_rate).to(u.s)
+            self.initialise_waveforms(waveform, dt, tel_id)
+        n_pixels, _ = waveform.shape
+        # The nsb_waveform array is randomised along the first axis,
+        # then the n_pixels first elements with correct gain are used for the injection
+        self.rng.shuffle(self.nsb_waveforms[tel_id])
+        waveform += ((is_high_gain == 0)[:, np.newaxis] * self.nsb_waveforms[tel_id][:n_pixels, 0] +
+                     (is_high_gain == 1)[:, np.newaxis] * self.nsb_waveforms[tel_id][:n_pixels, 1])
+
+    def _tune_nsb_on_waveform_direct(self, waveform, tel_id, is_high_gain, subarray):
+        """
+        Inject single photon pulses in existing R1 waveforms to increase NSB.
+
+        Parameters
+        ----------
+        waveform: ndarray
+            Charge (p.e. / ns) in each pixel and sampled time
+        tel_id: `int`
+            Telescope id associated to the waveform to tune
+        is_high_gain: ndarray of boolean
+            Gain channel used per pixel: True=hg, False=lg
+        subarray: `ctapipe.instrument.subarray.SubarrayDescription`
+
+        """
+        n_pixels, n_samples = waveform.shape
+        readout = subarray.tel[tel_id].camera.readout
+        sampling_rate = readout.sampling_rate
+        dt = (1.0 / sampling_rate).to(u.s)
+        duration = (self.extra_samples + n_samples) * dt
+        t = np.arange(-self.extra_samples, n_samples) * dt.to_value(u.ns)
+        mean_added_nsb = (self.added_nsb_fraction * self.original_nsb[tel_id] * duration).to_value("")
+        additional_nsb = self.rng.poisson(mean_added_nsb, n_pixels)
+        added_nsb_time = self.rng.uniform(-self.extra_samples * dt.to_value(u.ns),
+                                          -self.extra_samples * dt.to_value(u.ns) + duration.to_value(u.ns),
+                                          (n_pixels, max(additional_nsb)))
+        added_nsb_amp = self.charge_spe_cumulative_pdf(self.rng.uniform(size=(n_pixels, max(additional_nsb))))
+        baseline_correction = (self.added_nsb_fraction * self.original_nsb[tel_id] * dt).to_value("")
+        waveform -= baseline_correction
+        waveform += nsb_only_waveforms(
+            time=t[self.extra_samples:],
+            is_high_gain=is_high_gain,
+            additional_nsb=additional_nsb,
+            amplitude=added_nsb_amp,
+            t_0=added_nsb_time,
+            t0_template=self.pulse_template[tel_id].t0,
+            dt_template=self.pulse_template[tel_id].dt,
+            a_hg_template=self.pulse_template[tel_id].amplitude_HG,
+            a_lg_template=self.pulse_template[tel_id].amplitude_LG
+        )
diff --git a/lstchain/io/io.py b/lstchain/io/io.py
index 79a82ad5..d10498b0 100644
--- a/lstchain/io/io.py
+++ b/lstchain/io/io.py
@@ -1279,7 +1279,7 @@ def extract_simulation_nsb(filename):
         for _, line in f.history:
             line = line.decode('utf-8').strip().split(' ')
             if next_nsb and line[0] == 'NIGHTSKY_BACKGROUND':
-                nsb[tel_id] = float(line[1].strip('all:'))
+                nsb[tel_id] = float(line[1].strip('all:')) * u.GHz
                 tel_id = tel_id+1
             if line[0] == 'STORE_PHOTOELECTRONS':
                 next_nsb = True
@@ -1288,7 +1288,7 @@ def extract_simulation_nsb(filename):
     log.warning('Original MC night sky background extracted from the config history in the simtel file.\n'
                 'This is done for existing LST MC such as the one created using: '
                 'https://github.com/cta-observatory/lst-sim-config/tree/sim-tel_LSTProd2_MAGICST0316'
-                '\nExtracted values are: ' + str(np.asarray(nsb)) + 'GHz. Check that it corresponds to expectations.')
+                '\nExtracted values are: ' + str(np.asarray(nsb)) + '. Check that it corresponds to expectations.')
     return nsb
 
 
diff --git a/lstchain/reco/r0_to_dl1.py b/lstchain/reco/r0_to_dl1.py
index 3f2db65a..ce08cbaa 100644
--- a/lstchain/reco/r0_to_dl1.py
+++ b/lstchain/reco/r0_to_dl1.py
@@ -23,7 +23,7 @@ from ctapipe.image import (
     tailcuts_clean,
 )
 from ctapipe.image import number_of_islands, apply_time_delta_cleaning
-from ctapipe.io import EventSource, HDF5TableWriter, DataWriter 
+from ctapipe.io import EventSource, HDF5TableWriter, DataWriter
 from ctapipe.utils import get_dataset_path
 from traitlets.config import Config
 from ctapipe_io_lst.constants import (
@@ -37,7 +37,7 @@ from ..data import NormalizedPulseTemplate
 from ..calib.camera import load_calibrator_from_config
 from ..calib.camera.calibration_calculator import CalibrationCalculator
 from ..image.cleaning import apply_dynamic_cleaning
-from ..image.modifier import tune_nsb_on_waveform, calculate_required_additional_nsb
+from ..image.modifier import calculate_required_additional_nsb, WaveformNsbTunner
 from .reconstructor import TimeWaveformFitter
 from ..image.muon import analyze_muon_event, tag_pix_thr
 from ..image.muon import create_muon_table, fill_muon_event
@@ -59,7 +59,6 @@ from ..io import (
     write_subarray_tables
 )
 
-
 from ..io.io import add_column_table, extract_simulation_nsb, dl1_params_lstcam_key, get_resource_path
 from ..io.lstcontainers import ExtraImageInfo, DL1MonitoringEventIndexContainer
 from ..paths import parse_r0_filename, run_to_dl1_filename, r0_to_dl1_filename
@@ -67,7 +66,6 @@ from ..visualization.plot_reconstructor import plot_debug
 
 logger = logging.getLogger(__name__)
 
-
 __all__ = [
     'add_disp_to_parameters_table',
     'get_dl1',
@@ -75,7 +73,6 @@ __all__ = [
     'r0_to_dl1',
 ]
 
-
 cleaning_method = tailcuts_clean
 
 
@@ -169,11 +166,11 @@ def parametrize_image(image, peak_time, signal_pixels, camera_geometry, focal_le
 
 
 def get_dl1(
-    calibrated_event,
-    subarray,
-    telescope_id,
-    dl1_container=None,
-    custom_config={}
+        calibrated_event,
+        subarray,
+        telescope_id,
+        dl1_container=None,
+        custom_config={}
 ):
     """
     Return a DL1ParametersContainer of extracted features from a calibrated event.
@@ -312,9 +309,9 @@ def apply_lh_fit(
 
 
 def r0_to_dl1(
-    input_filename=None,
-    output_filename=None,
-    custom_config={},
+        input_filename=None,
+        output_filename=None,
+        custom_config={},
 ):
     """
     Chain r0 to dl1
@@ -375,11 +372,10 @@ def r0_to_dl1(
     )
 
     if not is_simu:
-
         # Pulse extractor for muon ring analysis. Same parameters (window_width and _shift) as the one for showers, but
         # using GlobalPeakWindowSum, since the signal for the rings is expected to be very isochronous
-        r1_dl1_calibrator_for_muon_rings = CameraCalibrator(image_extractor_type = config['image_extractor_for_muons'],
-                                                            config=Config(config),subarray = subarray)
+        r1_dl1_calibrator_for_muon_rings = CameraCalibrator(image_extractor_type=config['image_extractor_for_muons'],
+                                                            config=Config(config), subarray=subarray)
 
         # Component to process interleaved pedestal and flat-fields
         calib_config = Config({config['calibration_product']: config[config['calibration_product']]})
@@ -407,15 +403,15 @@ def r0_to_dl1(
             metadata=metadata,
         )
     nsb_tuning = False
-    nsb_tuning_args = None
+    nsb_tunner = None
     if 'waveform_nsb_tuning' in config.keys():
         nsb_tuning = config['waveform_nsb_tuning']['nsb_tuning']
         if nsb_tuning:
             if is_simu:
                 nsb_original = extract_simulation_nsb(input_filename)
-                pulse_template = NormalizedPulseTemplate.load_from_eventsource(
-                    subarray.tel[1].camera.readout, resample=True
-                )
+                pulse_templates = {tel_id: NormalizedPulseTemplate.load_from_eventsource(
+                    subarray.tel[tel_id].camera.readout, resample=True)
+                    for tel_id in config['source_config']['LSTEventSource']['allowed_tels']}
                 if 'nsb_tuning_ratio' in config['waveform_nsb_tuning'].keys():
                     # get value from config to possibly extract it beforehand on multiple files for averaging purposes
                     # or gain time
@@ -426,18 +422,23 @@ def r0_to_dl1(
                     nsb_tuning_ratio = calculate_required_additional_nsb(input_filename,
                                                                          config['waveform_nsb_tuning']['target_data'],
                                                                          config=config)[0]
-
-                spe_location = config['waveform_nsb_tuning']['spe_location'] if 'spe_location' in config['waveform_nsb_tuning'] and config['waveform_nsb_tuning']['spe_location'] is not None else get_resource_path("data/SinglePhE_ResponseInPhE_expo2Gaus.dat")
+                spe_location = (config['waveform_nsb_tuning']['spe_location']
+                                or get_resource_path("data/SinglePhE_ResponseInPhE_expo2Gaus.dat"))
                 spe = np.loadtxt(spe_location).T
                 spe_integral = np.cumsum(spe[1])
                 charge_spe_cumulative_pdf = interp1d(spe_integral, spe[0], kind='cubic',
                                                      bounds_error=False, fill_value=0.,
                                                      assume_sorted=True)
+                pre_computed_multiplicity = config['waveform_nsb_tuning'].get('pre_computed_multiplicity', 10)
                 logger.info('Tuning NSB on MC waveform from '
                             + str(np.asarray(nsb_original))
-                            + 'GHz to {0:d}%'.format(int(nsb_tuning_ratio * 100 + 100.5))
+                            + ' to {0:d}%'.format(int(nsb_tuning_ratio * 100 + 100.5))
                             + ' for telescopes ids ' + str(config['source_config']['LSTEventSource']['allowed_tels']))
-                nsb_tuning_args = [nsb_tuning_ratio, nsb_original, pulse_template, charge_spe_cumulative_pdf]
+                nsb_tunner = WaveformNsbTunner(nsb_tuning_ratio,
+                                               nsb_original,
+                                               pulse_templates,
+                                               charge_spe_cumulative_pdf,
+                                               pre_computed_multiplicity)
             else:
                 logger.warning('NSB tuning on waveform active in config but file is real data, option will be ignored')
                 nsb_tuning = False
@@ -450,7 +451,7 @@ def r0_to_dl1(
             tmp_source = EventSource(input_url=input_filename,
                                      config=Config(config["source_config"]))
             if is_simu:
-                lhfit_fitter.get_ped_from_true_signal_less(tmp_source, nsb_tuning_args)
+                lhfit_fitter.get_ped_from_true_signal_less(tmp_source, nsb_tunner)
             else:
                 lhfit_fitter.get_ped_from_interleaved(tmp_source)
             del tmp_source
@@ -464,21 +465,22 @@ def r0_to_dl1(
         # create output dir in the data-tree if necessary
         dir = f"{dir}/interleaved"
         os.makedirs(dir, exist_ok=True)
-        if 'dl1' in name: 
+        if 'dl1' in name:
             name = name.replace('dl1', 'interleaved').replace('LST-1.1', 'LST-1')
         else:
             name = f"interleaved_{name}"
         interleaved_output_file = Path(dir, name)
-        interleaved_writer = DataWriter(event_source=source,output_path=interleaved_output_file,config=interleaved_writer_config)
+        interleaved_writer = DataWriter(event_source=source, output_path=interleaved_output_file,
+                                        config=interleaved_writer_config)
         interleaved_writer._writer.exclude("/r1/event/telescope/.*", "selected_gain_channel")
 
-    with HDF5TableWriter( 
-        filename=output_filename,
-        group_name='dl1/event',
-        mode='a',
-        filters=HDF5_ZSTD_FILTERS,
-        add_prefix=True,
-        # overwrite=True,
+    with HDF5TableWriter(
+            filename=output_filename,
+            group_name='dl1/event',
+            mode='a',
+            filters=HDF5_ZSTD_FILTERS,
+            add_prefix=True,
+            # overwrite=True,
     ) as writer:
 
         setup_writer(writer, source.subarray, is_simulation=is_simu)
@@ -512,10 +514,10 @@ def r0_to_dl1(
 
                     tel_id = calibration_calculator.tel_id
 
-
-                    #initialize the event monitoring data
+                    # initialize the event monitoring data
                     event.mon = deepcopy(source.r0_r1_calibrator.mon_data)
-                    for container in [event.mon.tel[tel_id].pedestal, event.mon.tel[tel_id].flatfield, event.mon.tel[tel_id].calibration]:
+                    for container in [event.mon.tel[tel_id].pedestal, event.mon.tel[tel_id].flatfield,
+                                      event.mon.tel[tel_id].calibration]:
                         add_global_metadata(container, metadata)
                         add_config_metadata(container, config)
 
@@ -537,42 +539,38 @@ def r0_to_dl1(
                     # these data a supposed to be replaced by the Cat_B data in a short future
                     if new_ped_event or new_ff_event:
                         write_calibration_data(writer,
-                                           calibration_index,
-                                           event.mon.tel[tel_id],
-                                           new_ped=new_ped_event, new_ff=new_ff_event)
-                    
+                                               calibration_index,
+                                               event.mon.tel[tel_id],
+                                               new_ped=new_ped_event, new_ff=new_ff_event)
+
                     # write the calibrated R1 waveform without gain selection
                     source.r0_r1_calibrator.select_gain = False
                     source.r0_r1_calibrator.calibrate(event)
-                
+
                     if interleaved_writer is not None:
                         interleaved_writer(event)
 
                     # gain select the events
                     source.r0_r1_calibrator.select_gain = True
 
-                    r1 = event.r1.tel[tel_id]                   
+                    r1 = event.r1.tel[tel_id]
                     r1.selected_gain_channel = source.r0_r1_calibrator.gain_selector(event.r0.tel[tel_id].waveform)
                     r1.waveform = r1.waveform[r1.selected_gain_channel, PIXEL_INDEX]
 
                     event.calibration.tel[tel_id].dl1.time_shift = \
-                    event.calibration.tel[tel_id].dl1.time_shift[r1.selected_gain_channel, PIXEL_INDEX]
-                    
+                        event.calibration.tel[tel_id].dl1.time_shift[r1.selected_gain_channel, PIXEL_INDEX]
+
                     event.calibration.tel[tel_id].dl1.relative_factor = \
-                    event.calibration.tel[tel_id].dl1.relative_factor[r1.selected_gain_channel, PIXEL_INDEX]
-                    
+                        event.calibration.tel[tel_id].dl1.relative_factor[r1.selected_gain_channel, PIXEL_INDEX]
+
             # Option to add nsb in waveforms
             if nsb_tuning:
                 # FIXME? assumes same correction ratio for all telescopes
                 for tel_id in config['source_config']['LSTEventSource']['allowed_tels']:
                     waveform = event.r1.tel[tel_id].waveform
-                    readout = subarray.tel[tel_id].camera.readout
-                    sampling_rate = readout.sampling_rate.to_value(u.GHz)
-                    dt = (1.0 / sampling_rate)
                     selected_gains = event.r1.tel[tel_id].selected_gain_channel
-                    mask_high = (selected_gains == 0)
-                    tune_nsb_on_waveform(waveform, nsb_tuning_ratio, nsb_original[tel_id] * u.GHz,
-                                         dt * u.ns, pulse_template, mask_high, charge_spe_cumulative_pdf)
+                    mask_high = selected_gains == 0
+                    nsb_tunner.tune_nsb_on_waveform(waveform, tel_id, mask_high, subarray)
 
             # create image for all events
             r1_dl1_calibrator(event)
@@ -604,7 +602,6 @@ def r0_to_dl1(
 
                 dl1_container.fill_event_info(event)
 
-
                 # Will determine whether this event has to be written to the
                 # DL1 output or not.
                 if is_simu:
@@ -629,7 +626,7 @@ def r0_to_dl1(
                 if not is_simu:
 
                     calibration_mon = source.r0_r1_calibrator.mon_data.tel[telescope_id].calibration
- 
+
                     dl1_container.ucts_time = 0
                     # convert Time to unix timestamp in (UTC) to keep compatibility
                     # with older lstchain
@@ -644,7 +641,6 @@ def r0_to_dl1(
                 else:
                     dl1_container.trigger_type = event.trigger.event_type
 
-
                 dl1_container.az_tel = event.pointing.tel[telescope_id].azimuth
                 dl1_container.alt_tel = event.pointing.tel[telescope_id].altitude
                 dl1_container.sin_az_tel = np.sin(dl1_container.az_tel)
@@ -682,7 +678,7 @@ def r0_to_dl1(
                     bad_pixels = calibration_mon.unusable_pixels[0]
 
                     # Set to 0 unreliable pixels:
-                    image = dl1_tel.image*(~bad_pixels)
+                    image = dl1_tel.image * (~bad_pixels)
 
                     # process only promising events, in terms of # of pixels with large signals:
                     if tag_pix_thr(image):
@@ -696,7 +692,7 @@ def r0_to_dl1(
                         # integrator in case of crazy pixels!  TBD: can this be done in a simpler
                         # way?
                         bad_pixels = bad_pixels_hg | bad_pixels_lg
-                        bad_waveform = np.transpose(np.array(numsamples*[bad_pixels]))
+                        bad_waveform = np.transpose(np.array(numsamples * [bad_pixels]))
 
                         # print('hg bad pixels:',np.where(bad_pixels_hg))
                         # print('lg bad pixels:',np.where(bad_pixels_lg))
@@ -714,10 +710,10 @@ def r0_to_dl1(
                             good_ring = False
                         else:
                             muonintensityparam, dist_mask, \
-                            ring_size, size_outside_ring, muonringparam, \
-                            good_ring, radial_distribution, \
-                            mean_pixel_charge_around_ring,\
-                            muonpars = \
+                                ring_size, size_outside_ring, muonringparam, \
+                                good_ring, radial_distribution, \
+                                mean_pixel_charge_around_ring, \
+                                muonpars = \
                                 analyze_muon_event(subarray,
                                                    tel_id, event.index.event_id,
                                                    image, good_ring_config=None,
@@ -758,10 +754,10 @@ def r0_to_dl1(
 
                 # writes mc information per telescope, including photo electron image
                 if (
-                    is_simu
-                    and config['write_pe_image']
-                    and event.simulation.tel[telescope_id].true_image is not None
-                    and event.simulation.tel[telescope_id].true_image.any()
+                        is_simu
+                        and config['write_pe_image']
+                        and event.simulation.tel[telescope_id].true_image is not None
+                        and event.simulation.tel[telescope_id].true_image.any()
                 ):
                     event.simulation.tel[telescope_id].prefix = ''
                     writer.write(
@@ -799,13 +795,12 @@ def r0_to_dl1(
         muon_output_filename = Path(dir, name)
         table = Table(muon_parameters)
         table.write(muon_output_filename, format='fits', overwrite=True)
-        
+
         # close the interleaved output file and write metadata
         if interleaved_writer is not None:
             interleaved_writer.finish()
 
 
-
 def add_disp_to_parameters_table(dl1_file, table_path, focal):
     """
     Reconstruct the disp parameters and source position from a DL1 parameters table and write the result in the file
diff --git a/lstchain/reco/reconstructor.py b/lstchain/reco/reconstructor.py
index 492694af..31335f3f 100644
--- a/lstchain/reco/reconstructor.py
+++ b/lstchain/reco/reconstructor.py
@@ -9,7 +9,6 @@ from ctapipe.core import TelescopeComponent
 from ctapipe.core.traits import Bool, Float, FloatTelescopeParameter, Int, Unicode
 
 from lstchain.data.normalised_pulse_template import NormalizedPulseTemplate
-from lstchain.image.modifier import tune_nsb_on_waveform
 from lstchain.io.lstcontainers import DL1LikelihoodParametersContainer
 from lstchain.reco.reconstructorCC import log_pdf as log_pdf
 
@@ -140,20 +139,15 @@ class TimeWaveformFitter(TelescopeComponent):
             self.error[tel_id] = std
             self.allowed_pixels = (std > 0.5 * np.median(std))
 
-    def get_ped_from_true_signal_less(self, source, nsb_tuning_args):
+    def get_ped_from_true_signal_less(self, source, nsb_tunner):
         """
         Use pixels with no true signal in MC to extract the pedestal variance, assumed uniform in the camera.
         """
         self.error = {}
         waveforms = {}
         n_pix = {}
-        dt = {}
         for tel_id in self.subarray.tel:
             waveforms[tel_id] = []
-            readout = self.subarray.tel[tel_id].camera.readout
-            sampling_rate = readout.sampling_rate.to_value(u.GHz)
-            dt[tel_id] = (1.0 / sampling_rate)
-            waveforms[tel_id] = []
         for i, event in enumerate(source):
             if i > 50 and len(waveforms) > 1000:
                 break
@@ -163,11 +157,10 @@ class TimeWaveformFitter(TelescopeComponent):
                         n_pix[tel_id] = event.r1.tel[tel_id].waveform.shape[0]
                     mask = event.simulation.tel[tel_id].true_image == 0
                     wave = event.r1.tel[tel_id].waveform
-                    if nsb_tuning_args is not None:
+                    if nsb_tunner is not None:
                         selected_gains = event.r1.tel[tel_id].selected_gain_channel
                         mask_high = (selected_gains == 0)
-                        tune_nsb_on_waveform(wave, nsb_tuning_args[0], nsb_tuning_args[1][tel_id] * u.GHz,
-                                             dt[tel_id] * u.ns, nsb_tuning_args[2], mask_high, nsb_tuning_args[3])
+                        nsb_tunner.tune_nsb_on_waveform(wave, tel_id, mask_high, self.subarray)
                     waveforms[tel_id].append(wave[mask].flatten())
         for tel_id, tel_waveforms in waveforms.items():
             self.error[tel_id] = np.full(n_pix[tel_id], np.nanstd(np.concatenate(tel_waveforms)))


Test patch:
diff --git a/lstchain/image/tests/test_modifier.py b/lstchain/image/tests/test_modifier.py
index d48596de..a23c25ef 100644
--- a/lstchain/image/tests/test_modifier.py
+++ b/lstchain/image/tests/test_modifier.py
@@ -77,17 +77,18 @@ def test_calculate_required_additional_nsb(mc_gamma_testfile, observed_dl1_files
 
 def test_tune_nsb_on_waveform():
     import astropy.units as u
+    from ctapipe_io_lst import LSTEventSource
     from scipy.interpolate import interp1d
     from lstchain.io.io import get_resource_path
-    from lstchain.image.modifier import tune_nsb_on_waveform
+    from lstchain.image.modifier import WaveformNsbTunner
     from lstchain.data.normalised_pulse_template import NormalizedPulseTemplate
 
     waveform = np.array(
-        [[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0],
-         [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]]
+        [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+         [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]
     )
-    added_nsb_fraction, original_nsb = 1.0, 0.1 * u.GHz
-    dt = 1 * u.ns
+    added_nsb_fraction, original_nsb = 1.0, {1: 0.1 * u.GHz}
+    subarray = LSTEventSource.create_subarray(1)
     amplitude_HG = np.zeros(40)
     amplitude_LG = np.zeros(40)
     amplitude_HG[19] = 0.25
@@ -96,17 +97,32 @@ def test_tune_nsb_on_waveform():
     amplitude_LG[19] = 0.4
     amplitude_LG[20] = 1.0
     amplitude_LG[21] = 0.4
-    time = np.linspace(-10,30,40)
-    pulse_templates = NormalizedPulseTemplate(amplitude_HG, amplitude_LG, time, amplitude_HG_err=None,
-                                              amplitude_LG_err=None)
-    gain = np.array(['HG', 'LG'])
+    time = np.linspace(-10, 30, 40)
+    pulse_templates = {1: NormalizedPulseTemplate(amplitude_HG, amplitude_LG, time, amplitude_HG_err=None,
+                                                  amplitude_LG_err=None)}
+    gain = np.array([1, 0])
     spe = np.loadtxt(get_resource_path("data/SinglePhE_ResponseInPhE_expo2Gaus.dat")).T
     spe_integral = np.cumsum(spe[1])
     charge_spe_cumulative_pdf = interp1d(spe_integral, spe[0], kind='cubic',
                                          bounds_error=False, fill_value=0.,
                                          assume_sorted=True)
-    tune_nsb_on_waveform(waveform, added_nsb_fraction, original_nsb,
-                         dt, pulse_templates, gain, charge_spe_cumulative_pdf)
-    #  assert may be randomly wrong in very unusual cases
+    nsb_tunner = WaveformNsbTunner(added_nsb_fraction,
+                                   original_nsb,
+                                   pulse_templates,
+                                   charge_spe_cumulative_pdf,
+                                   pre_computed_multiplicity=0)
+    nsb_tunner.tune_nsb_on_waveform(waveform, 1, gain, subarray)
+    assert np.any(waveform != 0)
+    assert np.isclose(np.mean(waveform), 0.0, atol=0.2)
+    waveform = np.array(
+        [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+         [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]
+    )
+    nsb_tunner = WaveformNsbTunner(added_nsb_fraction,
+                                   original_nsb,
+                                   pulse_templates,
+                                   charge_spe_cumulative_pdf,
+                                   pre_computed_multiplicity=10)
+    nsb_tunner.tune_nsb_on_waveform(waveform, 1, gain, subarray)
     assert np.any(waveform != 0)
     assert np.isclose(np.mean(waveform), 0.0, atol=0.2)
diff --git a/lstchain/io/tests/test_io.py b/lstchain/io/tests/test_io.py
index 6442eff6..9db95dea 100644
--- a/lstchain/io/tests/test_io.py
+++ b/lstchain/io/tests/test_io.py
@@ -154,9 +154,10 @@ def test_trigger_type_in_dl1_params(simulated_dl1_file):
 
 def test_extract_simulation_nsb(mc_gamma_testfile):
     from lstchain.io.io import extract_simulation_nsb
+    import astropy.units as u
 
     nsb = extract_simulation_nsb(mc_gamma_testfile)
-    assert np.isclose(nsb[1], 0.246, rtol=0.1)
+    assert np.isclose(nsb[1].to_value(u.GHz), 0.246, rtol=0.1)
 
 
 def test_remove_duplicated_events():
diff --git a/lstchain/tests/test_lstchain.py b/lstchain/tests/test_lstchain.py
index d88bab4b..771204f8 100644
--- a/lstchain/tests/test_lstchain.py
+++ b/lstchain/tests/test_lstchain.py
@@ -129,6 +129,10 @@ def test_r0_to_dl1_lhfit_mc(tmp_path, mc_gamma_testfile):
     config['lh_fit_config']["spatial_selection"] = 'dvr'
     config['lh_fit_config']["use_interleaved"] = True
     config['waveform_nsb_tuning']['nsb_tuning'] = True
+    config['waveform_nsb_tuning']['pre_computed_multiplicity'] = 10  # default value
+    r0_to_dl1(mc_gamma_testfile, custom_config=config, output_filename=tmp_path / "tmp.h5")
+    os.remove(tmp_path / "tmp.h5")
+    config['waveform_nsb_tuning']['pre_computed_multiplicity'] = 0
     r0_to_dl1(mc_gamma_testfile, custom_config=config, output_filename=tmp_path / "tmp.h5")