sed.py

#!/usr/bin/env python
# -*- coding:utf-8 -*-

# @Author: Shuai Feng (hebtu.edu.cn)
# @Time:   2022-09-25

import numpy as np
from scipy.interpolate import interp1d
from astropy.io import fits
from astropy.io import ascii
from astropy import units as u

# ----------------
# Magnitude Module

from scipy.interpolate import interp1d

def Calzetti_Law(wave, Rv = 4.05):
    """Dust Extinction Curve by Calzetti et al. (2000)

    Args:
        wave (float): Wavelength
        Rv (float, optional): Extinction curve. Defaults to 4.05.

    Returns:
        float: Extinction value E(B-V)
    """
    
    wave_number = 1./(wave * 1e-4)
    reddening_curve = np.zeros(len(wave))
    
    idx = np.logical_and(wave >= 1200, wave <= 6300)
    reddening_curve[idx] = 2.659 * ( -2.156 + 1.509 * wave_number[idx] - 0.198 * \
                    (wave_number[idx] ** 2)) + 0.011 * (wave_number[idx] **3 ) + Rv
                                    
    idx = np.logical_and(wave >= 6300, wave <= 22000)
    reddening_curve[idx] = 2.659 * ( -1.857 + 1.040 * wave_number[idx]) + Rv
    return reddening_curve

def reddening(wave, flux, ebv = 0.0, law = 'calzetti', Rv = 4.05):
    """
    Reddening an input spectra through a given reddening curve.

    Args:
        wave (float): Wavelength of input spectra
        flux (float): Flux of input spectra
        ebv (float, optional): Extinction value. Defaults to 0.
        law (str, optional): Extinction curve. Defaults to 'calzetti'.
        Rv (float, optional): _description_. Defaults to 4.05.

    Returns:
        float: Flux of spectra after reddening.
    """
    if law == 'calzetti':
        curve = Calzetti_Law(wave, Rv = Rv)
        fluxNew = flux / (10. ** (0.4 * ebv * curve))
    return fluxNew

def flux_to_mag(wave, flux, path, band='GAIA_bp'):
    """Convert flux of given spectra to magnitude

    Args:
        wave (float): Wavelength
        flux (float): Flux of spectra
        band (str, optional): Filter band name. Defaults to 'GAIA_bp'.

    Returns:
        float: value of magnitude
    """
    # /home/yan/MCI_sim/MCI_input/SED_Code/data
    ##import os
    
    ###parent = os.path.dirname(os.path.realpath(__file__))
        
    band = ascii.read(path+'MCI_inputData/SED_Code/seddata/' + band + '.dat')
    wave0= band['col1']
    curv0= band['col2']
    
    # Setting the response
    func = interp1d(wave0, curv0)
    response = np.copy(wave)
    ind_extra = (wave > max(wave0)) | (wave < min(wave0))
    response[ind_extra] = 0
    ind_inside = (wave < max(wave0)) & (wave > min(wave0))
    response[ind_inside] = func(wave[ind_inside])

    # Total Flux
    Tflux = np.trapz(flux * response, wave) / np.trapz(response, wave)
    
    return -2.5 * np.log10(Tflux)

def calibrate(wave, flux, mag, path,band='GAIA_bp'):
    """
    Calibrate the spectra according to the magnitude.

    Args:
        wave (float): Wavelength
        flux (float): Flux of spectra
        mag (float): Input magnitude.
        band (str, optional): Filter band name. Defaults to 'GAIA_bp'.

    Returns:
        float: Flux of calibrated spectra. Units: 1e-17 erg/s/A/cm^2
    """
    inst_mag = flux_to_mag(wave, flux, path,band = band)
    instflux = 10 ** (-0.4 * inst_mag)
    realflux = (mag * u.STmag).to(u.erg/u.s/u.cm**2/u.AA).value

    # Normalization
    flux_ratio = realflux / instflux
    flux_calibrate = flux * flux_ratio * 1e17                        # Units: 10^-17 erg/s/A/cm^2
    
    return flux_calibrate

# ------------
# SED Template

