target.py

import os
import re
import json
import numpy as np
from scipy import constants
from astropy.io import fits
from astropy.coordinates import SkyCoord

from CpicImgSim.config import cpism_refdata, MAG_SYSTEM
from CpicImgSim.config import S  # pysynphot
from CpicImgSim.optics import filter_throughput
from CpicImgSim.io import log

PLANK_H = constants.h  # ~6.62607015e-34
LIGHT_SPEED = constants.c  # ~3e8
DEG2RAD = np.pi / 180
R_JUPITER_METER = 6.99e4
AU_METER = 1.49e8

DEFAULT_FILTER_FILE = f'{cpism_refdata}\\throughtput\\f661_total.fits'
HYBRID_MODEL_FILE = f'{cpism_refdata}\\target_model\\hybrid_model.fits'
BCC_MODEL_FOLDER = f'{cpism_refdata}\\target_model\\bccmodels'

CATALOG_CACHE = {}

class AlbedoCat(S.spectrum.TabularSpectralElement, S.catalog.Icat):
    """
    This class is used to generate a spectrum from the planet reflection model.
    
    Parameters
    ----------
    phase : float
        [degree], The phase angle of the planet, range from 0 to 180.
    metallicity : float
        The metallicity of the planet. log(Z/Zsun), range from 0 to 2.
    f_sed : float
        The sedimentation efficiency of the cloud. log(f_sed), range from -2 to 2. 2 is for cloud free case.

    Notes
    -----
    f_sed is only reliable between -2 and log10(6) and the cloud free case (2).
    Values between log10(6) and 2 are interpolated from the cloud free case (2) and log10(6).
    
    """
    
    def __init__(self, 
                phase: float, 
                metallicity: float,
                f_sed: float,
            ):
        catname = 'BCCalbedo'
        self.isAnalytic=False
        self.name = f"{catname}_{phase}_{metallicity}_{f_sed}"
        self.parameter_names = ['phase', 'metallicity', 'f_sed']
        self.catalog_folder = BCC_MODEL_FOLDER
        filename = os.path.join(self.catalog_folder, 'catalog.fits')

        if filename in CATALOG_CACHE:
            indices = CATALOG_CACHE[filename]
        else:
            with fits.open(filename) as table:
                indexList = table[1].data.field('INDEX')
                filenameList = table[1].data.field('FILENAME')
            
            indices = self._getArgs(indexList, filenameList)
            CATALOG_CACHE[filename] = indices

        list0,list1 = self._breakList(indices, 0, phase)

        list2,list3 = self._breakList(list0, 1, metallicity)
        list4,list5 = self._breakList(list1, 1, metallicity)

        list6,list7   = self._breakList(list2, 2, f_sed)
        list8,list9   = self._breakList(list3, 2, f_sed)
        list10,list11 = self._breakList(list4, 2, f_sed)
        list12,list13 = self._breakList(list5, 2, f_sed)

        sp1, wave, waveunits = self._getSpectrum(list6[0],  catname, wave_output=True)
        sp2 = self._getSpectrum(list7[0],  catname)
        sp3 = self._getSpectrum(list8[0],  catname)
        sp4 = self._getSpectrum(list9[0],  catname)
        sp5 = self._getSpectrum(list10[0], catname)
        sp6 = self._getSpectrum(list11[0], catname)
        sp7 = self._getSpectrum(list12[0], catname)
        sp8 = self._getSpectrum(list13[0], catname)

        spa1 = self._interpolateSpectrum(sp1, sp2, f_sed)
        spa2 = self._interpolateSpectrum(sp3, sp4, f_sed)
        spa3 = self._interpolateSpectrum(sp5, sp6, f_sed)
        spa4 = self._interpolateSpectrum(sp7, sp8, f_sed)

        spa5 = self._interpolateSpectrum(spa1, spa2, metallicity)
        spa6 = self._interpolateSpectrum(spa3, spa4, metallicity)

        spa7 = self._interpolateSpectrum(spa5, spa6, phase)

        sp = spa7[0]

        self._wavetable = wave * 1e4
        self._throughputtable = sp
        self.waveunits = S.units.Units(waveunits.lower())
        self.warnings = {}

    def _getSpectrum(self, parlist, catdir, wave_output=False):
        name = parlist[3]

        filename = name.split('[')[0]
        column = name.split('[')[1][:-1]

        filename = f"{self.catalog_folder}/{filename}"
        
        # sp = S.spectrum.TabularSpectralElement(filename, thrucol=column)

        with fits.open(filename) as td:
            sp = td[1].data.field(column)
            wave = td[1].data.field('wavelength')
            # waveunits = td[1].header['tunit1']
            waveunits = 'micron'

        result = []
        for member in parlist:
            result.append(member)

        result.pop()
        result.append(sp)

        if wave_output:
            return result, wave, waveunits

        return result



