Constructing an oblate stellar model in PyMC3 w/ starry

oblate_sketch.py

import pymc3 as pm
import pymc3_ext as pmx
import numpy as np
import starry
import theano.tensor as tt
from astropy import constants, units
import matplotlib.pyplot as plt


def get_stellar_inclination_prior(incstar):
    """
    Applies a `pm.Potential` likelihood penalty for the inclination.

    Args:
        incstar (scalar): The stellar inclination in radians
    """
    return pm.Potential("isotropy", tt.log(tt.sin(incstar)))


def get_orbital_inclination(b, mstar, porb):
    """
    Returns the orbital inclination in radians.

    Args:
        b (scalar): The impact parameter
        mstar (scalar): The stellar mass in Solar masses
        porb (scalar): The orbital period in days
    """
    G = constants.G.to(units.R_sun ** 3 / units.M_sun / units.day ** 2).value
    aRs = (np.sqrt(G * mstar) * porb / (2 * np.pi)) ** (2.0 / 3.0)
    return tt.arccos(b / aRs)


def get_star(omega, beta, tpole, u, incstar, mstar, rstar, wav, tput):
    """
    Returns a `starry.Primary` instance.

    Args:
        omega (scalar): The unitless angular velocity
        beta (scalar): The von Zeipel law coefficient
        tpole (scalar): The polar temperature in K
        u (vector): The limb darkening coefficients
        incstar (scalar): The stellar inclination in radians
        mstar (scalar): The stellar mass in Solar masses
        rstar (scalar): The stellar radius in Solar radii
        wav (vector): The wavelength array in nm
        tput (vector): The TESS bandpass throughput evaluated on the wav array
    """
    map = starry.Map(udeg=len(u), gdeg=4, oblate=True, nw=wav.shape[0])
    map.omega = omega
    map.beta = beta
    map.wav = wav
    map.tpole = tpole
    map.f = 1 - 2 / (omega ** 2 + 2)
    map[1:] = u
    map.inc = incstar * 180 / np.pi
    map.amp = tput
    return starry.Primary(map, m=mstar, r=rstar)


def get_planet(porb, r, t0, Omega, incorb, wav):
    """
    Returns a `starry.Secondary` instance.

    Args:
        porb (scalar): The orbital period in days
        r (scalar): The planetary radius in Solar radii
        t0 (scalar): The time of transit in days
        Omega (scalar: The longitude of ascending node of the orbit in radians
        incorb (scalar): The orbital inclination in radians
        wav (vector): The wavelength array in nm
    """
    return starry.Secondary(
        starry.Map(amp=0, nw=wav.shape[0]),
        porb=porb,
        r=r,
        t0=0,
        Omega=Omega,
        omega=0,
        m=0,
        inc=incorb * 180 / np.pi,
    )


with pm.Model() as model:

    # Star
    omega = pm.Uniform("omega", 0, 1)
    beta = 1.23
    nw = 50
    wav = np.linspace(600, 900, nw)
    tput = np.ones_like(nw)
    tpole = 7340
    u = [0.5, 0.25]
    incstar = pmx.Periodic("incstar", lower=0, upper=np.pi)
    mstar = 1.59
    rstar = 1.561
    star = get_star(omega, beta, tpole, u, incstar, mstar, rstar, wav, tput)

    # Planet
    porb = 1.2198681
    rprstar = 0.1087
    r = rprstar * rstar
    t0 = 0
    Omega = pmx.Angle("Omega")
    b = pm.Uniform("b", lower=0, upper=1)
    incorb = get_orbital_inclination(b, mstar, porb)
    planet = get_planet(porb, r, t0, Omega, incorb, wav)

    # System
    sys = starry.System(star, planet)

    # Flux
    time = np.linspace(-0.1, 0.1, 500)
    flux_model = sys.flux(time, integrated=True)

    plt.plot(time, pmx.eval_in_model(flux_model, point=model.test_point))
    plt.show()