class Gal_Temp():
    """
    Template of Galaxy SED
    """
    
    def __init__(self,path):
        
        ###import os
        
        ###parent = os.path.dirname(os.path.realpath(__file__))
        self.path=path
        hdulist = fits.open(self.path+'MCI_inputData/SED_Code/seddata/galaxy_temp.fits')
        self.wave = hdulist[1].data['wave']
        self.flux = hdulist[2].data
        self.age_grid = hdulist[3].data['logAge']
        self.feh_grid = hdulist[3].data['FeH']
        
    def toMag(self, redshift = 0):
        """Calculating magnitude

        Args:
            redshift (float, optional): redshift of spectra. Defaults to 0.
        """
        wave = self.wave * (1 + redshift)
        self.umag = flux_to_mag(wave, self.flux, self.path,band='SDSS_u')
        self.gmag = flux_to_mag(wave, self.flux, self.path,band='SDSS_g')
        self.rmag = flux_to_mag(wave, self.flux, self.path,band='SDSS_r')
        self.imag = flux_to_mag(wave, self.flux, self.path,band='SDSS_i')
        self.zmag = flux_to_mag(wave, self.flux, self.path,band='SDSS_z')
        
class Star_Temp():
    """
    Template of Stellar SED
    """
    
    def __init__(self,path):
        
        ##import os
        self.path=path
        ####parent = os.path.dirname(os.path.realpath(__file__))
        ###print("获取其父目录——" + parent)  # 从当前文件路径中获取目录
        
        hdulist = fits.open(path+'MCI_inputData/SED_Code/seddata/stellar_temp.fits')
        self.wave = hdulist[1].data['wave']
        self.flux = hdulist[2].data
        self.Teff_grid = hdulist[3].data['Teff']
        self.FeH_grid = hdulist[3].data['FeH']
        
        self.bpmag = flux_to_mag(self.wave, self.flux, path,band='GAIA_bp')
        self.rpmag = flux_to_mag(self.wave, self.flux, path,band='GAIA_rp')
        
        
    def toMag(self):
        wave = self.wave
        self.bpmag = flux_to_mag(wave, self.flux, self.path,band='GAIA_bp')
        self.rpmag = flux_to_mag(wave, self.flux, self.path,band='GAIA_rp')
        
# -------------
# SED Modelling

def Model_Stellar_SED(wave, bp, rp, temp):
    """Modelling stellar SED based on bp, rp magnitude

    Args:
        wave (float): Wavelength
        bp (float): Magnitude of GAIA BP band
        rp (float): Magnitude of GAIA RP band
        temp (class): Class of stellar template

    Returns:
        float array: Spectral energy distribution of stellar SED, 
        which have the same length to the input wave
    """

    color0 = bp - rp
    colors = temp.bpmag - temp.rpmag
    
    idx = np.argmin(np.abs(colors - color0))
    flux0 = temp.flux[idx]
    flux1 = np.interp(wave, temp.wave, flux0)
    flux  = calibrate(wave, flux1, rp, band = 'GAIA_rp')
    return flux

def Model_Galaxy_SED(wave, ugriz, z, temp):
    """Modelling galaxy SED based on u,g,r,i,z magnitude

    Args:
        wave (float): Wavelength
        ugriz (float, array): The array of magnitude of SDSS ugriz band
        z (float): Redshift
        temp (class): Class of gaalxy template

    Returns:
        float array: Spectral energy distribution of stellar SED, 
        which have the same length to the input wave
    """
    sed = 10. ** (-0.4 * ugriz)
    sed = sed / sed[2]
    ntemp = len(temp.rmag)
    dmag  = np.zeros(ntemp)
    for j in range(ntemp):
        ugriz0 = np.array([temp.umag[j], temp.gmag[j], temp.rmag[j], 
                           temp.imag[j], temp.zmag[j]])
        sed0 = 10. ** (-0.4 * ugriz0)
        sed0 = sed0 / sed0[2]
        dmag[j] = np.sum(np.abs(sed - sed0))
    
    idx = np.argmin(dmag)
    flux0 = temp.flux[idx]
    
    # Effect of E(B-V)
    ri0 = ugriz[2] - ugriz[3]
    ri  = temp.rmag - temp.imag
    dri = ri0 - ri[idx]
    Alambda = Calzetti_Law(np.array([6213 / (1 + z), 7625 / (1 + z)]))
    eri0 = (Alambda[0] - Alambda[1])
    ebv = dri/eri0
    if ebv<0:
        ebv=0
    if ebv>0.5:
        ebv=0.5
    flux1 = reddening(temp.wave, flux0, ebv = ebv)
    
    flux2 = np.interp(wave, temp.wave * (1 + z), flux1)
    flux  = calibrate(wave, flux2, ugriz[2], band = 'SDSS_r')
    return flux