Source code for nectarchain.dqm.mean_waveforms

import os

import numpy as np
from matplotlib import pyplot as plt

from .dqm_summary_processor import DQMSummary

__all__ = ["MeanWaveFormsHighLowGain"]




class MeanWaveFormsHighLowGain(DQMSummary):
    def __init__(self, gaink, r0=False):
        self.k = gaink
        self.Pix = None
        self.Samp = None
        self.tel_id = None
        self.Mwf = None
        self.Mwf_ped = None
        self.counter_evt = None
        self.counter_ped = None
        self.Mwf_average = None
        self.Mwf_ped_average = None
        self.Mwf_Mean_overPix = []
        self.Mwf_ped_Mean_overPix = []
        self.MeanWaveForms_Results_Dict = {}
        self.MeanWaveForms_Figures_Dict = {}
        self.MeanWaveForms_Figures_Names_Dict = {}

        self.wf_list_plot = None

        gain_c = "High" if self.k == 0 else "Low"
        self.gain_c = gain_c

        self.figure_keys = {
            "physical": f"FIGURE-WF-Physical-{gain_c}-GAIN",
            "pedestals": f"FIGURE-WF-Pedestals-{gain_c}-GAIN",
            "combined": f"FIGURE-WF-COMBINED-{gain_c}-GAIN",
        }

        self.figure_filenames = {
            "physical": f"_MeanWaveforms_Physical_{gain_c}Gain.png",
            "pedestals": f"_MeanWaveforms_Pedestals_{gain_c}Gain.png",
            "combined": f"_MeanWaveforms_CombinedPlot_{gain_c}Gain.png",
        }

        super().__init__(r0)



    def configure_for_run(self, path, Pix, Samp, Reader1, **kwargs):
        """Initialize waveform buffers and counters for a new run."""
        self.Pix = Pix
        self.Samp = Samp

        self.tel_id = Reader1.subarray.tel_ids[0]

        self.Mwf = np.zeros((Pix, Samp))
        self.Mwf_ped = np.zeros((Pix, Samp))

        self.counter_evt = 0
        self.counter_ped = 0

        self.Mwf_average = np.zeros((Pix, Samp))
        self.Mwf_ped_average = np.zeros((Pix, Samp))

        self.wf_list_plot = list(np.arange(1, Samp + 1))


    def process_event(self, evt, noped):
        if evt.trigger.event_type.value == 32:  # only peds now
            self.counter_ped += 1
        else:
            self.counter_evt += 1

        for ipix in range(self.Pix):
            if self.r0:
                waveform = evt.r0.tel[self.tel_id].waveform[self.k][ipix]
            else:
                # This should accommodate cases were the shape of waveforms is 2D
                # (1855,60), or 3D (2, 1855, 60) for 2-gain channels or
                # (1, 1855, 60) for single-gain channel
                waveform = evt.r1.tel[self.tel_id].waveform[..., ipix, :]
            if evt.trigger.event_type.value == 32:  # only peds now
                self.Mwf_ped[ipix, :] += waveform
            else:
                self.Mwf[ipix, :] += waveform
        return None



    def finish_run(self):
        """Compute mean waveforms over events and pixels."""

        if self.counter_evt > 0:
            self.Mwf_average = self.Mwf / self.counter_evt  # get average
            # get average over pixels
            self.Mwf_Mean_overPix = np.mean(self.Mwf_average, axis=0)

        if self.counter_ped > 0:
            # get average pedestals
            self.Mwf_ped_average = self.Mwf_ped / self.counter_ped
            self.Mwf_ped_Mean_overPix = np.mean(self.Mwf_ped_average, axis=0)

        return None




    def get_results(self):
        """Store waveform statistics in results dictionary by gain and type."""

        self.MeanWaveForms_Results_Dict[
            f"WF-PHY-AVERAGE-PIX-{self.gain_c.upper()}-GAIN"
        ] = self.Mwf_Mean_overPix
        if self.counter_ped > 0:
            self.MeanWaveForms_Results_Dict[
                f"WF-AVERAGE-PED-PIX-{self.gain_c.upper()}-GAIN"
            ] = self.Mwf_ped_Mean_overPix

        return self.MeanWaveForms_Results_Dict


    def plot_results(self, name, fig_path):
        wf_list = np.array(self.wf_list_plot)

        colors = ["blue", "red"]
        titles = ["physical", "pedestals"]

        full_fig, full_ax = plt.subplots()
        array_plot = [self.Mwf_average]

        if self.counter_ped > 0:
            array_plot.append(self.Mwf_ped_average)

        for i, x in enumerate(array_plot):
            key = titles[i]
            fig_key = self.figure_keys[key]
            full_name = name + self.figure_filenames[key]
            fig_name = os.path.join(fig_path, full_name)

            part_fig, part_ax = plt.subplots()

            for ipix in range(self.Pix):
                part_ax.plot(
                    wf_list, x[ipix, :], color=colors[i], alpha=0.005, linewidth=1
                )
                full_ax.plot(
                    wf_list, x[ipix, :], color=colors[i], alpha=0.005, linewidth=1
                )

            part_ax.set_title(f"Mean Waveforms {key.capitalize()} ({self.gain_c} Gain)")
            part_ax.set_xlabel("Samples")
            part_ax.set_ylabel("Amplitude (DC)")
            part_ax.grid()

            self.MeanWaveForms_Figures_Dict[fig_key] = part_fig
            self.MeanWaveForms_Figures_Names_Dict[fig_key] = fig_name

            plt.close(part_fig)

        # Combined figure setup
        full_ax.set_title(f"Mean Waveforms Combined Plot ({self.gain_c} Gain)")
        full_ax.set_xlabel("Samples")
        full_ax.set_ylabel("Amplitude (DC)")
        full_ax.grid()

        combined_key = self.figure_keys["combined"]
        combined_path = os.path.join(fig_path, name + self.figure_filenames["combined"])

        self.MeanWaveForms_Figures_Dict[combined_key] = full_fig
        self.MeanWaveForms_Figures_Names_Dict[combined_key] = combined_path

        plt.close(full_fig)

        return self.MeanWaveForms_Figures_Dict, self.MeanWaveForms_Figures_Names_Dict