target.py

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

from .config import S  # pysynphot
from .config import config
from .optics import filter_throughput
from .io import log
from pysynphot.renorm import StdRenorm


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 = config['default_band']
HYBRID_MODEL_FILE = config['hybrid_model']
BCC_MODEL_FOLDER = config['bcc_model']
MAG_SYSTEM = config['mag_system']
CATALOG_CACHE = {}

class AlbedoCat(S.spectrum.TabularSpectralElement, S.catalog.Icat):
    """Generate albedo spectrum from the planet reflection model in Batalha et al. 2018
    
    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).

    Reference
    ---------
    Color Classification of Extrasolar Giant Planets: Prospects and Cautions
    Natasha E. Batalha et al 2018 AJ 156 158
    
    """
    
    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: str) -> tuple:
    """
    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: int,
        stn: int) -> tuple:
    """
    interpret of spectral number (by __sptype2num) to get t_eff and log_g
    look up table from the document of ck04model
    """
    with open(config['sp2teff_model'], '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') -> S.ArraySpectrum:
    """
    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 standard_spectrum(
        magnitude: float) -> S.ArraySpectrum:
    """Standard spectrum of magnitude system. 
    For example, the standard_spectrum of vega megnitude is vega spectrum. 
    The standard spectrum of ab magnitude is a flat spectrum.

    Parameters
    -----------
    magnitude : float
        magnitude of the standard spectrum
    
    Returns
    ----------
    star_sp : S.spectrum.Spectrum
    
    """
    inner_std = S.units.Units(MAG_SYSTEM).StdSpectrum
    std_spectrum = S.ArraySpectrum(
        inner_std.wave,
        inner_std.flux,
        inner_std.waveunits,
        inner_std.fluxunits
    )
    filter = filter_throughput(DEFAULT_FILTER_FILE)
    std_spectrum = StdRenorm(std_spectrum, filter, magnitude, MAG_SYSTEM)
    std_spectrum.convert('photlam')
    return std_spectrum
    
def bcc_spectrum(
        coe_cloud: float,
        coe_metalicity: float) -> S.spectrum.ArraySpectralElement:
    """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) -> S.spectrum.ArraySpectralElement:
    """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) -> tuple:
    """
    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


def detect_template_path(template: str) -> str:
    """Find template file in catalog folder or current folder.

    Parameters
    ----------
    template: str
        template file name
    
    Returns
    -------
    str
        absolute path of template file

    """

    if os.path.isfile(template):
        return os.path.abspath(template)
    
    catalog = config['catalog_folder']
    cat_temp = os.path.join(catalog, template)
    if os.path.isfile(cat_temp):
        return cat_temp

    raise FileExistsError(f'cant find {template} in ./ or catalog folder')
    

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):
        """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 projection x, y of each target according to its ra and dec
            also center star's distance is used to calculate seperation of planet

        Examples: 
        --------

        cstar = {
            'magnitude': 0,
            'ra': '120d',
            'dec': '40d',
            'distance': 10,
            'sptype': 'G0III'
            'sp_model': 'star'
        }
             
        stars = {
            'magnitude': 15,
            'ra': '120.001d',
            'dec': '40.001d',
            'sptype': 'F0III',
            'sp_model': 'blackbody'
        }
            
        planets = {
            'radius': 2,
            'pangle': 60,
            'coe_cloud': 0.3,
            'coe_metal': 0.7,
            'separation': 0.5,
            'phase_angle': 90,                 
            'sp_model': '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'
        }
        """
        self.sp_model = info.get('sp_model', 'star')
        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():
            self.image = fits.getdata(detect_template_path(info['image']))

        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 == 'reference':
            self.spectrum = standard_spectrum(info['magnitude'])

        if self.sp_model == 'template_star':
            self.spectrum = S.FileSpectrum(detect_template_path(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', 'template_planet']:
            planet = info
            phase_angle = planet.get('phase_angle', 90)
            sp_model = self.sp_model

            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)
            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)
            else: #sp_model == 'template_planet'
                albedo_spect = S.FileBandpass(detect_template_path(planet['template']))
                contrast = 1

            self.spectrum = cstar.spectrum * albedo_spect * contrast

def planet_contrast(
        planet_x_au: float,
        planet_y_au: float,
        phase_angle: float,
        radius: float) -> 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) -> list:
    """
    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
        - 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']
    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 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


def target_file_load(
        target: Union[dict, str]) -> dict:
    """Generate target dict from file, string or dict.

    Parameters
    ----------
    target: Union[dict, str]
        target file name, formated string or target dict.
    
    Outputs
    --------
    target: dict
        dictionary of target. Format of input 

    Note
    -------
    If target is a string start with *, it will be treated as a formatted string.
    e.g. "*5.1/25.3(1.3,1.5)/22.1(2.3,-4.5)" which means a central object
    with magnitude 5.1, and two substellar with magnitude 25.3 and 22.1, respectively. 
    The spectrum of each object is standard refernece spectrum of ABmag.
    The first number in the parenthesis is the x position in arcsec, and the second is the y position.

    If target is a string without *, it will be treated as file name. And find the file in catalog folder.
    The file need to be in yaml format.
    And end with .yaml (note .yml not work). If not .yaml will be added.

    If target is a dict, it will be returned directly.

    If all the above conditions are not met, an empty dict will be returned.
    """
    if isinstance(target, dict):
        return target

    if not target: # None or empty string
        return {}
    
    if isinstance(target, str): #filename or formatted string
        target = target.strip()
        if not target:
            return {}
        
        catalog_folder = config['catalog_folder']
        target_file = target
        target_file += '.yaml' if target_file[-5:].lower() != '.yaml' else ""
        target_name = os.path.basename(target_file)[:-5]
        file_search = [target_file, os.path.join(catalog_folder, target_file)]
        
        for file in file_search:
            if os.path.isfile(file):
                with open(file) as fid:
                    target = yaml.load(fid, Loader=yaml.FullLoader)
                    target['name'] = target_name
                    return target

        target_str = target
        if (target_str[0] == '*'):
            objects = target_str[1:].split('/')
            cstar_mag = float(objects[0])
            cstar = {
                'magnitude': cstar_mag,
                'ra': '0d',
                'dec': '0d',
                'sp_model': 'reference',
                'distance': 10,
            }

            stars = []
            for sub_stellar in objects[1:]:
                float_regex = R"[+-]?\d+(?:\.\d+)?"
                match = re.match(
                    rf"({float_regex})\(({float_regex}),({float_regex})\)", sub_stellar)
                if not match:
                    log.error(f'Wrong format for sub stellar: {sub_stellar}, Skip it')
                    continue
                mag = float(match.group(1))
                x = float(match.group(2))
                y = float(match.group(3))
                pangle = np.arctan2(x, y) * 180 / np.pi
                separation = np.sqrt(x**2 + y**2)
                stars.append({
                    'magnitude': mag,
                    'pangle': pangle,
                    'separation': separation,
                    'sp_model': 'reference',
                })
            target_dict = {
                'name': target_str[1:],
                'cstar': cstar,
                'objects': stars,
            }
            return target_dict

    log.error(f'Wrong format for target: {target}, using blank target instead')
    return {}