def _sptype2num(spectral_type):
    """
    convert spectral type string to number, for interpretation

    case0: normal case
    - M1V: 6.1, 5
    - O5IV: 0.5, 4
    - F3V: 3.3, 5
    - K4.5II: 5.45, 2

    case 1: star type or subtype missing
    zero or V will return
    - M1: 6.1, 5
    - M: 6.0, 5

    case 2: spectral type + subtype
    subtype will be ignored
    - K3Vvar: 5.3, 5
    - F6Vbwvar: 3.6, 5
    - K0IV SB: 5.0, 4
    - F5V+: 3.5, 5

    case 3: multi spectral type
    only the first sptype is used
    - G5IV/V +K1IV: 4.5, 4
    - F7IV-V: 3.7, 4
    - O4/O5IV: 0.4, 0
    - G3/G5V: 4.3, 0

    case 4: illegal input
    ValueError will be raised
    """
    obafgkm = 'OBAFGKML'
    spectral_type = spectral_type.upper()
    # match spectral type such as M1V, O5IV, F3V, K4.5II
    matched = re.match(
        RF'^([{obafgkm}])([0-9]\d*\.?\d*)*([IV]*)', spectral_type)

    if not matched:
        raise ValueError(f"illegal spectral type input: {spectral_type}")

    shorttype = obafgkm.find(matched.group(1))

    subtype = 0.0
    if matched.group(2):
        subtype = float(matched.group(2))

    stlist = ['O', 'I', 'II', 'III', 'IV', 'V']
    startype = dict(zip(stlist, range(len(stlist)))).get(matched.group(3), 5)

    return shorttype + subtype / 10, startype


def _spnum2teff(spn, stn):
    """
    interpret of spectral number (by __sptype2num) to get t_eff and log_g
    look up table from the document of ck04model
    """
    with open(f'{cpism_refdata}\\target_model\\sptype2teff_lut.json', 'r') as fid:
        sptype_teff_lut = json.load(fid)

    def _interp(spn, stn):
        lut = sptype_teff_lut[f'{stn}']
        teff = np.interp(spn, lut[0], lut[1])
        logg = np.interp(spn, lut[0], lut[2])
        return teff, logg

    stn = 5 if stn not in [1, 2, 3, 4, 5] else stn

    if stn in [1, 3, 5]:
        return _interp(spn, stn)
    else:
        teff_low, logg_low = _interp(spn, stn-1)
        teff_high, logg_high = _interp(spn, stn+1)
        return (teff_high + teff_low)/2, (logg_high + logg_low)/2


def star_photlam(
        magnitude: float,
        sptype: str,
        is_blackbody: bool = False,
        mag_input_band: str = 'f661'):
    """
    genrate flux spectrum of a star by its observal magnitude and spectral type

    Parameters
    ----------
    magnitude: float
        magnitude of the star
    sptype: str
        spectral type of the star
    is_blackbody: bool
        if True, use blackbody spectrum instead of ck04model
    mag_input_band: str
        bandpass of the input magnitude, default is f661

    Returns
    -------
    pysynphot.spectrum.ArraySpectrum
        spectrum of the star in photlam unit
    """
    spn, stn = _sptype2num(sptype)
    t_eff, log_g = _spnum2teff(spn, stn)

    log.debug(f"{sptype} star => [{t_eff=:}, {log_g=:}]; {is_blackbody=:}")
    filter = filter_throughput(mag_input_band)

    if not is_blackbody:
        METALLICITY = 0
        spectrum = S.Icat('ck04models', t_eff, METALLICITY, log_g)
    else:
        spectrum = S.BlackBody(t_eff)

    star_sp = spectrum.renorm(magnitude, MAG_SYSTEM, filter)
    star_sp.convert('photlam')
    return star_sp


