diff --git a/thejoker/__init__.py b/thejoker/__init__.py index 67880d0..6e1b589 100644 --- a/thejoker/__init__.py +++ b/thejoker/__init__.py @@ -31,6 +31,10 @@ "phase_coverage_per_period", ] + # SB2: + from .thejoker_sb2 import * + from .prior_sb2 import JokerSB2Prior + __bibtex__ = __citation__ = """@ARTICLE{thejoker, author = {{Price-Whelan}, Adrian M. and {Hogg}, David W. and @@ -55,3 +59,12 @@ adsnote = {Provided by the SAO/NASA Astrophysics Data System} } """ + +__all__ = [ + 'TheJoker', + 'RVData', + 'JokerSamples', + 'JokerPrior', + 'plot_rv_curves', + 'TheJokerSB2' +] diff --git a/thejoker/data.py b/thejoker/data.py index 5dc82ca..3d0335b 100644 --- a/thejoker/data.py +++ b/thejoker/data.py @@ -33,11 +33,13 @@ class RVData: (days). Set to ``False`` to disable subtracting the reference time. clean : bool (optional) Filter out any NaN or Inf data points. + sort : bool (optional) + Whether or not to sort on time. """ @u.quantity_input(rv=u.km / u.s, rv_err=[u.km / u.s, (u.km / u.s) ** 2]) - def __init__(self, t, rv, rv_err, t_ref=None, clean=True): + def __init__(self, t, rv, rv_err, t_ref=None, clean=True, sort=True): # For speed, time is saved internally as BMJD: if isinstance(t, Time): _t_bmjd = t.tcb.mjd @@ -94,15 +96,16 @@ def __init__(self, t, rv, rv_err, t_ref=None, clean=True): else: self.rv_err = self.rv_err[idx] - # sort on times - idx = self._t_bmjd.argsort() - self._t_bmjd = self._t_bmjd[idx] - self.rv = self.rv[idx] - if self._has_cov: - self.rv_err = self.rv_err[idx] - self.rv_err = self.rv_err[:, idx] - else: - self.rv_err = self.rv_err[idx] + if sort: + # sort on times + idx = self._t_bmjd.argsort() + self._t_bmjd = self._t_bmjd[idx] + self.rv = self.rv[idx] + if self._has_cov: + self.rv_err = self.rv_err[idx] + self.rv_err = self.rv_err[:, idx] + else: + self.rv_err = self.rv_err[idx] if t_ref is False: self.t_ref = None diff --git a/thejoker/likelihood_helpers.py b/thejoker/likelihood_helpers.py index 257a918..e6fcca6 100644 --- a/thejoker/likelihood_helpers.py +++ b/thejoker/likelihood_helpers.py @@ -66,8 +66,13 @@ def marginal_ln_likelihood_inmem(joker_helper, prior_samples_batch): return np.array(ll) -def make_full_samples_inmem(joker_helper, prior_samples_batch, rng, n_linear_samples=1): - from .samples import JokerSamples +def make_full_samples_inmem( + joker_helper, prior_samples_batch, rng, n_linear_samples=1, SamplesCls=None +): + if SamplesCls is None: + from .samples import JokerSamples + + SamplesCls = JokerSamples if prior_samples_batch.dtype != np.float64: prior_samples_batch = prior_samples_batch.astype(np.float64) @@ -77,7 +82,7 @@ def make_full_samples_inmem(joker_helper, prior_samples_batch, rng, n_linear_sam ) # unpack the raw samples - samples = JokerSamples.unpack( + samples = SamplesCls.unpack( raw_samples, joker_helper.internal_units, t_ref=joker_helper.data.t_ref, @@ -96,6 +101,7 @@ def rejection_sample_inmem( max_posterior_samples=None, n_linear_samples=1, return_all_logprobs=False, + SamplesCls=None, ): if max_posterior_samples is None: max_posterior_samples = len(prior_samples_batch) @@ -114,6 +120,7 @@ def rejection_sample_inmem( prior_samples_batch[good_samples_idx], rng, n_linear_samples=n_linear_samples, + SamplesCls=SamplesCls, ) if ln_prior is not None and ln_prior is not False: @@ -136,6 +143,7 @@ def iterative_rejection_inmem( init_batch_size=None, growth_factor=128, n_linear_samples=1, + SamplesCls=None, ): n_total_samples = len(prior_samples_batch) @@ -219,6 +227,7 @@ def iterative_rejection_inmem( prior_samples_batch[full_samples_idx], rng, n_linear_samples=n_linear_samples, + SamplesCls=SamplesCls, ) # FIXME: copy-pasted from function above diff --git a/thejoker/multiproc_helpers.py b/thejoker/multiproc_helpers.py index 1d242f6..73d0448 100644 --- a/thejoker/multiproc_helpers.py +++ b/thejoker/multiproc_helpers.py @@ -155,6 +155,7 @@ def make_full_samples( samples_idx, n_linear_samples=1, n_batches=None, + SamplesCls=JokerSamples, ): task_args = (prior_samples_file, joker_helper, n_linear_samples) results = run_worker( @@ -164,14 +165,14 @@ def make_full_samples( task_args=task_args, n_batches=n_batches, samples_idx=samples_idx, - rng=rng, + random_state=rng, ) # Concatenate all of the raw samples arrays raw_samples = np.concatenate(results) # unpack the raw samples - samples = JokerSamples.unpack( + samples = SamplesCls.unpack( raw_samples, joker_helper.internal_units, t_ref=joker_helper.data.t_ref, @@ -195,6 +196,7 @@ def rejection_sample_helper( n_batches=None, randomize_prior_order=False, return_all_logprobs=False, + SamplesCls=None, ): # Total number of samples in the cache: with tb.open_file(prior_samples_file, mode="r") as f: @@ -271,6 +273,7 @@ def rejection_sample_helper( full_samples_idx, n_linear_samples=n_linear_samples, n_batches=n_batches, + SamplesCls=SamplesCls, ) if return_logprobs: @@ -300,6 +303,7 @@ def iterative_rejection_helper( return_logprobs=False, n_batches=None, randomize_prior_order=False, + SamplesCls=None, ): # Total number of samples in the cache: with tb.open_file(prior_samples_file, mode="r") as f: @@ -412,6 +416,7 @@ def iterative_rejection_helper( full_samples_idx, n_linear_samples=n_linear_samples, n_batches=n_batches, + SamplesCls=SamplesCls, ) # FIXME: copy-pasted from function above diff --git a/thejoker/prior.py b/thejoker/prior.py index 3d2eaa6..4e29768 100644 --- a/thejoker/prior.py +++ b/thejoker/prior.py @@ -38,6 +38,8 @@ def _validate_model(model): class JokerPrior: + _sb2 = False + def __init__(self, pars=None, poly_trend=1, v0_offsets=None, model=None): """ This class controls the prior probability distributions for the @@ -121,7 +123,9 @@ def __init__(self, pars=None, poly_trend=1, v0_offsets=None, model=None): # are only used to validate that the units for each parameter are # equivalent to these self._nonlinear_equiv_units = get_nonlinear_equiv_units() - self._linear_equiv_units = get_linear_equiv_units(self.poly_trend) + self._linear_equiv_units = get_linear_equiv_units( + self.poly_trend, sb2=self._sb2 + ) self._v0_offsets_equiv_units = get_v0_offsets_equiv_units(self.n_offsets) self._all_par_unit_equiv = { **self._nonlinear_equiv_units, @@ -291,10 +295,7 @@ def __repr__(self): def __str__(self): return ", ".join(self.par_names) - @deprecated_renamed_argument( - "random_state", "rng", since="v1.3", warning_type=DeprecationWarning - ) - def sample( + def _get_raw_samples( self, size=1, generate_linear=False, @@ -303,29 +304,6 @@ def sample( dtype=None, **kwargs, ): - """ - Generate random samples from the prior. - - Parameters - ---------- - size : int (optional) - The number of samples to generate. - generate_linear : bool (optional) - Also generate samples in the linear parameters. - return_logprobs : bool (optional) - Generate the log-prior probability at the position of each sample. - **kwargs - Additional keyword arguments are passed to the - `~thejoker.JokerSamples` initializer. - - Returns - ------- - samples : `thejoker.Jokersamples` - The random samples. - - """ - from .samples import JokerSamples - if dtype is None: dtype = np.float64 @@ -339,11 +317,6 @@ def sample( ) } - if generate_linear: - par_names = self.par_names - else: - par_names = list(self._nonlinear_equiv_units.keys()) - # MAJOR HACK RELATED TO UPSTREAM ISSUES WITH pymc3: # init_shapes = {} # for name, par in sub_pars.items(): @@ -374,12 +347,68 @@ def sample( logp.append(_logp) log_prior = np.sum(logp, axis=0) + else: + log_prior = None # CONTINUED MAJOR HACK RELATED TO UPSTREAM ISSUES WITH pymc3: # for name, par in sub_pars.items(): # if hasattr(par, "distribution"): # par.distribution.shape = init_shapes[name] + return raw_samples, sub_pars, log_prior + + @deprecated_renamed_argument( + "random_state", "rng", since="v1.3", warning_type=DeprecationWarning + ) + def sample( + self, + size=1, + generate_linear=False, + return_logprobs=False, + rng=None, + dtype=None, + **kwargs, + ): + """ + Generate random samples from the prior. + + .. note:: + + Right now, generating samples with the prior values is slow (i.e. + with ``return_logprobs=True``) because of pymc3 issues (see + discussion here: + https://discourse.pymc.io/t/draw-values-speed-scaling-with-transformed-variables/4076). + This will hopefully be resolved in the future... + + Parameters + ---------- + size : int (optional) + The number of samples to generate. + generate_linear : bool (optional) + Also generate samples in the linear parameters. + return_logprobs : bool (optional) + Generate the log-prior probability at the position of each sample. + **kwargs + Additional keyword arguments are passed to the + `~thejoker.JokerSamples` initializer. + + Returns + ------- + samples : `thejoker.Jokersamples` + The random samples. + + """ + from thejoker.samples import JokerSamples + + raw_samples, sub_pars, log_prior = self._get_raw_samples( + size, generate_linear, return_logprobs, rng, dtype, **kwargs + ) + + if generate_linear: + par_names = self.par_names + else: + par_names = list(self._nonlinear_equiv_units.keys()) + # Apply units if they are specified: prior_samples = JokerSamples( poly_trend=self.poly_trend, n_offsets=self.n_offsets, **kwargs @@ -448,9 +477,8 @@ def default_nonlinear_prior(P_min=None, P_max=None, s=None, model=None, pars=Non if isinstance(s, pt.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") + elif 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 = {} diff --git a/thejoker/prior_helpers.py b/thejoker/prior_helpers.py index ff9451e..a02e13b 100644 --- a/thejoker/prior_helpers.py +++ b/thejoker/prior_helpers.py @@ -31,12 +31,19 @@ def validate_poly_trend(poly_trend): return poly_trend, vtrend_names -def get_linear_equiv_units(poly_trend): +def get_linear_equiv_units(poly_trend, sb2=False): poly_trend, v_names = validate_poly_trend(poly_trend) - return { - 'K': u.m/u.s, - **{name: u.m/u.s/u.day**i for i, name in enumerate(v_names)} - } + if sb2: + return { + 'K1': u.m/u.s, + 'K2': u.m/u.s, + **{name: u.m/u.s/u.day**i for i, name in enumerate(v_names)} + } + else: + return { + 'K': u.m/u.s, + **{name: u.m/u.s/u.day**i for i, name in enumerate(v_names)} + } def validate_sigma_v(sigma_v, poly_trend, v_names): diff --git a/thejoker/prior_sb2.py b/thejoker/prior_sb2.py new file mode 100644 index 0000000..c627047 --- /dev/null +++ b/thejoker/prior_sb2.py @@ -0,0 +1,321 @@ +# Third-party +import astropy.units as u +import numpy as np + +# Project +from .prior_helpers import validate_poly_trend, validate_sigma_v +from .prior import JokerPrior, _validate_model, default_nonlinear_prior +from .samples import JokerSamples + +__all__ = ['JokerSB2Prior'] + + +class JokerSB2Prior(JokerPrior): + _sb2 = True + + def __init__(self, pars=None, poly_trend=1, model=None): + """ + This class controls the prior probability distributions for the + parameters used in The Joker. + + TODO: + + This initializer is meant to be flexible, allowing you to specify the + prior distributions on the linear and nonlinear parameters used in The + Joker. However, for many use cases, you may want to just use the + default prior: To initialize this object using the default prior, see + the alternate initializer `JokerPrior.default()`. + + Parameters + ---------- + pars : dict, list (optional) + Either a list of pymc3 variables, or a dictionary of variables with + keys set to the variable names. If any of these variables are + defined as deterministic transforms from other variables, see the + next parameter below. + poly_trend : int (optional) + Specifies the number of coefficients in an additional polynomial + velocity trend, meant to capture long-term trends in the data. The + default here is ``polytrend=1``, meaning one term: the (constant) + systemtic velocity. For example, ``poly_trend=3`` will sample over + parameters of a long-term quadratic velocity trend. + model : `pymc3.Model` + This is either required, or this function must be called within a + pymc3 model context. + + """ + super().__init__(pars=pars, poly_trend=poly_trend, model=model) + + @classmethod + def default(cls, P_min=None, P_max=None, + sigma_K0_1=None, P0_1=1*u.year, + sigma_K0_2=None, P0_2=1*u.year, + sigma_v=None, s=None, poly_trend=1, + model=None, pars=None): + r""" + An alternative initializer to set up the default prior for The Joker. + + The default prior is: + + .. math:: + + p(P) \propto \frac{1}{P} \quad ; \quad P \in (P_{\rm min}, P_{\rm max})\\ + p(e) = B(a_e, b_e)\\ + p(\omega) = \mathcal{U}(0, 2\pi)\\ + p(M_0) = \mathcal{U}(0, 2\pi)\\ + p(s) = 0\\ + p(K) = \mathcal{N}(K \,|\, \mu_K, \sigma_K)\\ + \sigma_K = \sigma_{K, 0} \, \left(\frac{P}{P_0}\right)^{-1/3} \, \left(1 - e^2\right)^{-1/2} + + and the priors on any polynomial trend parameters are assumed to be + independent, univariate Normals. + + This prior has sensible choices for typical binary star or exoplanet + use cases, but if you need more control over the prior distributions + you might need to use the standard initializer (i.e. + ``JokerPrior(...)```) and specify all parameter distributions manually. + See `the documentation `_ for tutorials + that demonstrate this functionality. + + Parameters + ---------- + P_min : `~astropy.units.Quantity` [time] + Minimum period for the default period prior. + P_max : `~astropy.units.Quantity` [time] + Maximum period for the default period prior. + sigma_K0 : `~astropy.units.Quantity` [speed] + The scale factor, :math:`\sigma_{K, 0}` in the equation above that + sets the scale of the semi-amplitude prior at the reference period, + ``P0``. + P0 : `~astropy.units.Quantity` [time] + The reference period, :math:`P_0`, used in the prior on velocity + semi-amplitude (see equation above). + sigma_v : `~astropy.units.Quantity` (or iterable of) + The standard deviations of the velocity trend priors. + s : `~astropy.units.Quantity` [speed] + The jitter value, assuming it is constant. + poly_trend : int (optional) + Specifies the number of coefficients in an additional polynomial + velocity trend, meant to capture long-term trends in the data. The + default here is ``polytrend=1``, meaning one term: the (constant) + systemtic velocity. For example, ``poly_trend=3`` will sample over + parameters of a long-term quadratic velocity trend. + v0_offsets : list (optional) + A list of additional Gaussian parameters that set systematic offsets + of subsets of the data. TODO: link to tutorial here + model : `pymc3.Model` (optional) + If not specified, this will create a model instance and store it on + the prior object. + pars : dict, list (optional) + Either a list of pymc3 variables, or a dictionary of variables with + keys set to the variable names. If any of these variables are + defined as deterministic transforms from other variables, see the + next parameter below. + """ + + model = _validate_model(model) + + nl_pars = default_nonlinear_prior(P_min, P_max, s=s, + model=model, pars=pars) + l_pars = default_linear_prior_sb2(sigma_K0_1=sigma_K0_1, P0_1=P0_1, + sigma_K0_2=sigma_K0_2, P0_2=P0_2, + sigma_v=sigma_v, + poly_trend=poly_trend, model=model, + pars=pars) + + pars = {**nl_pars, **l_pars} + obj = cls(pars=pars, model=model, poly_trend=poly_trend) + + return obj + + def __repr__(self): + return f'' + + def sample(self, size=1, generate_linear=False, return_logprobs=False, + random_state=None, dtype=None, **kwargs): + """ + Generate random samples from the prior. + + .. note:: + + Right now, generating samples with the prior values is slow (i.e. + with ``return_logprobs=True``) because of pymc3 issues (see + discussion here: + https://discourse.pymc.io/t/draw-values-speed-scaling-with-transformed-variables/4076). + This will hopefully be resolved in the future... + + Parameters + ---------- + size : int (optional) + The number of samples to generate. + generate_linear : bool (optional) + Also generate samples in the linear parameters. + return_logprobs : bool (optional) + Generate the log-prior probability at the position of each sample. + **kwargs + Additional keyword arguments are passed to the + `~thejoker.JokerSamples` initializer. + + Returns + ------- + samples : `thejoker.Jokersamples` + The random samples. + + """ + import exoplanet.units as xu + + raw_samples, sub_pars, log_prior = self._get_raw_samples( + size, generate_linear, return_logprobs, random_state, dtype, + **kwargs) + + if generate_linear: + raw_samples['K2'] = (np.sign(raw_samples['K1']) * + np.sign(raw_samples['K2']) * + raw_samples['K2']) + + if generate_linear: + par_names = self.par_names + else: + par_names = list(self._nonlinear_equiv_units.keys()) + + # Split into primary / secondary: + prior_samples = [] + for k in range(2): + # Apply units if they are specified: + s = JokerSamples(poly_trend=self.poly_trend, + n_offsets=self.n_offsets, + **kwargs) + for name in par_names: + p = sub_pars[name] + unit = getattr(p, xu.UNIT_ATTR_NAME, u.one) + + if (p.name == 'K1' and k == 0) or (p.name == 'K2' and k == 1): + name = 'K' + + elif p.name not in s._valid_units.keys(): + continue + + s[name] = np.atleast_1d(raw_samples[p.name]) * unit + + if return_logprobs: + s['ln_prior'] = log_prior + + prior_samples.append(s) + + # Secondary argument of pericenter: + prior_samples[1]['omega'] = prior_samples[1]['omega'] - 180*u.deg + + # TODO: right now, elsewhere, we assume the log_prior is a single value + # for each sample (i.e. the total prior value). In principle, we could + # store all of the individual log-prior values (for each parameter), + # like here: + # log_prior = {k: np.atleast_1d(v) + # for k, v in log_prior.items()} + # log_prior = Table(log_prior)[par_names] + + return prior_samples + + +@u.quantity_input(sigma_K0=u.km/u.s, P0=u.day) +def default_linear_prior_sb2(sigma_K0_1=None, P0_1=None, + sigma_K0_2=None, P0_2=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: + + * ``K1``, velocity semi-amplitude for primary: Normal distribution, but with + a variance that scales with period and eccentricity. + * ``K2``, velocity semi-amplitude for secondary: 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_1 : `~astropy.units.Quantity` [speed] + P0_1 : `~astropy.units.Quantity` [time] + sigma_K0_2 : `~astropy.units.Quantity` [speed] + P0_2 : `~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 'K1' not in pars: + if sigma_K0_1 is None or P0_1 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) for both the " + "primary and secondary") + + # Default prior on semi-amplitude: scales with period and + # eccentricity such that it is flat with companion mass + v_unit = sigma_K0_1.unit + out_pars['K1'] = xu.with_unit( + FixedCompanionMass('K1', P=P, e=e, + sigma_K0=sigma_K0_1, P0=P0_1), + v_unit) + else: + v_unit = getattr(pars['K1'], xu.UNIT_ATTR_NAME, u.one) + + if 'K2' not in pars: + if sigma_K0_2 is None or P0_2 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) for both the " + "primary and secondary") + + # Default prior on semi-amplitude: scales with period and + # eccentricity such that it is flat with companion mass + v_unit = sigma_K0_1.unit + out_pars['K2'] = xu.with_unit( + FixedCompanionMass('K2', P=P, e=e, + sigma_K0=sigma_K0_2, P0=P0_2), + v_unit) + else: + v_unit = getattr(pars['K2'], 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 diff --git a/thejoker/samples_helpers.py b/thejoker/samples_helpers.py index d408aa1..d49e4e7 100644 --- a/thejoker/samples_helpers.py +++ b/thejoker/samples_helpers.py @@ -25,9 +25,8 @@ def _custom_tbl_dtype_compare(dtype1, dtype2): return False if d1.get(k, "") != d2.get(k, ""): return False - else: - if d1.get(k, "1") != d2.get(k, "2"): - return False + elif d1.get(k, "1") != d2.get(k, "2"): + return False return True diff --git a/thejoker/src/fast_likelihood.pyx b/thejoker/src/fast_likelihood.pyx index 2c1e9f4..9f5ed5a 100644 --- a/thejoker/src/fast_likelihood.pyx +++ b/thejoker/src/fast_likelihood.pyx @@ -580,3 +580,280 @@ cdef class CJokerHelper: ll = self.likelihood_worker(1) # the 1 is "True" return ll + + +cdef class CJokerSB2Helper(CJokerHelper): + cdef: + double sigma_K0_1 # TODO: total HACK + double P0_1 # TODO: total HACK + double max_K_1 # TODO: total HACK + double sigma_K0_2 # TODO: total HACK + double P0_2 # TODO: total HACK + double max_K_2 # TODO: total HACK + + def __reduce__(self): + return (CJokerSB2Helper, (self.data, self.prior, + np.array(self.trend_M))) + + def __init__(self, data, prior, double[:, ::1] trend_M): + cdef int i, j, n + + self.prior = prior + self.data = data + + # Internal units needed for calculations below. + # Note: order here matters! This is the order in which prior samples + # will be unpacked externally! + self.internal_units = {} + self.internal_units['P'] = _nonlinear_internal_units['P'] + self.internal_units['e'] = _nonlinear_internal_units['e'] + self.internal_units['omega'] = _nonlinear_internal_units['omega'] + self.internal_units['M0'] = _nonlinear_internal_units['M0'] + self.internal_units['s'] = self.data.rv.unit + self.internal_units['K1'] = self.data.rv.unit + self.internal_units['K2'] = self.data.rv.unit + self.internal_units['v0'] = self.data.rv.unit + + # v0 offsets must be between v0 and v1, if v1 is present + for offset in prior.v0_offsets: + self.internal_units[offset.name] = self.data.rv.unit + + for i, name in enumerate(prior._v_trend_names): + self.internal_units[name] = self.data.rv.unit / u.day ** i + + # The assumed order of the nonlinear parameters used below to read from + # packed samples array + self.packed_order = _nonlinear_packed_order + + import exoplanet.units as xu + + # Counting: + self.n_times = len(data) # number of data pints + self.n_poly = prior.poly_trend # polynomial trend terms + self.n_offsets = 1 # DISABLED FOR SB2 - v0 offsets + self.n_linear = 2 + self.n_poly # K1, K2, trend + self.n_pars = len(prior.par_names) + + # Data: + self.t0 = data._t_ref_bmjd + self.t = np.ascontiguousarray(data._t_bmjd, dtype='f8') + self.rv = np.ascontiguousarray(data.rv.value, dtype='f8') + self.ivar = np.ascontiguousarray( + data.ivar.to_value(1 / self.data.rv.unit**2), dtype='f8') + self.trend_M = trend_M + + # ivar with jitter included + self.s_ivar = np.zeros(self.n_times, dtype='f8') + + # Transpose of design matrix: Fill the columns for the linear part of M + # - trend shape: K1, K2, v0, poly_trend-1 + if (trend_M.shape[0] != self.n_times + or trend_M.shape[1] != self.n_linear): + raise ValueError("Invalid design matrix shape: {}, expected: {}" + .format(trend_M.shape, + (self.n_times, self.n_linear))) + + self.M_T = np.zeros((self.n_linear, self.n_times)) + for n in range(self.n_times): + for i in range(self.n_linear): + self.M_T[i, n] = trend_M[n, i] + + # Needed for temporary storage in likelihood_worker: + self.Btmp = np.zeros((self.n_times, self.n_times), dtype=np.float64) + self.Atmp = np.zeros((self.n_linear, self.n_linear), dtype=np.float64) + self.npar_ipiv = np.zeros(self.n_linear, dtype=np.int32) + self.ntime_ipiv = np.zeros(self.n_times, dtype=np.int32) + self.npar_work = np.zeros(self.n_linear, dtype=np.float64) + self.ntime_work = np.zeros(self.n_times, dtype=np.float64) + + # Needed for internal work / output from likelihood_worker: + self.B = np.zeros((self.n_times, self.n_times), dtype=np.float64) + self.Binv = np.zeros((self.n_times, self.n_times), dtype=np.float64) + self.b = np.zeros(self.n_times, dtype=np.float64) + self.Ainv = np.zeros((self.n_linear, self.n_linear), dtype=np.float64) + self.A = np.zeros((self.n_linear, self.n_linear), dtype=np.float64) + self.a = np.zeros(self.n_linear, dtype=np.float64) + + # TODO: Lambda should be a matrix, but we currently only support + # diagonal variance on Lambda + self.mu = np.zeros(self.n_linear) + self.Lambda = np.zeros(self.n_linear) + + # --------------------------------------------------------------------- + # TODO: This is a bit of a hack: assumes FixedCompanionMass prior + self.fixed_K_prior = 0 + + for i, name in enumerate(prior._linear_equiv_units.keys()): + dist = prior.model[name].distribution + _unit = getattr(prior.model[name], xu.UNIT_ATTR_NAME) + to_unit = self.internal_units[name] + mu = (dist.mean.eval() * _unit).to_value(to_unit) + + if name == 'K1' and self.fixed_K_prior == 0: + # TODO: here's the major hack + self.sigma_K0_1 = dist._sigma_K0.to_value(to_unit) + self.P0_1 = dist._P0.to_value(getattr(prior.pars['P'], + xu.UNIT_ATTR_NAME)) + self.max_K_1 = dist._max_K.to_value(to_unit) + self.mu[i] = mu + + elif name == 'K2' and self.fixed_K_prior == 0: + # TODO: here's the major hack + self.sigma_K0_2 = dist._sigma_K0.to_value(to_unit) + self.P0_2 = dist._P0.to_value(getattr(prior.pars['P'], + xu.UNIT_ATTR_NAME)) + self.max_K_2 = dist._max_K.to_value(to_unit) + self.mu[i] = mu + + elif name == 'v0': + self.Lambda[i] = (dist.sd.eval() * _unit).to_value(to_unit) ** 2 + self.mu[i] = mu + + # --------------------------------------------------------------------- + + # TODO: instead of copy-pasting these, I could modify c_rv_from_elements to + # multiple the input "rv" by the solved rv (instead of just replacing the + # value) + cpdef batch_marginal_ln_likelihood(self, double[:, ::1] chunk): + """Compute the marginal log-likelihood for a batch of prior samples. + + Parameters + ---------- + chunk : numpy.ndarray + A chunk of nonlinear parameter prior samples. + Expected order: P, e, omega, M0, s (jitter). + """ + cdef: + int n, i, j + int n_samples = chunk.shape[0] + double P, e, om, M0 + + # the log-likelihood values + double[::1] ll = np.full(n_samples, np.nan) + + for n in range(n_samples): + # NOTE: need to make sure the chunk is always in this order! If this + # is changed, change "packed_order" above + P = chunk[n, 0] + e = chunk[n, 1] + om = chunk[n, 2] + M0 = chunk[n, 3] + + c_rv_from_elements(&self.t[0], &self.M_T[0, 0], self.n_times, + P, 1., e, om, M0, self.t0, + anomaly_tol, anomaly_maxiter) + + # Copy the unscaled rv sequence from K1 to K2: + for j in range(self.n_times): + self.M_T[1, j] = self.M_T[0, j] + + for j in range(self.n_times): + for i in range(2): + self.M_T[i, j] = self.M_T[i, j] * self.trend_M[j, i] + + # Note: jitter must be in same units as the data RV's / ivar + get_ivar(self.ivar, chunk[n, 4], self.s_ivar) + + # TODO: this is a continuation of the massive hack introduced above. + if self.fixed_K_prior == 0: + self.Lambda[0] = (self.sigma_K0_1**2 / (1 - e**2) + * (P / self.P0_1)**(-2/3.)) + self.Lambda[0] = min(self.max_K_1**2, self.Lambda[0]) + + self.Lambda[1] = (self.sigma_K0_2**2 / (1 - e**2) + * (P / self.P0_2)**(-2/3.)) + self.Lambda[1] = min(self.max_K_2**2, self.Lambda[1]) + + # compute things needed for the ln(likelihood) + ll[n] = self.likelihood_worker(0) + + return ll + + cpdef batch_get_posterior_samples(self, double[:, ::1] chunk, + int n_linear_samples_per, + object random_state): + """TODO: + + Parameters + ---------- + chunk : numpy.ndarray + A chunk of nonlinear parameter prior samples. + Expected order: P, e, omega, M0, s (jitter). + """ + + cdef: + int n, i, j, k + int n_samples = chunk.shape[0] + double P, e, om, M0 + + # Transpose of design matrix + double[:, ::1] M_T = np.zeros((self.n_linear, self.n_times)) + + # the log-likelihood values + double[:, ::1] ll = np.full((n_samples, n_linear_samples_per), + np.nan) + double _ll + + # The samples + double[:, :, ::1] samples = np.zeros((n_samples, + n_linear_samples_per, + self.n_pars)) + double[:, ::1] linear_pars = np.zeros((n_linear_samples_per, + self.n_linear)) + + for n in range(n_samples): + # NOTE: need to make sure the chunk is always in this order! If this + # is changed, change "packed_order" above + P = chunk[n, 0] + e = chunk[n, 1] + om = chunk[n, 2] + M0 = chunk[n, 3] + + c_rv_from_elements(&self.t[0], &self.M_T[0, 0], self.n_times, + P, 1., e, om, M0, self.t0, + anomaly_tol, anomaly_maxiter) + + # Copy the unscaled rv sequence from K1 to K2: + for j in range(self.n_times): + self.M_T[1, j] = self.M_T[0, j] + + for j in range(self.n_times): + for i in range(2): + self.M_T[i, j] = self.M_T[i, j] * self.trend_M[j, i] + + # Note: jitter must be in same units as the data RV's / ivar + get_ivar(self.ivar, chunk[n, 4], self.s_ivar) + + # TODO: this is a continuation of the massive hack introduced above. + if self.fixed_K_prior == 0: + self.Lambda[0] = (self.sigma_K0_1**2 / (1 - e**2) + * (P / self.P0_1)**(-2/3.)) + self.Lambda[0] = min(self.max_K_1**2, self.Lambda[0]) + + self.Lambda[1] = (self.sigma_K0_2**2 / (1 - e**2) + * (P / self.P0_2)**(-2/3.)) + self.Lambda[1] = min(self.max_K_2**2, self.Lambda[1]) + + # compute likelihood, but also generate a, Ainv + _ll = self.likelihood_worker(1) # the 1 is "True" + + # TODO: FIXME: this calls back to numpy at the Python layer + # - use https://github.com/bashtage/randomgen instead? + # a and Ainv are populated by the likelihood_worker() + linear_pars = random_state.multivariate_normal( + self.a, np.linalg.inv(self.Ainv), size=n_linear_samples_per) + + for j in range(n_linear_samples_per): + ll[n, j] = _ll + + samples[n, j, 0] = P + samples[n, j, 1] = e + samples[n, j, 2] = om + samples[n, j, 3] = M0 + samples[n, j, 4] = chunk[n, 4] # s, jitter + + for k in range(self.n_linear): + samples[n, j, 5 + k] = linear_pars[j, k] + + return (np.array(samples).reshape(n_samples * n_linear_samples_per, -1), + np.array(ll).reshape(n_samples * n_linear_samples_per)) diff --git a/thejoker/thejoker.py b/thejoker/thejoker.py index b0ec243..23fccb0 100644 --- a/thejoker/thejoker.py +++ b/thejoker/thejoker.py @@ -41,6 +41,8 @@ class TheJoker: Default: ``~/.thejoker`` """ + _samples_cls = JokerSamples + @deprecated_renamed_argument( "random_state", "rng", since="v1.3", warning_type=DeprecationWarning ) @@ -237,6 +239,7 @@ def rejection_sample( max_posterior_samples=max_posterior_samples, n_linear_samples=n_linear_samples, return_all_logprobs=return_all_logprobs, + SamplesCls=self._samples_cls, ) else: @@ -252,6 +255,7 @@ def rejection_sample( n_batches=n_batches, randomize_prior_order=randomize_prior_order, return_all_logprobs=return_all_logprobs, + SamplesCls=self._samples_cls, ) return samples @@ -349,6 +353,7 @@ def iterative_rejection_sample( init_batch_size=init_batch_size, growth_factor=growth_factor, n_linear_samples=n_linear_samples, + SamplesCls=self._samples_cls, ) else: @@ -365,6 +370,7 @@ def iterative_rejection_sample( return_logprobs=return_logprobs, n_batches=n_batches, randomize_prior_order=randomize_prior_order, + SamplesCls=self._samples_cls, ) return samples diff --git a/thejoker/thejoker_sb2.py b/thejoker/thejoker_sb2.py new file mode 100644 index 0000000..73f23f3 --- /dev/null +++ b/thejoker/thejoker_sb2.py @@ -0,0 +1,145 @@ +from astropy.time import Time +import astropy.units as u +import numpy as np + +from .src.fast_likelihood import CJokerSB2Helper +from .samples import JokerSamples +from .likelihood_helpers import get_trend_design_matrix +from .thejoker import TheJoker + +__all__ = ['TheJokerSB2', 'JokerSB2Samples'] + + +def validate_prepare_data_sb2(data, poly_trend, t_ref=None): + """Internal function. + + Used to take an input ``RVData`` instance, or a list/dict of ``RVData`` + instances, and produce concatenated time, RV, and error arrays, along + with a consistent t_ref. + """ + from .data import RVData + + # If we've gotten here, data is dict-like: + rv_unit = None + t = [] + rv = [] + err = [] + ids = [] + for k in data.keys(): + d = data[k] + + if not isinstance(d, RVData): + raise TypeError(f"All data must be specified as RVData instances: " + f"Object at key '{k}' is a '{type(d)}' instead.") + + if d._has_cov: + raise NotImplementedError("We currently don't support " + "multi-survey data when a full " + "covariance matrix is specified. " + "Raise an issue in adrn/thejoker if " + "you want this functionality.") + + if rv_unit is None: + rv_unit = d.rv.unit + + t.append(d.t.tcb.mjd) + rv.append(d.rv.to_value(rv_unit)) + err.append(d.rv_err.to_value(rv_unit)) + ids.append([k] * len(d)) + + t = np.concatenate(t) + rv = np.concatenate(rv) * rv_unit + err = np.concatenate(err) * rv_unit + ids = np.concatenate(ids) + + all_data = RVData(t=Time(t, format='mjd', scale='tcb'), + rv=rv, rv_err=err, + t_ref=t_ref, sort=False) + K_M = np.zeros((len(all_data), 2)) + K_M[ids == '1', 0] = 1. + K_M[ids == '2', 1] = -1. + trend_M = get_trend_design_matrix(all_data, ids=None, poly_trend=poly_trend) + + return all_data, ids, np.hstack((K_M, trend_M)) + + +class JokerSB2Samples(JokerSamples): + + def __init__(self, samples=None, t_ref=None, poly_trend=None, + **kwargs): + """ + A dictionary-like object for storing prior or posterior samples from + The Joker, with some extra functionality. + + Parameters + ---------- + samples : `~astropy.table.QTable`, table-like (optional) + The samples data as an Astropy table object, or something + convertable to an Astropy table (e.g., a dictionary with + `~astropy.units.Quantity` object values). This is optional because + the samples data can be added later by setting keys on the resulting + instance. + poly_trend : int (optional) + Specifies the number of coefficients in an additional polynomial + velocity trend, meant to capture long-term trends in the data. See + the docstring for `thejoker.JokerPrior` for more details. + t_ref : `astropy.time.Time`, numeric (optional) + The reference time for the orbital parameters. + **kwargs + Additional keyword arguments are stored internally as metadata. + """ + super().__init__(t_ref=t_ref, poly_trend=poly_trend, **kwargs) + + self._valid_units['K1'] = self._valid_units.pop('K') + self._valid_units['K2'] = self._valid_units.get('K1') + + if samples is not None: + for col in samples.colnames: + self[col] = samples[col] + + def __repr__(self): + return (f'') + + @property + def primary(self): + new_samples = JokerSamples(t_ref=self.t_ref, poly_trend=self.poly_trend) + new_samples['K'] = self['K1'] + for name in new_samples._valid_units.keys(): + if name == 'K' or name not in self.tbl.colnames: + continue + + else: + new_samples[name] = self[name] + return new_samples + + @property + def secondary(self): + new_samples = JokerSamples(t_ref=self.t_ref, poly_trend=self.poly_trend) + new_samples['K'] = self['K2'] + for name in new_samples._valid_units.keys(): + if name == 'K' or name not in self.tbl.colnames: + continue + + elif name == 'omega': + new_samples[name] = self[name] - 180*u.deg + + else: + new_samples[name] = self[name] + return new_samples + + def get_orbit(self, index=None, which='1', **kwargs): + pass + + +class TheJokerSB2(TheJoker): + _samples_cls = JokerSB2Samples + + def _make_joker_helper(self, data): + assert len(data) == 2 + assert '1' in data.keys() and '2' in data.keys() + all_data, ids, M = validate_prepare_data_sb2( + data, self.prior.poly_trend, t_ref=data['1'].t_ref) + + joker_helper = CJokerSB2Helper(all_data, self.prior, M) + return joker_helper