"""
Algorithms to select correct gain channel
"""
from abc import abstractmethod
from enum import IntEnum
import numpy as np
from ctapipe.core import Component, traits
__all__ = [
"GainChannel",
"GainSelector",
"ManualGainSelector",
"ThresholdGainSelector",
]
[docs]class GainChannel(IntEnum):
"""
Possible gain channels
"""
HIGH = 0
LOW = 1
[docs]class GainSelector(Component):
"""
Base class for algorithms that decide on the gain channel to use
"""
[docs] def __call__(self, waveforms):
"""
Reduce the waveform to a single gain channel
Parameters
----------
waveforms : ndarray
Waveforms stored in a numpy array of shape
(n_chan, n_pix, n_samples).
Returns
-------
selected_gain_channel : ndarray
Gain channel to use for each pixel
Shape: n_pix
Dtype: int8
"""
if waveforms.ndim == 2: # Return None if already gain selected
return None
elif waveforms.ndim == 3:
n_channels, n_pixels, _ = waveforms.shape
if n_channels == 1: # Must be first channel if only one channel
return np.zeros(n_pixels, dtype=np.int8)
else:
return self.select_channel(waveforms)
else:
raise ValueError(f"Cannot handle waveform array of shape: {waveforms.ndim}")
[docs] @abstractmethod
def select_channel(self, waveforms):
"""
Abstract method to be defined by a GainSelector subclass.
Call the relevant functions to decide on the gain channel used for
each pixel.
Parameters
----------
waveforms : ndarray
Waveforms stored in a numpy array of shape
(n_chan, n_pix, n_samples).
Returns
-------
selected_gain_channel : ndarray
Gain channel to use for each pixel
Shape: n_pix
Dtype: int8
"""
[docs]class ManualGainSelector(GainSelector):
"""
Manually choose a gain channel.
"""
channel = traits.CaselessStrEnum(
["HIGH", "LOW"], default_value="HIGH", help="Which gain channel to retain"
).tag(config=True)
[docs] def select_channel(self, waveforms):
n_pixels = waveforms.shape[1]
return np.full(n_pixels, GainChannel[self.channel])
[docs]class ThresholdGainSelector(GainSelector):
"""
Select gain channel according to a maximum threshold value.
"""
threshold = traits.Float(
default_value=4000,
help="Threshold value in waveform sample units. If a waveform "
"contains a sample above this threshold, use the low gain "
"channel for that pixel.",
).tag(config=True)
[docs] def select_channel(self, waveforms):
return (waveforms[0] > self.threshold).any(axis=1).astype(np.int8)