def bcc_spectrum(
        coe_cloud: float,
        coe_metalicity: float
):
    """Albedo spectrum of planet using BCC model (Batalha et al. 2018),

    Parameters
    ----------
    coe_cloud: float
        The sedimentation efficiency of the cloud. log(f_sed), range from -2 to 2. 2 is for cloud free case.
    coe_metalicity: float
        The metallicity of the planet. log(Z/Zsun), range from 0 to 2.

    Returns
    -------
    pysynphot.spectrum.ArrayBandpass
        albedo spectrum of the planet
    """

    spec = AlbedoCat(0, coe_metalicity, coe_cloud)
    spec.convert('angstrom')
    return spec


def hybrid_albedo_spectrum(
        coe_b: float,
        coe_r: float):
    """Albedo spectrum of planet using hybrid-jupiter-neptune model (Lacy et al. 2018)
    jupiter and neptune spectrum is from Karkoschka’s 1994

    Parameters
    ----------
    coe_b: float
        coefficient of blue spectrum, 1 for jupiter, 0 for neptune
    coe_r: float
        coefficient of red spectrum, 1 for jupiter, 0 for neptune

    Returns
    -------
    pysynphot.spectrum.ArrayBandpass
        albedo spectrum of the planet
    """
    log.debug(f"planet hybrid spectrum with {coe_b=:}, {coe_r=:}")
    model = fits.getdata(HYBRID_MODEL_FILE)
    spec = model[1, :] * coe_r
    spec += model[2, :] * coe_b
    spec += model[3, :] * (1 - coe_r)
    spec += model[4, :] * (1 - coe_b)

    albedo = S.ArrayBandpass(
        wave=model[0, :],
        throughput=spec,
        waveunits='nm'
    )
    albedo.convert('angstrom')
    return albedo


def extract_target_x_y(
        target: dict,
        ra0: str = None,
        dec0: str = None):
    """
    extract x, y of target from target dict

    Parameters
    ----------
    target: dict
        target dict. must contain either (ra, dec) or (pangle, spearation)
    ra0: str
        ra of center star. must be provided if (ra, dec) of target is used
    dec0: str
        dec of center star. must be provided if (ra, dec) of target is used

    Returns
    -------
    x, y: float
        x, y of target in arcsec

    Raises
    ------
    ValueError
        if (ra, dec) of target is used but (ra, dec) of center star is not provided.

    ValueError
        one of (ra, dec) or (pangle, spearation) is not provided.
    """

    def _pa2xy(p_angle, separation):
        p_angle_rad = p_angle * DEG2RAD
        x = separation * np.sin(p_angle_rad)
        y = separation * np.cos(p_angle_rad)
        log.debug(f"({p_angle=:}, {separation=:}) => ({x=:}, {y=:})")
        return x, y

    if 'pangle' in target.keys() and 'separation' in target.keys():
        return _pa2xy(target['pangle'], target['separation'])

    if 'ra' not in target.keys() or 'dec' not in target.keys():
        raise ValueError(
            'either (ra, dec) or (pangle, separation) needed in target dict')

    if ra0 is None or dec0 is None:
        raise ValueError(
            '(ra, dec) of center star must be provided if (ra, dec) of bkstar is used'
        )

    ra, dec = target['ra'], target['dec']
    log.debug(f"target: {ra=:}, {dec=:}, center star: {ra0=:}, {dec0=:}")
    cstar = SkyCoord(ra0, dec0)
    bkstar = SkyCoord(ra, dec)
    separation = cstar.separation(bkstar).arcsec
    p_angle = cstar.position_angle(bkstar).degree
    x, y = _pa2xy(p_angle, separation)

    return x, y
    

