mors.spectrum

def CalcBandFluxes(self):
    """Calculate integrated fluxes for each band.

    Integrated fluxes will have units of [erg s-1 cm-2] scaled to 1 AU.
    """

    # For each spectral band
    idxs = []
    i_lo = 0
    for b in bands_ascending:

        # Get band indicies
        for i in range(i_lo, self.nbins, 1):
            wb = WhichBand(self.wl[i])

            # Out of range
            if wb == None:
                continue

            if b in wb:
                # In current band
                idxs.append(i)
            else:
                # End of band
                i_lo = i
                break

        # With idxs defined, integrate over band
        band_wl = self.wl[idxs]
        band_fl = self.fl[idxs]
        self.fl_integ[b] = np.trapezoid(band_fl,band_wl)

        # Reset idxs
        idxs = []

    # For bolometric "band"
    self.fl_integ["bo"] = np.trapezoid(self.fl,self.wl)

    return self.fl_integ

def ExtendPlanck(self, Teff:float, R_star:float, wl_max:float):
    """Extend spectrum to longer wavelengths using planck function.

    Parameters
    ----------
    Teff : float
        Effective temperature of star
    R_star : float
        Radius of star [m]
    wl_max : float
        New maximum wavelength [nm]
    """


    # Already extended
    if wl_max < self.wl[-1]:
        return

    # Calc wavelength extension
    wl_ext = np.logspace(np.log10(self.wl[-1]), np.log10(wl_max), 300)[1:]

    # Evalulate planck function
    fl_ext = PlanckFunction_surf(wl_ext, Teff)

    # Scale to 1 AU
    fl_ext = ScaleTo1AU(fl_ext, R_star)

    # Update Spectrum object data
    self.ext_long = len(self.wl)
    spec_wl = np.concatenate((self.wl, wl_ext))
    spec_fl = np.concatenate((self.fl, fl_ext))

    # Store
    self.LoadDirectly(spec_wl, spec_fl)

def ExtendShortwave(self, wl_min:float):
    """Extend spectrum to shorter wavelengths using constant value.

    Parameters
    ----------
    wl_min : float
        New minimum wavelength [nm]
    """

    # Already extended
    if wl_min > self.wl[0]:
        return

    # Calc wavelength extension
    wl_ext = np.logspace(np.log10(wl_min), np.log10(self.wl[0]), 200)[:-1]

    # Evalulate planck function
    fl_ext = np.ones(np.shape(wl_ext)) * self.fl[0]

    # Update Spectrum object data
    self.ext_short = len(wl_ext)
    spec_wl = np.concatenate((wl_ext,self.wl))
    spec_fl = np.concatenate((fl_ext,self.fl))

    # Store
    self.LoadDirectly(spec_wl, spec_fl)

def LoadDirectly(self, spec_wl:np.ndarray, spec_fl:np.ndarray):
    """Store spectral data in object.

    Scaled to 1 AU.

    Parameters
    ----------
    spec_wl : np.ndarray
        Array of wavelengths [nm]
    spec_fl : np.ndarray
        Array of fluxes [erg s-1 cm-2 nm-1]
    """


    # Check length
    if len(spec_wl) != len(spec_fl):
        raise Exception("Stellar spectrum size mismatch (%d and %d)"%(len(spec_wl), len(spec_fl)))
    if len(spec_wl) < 10:
        raise Exception("Stellar spectrum size too small (%d bins)"%len(spec_wl))

    # Check reversal (should be wl ascending)
    if spec_wl[4] < spec_wl[0]:
        spec_wl = spec_wl[::-1]
        spec_fl = spec_fl[::-1]

    # Replace NaN and zero values
    for i in range(len(spec_fl)):
        if not np.isfinite(spec_fl[i]):
            spec_fl[i] = 0.0
        spec_fl[i] = max(spec_fl[i], 1e-20)

    # Calculate bin width
    binwidth_wl = spec_wl[1:] - spec_wl[0:-1]

    # Store
    self.nbins = len(spec_wl)
    self.wl = spec_wl
    self.fl = spec_fl
    self.binwidth = np.array(binwidth_wl, dtype=float)

    self.loaded = True

    return self

def LoadTSV(self, fp:str):
    """Load stellar spectrum from TSV file into memory.

    Scaled to 1 AU. File should be whitespace delimited with
    units of [nm] and [erg s-1 cm-2 nm-1].

    Parameters
    ----------
    fp : str
        Path to file
    """

    log.debug("Loading stellar spectrum from TSV file")

    # Check path
    fp = os.path.abspath(fp)
    if not os.path.isfile(fp):
        raise Exception("Cannot find TSV file at '%s'"%fp)

    # Load file
    spec_data = np.loadtxt(fp).T
    spec_wl = spec_data[0]
    spec_fl = spec_data[1]

    # Process data
    spec_wl = np.array(spec_wl, dtype=float)
    spec_fl = np.array(spec_fl, dtype=float)

    # Store
    self.LoadDirectly(spec_wl, spec_fl)

    self.loaded = True

    return self

def WriteTSV(self, fp:str):
    """Write spectrum to file(s) on disk.

    Parameters
    ----------
    fp : str
        Path to file
    """

    log.debug("Writing stellar spectrum to TSV file")

    fp = os.path.abspath(fp)
    X = np.array([self.wl,self.fl]).T
    header = "WL(nm)\t Flux(ergs/cm**2/s/nm)    Stellar flux (1 AU)"

    np.savetxt(fp, X, header=header,fmt='%1.4e',delimiter='\t')

    return fp