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