class TargetOjbect(object):
    """A helper class to generate target spectrum and albedo spectrum

    Attributes
    ----------
    x: float
        x of target in arcsec
    y: fload
        y of target in arcsec
    ra: str
        ra string for center star, such as '15d23m05s'
    dec: str
        dec string for center star
    distance: float
        distance of center star in pc
    image: 2d np.array
        image of the target
    spectrum: pysynphot.spectrum.Spectrum
        spectrum of the target
    """

    def __init__(self, info, cstar=None, sp_model=None):
        """Initialize a target object

        Parameters
        ----------
        info: dict
            target info, see Example for more details
        cstar: TargetOjbect or None
            center star object bounded, if None, means the target is the center star itself
            center star object is used to calculate the x, y of each target according to its ra and dec
            also center star's distance is used to calculate seperation of planet
        sp_model: str
            keyword to compatible with V1.0 input. See Notes
            spectral type of the target, can be 'star' or 'planet'

        Notes
        -----
        keyargv 中的sp_model 完全用来用来兼容V1.0版的输入,执行逻辑如下:
         1. 如果没有设置sp_model，对应了V2.0的输入，则sp_model从info中获取
         2. 如果设置了sp_model，则区分为star和planet
         3. 如果sp_model为star, 则判断是否设置了isblackbody，如果是的话，则sp_model改为blackbody
         4. 如果sp_model为planet，则进一步区分为hybrid_planet 和 bcc_planet 和 template_planet
         注意：2.0版本的输入尽量不要使用sp_model关键字,逻辑还是有点复杂的，容易乱，而应该直接在info中设置

        Examples: 
        --------

        cstar = {
            'magnitude': 0,
            'ra': '120d',
            'dec': '40d',
            'distance': 10,
            'sptype': 'G0III'
        }
             
        stars = {
            'magnitude': 15,
            'ra': '120.001d',
            'dec': '40.001d',
            'sptype': 'F0III',
            'sp_spectrum': 'blackbody',
        }
            
        planets = {
            'radius': 2,
            'pangle': 60,
            'coe_cloud': 0.3,
            'coe_metal': 0.7,
            'separation': 0.5,
            'phase_angle': 90,                 
            'sp_spectrum': 'hybrid_planet',
            'image': 'extend_planet.fits'
        }

        # planet using input costum albedo spectrum!
        # Note that the albedo spectrum is not normalized! 
        # so the contrast of the planet sould be considered in the input file by user!
        # The input file is in pysynphot.FileSpectralElement format.
        # See the documentation of pysynphot for details. 

        planets = {
            'pangle': 60,
            'separation': 0.5,              
            'sp_spectrum': 'template_planet',
            'template': 'planet_albedo.fits'
        }
        """

        if sp_model is None:  
            sp_model = info.get('sp_type', 'star') # 如果不设置，默认为star
        elif sp_model == 'star' and info.get('isblackbody', False):
            #原先的star 现在分为star和blackbody两个类型
            sp_model = 'blackbody'
        elif sp_model == 'planet':
            # 原先的planet 现在分为hybrid_planet 和 bcc_planet两类
            sp_model = info.get('model', 'hybrid_planet')
            if sp_model not in ['bcc_planet', 'hybrid_planet', 'template_planet']:
                Warning(f'Unknown planet model {sp_model}, using default[hybrid] planet model' )

        self.sp_model = sp_model

        if cstar is None:
            self.x, self.y = 0, 0
            self.ra = info['ra']
            self.dec = info['dec']
            self.distance = info.get('distance', None)
        else:
            self.x, self.y = extract_target_x_y(info, cstar.ra, cstar.dec)


        self.image = None
        if 'image' in info.keys():
            file = info['image']
            if os.path.exists(file):
                self.image = fits.getdata(file)    
            else:
                self.image = None
                Warning(f"extend Ojbect File {file} not found, return None!")
        

        if self.sp_model == 'blackbody':
            self.spectrum = star_photlam(
                info['magnitude'],
                info['sptype'],
                is_blackbody=True,
                mag_input_band=info.get('mag_input_band', 'f661')
            )

        if self.sp_model == 'template_star':
            self.spectrum = S.FileSpectrum(info['template'])

        if self.sp_model == 'star':
            self.spectrum = star_photlam(
                info['magnitude'],
                info['sptype'],
                is_blackbody=False,
                mag_input_band=info.get('mag_input_band', 'f661')
            )

        if self.sp_model in ['hybrid_planet', 'bcc_planet']:
            planet = info
            phase_angle = planet.get('phase_angle', 90)

            radius = planet.get('radius', 1)

            if cstar.distance is None:
                raise ValueError('distance of center star must be provided if planet is added')


            if planet.get('contrast', None) is not None:
                contrast = planet['contrast']
            else:
                contrast = planet_contrast(
                    self.x * cstar.distance,
                    self.y * cstar.distance,
                    phase_angle,
                    radius,
                )

            if sp_model == 'hybrid_planet':
                coe_blue, coe_red = planet.get('coe_b', 1), planet.get('coe_r', 1)
                albedo_spect = hybrid_albedo_spectrum(coe_blue, coe_red) * contrast
            elif sp_model == 'bcc_planet':
                coe_c, coe_m = planet.get('coe_cloud', 2), planet.get('coe_metal', 0)
                albedo_spect = bcc_spectrum(coe_c, coe_m) * contrast
            if sp_model == 'template_planet':
                albedo_spect = S.FileBandpass(planet['template'])

            self.spectrum = cstar.spectrum * albedo_spect 



