ESBO-ETC/esbo_etc/classes/SpectralQty.py

from ..lib.helpers import isLambda, readCSV
from ..lib.logger import logger
from scipy.interpolate import interp1d
import astropy.units as u
from typing import Union, Callable
import os
from scipy.integrate import trapz
import numpy as np


# noinspection PyUnresolvedReferences
class SpectralQty:
    """
    A class to hold and work with spectral quantities
    """

    def __init__(self, wl: u.Quantity, qty: u.Quantity, fill_value: Union[bool, int, float] = 0):
        """
        Initialize a new spectral quantity

        Parameters
        ----------
        wl : Quantity
            The binned wavelengths
        qty : Quantity
            The quantity values corresponding to the binned wavelengths. If the values are supplied without a unit,
            they are assumed to be dimensionless.
        fill_value : Union[bool, int, float]
            How to treat missing values. True enables extrapolation, False disables extrapolation and the spectrum will
            be truncated. If a numeric value is given, the missing values will be filled with this value.
        Returns
        -------
        sqty : SpectralQty
            The created spectral quantity.
        """
        # Check if both lengths are equal
        if wl.value.size == qty.value.size:
            # check if units are given. If not, add a dimensionless unit
            if hasattr(wl, "unit"):
                self.wl = wl
            else:
                self.wl = wl * u.dimensionless_unscaled
            if hasattr(qty, "unit"):
                self.qty = qty
            else:
                self.qty = qty * u.dimensionless_unscaled
        else:
            logger.error("Lengths not matching")
        self._fill_value = fill_value

    @classmethod
    def fromFile(cls, file: str, wl_unit_default: u.Quantity = None, qty_unit_default: u.Quantity = None,
                 fill_value: Union[bool, int, float] = 0) -> "SpectralQty":
        """
        Initialize a new spectral quantity and read the values from a file

        Parameters
        ----------
        file : str
            Path to the file to read the values from. The file needs to provide two columns: wavelength
            and the corresponding spectral quantity. The format of the file will be guessed by
            `astropy.io.ascii.read()`. If the file doesn't provide units via astropy's enhanced CSV format, the units
            will be read from the column headers or otherwise assumed to be *wl_unit_default* and *qty_unit_default*.
        wl_unit_default : Quantity
            Default unit to be used for the wavelength column if no units are provided by the file.
        qty_unit_default : Quantity
            Default unit to be used for the quantity column if no units are provided by the file.
        fill_value : Union[bool, int, float]
            How to treat missing values. True enables extrapolation, False disables extrapolation and the spectrum will
            be truncated. If a numeric value is given, the missing values will be filled with this value.
        Returns
        -------
        sqty : SpectralQty
            The created spectral quantity.
        """
        # Read the file
        data = readCSV(file, [wl_unit_default, qty_unit_default] if wl_unit_default is not None and
                       qty_unit_default is not None else None)
        return cls(data[data.colnames[0]].quantity, data[data.colnames[1]].quantity, fill_value=fill_value)

    def __str__(self, precision: int = 4) -> str:
        """
        Convert a SpectralQty object to a string representation

        Parameters
        ----------
        precision : int
            Precision of the printed values

        Returns
        -------
        ret : str
            String representation of the object
        """
        wl_str = []
        qty_str = []
        for i in range(len(self.wl)):
            wl_str_temp = "%%.%dg" % precision % self.wl[i].value
            qty_str_temp = "%%.%dg" % precision % self.qty[i].value
            wl_str.append(wl_str_temp.ljust(max(len(wl_str_temp), len(qty_str_temp)), " "))
            qty_str.append(qty_str_temp.ljust(max(len(wl_str_temp), len(qty_str_temp)), " "))
        return "Wavelength: [" + ", ".join(wl_str) + "] " + self.wl.unit.to_string() + os.linesep +\
               "Quantitiy:  [" + ", ".join(qty_str) + "] " + self.qty.unit.to_string()

    def __eq__(self, other) -> bool:
        """
        Check if this object is equal to another object

        Parameters
        ----------
        other : SpectralQty
            Object to compare with

        Returns
        -------
        res : bool
            Result of the comparison
        """
        return self.wl.unit.is_equivalent(other.wl.unit) and self.qty.unit.is_equivalent(other.qty.unit) and \
            len(self.wl) == len(other.wl) and len(self.qty) == len(other.qty) and \
            np.allclose(self.wl.value, other.wl.to(self.wl.unit).value) and \
            np.allclose(self.qty.value, other.qty.to(self.qty.unit).value)

    def __add__(self, other: Union[int, float, u.Quantity, "SpectralQty", Callable[[u.Quantity], u.Quantity]]) ->\
            "SpectralQty":
        """
        Calculate the sum with another object

        Parameters
        ----------
        other : Union[int, float, u.Quantity, "SpectralQty", Callable]
            Addend to be added to this object. If the binning of the object on the right hand side differs
            from the binning of the left object, the object on the right hand side will be rebinned.

        Returns
        -------
        sum : SpectralQty
            The sum of both objects
        """
        # Summand is of type int or float, use same unit
        if isinstance(other, int) or isinstance(other, float):
            return SpectralQty(self.wl, self.qty + other * self.qty.unit)
        # Summand is of type Quantity
        elif isinstance(other, u.Quantity):
            if other.unit == self.qty.unit:
                return SpectralQty(self.wl, self.qty + other)
            else:
                raise TypeError("Units are not matching for addition.")
        # Summand is of type lambda
        elif isLambda(other):
            return SpectralQty(self.wl, self.qty + other(self.wl).value * other(self.wl[0]).unit)
        # Summand is of type SpectralQty
        else:
            if other.wl.unit.is_equivalent(self.wl.unit) and other.qty.unit.is_equivalent(self.qty.unit):
                if len(self.wl) == len(other.wl) and (self.wl == other.wl).all():
                    # Wavelengths are matching, just add the quantities
                    return SpectralQty(self.wl, self.qty + other.qty)
                else:
                    # Wavelengths are not matching, rebinning needed
                    other_rebinned = other.rebin(self.wl)
                    if len(self.wl) == len(other_rebinned.wl) and (self.wl == other_rebinned.wl).all():
                        return SpectralQty(self.wl, self.qty + other_rebinned.qty)
                    else:
                        # Wavelengths are still not matching as extrapolation is disabled, rebin this spectral quantity
                        return SpectralQty(other_rebinned.wl, self.rebin(other_rebinned.wl).qty + other_rebinned.qty)
            else:
                logger.error("Units are not matching for addition.")

    __radd__ = __add__

    def __sub__(self, other: Union[int, float, u.Quantity, "SpectralQty", Callable[[u.Quantity], u.Quantity]]) ->\
            "SpectralQty":
        """
        Calculate the difference to another object

        Parameters
        ----------
        other : Union[int, float, u.Quantity, "SpectralQty", Callable]
            Subtrahend to be subtracted from this object. If the binning of the object on the right hand side differs
            from the binning of the left object, the object on the right hand side will be rebinned.

        Returns
        -------
        diff : SpectralQty
            The difference of both objects
        """
        # Subtrahend is of type int or float, use same unit
        if isinstance(other, int) or isinstance(other, float):
            return SpectralQty(self.wl, self.qty - other * self.qty.unit)
        # Subtrahend is of type Quantity
        elif isinstance(other, u.Quantity):
            if other.unit == self.qty.unit:
                return SpectralQty(self.wl, self.qty - other)
            else:
                raise TypeError('Units are not matching for subtraction.')
        # Subtrahend is of type lambda
        elif isLambda(other):
            return SpectralQty(self.wl, self.qty - other(self.wl).value * other(self.wl[0]).unit)
        # Subtrahend is of type SpectralQty
        else:
            if other.wl.unit.is_equivalent(self.wl.unit) and other.qty.unit.is_equivalent(self.qty.unit):
                # Wavelengths are matching, just subtract the quantities
                if len(self.wl) == len(other.wl) and (self.wl == other.wl).all():
                    return SpectralQty(self.wl, self.qty - other.qty)
                # Wavelengths are not matching, rebinning needed
                else:
                    # Rebin subtrahend
                    other_rebinned = other.rebin(self.wl)
                    if len(self.wl) == len(other_rebinned.wl) and (self.wl == other_rebinned.wl).all():
                        return SpectralQty(self.wl, self.qty - other_rebinned.qty)
                    else:
                        # Wavelengths are still not matching as extrapolation is disabled, rebin this spectral quantity
                        return SpectralQty(other_rebinned.wl, self.rebin(other_rebinned.wl).qty - other_rebinned.qty)
            else:
                logger.error("Units are not matching for substraction.")

    def __mul__(self, other: Union[int, float, u.Quantity, "SpectralQty", Callable[[u.Quantity], u.Quantity]]) ->\
            "SpectralQty":
        """
        Calculate the product with another object

        Parameters
        ----------
        other : Union[int, float, u.Quantity, "SpectralQty", Callable]
            Factor to be multiplied with this object. If the binning of the object on the right hand side differs
            from the binning of the left object, the object on the right hand side will be rebinned.

        Returns
        -------
        prod : SpectralQty
            The product of both objects
        """
        # Factor is of type int, float or Quantity, just multiply
        if isinstance(other, int) or isinstance(other, float) or isinstance(other, u.Quantity):
            return SpectralQty(self.wl, self.qty * other)
        # Factor is of type lambda
        elif isLambda(other):
            return SpectralQty(self.wl, self.qty * other(self.wl).value * other(self.wl[0]).unit)
        # Factor is of type SpectralQty
        else:
            if other.wl.unit.is_equivalent(self.wl.unit):
                # Wavelengths are matching, just multiply the quantities
                if len(self.wl) == len(other.wl) and (self.wl == other.wl).all():
                    return SpectralQty(self.wl, self.qty * other.qty)
                # Wavelengths are not matching, rebinning needed
                else:
                    # Rebin factor
                    other_rebinned = other.rebin(self.wl)
                    if len(self.wl) == len(other_rebinned.wl) and (self.wl == other_rebinned.wl).all():
                        return SpectralQty(self.wl, self.qty * other_rebinned.qty)
                    else:
                        # Wavelengths are still not matching as extrapolation is disabled, rebin this spectral quantity
                        return SpectralQty(other_rebinned.wl, self.rebin(other_rebinned.wl).qty * other_rebinned.qty)
            else:
                logger.error("Units are not matching for multiplication.")

    __rmul__ = __mul__

    def __truediv__(self, other: Union[int, float, u.Quantity, "SpectralQty", Callable[[u.Quantity], u.Quantity]]) ->\
            "SpectralQty":
        """
        Calculate the quotient with another object

        Parameters
        ----------
        other : Union[int, float, u.Quantity, "SpectralQty", Callable]
            Divisor for this object. If the binning of the object on the right hand side differs
            from the binning of the left object, the object on the right hand side will be rebinned.

        Returns
        -------
        quot : SpectralQty
            The quotient of both objects
        """
        # Factor is of type int, float or Quantity, just multiply
        if isinstance(other, int) or isinstance(other, float) or isinstance(other, u.Quantity):
            return SpectralQty(self.wl, self.qty / other)
        # Factor is of type lambda
        elif isLambda(other):
            return SpectralQty(self.wl, self.qty / other(self.wl).value / other(self.wl[0]).unit)
        # Factor is of type SpectralQty
        else:
            if other.wl.unit.is_equivalent(self.wl.unit):
                # Wavelengths are matching, just multiply the quantities
                if len(self.wl) == len(other.wl) and (self.wl == other.wl).all():
                    return SpectralQty(self.wl, self.qty / other.qty)
                # Wavelengths are not matching, rebinning needed
                else:
                    # Rebin factor
                    other_rebinned = other.rebin(self.wl)
                    if len(self.wl) == len(other_rebinned.wl) and (self.wl == other_rebinned.wl).all():
                        return SpectralQty(self.wl, self.qty / other_rebinned.qty)
                    else:
                        # Wavelengths are still not matching as extrapolation is disabled, rebin this spectral quantity
                        return SpectralQty(other_rebinned.wl, self.rebin(other_rebinned.wl).qty / other_rebinned.qty)
            else:
                logger.error("Units are not matching for division.")

    def __pow__(self, other: Union[int, float]) -> "SpectralQty":
        """
        Raise to spectral quantity to the given power.

        Parameters
        ----------
        other : Union[int, float]
            The exponent for this object.

        Returns
        -------
        prod : SpectralQty
            The spectral quantity raised to the given power.
        """
        # Factor is of type int, float or Quantity, just multiply
        if isinstance(other, int) or isinstance(other, float):
            return SpectralQty(self.wl, self.qty ** other)

    def rebin(self, wl: u.Quantity) -> "SpectralQty":
        """
        Resample the spectral quantity sqty(wl) over the new grid wl, rebinning if necessary, otherwise interpolates.
        Copied from ExoSim (https://github.com/ExoSim/ExoSimPublic).

        Parameters
        ----------
        wl : Quantity
            new binned wavelengths

        Returns
        -------
        sqty : SpectralQty
            The rebinned spectral quantity
        """

        if not wl.unit.is_equivalent(self.wl.unit):
            logger.error("Mismatching units for rebinning: " + wl.unit + ", " + self.wl.unit)
        if (wl.value.size == 1 and (wl < min(self.wl) or wl > max(self.wl))) or (
                wl.value.size > 1 and (min(wl) < min(self.wl) or max(wl) > max(self.wl))):
            if isinstance(self._fill_value, bool):
                if not self._fill_value:
                    logger.warning("Extrapolation disabled, bandwidth will be reduced.")
                    # Remove new wavelengths where extrapolation would have been necessary
                    wl = [x.value for x in wl if min(self.wl) <= x <= max(self.wl)] * wl.unit
                f = interp1d(self.wl, self.qty.value, fill_value="extrapolate")
            else:
                f = interp1d(self.wl, self.qty.value, fill_value=self._fill_value, bounds_error=False)
        else:
            f = interp1d(self.wl, self.qty.value)
        return SpectralQty(wl, f(wl.to(self.wl.unit)) * self.qty.unit)

    def integrate(self) -> u.Quantity:
        """
        Integrate the spectral quantity over the given spectrum using

        Returns
        -------
        int : Quantity
            The integrated quantity
        """
        return u.Quantity(trapz(self.qty, self.wl))