Source code for skcomponents.processing.noise

import numpy as np
from skcomponents.processing.processor import SignalProcessor
from typing import Union


class GaussianNoise(SignalProcessor):
    """
    Add gaussian noise to a signal

    Parameters
    ----------
    noise_level : float
        noise level in fraction of signal (1% = 0.01)
    """

    def __init__(self, noise_level: float = 0.01, relative=True):

        self.noise_level = noise_level
        self._relative = relative

    def process_signal(self, signal: np.ndarray) -> np.ndarray:

        if self._relative:
            noise = signal * self.noise_level * np.random.standard_normal(signal.shape)
        else:
            noise = self.noise_level * np.random.standard_normal(signal.shape)

        return signal + noise

    def noise_estimate(self, signal: np.ndarray) -> np.ndarray:

        if self._relative:
            return self.noise_level * signal
        else:
            return self.noise_level * np.ones_like(signal)


class PoissonCounting(SignalProcessor):
    """
    Add poisson counting error to the signal
    """
    def process_signal(self, signal: np.ndarray) -> np.ndarray:

        try:
            mask = np.isnan(signal)
            if np.any(mask):
                signal[mask] = 0
            return np.random.poisson(signal)
        except ValueError as e:
            print(e)

    def noise_estimate(self, signal: np.ndarray) -> np.ndarray:

        return np.sqrt(signal)


class SimpleDarkCurrent(SignalProcessor):

    def __init__(self, dark_current: Union[float, np.ndarray],
                 exposure_time: float = 1.0,
                 temperature: np.ndarray = None,
                 sensor_temperature: float = None,
                 noisy: bool = True):
        """

        Parameters
        ----------
        dark_current : float
            dark current in Amperes. If a single number is provided it is assumed constant for all temperatures. If
            an array is provided a correspending `temperature` of the same size should be given.
        exposure_time : float
            exposure time of the image
        temperature :
            array of temperatures at which dark current is defined
        sensor_temperature :
            temperature of the sensor. Used only if an array of `dark_current`, `temperature` is provided.
        noisy : bool
            if True the dark current signal is assumed to be have a Poisson distribution. Otherwise a constant value is
            added.
        """
        self.dark_current = dark_current
        self.temperature = temperature
        self.sensor_temperature = sensor_temperature
        self.exposure_time = exposure_time
        self.noisy = noisy

    @property
    def electrons_per_second(self):
        if self.temperature is not None:
            # interpolate in log(DC), inverse temperature space
            tidx = np.argsort(1 / self.temperature)
            temp = 1 / self.temperature[tidx]
            dc_array = np.log(self.dark_current[tidx])
            dc = np.exp(np.interp(1 / self.sensor_temperature, temp, dc_array,
                        left=dc_array[0], right=dc_array[-1]))
        else:
            dc = self.dark_current

        return dc * 6.241509074e18

    def process_signal(self, signal: np.ndarray) -> np.ndarray:

        if self.noisy:
            if type(self.exposure_time) is np.ndarray:
                if signal.shape[-1] == self.exposure_time.shape[0]:
                    # return signal * self.exposure_time
                    dc = PoissonCounting().process_signal(np.ones_like(signal) * (self.electrons_per_second * self.exposure_time))
                elif signal.shape[0] == self.exposure_time.shape[0]:
                    # return signal * self.exposure_time[:, np.newaxis, np.newaxis]
                    dc = PoissonCounting().process_signal(np.ones_like(signal) * (self.electrons_per_second * self.exposure_time)[:, np.newaxis, np.newaxis])
                else:
                    raise ValueError('Cannot broadcast signal and exposure_time, '
                                     'wavelength should be the first or last dimension')
            else:
                dc = PoissonCounting().process_signal(np.ones_like(signal) * (self.electrons_per_second * self.exposure_time))
        else:
            dc = self.electrons_per_second * self.exposure_time
        return signal + dc

    def noise_estimate(self, signal: np.ndarray) -> np.ndarray:

        return np.ones_like(signal) * np.sqrt(self.electrons_per_second * self.exposure_time)


class DarkCurrent(SignalProcessor):
    """
    Add dark current to a signal

    Parameters
    ----------
    exposure_time : float
        exposure time of the measurement [s]
    temperature : float
        ccd temperature [K]
    tref : float
        reference temperature of the dark current parameters [K]
    e_pixel_s : float
        electrons per 
    """
    def __init__(self, exposure_time: float = None,
                 temperature: float = 273.15,
                 tref: float = 273.15,
                 e_pixel_s: float = 200,
                 c1: float = 1.14e6,
                 c2: float = 9080,
                 ratio_exp: float = 3.0):

        self.exposure_time = exposure_time
        self.readout_time = 0
        self.temperature = temperature
        self.tref = tref
        self.e_pixel_s = e_pixel_s
        self.c1 = c1
        self.c2 = c2
        self.ratio_exp = ratio_exp

    def process_signal(self, signal: np.ndarray, temperature: float = None, exposure_time=None) -> np.ndarray:

        if exposure_time is None:
            exposure_time = self.exposure_time

        if exposure_time is None:
            raise ValueError('exposure not set')

        dark_current = self.dark_signal_rate(temperature) * (self.readout_time + self.exposure_time)
        return signal + dark_current

    def dark_signal_rate(self, ccd_temperature: float = None):
        if ccd_temperature is None:
            ccd_temperature = self.temperature

        # Modified Doug's equation
        dark_signal_ratio_at_Tref = self.c1 * (self.tref ** self.ratio_exp) * np.exp(-self.c2 / self.tref)
        dark_signal_ratio_TCCD = self.c1 * (ccd_temperature ** self.ratio_exp) * np.exp(-self.c2 / ccd_temperature)
        dark_signal_ratio_TCCD /= dark_signal_ratio_at_Tref
        dark_signal_rate_TCCD = dark_signal_ratio_TCCD * self.e_pixel_s
        return dark_signal_rate_TCCD