From 74c9c4636cafe6d0f76a610ed88fb3a668b8141f Mon Sep 17 00:00:00 2001
From: Gabriel Emery <gabriel-emery@hotmail.fr>
Date: Tue, 9 Jul 2024 16:21:05 +0200
Subject: [PATCH 1/6] Implement a new class to handle the NSB tuning on
 waveform. Adds the possibility to pre-generate NSB only waveforms to randomly
 apply to events.

---
 lstchain/data/lstchain_lhfit_config.json |   3 +-
 lstchain/image/modifier.py               | 202 ++++++++++++++++-------
 lstchain/image/tests/test_modifier.py    |  40 +++--
 lstchain/io/io.py                        |   4 +-
 lstchain/io/tests/test_io.py             |   3 +-
 lstchain/reco/r0_to_dl1.py               | 131 +++++++--------
 lstchain/reco/reconstructor.py           |  14 +-
 lstchain/tests/test_lstchain.py          |   3 +
 8 files changed, 242 insertions(+), 158 deletions(-)

diff --git a/lstchain/data/lstchain_lhfit_config.json b/lstchain/data/lstchain_lhfit_config.json
index 0536811083..4b876ab25e 100644
--- a/lstchain/data/lstchain_lhfit_config.json
+++ b/lstchain/data/lstchain_lhfit_config.json
@@ -280,7 +280,8 @@
   "waveform_nsb_tuning":{
     "nsb_tuning": false,
     "nsb_tuning_ratio": 0.52,
-    "spe_location": null
+    "spe_location": null,
+    "pre_computed_multiplicity": 0
   },
   "lh_fit_config": {
     "sigma_s": [
diff --git a/lstchain/image/modifier.py b/lstchain/image/modifier.py
index 1c0f28f853..dba5881509 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 @@
     '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,54 +398,14 @@ 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'
@@ -577,3 +535,119 @@ 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:
+    def __init__(self, added_nsb_fraction, original_nsb, pulse_template, charge_spe_cumulative_pdf,
+                 pre_computed_multiplicity=10):
+        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
+        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):
+        log.info(f"Pre-generating nsb waveforms for nsb tuning and telescope id {tel_id}.")
+        n = 25
+        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 = (n + n_samples) * dt  # TODO check needed time window, effect of edges
+        t = np.arange(-n, 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(-n * dt.to_value(u.ns), -n * 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[n:],
+            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[n:],
+            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: charge (p.e. / ns) in each pixel and sampled time
+        the single p.e. pulse template used for the injection
+        tel_id: Telescope id associated to the waveform to tune
+        is_high_gain: : boolean 1D array
+        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
+        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: charge (p.e. / ns) in each pixel and sampled time
+        the single p.e. pulse template used for the injection
+        tel_id: Telescope id associated to the waveform to tune
+        is_high_gain: : boolean 1D array
+        Gain channel used per pixel: True=hg, False=lg
+        subarray: `ctapipe.instrument.subarray.SubarrayDescription`
+        """
+        n = 25
+        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 = (n + n_samples) * dt  # TODO check needed time window, effect of edges
+        t = np.arange(-n, 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(-n * dt.to_value(u.ns), -n * 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[n:],
+            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/image/tests/test_modifier.py b/lstchain/image/tests/test_modifier.py
index d48596de52..a23c25eff2 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/io.py b/lstchain/io/io.py
index 79a82ad556..d10498b0be 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/io/tests/test_io.py b/lstchain/io/tests/test_io.py
index 6442eff66e..9db95deaa3 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/reco/r0_to_dl1.py b/lstchain/reco/r0_to_dl1.py
index 3f2db65a45..ce08cbaa64 100644
--- a/lstchain/reco/r0_to_dl1.py
+++ b/lstchain/reco/r0_to_dl1.py
@@ -23,7 +23,7 @@
     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 ..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 @@
     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 @@
 
 logger = logging.getLogger(__name__)
 
-
 __all__ = [
     'add_disp_to_parameters_table',
     'get_dl1',
@@ -75,7 +73,6 @@
     '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 492694af50..0a7606c24b 100644
--- a/lstchain/reco/reconstructor.py
+++ b/lstchain/reco/reconstructor.py
@@ -9,7 +9,7 @@
 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.image.modifier import WaveformNsbTunner
 from lstchain.io.lstcontainers import DL1LikelihoodParametersContainer
 from lstchain.reco.reconstructorCC import log_pdf as log_pdf
 
@@ -140,20 +140,15 @@ def get_ped_from_interleaved(self, source):
             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 +158,10 @@ def get_ped_from_true_signal_less(self, source, nsb_tuning_args):
                         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)))
diff --git a/lstchain/tests/test_lstchain.py b/lstchain/tests/test_lstchain.py
index d88bab4b76..1d02a054a9 100644
--- a/lstchain/tests/test_lstchain.py
+++ b/lstchain/tests/test_lstchain.py
@@ -130,6 +130,9 @@ def test_r0_to_dl1_lhfit_mc(tmp_path, mc_gamma_testfile):
     config['lh_fit_config']["use_interleaved"] = True
     config['waveform_nsb_tuning']['nsb_tuning'] = True
     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")
 
 
 @pytest.mark.private_data

From 7082efd3f8d4437ae14d64061f2794d141f6a639 Mon Sep 17 00:00:00 2001
From: Gabriel Emery <gabriel-emery@hotmail.fr>
Date: Tue, 9 Jul 2024 17:15:02 +0200
Subject: [PATCH 2/6] Remove unused import

---
 lstchain/reco/reconstructor.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/lstchain/reco/reconstructor.py b/lstchain/reco/reconstructor.py
index 0a7606c24b..31335f3fc3 100644
--- a/lstchain/reco/reconstructor.py
+++ b/lstchain/reco/reconstructor.py
@@ -9,7 +9,6 @@
 from ctapipe.core.traits import Bool, Float, FloatTelescopeParameter, Int, Unicode
 
 from lstchain.data.normalised_pulse_template import NormalizedPulseTemplate
-from lstchain.image.modifier import WaveformNsbTunner
 from lstchain.io.lstcontainers import DL1LikelihoodParametersContainer
 from lstchain.reco.reconstructorCC import log_pdf as log_pdf
 

From c06d74296c0e576e82de2e78680ab2c14a08937b Mon Sep 17 00:00:00 2001
From: Gabriel Emery <gabriel-emery@hotmail.fr>
Date: Thu, 11 Jul 2024 11:18:59 +0200
Subject: [PATCH 3/6] Improve docstring and comments.

---
 lstchain/image/modifier.py | 55 ++++++++++++++++++++++++++++++--------
 1 file changed, 44 insertions(+), 11 deletions(-)

diff --git a/lstchain/image/modifier.py b/lstchain/image/modifier.py
index dba5881509..62b7b9c1cd 100644
--- a/lstchain/image/modifier.py
+++ b/lstchain/image/modifier.py
@@ -407,7 +407,6 @@ def calculate_noise_parameters(simtel_filename, data_dl1_filename,
 
 
 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
@@ -538,8 +537,26 @@ def calculate_required_additional_nsb(simtel_filename, data_dl1_filename, config
 
 
 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
@@ -554,6 +571,18 @@ def __init__(self, added_nsb_fraction, original_nsb, pulse_template, charge_spe_
             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 = 25
         n_pixels, n_samples = waveform.shape
@@ -597,11 +626,12 @@ def _tune_nsb_on_waveform_precomputed(self, waveform, tel_id, is_high_gain, suba
         Inject single photon pulses in existing R1 waveforms to increase NSB using pre-computed pure nsb waveforms.
         Parameters
         ----------
-        waveform: charge (p.e. / ns) in each pixel and sampled time
-        the single p.e. pulse template used for the injection
-        tel_id: Telescope id associated to the waveform to tune
-        is_high_gain: : boolean 1D array
-        Gain channel used per pixel: True=hg, False=lg
+        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:
@@ -610,6 +640,8 @@ def _tune_nsb_on_waveform_precomputed(self, waveform, tel_id, is_high_gain, suba
             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])
@@ -619,11 +651,12 @@ 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: charge (p.e. / ns) in each pixel and sampled time
-        the single p.e. pulse template used for the injection
-        tel_id: Telescope id associated to the waveform to tune
-        is_high_gain: : boolean 1D array
-        Gain channel used per pixel: True=hg, False=lg
+        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 = 25

From 2a4e061b763f678a5ba06dda615e3a2e69c376e5 Mon Sep 17 00:00:00 2001
From: Gabriel Emery <gabriel-emery@hotmail.fr>
Date: Thu, 11 Jul 2024 11:33:48 +0200
Subject: [PATCH 4/6] Improve docstring

---
 lstchain/image/modifier.py | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/lstchain/image/modifier.py b/lstchain/image/modifier.py
index 62b7b9c1cd..5bda1a5808 100644
--- a/lstchain/image/modifier.py
+++ b/lstchain/image/modifier.py
@@ -411,6 +411,7 @@ def calculate_required_additional_nsb(simtel_filename, data_dl1_filename, config
     """
     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
@@ -424,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)
@@ -556,6 +559,7 @@ def __init__(self, added_nsb_fraction, original_nsb, pulse_template, charge_spe_
         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
@@ -574,14 +578,16 @@ 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
+            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
+            Telescope identifier for which nsb waveforms are generated
+
         """
         log.info(f"Pre-generating nsb waveforms for nsb tuning and telescope id {tel_id}.")
         n = 25
@@ -624,15 +630,17 @@ def initialise_waveforms(self, waveform, dt, tel_id):
     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
+            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
@@ -649,15 +657,17 @@ def _tune_nsb_on_waveform_precomputed(self, waveform, tel_id, is_high_gain, suba
     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
+            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 = 25
         n_pixels, n_samples = waveform.shape

From a2944185651c6cd134232e6315223287da0a3933 Mon Sep 17 00:00:00 2001
From: Gabriel Emery <gabriel-emery@hotmail.fr>
Date: Thu, 11 Jul 2024 13:29:47 +0200
Subject: [PATCH 5/6] Update waveform nsb tunning section of config files. Move
 parameter definition done twice in WaveformNsbTunner to init. Clarify
 parameter naming and add comment.

---
 lstchain/data/lstchain_lhfit_config.json    |  4 +--
 lstchain/data/lstchain_standard_config.json |  5 ++--
 lstchain/image/modifier.py                  | 27 ++++++++++++---------
 3 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/lstchain/data/lstchain_lhfit_config.json b/lstchain/data/lstchain_lhfit_config.json
index 4b876ab25e..c761c19937 100644
--- a/lstchain/data/lstchain_lhfit_config.json
+++ b/lstchain/data/lstchain_lhfit_config.json
@@ -279,9 +279,9 @@
   },
   "waveform_nsb_tuning":{
     "nsb_tuning": false,
-    "nsb_tuning_ratio": 0.52,
+    "nsb_tuning_ratio": 0.5,
     "spe_location": null,
-    "pre_computed_multiplicity": 0
+    "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 d63c9e0f8e..74312da664 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 5bda1a5808..9f52390238 100644
--- a/lstchain/image/modifier.py
+++ b/lstchain/image/modifier.py
@@ -566,6 +566,11 @@ def __init__(self, added_nsb_fraction, original_nsb, pulse_template, charge_spe_
         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()
@@ -590,20 +595,20 @@ def initialise_waveforms(self, waveform, dt, tel_id):
 
         """
         log.info(f"Pre-generating nsb waveforms for nsb tuning and telescope id {tel_id}.")
-        n = 25
         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 = (n + n_samples) * dt  # TODO check needed time window, effect of edges
-        t = np.arange(-n, n_samples) * dt.to_value(u.ns)
+        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(-n * dt.to_value(u.ns), -n * dt.to_value(u.ns) + duration.to_value(u.ns),
+        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[n:],
+            time=t[self.extra_samples:],
             is_high_gain=np.zeros(self.multiplicity * n_pixels),
             additional_nsb=additional_nsb,
             amplitude=added_nsb_amp,
@@ -614,7 +619,7 @@ def initialise_waveforms(self, waveform, dt, tel_id):
             a_lg_template=self.pulse_template[tel_id].amplitude_LG
         )
         nsb_waveforms[:, 1, :] += nsb_only_waveforms(
-            time=t[n:],
+            time=t[self.extra_samples:],
             is_high_gain=np.ones(self.multiplicity * n_pixels),
             additional_nsb=additional_nsb,
             amplitude=added_nsb_amp,
@@ -669,22 +674,22 @@ def _tune_nsb_on_waveform_direct(self, waveform, tel_id, is_high_gain, subarray)
         subarray: `ctapipe.instrument.subarray.SubarrayDescription`
 
         """
-        n = 25
         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 = (n + n_samples) * dt  # TODO check needed time window, effect of edges
-        t = np.arange(-n, n_samples) * dt.to_value(u.ns)
+        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(-n * dt.to_value(u.ns), -n * dt.to_value(u.ns) + duration.to_value(u.ns),
+        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[n:],
+            time=t[self.extra_samples:],
             is_high_gain=is_high_gain,
             additional_nsb=additional_nsb,
             amplitude=added_nsb_amp,

From 15e909351421fc6042cdfbc658ebd1a442938666 Mon Sep 17 00:00:00 2001
From: Gabriel Emery <gabriel-emery@hotmail.fr>
Date: Thu, 11 Jul 2024 13:36:14 +0200
Subject: [PATCH 6/6] Explicit config value in test

---
 lstchain/tests/test_lstchain.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/lstchain/tests/test_lstchain.py b/lstchain/tests/test_lstchain.py
index 1d02a054a9..771204f829 100644
--- a/lstchain/tests/test_lstchain.py
+++ b/lstchain/tests/test_lstchain.py
@@ -129,6 +129,7 @@ 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