target.py

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

from .config import cpism_refdata, MAG_SYSTEM
from .config import S  # pysynphot
from .optics import filter_throughput
from .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'


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 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


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
            - list of background stars. each dict must contain ra, dec, magnitude, sptype
        - planets: list of dict, optional
            - list of planets. each dict must contain pangle, separation, magnitude, radius

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

    Examples
    --------
    >>> targets = {
    ...     target = {
    ...         'name': 'TEST1',
    ...         'cstar': {
    ...             'magnitude': 0,
    ...             'ra': '120d',
    ...             'dec': '40d',
    ...             'distance': 10,
    ...             'sptype': 'G0III',
    ...         },
    ...         'stars': [
    ...             {
    ...                 'magnitude': 15,
    ...                 'ra': '120.001d',
    ...                 'dec': '40.001d',
    ...                 'sptype': 'F0III',
    ...                 'isblackbody': True
    ...             },
    ...             {
    ...                 'magnitude': 10,
    ...                 'pangle': -60,
    ...                 'separation': 2,
    ...                 'sptype': 'A2.5II',
    ...                 'isblackbody': False
    ...             },
    ...     
    ...         ],
    ...         'planets': [
    ...             {
    ...                 'radius': 2,
    ...                 'pangle': 60,
    ...                 'coe_b': 0.3,
    ...                 'coe_r': 0.7,
    ...                 'separation': 0.5,
    ...                 'phase_angle': 90,
    ...             }
    ...         ]
    ...     }
    >>> spectrum_generator(targets)
    """

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

    obj_sp_list = []

    cstar_spectrum = star_photlam(
        cstar['magnitude'],
        cstar['sptype'],
        is_blackbody=cstar.get('isblackbody', False),
        mag_input_band=cstar.get('mag_input_band', 'f661')
    )

    cstar_ra = cstar['ra']
    cstar_dec = cstar['dec']
    cstar_dist = cstar['distance']
    obj_sp_list.append([0, 0, cstar_spectrum])

    for star in stars:
        star_x, star_y = extract_target_x_y(star, cstar_ra, cstar_dec)

        spectrum = star_photlam(
            star['magnitude'],
            star['sptype'],
            is_blackbody=star.get('isblackbody', False),
            mag_input_band=cstar.get('mag_input_band', 'f661')
        )

        obj_sp_list.append([star_x, star_y, spectrum])

    for planet in planets:
        planet_x, planet_y = extract_target_x_y(planet)

        # angle between observation direction and star-planet direction
        phase_angle = planet.get('phase_angle', 90)

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

        contrast = planet_contrast(
            planet_x * cstar_dist,
            planet_y * cstar_dist,
            phase_angle,
            radius,
        )

        coe_blue, coe_red = planet.get('coe_b', 1), planet.get('coe_r', 1)
        albedo_spect = hybrid_albedo_spectrum(coe_blue, coe_red)

        spectrum = albedo_spect * cstar_spectrum * contrast
        obj_sp_list.append([planet_x, planet_y, spectrum])

    return obj_sp_list