Merge pull request #1120 from cta-observatory/njit_nsb_tuning_on_wave…

…form

Create pure NSB R1 waveforms

lstchain/image/modifier.py

@@ -13,6 +13,7 @@
 
 from lstchain.io import standard_config
 from lstchain.io.config import read_configuration_file
+from lstchain.reco.reconstructorCC import nsb_only_waveforms
 
 __all__ = [
     'add_noise_in_pixels',
@@ -406,12 +407,17 @@ def tune_nsb_on_waveform(waveform, added_nsb_fraction, original_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
-    for i in range(n_pixels):
-        for j in range(additional_nsb[i]):
-            waveform[i] += pulse_templates(t[n:], 'HG' if gain[i] else 'LG',
-                                           amplitude=added_nsb_amp[i][j],
-                                           t_0=added_nsb_time[i][j])
-
+    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

lstchain/reco/reconstructorCC.py

@@ -227,6 +227,72 @@ def template_interpolation(gain, times, t0, dt, a_hg, a_lg):
     return out
 
 
+@njit(cache=True)
+def nsb_only_waveforms(time, is_high_gain, additional_nsb, amplitude, t_0,
+                        t0_template, dt_template, a_hg_template, a_lg_template):
+    """
+    Generate waveforms of pure NSB. NSB photons injected through a fast interpolator using a
+    uniformly sampled normalised template with known time of first value and time step.
+    The method requires as input the number of NSB events to inject per pixel,
+    the times of injection and the events normalisations.
+
+    Interpolation code duplicated from function template_interpolation.
+
+    Parameters
+    ----------
+    time: float64 1D array
+        Times of each waveform samples
+    is_high_gain: boolean 1D array
+        Gain channel used per pixel: True=hg, False=lg
+    additional_nsb: float64 1D array
+        Number of NSB photons to inject per pixel
+    amplitude: float 2D array
+        Normalisation factor to apply to the template per photon in each pixel
+    t_0: float 2D array
+        Shift in the origin of time per photon in each pixel
+    t0_template: float64
+        Time of the first value of the pulse template
+    dt_template: float 64
+        Time step between template values
+    a_hg_template: float64 1D array
+        Template values for the high gain channel
+    a_lg_template: float64 1D array
+        Template values for the low gain channel
+
+    Returns
+    -------
+    nsb_waveform: float64 2D array
+        Charge (p.e. / ns) in each pixel and sampled time of the injected NSB photons
+    """
+    n_pixels = additional_nsb.shape[0]
+    m = time.shape[0]
+    nsb_waveform = np.zeros((n_pixels, m), dtype=np.float64)
+    size = a_hg_template.shape[0]
+    single_spe_waveform = np.empty(m)
+
+    times = time / dt_template
+    t0 = (t_0 + t0_template) / dt_template
+    for i in range(n_pixels):
+        for j in range(additional_nsb[i]):
+            # Find the index before the requested time
+            a = times - t0[i, j]
+
+            for k in range(m):
+                t = int(a[k])
+
+                if 0 < t + 1 < size:
+                    # Select the gain and interpolate the pulse template at the requested time
+                    single_spe_waveform[k] = a_hg_template[t] * (1. - a[k] + t) + a_hg_template[t + 1] * (a[k] - t) if is_high_gain[i] else\
+                        a_lg_template[t] * (1. - a[k] + t) + a_lg_template[t + 1] * (a[k] - t)
+                else:
+                    # Assume 0 if outside the recorded range
+                    single_spe_waveform[k] = 0.0
+
+            nsb_waveform[i] += amplitude[i, j] * single_spe_waveform
+
+    return nsb_waveform
+
+
 @njit(cache=True)
 def log_pdf(charge, t_cm, x_cm, y_cm, length, wl, psi, v, rl,
             data, error, is_high_gain, sig_s, crosstalks, times, time_shift,