def test_against_py():
with pm.Model():
K = xu.with_unit(pm.Normal('K', 0, 10.), u.km / u.s)
prior = JokerPrior.default(P_min=8 * u.day,
P_max=32768 * u.day,
s=0 * u.m / u.s,
sigma_v=100 * u.km / u.s,
pars={'K': K})
# t = np.random.uniform(0, 250, 16) + 56831.324
t = np.sort(np.random.uniform(0, 250, 3)) + 56831.324
rv = np.cos(t)
rv_err = np.random.uniform(0.1, 0.2, t.size)
data = RVData(t=t, rv=rv * u.km / u.s, rv_err=rv_err * u.km / u.s)
trend_M = get_constant_term_design_matrix(data)
samples = prior.sample(size=8192)
chunk, _ = samples.pack()
chunk = np.ascontiguousarray(chunk)
helper = CJokerHelper(data, prior, trend_M)
t0 = time.time()
cy_ll = helper.batch_marginal_ln_likelihood(chunk)
print("Cython:", time.time() - t0)
t0 = time.time()
py_ll = marginal_ln_likelihood(samples, prior, data)
print("Python:", time.time() - t0)
assert np.allclose(np.array(cy_ll), py_ll)
def test_mass_units():
P_earth = 365.256
Tper_earth = 2454115.5208333
inclination_earth = np.radians(45.0)
orbit1 = KeplerianOrbit(
period=P_earth,
t_periastron=Tper_earth,
incl=inclination_earth,
m_planet=units.with_unit(1.0, u.M_earth),
)
orbit2 = KeplerianOrbit(
period=P_earth,
t_periastron=Tper_earth,
incl=inclination_earth,
m_planet=1.0,
m_planet_units=u.M_earth,
)
t = np.linspace(Tper_earth, Tper_earth + 1000, 1000)
rv1 = orbit1.get_radial_velocity(t).eval()
rv_diff = np.max(rv1) - np.min(rv1)
assert rv_diff < 1.0, "with_unit"
rv2 = orbit2.get_radial_velocity(t).eval()
rv_diff = np.max(rv2) - np.min(rv2)
assert rv_diff < 1.0, "m_planet_units"
np.testing.assert_allclose(rv2, rv1)
def test_likelihood_helpers():
with pm.Model():
K = xu.with_unit(pm.Normal('K', 0, 1.), u.km / u.s)
prior = JokerPrior.default(P_min=8 * u.day,
P_max=32768 * u.day,
s=0 * u.m / u.s,
sigma_v=1 * u.km / u.s,
pars={'K': K})
# t = np.random.uniform(0, 250, 16) + 56831.324
t = np.sort(np.random.uniform(0, 250, 3)) + 56831.324
rv = np.cos(t)
rv_err = np.random.uniform(0.1, 0.2, t.size)
data = RVData(t=t, rv=rv * u.km / u.s, rv_err=rv_err * u.km / u.s)
trend_M = get_constant_term_design_matrix(data)
samples = prior.sample(size=16) # HACK: MAGIC NUMBER 16!
chunk, _ = samples.pack()
helper = CJokerHelper(data, prior, trend_M)
py_vals = get_aAbB(samples, prior, data)
for i in range(len(samples)):
ll = helper.test_likelihood_worker(np.ascontiguousarray(chunk[i]))
assert np.abs(ll) < 1e8
cy_vals = {}
for k in py_vals.keys():
cy_vals[k] = np.array(getattr(helper, k))
for k in py_vals:
# print(i, k)
# print('cy', cy_vals[k])
# print('py', py_vals[k][i])
# print(np.allclose(cy_vals[k], py_vals[k][i]))
# print()
assert np.allclose(cy_vals[k], py_vals[k][i])
def test_multi_data():
import exoplanet.units as xu
import pymc3 as pm
rnd = np.random.default_rng(42)
# Set up mulitple valid data objects:
_, raw1 = get_valid_input(rnd=rnd)
data1 = RVData(raw1['t_obj'], raw1['rv'], raw1['err'])
_, raw2 = get_valid_input(rnd=rnd, size=8)
data2 = RVData(raw2['t_obj'], raw2['rv'], raw2['err'])
_, raw3 = get_valid_input(rnd=rnd, size=4)
data3 = RVData(raw3['t_obj'], raw3['rv'], raw3['err'])
prior1 = JokerPrior.default(1*u.day, 1*u.year,
25*u.km/u.s,
sigma_v=100*u.km/u.s)
# Object should return input:
multi_data, ids, trend_M = validate_prepare_data(data1,
prior1.poly_trend,
prior1.n_offsets)
assert np.allclose(multi_data.rv.value, data1.rv.value)
assert np.all(ids == 0)
assert np.allclose(trend_M[:, 0], 1.)
# Three valid objects as a list:
with pm.Model():
dv1 = xu.with_unit(pm.Normal('dv0_1', 0, 1.),
u.km/u.s)
dv2 = xu.with_unit(pm.Normal('dv0_2', 4, 5.),
u.km/u.s)
prior2 = JokerPrior.default(1*u.day, 1*u.year,
25*u.km/u.s,
sigma_v=100*u.km/u.s,
v0_offsets=[dv1, dv2])
datas = [data1, data2, data3]
multi_data, ids, trend_M = validate_prepare_data(datas,
prior2.poly_trend,
prior2.n_offsets)
assert len(np.unique(ids)) == 3
assert len(multi_data) == sum([len(d) for d in datas])
assert 0 in ids and 1 in ids and 2 in ids
assert np.allclose(trend_M[:, 0], 1.)
# Three valid objects with names:
datas = {'apogee': data1, 'lamost': data2, 'weave': data3}
multi_data, ids, trend_M = validate_prepare_data(datas,
prior2.poly_trend,
prior2.n_offsets)
assert len(np.unique(ids)) == 3
assert len(multi_data) == sum([len(d) for d in datas.values()])
assert 'apogee' in ids and 'lamost' in ids and 'weave' in ids
assert np.allclose(trend_M[:, 0], 1.)
# Check it fails if n_offsets != number of data sources
with pytest.raises(ValueError):
validate_prepare_data(datas,
prior1.poly_trend,
prior1.n_offsets)
with pytest.raises(ValueError):
validate_prepare_data(data1,
prior2.poly_trend,
prior2.n_offsets)
# Check that this fails if one has a covariance matrix
data_cov = RVData(raw3['t_obj'], raw3['rv'], raw3['cov'])
with pytest.raises(NotImplementedError):
validate_prepare_data({'apogee': data1, 'test': data2,
'weave': data_cov},
prior2.poly_trend, prior2.n_offsets)
with pytest.raises(NotImplementedError):
validate_prepare_data([data1, data2, data_cov],
prior2.poly_trend,
prior2.n_offsets)
# Imports we need for defining the prior:
import astropy.units as u
import pymc3 as pm
import exoplanet.units as xu
import thejoker as tj
# The prior used to run The Joker:
with pm.Model() as model:
# Prior on the extra variance parameter:
s = xu.with_unit(pm.Lognormal('s', 0, 0.5), u.km / u.s)
# Set up the default Joker prior:
prior = tj.JokerPrior.default(P_min=2 * u.day,
P_max=1024 * u.day,
sigma_K0=30 * u.km / u.s,
sigma_v=100 * u.km / u.s,
s=s)
if case == 0:
# Default prior with standard parameters:
with pm.Model() as model:
default_nonlinear_prior(P_min=1*u.day, P_max=1*u.year)
default_linear_prior(sigma_K0=25*u.km/u.s,
P0=1*u.year,
sigma_v=10*u.km/u.s)
prior = JokerPrior(model=model)
return prior, default_expected_units
elif case == 1:
# Replace a nonlinear parameter
units = default_expected_units.copy()
with pm.Model() as model:
P = xu.with_unit(pm.Normal('P', 10, 0.5),
u.year)
units['P'] = u.year
default_nonlinear_prior(pars={'P': P})
default_linear_prior(sigma_K0=25*u.km/u.s,
P0=1*u.year,
sigma_v=10*u.km/u.s)
prior = JokerPrior(model=model)
return prior, units
elif case == 2:
# Replace a linear parameter
units = default_expected_units.copy()
with pm.Model() as model:
def test_get_consistent_inputs():
period0 = np.array([12.567, 45.132])
r_star0 = 1.235
m_star0 = 0.986
m_planet0 = with_unit(np.array([1.543, 2.354]), u.M_earth)
(
a1,
period1,
rho_star1,
r_star1,
m_star1,
m_planet1,
) = _get_consistent_inputs(None, period0, None, r_star0, m_star0,
m_planet0)
assert np.allclose(period0, period1.eval())
assert np.allclose(r_star0, r_star1.eval())
assert np.allclose(m_star0, m_star1.eval())
assert np.allclose(m_planet0.eval(),
m_planet1.eval() * u.M_sun.to(u.M_earth))
(
a2,
period2,
rho_star2,
r_star2,
m_star2,
m_planet2,
) = _get_consistent_inputs(a1, None, rho_star1, r_star0, None, m_planet1)
assert np.allclose(period1.eval(), period2.eval())
assert np.allclose(rho_star1.eval(), rho_star2.eval())
assert np.allclose(r_star1.eval(), r_star2.eval())
assert np.allclose(m_star1.eval(), m_star2.eval())
assert np.allclose(m_planet1.eval(), m_planet2.eval())
(
a3,
period3,
rho_star3,
r_star3,
m_star3,
m_planet3,
) = _get_consistent_inputs(a2, None, rho_star2, None, m_star2, m_planet2)
assert np.allclose(period1.eval(), period3.eval())
assert np.allclose(rho_star1.eval(), rho_star3.eval())
assert np.allclose(r_star1.eval(), r_star3.eval())
assert np.allclose(m_star1.eval(), m_star3.eval())
assert np.allclose(m_planet1.eval(), m_planet3.eval())
(
a4,
period4,
rho_star4,
r_star4,
m_star4,
m_planet4,
) = _get_consistent_inputs(a3, period3, None, r_star3, None, m_planet3)
assert np.allclose(period1.eval(), period4.eval())
assert np.allclose(rho_star1.eval(), rho_star4.eval())
assert np.allclose(r_star1.eval(), r_star4.eval())
assert np.allclose(m_star1.eval(), m_star4.eval())
assert np.allclose(m_planet1.eval(), m_planet4.eval())
(
a5,
period5,
rho_star5,
r_star5,
m_star5,
m_planet5,
) = _get_consistent_inputs(
a3,
None,
with_unit(rho_star3, u.g / u.cm**3),
r_star3,
None,
m_planet3,
)
assert np.allclose(period1.eval(), period5.eval())
assert np.allclose(rho_star1.eval(), rho_star5.eval())
assert np.allclose(r_star1.eval(), r_star5.eval())
assert np.allclose(m_star1.eval(), m_star5.eval())
assert np.allclose(m_planet1.eval(), m_planet5.eval())
with pytest.raises(ValueError):
_get_consistent_inputs(None, None, None, r_star3, m_star3, None)
with pytest.raises(ValueError):
_get_consistent_inputs(a3, period3, None, r_star3, m_star3, None)
with pytest.raises(ValueError):
_get_consistent_inputs(a3, None, rho_star3, r_star3, m_star3, None)
def default_linear_prior(sigma_K0=None,
P0=None,
sigma_v=None,
poly_trend=1,
model=None,
pars=None):
r"""
Retrieve pymc3 variables that specify the default prior on the linear
parameters of The Joker. See docstring of `JokerPrior.default()` for more
information.
The linear parameters an default prior forms are:
* ``K``, velocity semi-amplitude: Normal distribution, but with a variance
that scales with period and eccentricity.
* ``v0``, ``v1``, etc. polynomial velocity trend parameters: Independent
Normal distributions.
Parameters
----------
sigma_K0 : `~astropy.units.Quantity` [speed]
P0 : `~astropy.units.Quantity` [time]
sigma_v : iterable of `~astropy.units.Quantity`
model : `pymc3.Model`
This is either required, or this function must be called within a pymc3
model context.
"""
import pymc3 as pm
import exoplanet.units as xu
from .distributions import FixedCompanionMass
model = pm.modelcontext(model)
if pars is None:
pars = dict()
# dictionary of parameters to return
out_pars = dict()
# set up poly. trend names:
poly_trend, v_names = validate_poly_trend(poly_trend)
# get period/ecc from dict of nonlinear parameters
P = model.named_vars.get('P', None)
e = model.named_vars.get('e', None)
if P is None or e is None:
raise ValueError("Period P and eccentricity e must both be defined as "
"nonlinear parameters on the model.")
if v_names and 'v0' not in pars:
sigma_v = validate_sigma_v(sigma_v, poly_trend, v_names)
with model:
if 'K' not in pars:
if sigma_K0 is None or P0 is None:
raise ValueError("If using the default prior form on K, you "
"must pass in a variance scale (sigma_K0) "
"and a reference period (P0)")
# Default prior on semi-amplitude: scales with period and
# eccentricity such that it is flat with companion mass
v_unit = sigma_K0.unit
out_pars['K'] = xu.with_unit(
FixedCompanionMass('K', P=P, e=e, sigma_K0=sigma_K0, P0=P0),
v_unit)
else:
v_unit = getattr(pars['K'], xu.UNIT_ATTR_NAME, u.one)
for i, name in enumerate(v_names):
if name not in pars:
# Default priors are independent gaussians
# FIXME: make mean, mu_v, customizable
out_pars[name] = xu.with_unit(
pm.Normal(name, 0., sigma_v[name].value),
sigma_v[name].unit)
for k in pars.keys():
out_pars[k] = pars[k]
return out_pars
def default_nonlinear_prior(P_min=None,
P_max=None,
s=None,
model=None,
pars=None):
r"""
Retrieve pymc3 variables that specify the default prior on the nonlinear
parameters of The Joker. See docstring of `JokerPrior.default()` for more
information.
The nonlinear parameters an default prior forms are:
* ``P``, period: :math:`p(P) \propto 1/P`, over the domain
:math:`(P_{\rm min}, P_{\rm max})`
* ``e``, eccentricity: the short-period form from Kipping (2013)
* ``M0``, phase: uniform over the domain :math:`(0, 2\pi)`
* ``omega``, argument of pericenter: uniform over the domain
:math:`(0, 2\pi)`
* ``s``, additional extra variance added in quadrature to data
uncertainties: delta-function at 0
Parameters
----------
P_min : `~astropy.units.Quantity` [time]
P_max : `~astropy.units.Quantity` [time]
s : `~pm.model.TensorVariable`, ~astropy.units.Quantity` [speed]
model : `pymc3.Model`
This is either required, or this function must be called within a pymc3
model context.
"""
import theano.tensor as tt
import pymc3 as pm
from exoplanet.distributions import Angle
import exoplanet.units as xu
from .distributions import UniformLog, Kipping13Global
model = pm.modelcontext(model)
if pars is None:
pars = dict()
if s is None:
s = 0 * u.m / u.s
if isinstance(s, pm.model.TensorVariable):
pars['s'] = pars.get('s', s)
else:
if not hasattr(s, 'unit') or not s.unit.is_equivalent(u.km / u.s):
raise u.UnitsError(
"Invalid unit for s: must be equivalent to km/s")
# dictionary of parameters to return
out_pars = dict()
with model:
# Set up the default priors for parameters with defaults
# Note: we have to do it this way (as opposed to with .get(..., default)
# because this can only get executed if the param is not already
# defined, otherwise variables are defined twice in the model
if 'e' not in pars:
out_pars['e'] = xu.with_unit(Kipping13Global('e'), u.one)
# If either omega or M0 is specified by user, default to U(0,2π)
if 'omega' not in pars:
out_pars['omega'] = xu.with_unit(Angle('omega'), u.rad)
if 'M0' not in pars:
out_pars['M0'] = xu.with_unit(Angle('M0'), u.rad)
if 's' not in pars:
out_pars['s'] = xu.with_unit(
pm.Deterministic('s', tt.constant(s.value)), s.unit)
if 'P' not in pars:
if P_min is None or P_max is None:
raise ValueError(
"If you are using the default period prior, "
"you must pass in both P_min and P_max to set "
"the period prior domain.")
out_pars['P'] = xu.with_unit(
UniformLog('P', P_min.value, P_max.to_value(P_min.unit)),
P_min.unit)
for k in pars.keys():
out_pars[k] = pars[k]
return out_pars