Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_ylm():
map = starry.Map(ydeg=1)
map[1, :] = [0.3, 0.2, 0.1]
orbit = exo.orbits.KeplerianOrbit(period=1.0, m_star=1.0, r_star=1.0)
t = np.linspace(-0.25, 0.25, 100)
theta = 30.
# Compute the whole light curve with Theano
f1 = map.flux(t=t, orbit=orbit, ro=0.1, theta=theta, use_in_transit=False).eval()
# Compute just the transit with Theano
f2 = map.flux(t=t, orbit=orbit, ro=0.1, theta=theta, use_in_transit=True).eval()
# Compute the whole light curve without Theano
coords = orbit.get_relative_position(t)
xo = (coords[0] / orbit.r_star).eval()
yo = (coords[1] / orbit.r_star).eval()
zo = -(coords[2] / orbit.r_star).eval()
f3 = map.flux(xo=xo, yo=yo, zo=zo, ro=0.1, theta=theta)
def test_ylm_phase():
texp = 0.05
map = starry.Map(ydeg=2)
np.random.seed(11)
map[1:, :] = 0.1 * np.random.randn(8)
theta = np.linspace(0, 360, 10000)
t = np.linspace(-0.2, 0.2, 10000)
orbit = exo.orbits.KeplerianOrbit(period=1.0)
flux = map.flux(theta=theta)
window = int(texp / (t[1] - t[0]))
fluence_mavg = moving_average(flux, window)
fluence_starry = map.flux(t=t, orbit=orbit, theta=theta,
texp=texp, oversample=50).eval()
# The error is primarily coming from our moving average
# integrator, so let's be lenient
f1 = fluence_mavg[window:-window]
f2 = fluence_starry[window:-window]
assert np.allclose(f1, f2, atol=1e-4, rtol=1e-4)
def test_ylm_occ():
texp = 0.05
map = starry.Map(ydeg=2)
np.random.seed(11)
map[1:, :] = 0.1 * np.random.randn(8)
orbit = exo.orbits.KeplerianOrbit(period=1.0, m_star=1.0, r_star=1.0)
t = np.linspace(-0.2, 0.2, 10000)
flux = map.flux(t=t, orbit=orbit, ro=0.1).eval()
xo = orbit.get_relative_position(t)[0].eval()
yo = orbit.get_relative_position(t)[1].eval()
flux = map.flux(xo=xo, yo=yo, ro=0.1)
fluence_mavg = moving_average(flux, int(texp / (t[1] - t[0])))
fluence_starry = map.flux(t=t, orbit=orbit, ro=0.1,
texp=texp, oversample=30).eval()
fluence_starry_vec = map.flux(t=t, orbit=orbit, ro=0.1,
texp=np.ones_like(t) * texp, oversample=30).eval()
assert np.allclose(fluence_mavg, fluence_starry, fluence_starry_vec)
w=60,
length_unit=u.Rsun,
mass_unit=u.Msun,
angle_unit=u.degree,
time_unit=u.day,
)
# Define the system
sys = starry.System(A, b)
# Compute with starry
time = np.linspace(-0.5, 0.5, 1000)
rv1 = sys.rv(time, keplerian=True)
# Compute with exoplanet
orbit = exoplanet.orbits.KeplerianOrbit(
period=1.0,
t0=0.0,
incl=86.0 * np.pi / 180,
ecc=0.3,
omega=60 * np.pi / 180,
m_planet=0.01,
m_star=1.0,
r_star=1.0,
)
rv2 = orbit.get_radial_velocity(time).eval()
assert np.allclose(rv1, rv2)
def test_ld():
texp = 0.05
map = starry.Map(udeg=2)
map[1:] = [0.4, 0.26]
orbit = exo.orbits.KeplerianOrbit(period=1.0, m_star=1.0, r_star=1.0)
t = np.linspace(-0.2, 0.2, 10000)
flux = map.flux(t=t, orbit=orbit, ro=0.1).eval()
fluence_mavg = moving_average(flux, int(texp / (t[1] - t[0])))
fluence_starry = map.flux(t=t, orbit=orbit, ro=0.1,
texp=texp, oversample=30).eval()
fluence_starry_vec = map.flux(t=t, orbit=orbit, ro=0.1,
texp=np.ones_like(t) * texp, oversample=30).eval()
assert np.allclose(fluence_mavg, fluence_starry, fluence_starry_vec)
def test_ld():
map = starry.Map(udeg=2)
map[1:] = [0.4, 0.26]
orbit = exo.orbits.KeplerianOrbit(period=1.0, m_star=1.0, r_star=1.0)
t = np.linspace(-0.25, 0.25, 100)
# Compute the whole light curve with Theano
f1 = map.flux(t=t, orbit=orbit, ro=0.1, use_in_transit=False).eval()
# Compute just the transit with Theano
f2 = map.flux(t=t, orbit=orbit, ro=0.1, use_in_transit=True).eval()
# Compute the whole light curve without Theano
coords = orbit.get_relative_position(t)
xo = (coords[0] / orbit.r_star).eval()
yo = (coords[1] / orbit.r_star).eval()
b = np.sqrt(xo * xo + yo * yo)
zo = -(coords[2] / orbit.r_star).eval()
f3 = map.flux(b=b, zo=zo, ro=0.1)
dt = np.linspace(-0.5, 0.5, oversample)
stencil[1:-1:2] = 4
stencil[2:-1:2] = 2
else:
raise ValueError("Parameter `order` must be <= 2")
stencil /= np.sum(stencil)
if texp.ndim == 0:
dt = texp * dt
else:
dt = tt.shape_padright(texp) * dt
t = tt.shape_padright(t) + dt
t = tt.reshape(t, (-1,))
# Compute the relative positions of all bodies
orbit = exoplanet.orbits.KeplerianOrbit(
period=sec_porb,
t0=sec_t0,
incl=sec_iorb,
ecc=sec_ecc,
omega=sec_w,
Omega=sec_Omega,
m_planet=sec_m,
m_star=pri_m,
r_star=pri_r,
)
try:
x, y, z = orbit.get_relative_position(
t, light_delay=self.light_delay
)
except TypeError:
if self.light_delay:
import exoplanet
from packaging import version
import theano.tensor as tt
import numpy as np
from ..ops import autocompile
# NOTE: In version 0.1.7, DFM changed the coordinates
# so that the z-axis points TOWARD the observer!
if version.parse(exoplanet.__version__) > version.parse('0.1.7.dev0'):
z_sign = 1
else:
z_sign = -1
class KeplerianOrbit(exoplanet.orbits.KeplerianOrbit):
"""
A wrapper around `exoplanet.orbits.KeplerianOrbit` that
plays nice with `starry`. Refer to the docs of that class
for all accepted keywords. In addition to those, this class
accepts the following keyword arguments:
Args:
r_planet: The radius of the planet in ``R_sun``. Default is
the radius of the Earth.
rot_period: The period of rotation of the planet in days.
Default ``1.0``. Set to ``None`` to disable rotation.
theta0: The rotational phase in degrees at ``t=t0``.
Default ``0.0``
lazy:
"""
# The time of a reference transit for each planet
t0 = pm.Normal("t0", mu=t0_true, sd=1.0)
# The log period; also tracking the period itself
logP = pm.Normal("logP", mu=np.log(period_true), sd=0.1)
period = pm.Deterministic("period", pm.math.exp(logP))
# Normal distributions for the map coeffs
y = pm.Normal("y", mu=y_true, sd=1.0, shape=len(y_true))
# Normal distributions for r and b
r = pm.Normal("r", mu=0.06, sd=0.001)
b = pm.Normal("b", mu=0.4, sd=0.03)
# Set up a Keplerian orbit for the planets
orbit = xo.orbits.KeplerianOrbit(period=period, t0=t0, b=b)
# Compute the model light curve using starry
light_curve = op.get_light_curve(orbit=orbit, r=r, t=t, y=y) + mean
# Here we track the value of the model light curve for plotting
# purposes
pm.Deterministic("light_curve", light_curve)
# In this line, we simulate the dataset that we will fit
flux = xo.eval_in_model(light_curve)
flux += ferr * np.random.randn(len(flux))
# The likelihood function assuming known Gaussian uncertainty
pm.Normal("obs", mu=light_curve, sd=ferr, observed=flux)
# Fit for the maximum a posteriori parameters given the simuated
# The amlitudes should be sorted
pm.Potential("logK_order", tt.switch(logK[1:] > logK[:-1], -np.inf, 0.0))
# We also want to keep period physical but this probably won't be hit
pm.Potential("P_bound", tt.switch(P <= 0, -np.inf, 0.0))
# Eccentricity & argument of periasteron
ecc = pm.Uniform("ecc", lower=0, upper=0.99, shape=N_pl, testval=eccs)
omega = xo.distributions.Angle("omega", shape=N_pl, testval=omegas)
# Jitter & a quadratic RV trend
# logs = pm.Normal("logs", mu=np.log(np.median(yerr)), sd=5.0)
trend = pm.Normal("trend", mu=0, sd=10.0 ** -np.arange(3)[::-1], shape=3)
# Set up the orbit
orbit = xo.orbits.KeplerianOrbit(period=P, t0=t0, ecc=ecc, omega=omega)
# Set up the RV model and save it as a deterministic
# for plotting purposes later
vrad = orbit.get_radial_velocity(x, K=tt.exp(logK))
if N_pl == 1:
vrad = vrad[:, None]
# Define the background model
A = np.vander(x - 0.5 * (x.min() + x.max()), 3)
bkg = tt.dot(A, trend)
# Sum over planets and add the background to get the full model
rv_model = tt.sum(vrad, axis=-1) + bkg
# Simulate the data
y_true = xo.eval_in_model(rv_model)