def planet_contrast(
        planet_x_au: float,
        planet_y_au: float,
        phase_angle: float,
        radius: float):
    """
    calculate the contrast of a planet

    Parameters
    ----------
    planet_x_au: float
        x position of the planet in au
    planet_y_au: float
        y position of the planet in au
    phase_angle: float
        phase angle of the planet in degree
    radius: float
        radius of the planet in jupiter radius

    Returns
    -------
    contrast: float
        contrast of the planet
    """
    separation = np.sqrt(planet_x_au**2 + planet_y_au**2)
    phase_angle = phase_angle * DEG2RAD

    if np.sin(phase_angle) < 1e-9:
        raise ValueError('sin(phase_angle) can not be 0')

    sep_3d = separation / np.sin(phase_angle)

    # Lambert Scattering phase function
    # from Madhusudhan and Burrows 2012 equation 33.
    phase = (np.sin(phase_angle) + (np.pi - phase_angle)
             * np.cos(phase_angle)) / np.pi
    log.debug(f'alpha: {phase_angle/np.pi*180} {phase=}')

    contrast = (radius / sep_3d * R_JUPITER_METER / AU_METER)**2 * phase
    return contrast

def spectrum_generator(
        targets: dict):
    """
    generate the spectrum due to the input target list

    Parameters
    ----------
    targets: dict
        target dictionary which contains keyword 'cstar' (necessary), 'stars'(optional), 'planets' (optional).
        The values are: 
        
        - cstar: dict
            - center star information. must contain keywords ra, dec, distance, magnitude, sptype
        - stars: list of dict, optional, not preferred, compatible with old version
            - list of background stars. each dict must contain ra, dec, magnitude, sptype
        - planets: list of dict, optional, not preferred, compatible with old version
            - list of planets. each dict must contain pangle, separation, magnitude, radius
        - objects: list of dict, optional, recommended! V2.0 new!
            - list of targets. each dict must contain ra, dec, magnitude, sptype

    Returns
    -------
    obj_sp_list: list
        list of [x, y, spectrum, image] of each target
    """

    cstar = targets['cstar']
    stars = targets.get('stars', [])
    planets = targets.get('planets', [])
    objects = targets.get('objects', [])

    obj_sp_list = []

    cstar_obj = TargetOjbect(cstar)
    obj_sp_list.append([cstar_obj.x, cstar_obj.y, cstar_obj.spectrum, cstar_obj.image])

    for star in stars:
        star_obj = TargetOjbect(star, sp_model='star', cstar=cstar_obj)
        obj_sp_list.append([star_obj.x, star_obj.y, star_obj.spectrum, star_obj.image])

    for planet in planets:
        planet_obj = TargetOjbect(planet, sp_model='planet', cstar=cstar_obj)
        obj_sp_list.append([planet_obj.x, planet_obj.y, planet_obj.spectrum, planet_obj.image])

    for target in objects:
        target_obj = TargetOjbect(target, cstar=cstar_obj)
        obj_sp_list.append([target_obj.x, target_obj.y, target_obj.spectrum, target_obj.image])

    return obj_sp_list