import numpy as np
import astropy.units as u
import astropy.constants as consts
import abc
import time
from astropy.time import Time
import dynesty
import emcee
import ptemcee
import multiprocessing as mp
from multiprocessing import Pool
import orbitize.lnlike
import orbitize.priors
import orbitize.kepler
from orbitize import cuda_ext
import orbitize.results
import matplotlib.pyplot as plt
[docs]class Sampler(abc.ABC):
"""
Abstract base class for sampler objects.
All sampler objects should inherit from this class.
Written: Sarah Blunt, 2018
"""
def __init__(
self, system, like="chi2_lnlike", custom_lnlike=None, chi2_type="standard"
):
self.system = system
# check if `like` is a string or a function
if callable(like):
self.lnlike = like
else:
self.lnlike = getattr(orbitize.lnlike, like)
self.custom_lnlike = custom_lnlike
self.chi2_type = chi2_type
# check if need to handle covariances
self.has_corr = np.any(~np.isnan(self.system.data_table["quant12_corr"]))
@abc.abstractmethod
def run_sampler(self, total_orbits):
pass
def _logl(self, params):
"""
log likelihood function that interfaces with the orbitize objects
Comptues the sum of the log likelihoods of the data given the input model
Args:
params (np.array of float): RxM array
of fitting parameters, where R is the number of
parameters being fit, and M is the number of orbits
we need model predictions for. Must be in the same order
documented in System() above. If M=1, this can be a 1d array.
Returns:
float: sum of all log likelihoods of the data given input model
"""
# compute the model based on system params
model, jitter = self.system.compute_model(params)
# fold data/errors to match model output shape. In particualr, quant1/quant2 are interleaved
data = np.array(
[self.system.data_table["quant1"], self.system.data_table["quant2"]]
).T
# errors below required for lnlike function below
errs = np.array(
[self.system.data_table["quant1_err"], self.system.data_table["quant2_err"]]
).T
# covariances/correlations, if applicable
# we're doing this check now because the likelihood computation is much faster if we can skip it.
if self.has_corr:
corrs = self.system.data_table["quant12_corr"]
else:
corrs = None
# grab all seppa indices
seppa_indices = self.system.all_seppa
# compute lnlike
lnlikes = self.lnlike(
data, errs, corrs, model, jitter, seppa_indices, chi2_type=self.chi2_type
)
# return sum of lnlikes (aka product of likeliehoods)
lnlikes_sum = np.nansum(lnlikes, axis=(0, 1))
if self.custom_lnlike is not None:
lnlikes_sum += self.custom_lnlike(params)
if self.system.hipparcos_IAD is not None:
# compute Ra/Dec predictions at the Hipparcos IAD epochs
raoff_model, deoff_model, _ = self.system.compute_all_orbits(
params, epochs=self.system.hipparcos_IAD.epochs_mjd
)
(
raoff_model_hip_epoch,
deoff_model_hip_epoch,
_,
) = self.system.compute_all_orbits(
params, epochs=Time([1991.25], format="decimalyear").mjd
)
# subtract off position of star at reference Hipparcos epoch
raoff_model[:, 0, :] -= raoff_model_hip_epoch[:, 0, :]
deoff_model[:, 0, :] -= deoff_model_hip_epoch[:, 0, :]
# select body 0 raoff/deoff predictions & feed into Hip IAD lnlike fn
lnlikes_sum += self.system.hipparcos_IAD.compute_lnlike(
raoff_model[:, 0, :],
deoff_model[:, 0, :],
params,
self.system.param_idx,
)
if self.system.gaia is not None:
gaiahip_epochs = Time(
np.append(
self.system.gaia.hipparcos_epoch, self.system.gaia.gaia_epoch
),
format="decimalyear",
).mjd
# compute Ra/Dec predictions at the Gaia epoch
raoff_model, deoff_model, _ = self.system.compute_all_orbits(
params, epochs=gaiahip_epochs
)
# select body 0 raoff/deoff predictions & feed into Gaia module lnlike fn
lnlikes_sum += self.system.gaia.compute_lnlike(
raoff_model[:, 0, :],
deoff_model[:, 0, :],
params,
self.system.param_idx,
)
return lnlikes_sum
[docs]class OFTI(
Sampler,
):
"""
OFTI Sampler
Args:
system (system.System): ``system.System`` object
like (string): name of likelihood function in ``lnlike.py``
custom_lnlike (func): ability to include an addition custom likelihood
function in the fit. The function looks like
``clnlikes = custon_lnlike(params)`` where ``params`` is a RxM array
of fitting parameters, where R is the number of orbital paramters
(can be passed in system.compute_model()),
and M is the number of orbits we need model predictions for.
It returns ``clnlikes`` which is an array of
length M, or it can be a single float if M = 1.
Written: Isabel Angelo, Sarah Blunt, Logan Pearce, 2018
"""
def __init__(
self, system, like="chi2_lnlike", custom_lnlike=None, chi2_type="standard"
):
super(OFTI, self).__init__(
system, like=like, chi2_type=chi2_type, custom_lnlike=custom_lnlike
)
if (self.system.hipparcos_IAD is not None) or (len(self.system.rv[0] > 0)):
raise NotImplementedError(
"""
You can only use OFTI with relative astrometry measurements
(no Hipparcos IAD or RVs... yet). Use MCMC, you overachiever, and
settle in for a nice long orbit-fit. (But seriously, if you want
this functionality, let us know!)
"""
)
# compute priors and columns containing ra/dec and sep/pa
self.priors = self.system.sys_priors
# convert RA/Dec rows to sep/PA
convert_warning_print = False
for body_num in np.arange(self.system.num_secondary_bodies) + 1:
if len(self.system.radec[body_num]) > 0:
# only print the warning once.
if not convert_warning_print:
print(
"Converting ra/dec data points in data_table to sep/pa. Original data are stored in input_table."
)
convert_warning_print = True
self.system.convert_data_table_radec2seppa(body_num=body_num)
# these are of type astropy.table.column
self.sep_observed = self.system.data_table[
np.where(self.system.data_table["quant_type"] == "seppa")
]["quant1"].copy()
self.pa_observed = self.system.data_table[
np.where(self.system.data_table["quant_type"] == "seppa")
]["quant2"].copy()
self.sep_err = self.system.data_table[
np.where(self.system.data_table["quant_type"] == "seppa")
]["quant1_err"].copy()
self.pa_err = self.system.data_table[
np.where(self.system.data_table["quant_type"] == "seppa")
]["quant2_err"].copy()
self.meas_object = self.system.data_table[
np.where(self.system.data_table["quant_type"] == "seppa")
]["object"].copy()
# this is OK, ONLY IF we are only using self.epochs for computing RA/Dec from Keplerian elements
self.epochs = (
np.array(self.system.data_table["epoch"]) - self.system.tau_ref_epoch
)
# distinguishing all epochs from sep/pa epochs
self.epochs_seppa = (
np.array(
self.system.data_table[
np.where(self.system.data_table["quant_type"] == "seppa")
]["epoch"]
)
- self.system.tau_ref_epoch
)
self.epochs_rv = (
np.array(
self.system.data_table[
np.where(self.system.data_table["quant_type"] == "rv")
]["epoch"]
)
- self.system.tau_ref_epoch
)
# choose scale-and-rotate epoch
# for multiplanet support, this is now a list.
# For each planet, we find the measurment for it that corresponds to the smallest astrometric error
self.epoch_idx = []
min_sep_indices = np.argsort(
self.sep_err
) # indices of sep err sorted from smallest to higheset
min_sep_indices_body = self.meas_object[
min_sep_indices
] # the corresponding body_num that these sorted measurements correspond to
for i in range(self.system.num_secondary_bodies):
body_num = i + 1
this_object_meas = np.where(min_sep_indices_body == body_num)[0]
if np.size(this_object_meas) == 0:
# no data, no scaling
self.epoch_idx.append(None)
continue
# get the smallest measurement belonging to this body
best_epoch = min_sep_indices[this_object_meas][
0
] # already sorted by argsort
self.epoch_idx.append(best_epoch)
if (
len(self.system.rv[0]) > 0 and self.system.fit_secondary_mass
): # checking for RV data
self.rv_observed = self.system.data_table[
np.where(self.system.data_table["quant_type"] == "rv")
]["quant1"].copy()
self.rv_err = self.system.data_table[
np.where(self.system.data_table["quant_type"] == "rv")
]["quant1_err"].copy()
self.epoch_rv_idx = [
np.argmin(self.rv_observed),
np.argmax(self.rv_observed),
]
# create an empty results object
self.results = orbitize.results.Results(
self.system,
sampler_name=self.__class__.__name__,
post=None,
lnlike=None,
version_number=orbitize.__version__,
)
[docs] def prepare_samples(self, num_samples):
"""
Prepare some orbits for rejection sampling. This draws random orbits
from priors, and performs scale & rotate.
Args:
num_samples (int): number of orbits to draw and scale & rotate for
OFTI to run rejection sampling on
Return:
np.array: array of prepared samples. The first dimension has size of
num_samples. This should be passed into ``OFTI.reject()``
"""
# generate sample orbits
samples = np.empty([len(self.priors), num_samples])
for i in range(len(self.priors)):
if hasattr(self.priors[i], "draw_samples"):
samples[i, :] = self.priors[i].draw_samples(num_samples)
else: # param is fixed & has no prior
samples[i, :] = self.priors[i] * np.ones(num_samples)
# Make Converison to Standard Basis:
samples = self.system.basis.to_standard_basis(samples)
for body_num in np.arange(self.system.num_secondary_bodies) + 1:
sma = samples[
self.system.basis.standard_basis_idx["sma{}".format(body_num)], :
]
ecc = samples[
self.system.basis.standard_basis_idx["ecc{}".format(body_num)], :
]
inc = samples[
self.system.basis.standard_basis_idx["inc{}".format(body_num)], :
]
argp = samples[
self.system.basis.standard_basis_idx["aop{}".format(body_num)], :
]
lan = samples[
self.system.basis.standard_basis_idx["pan{}".format(body_num)], :
]
tau = samples[
self.system.basis.standard_basis_idx["tau{}".format(body_num)], :
]
plx = samples[self.system.basis.standard_basis_idx["plx"], :]
if self.system.fit_secondary_mass:
m0 = samples[self.system.basis.standard_basis_idx["m0"], :]
m1 = samples[
self.system.basis.standard_basis_idx["m{}".format(body_num)], :
]
mtot = m0 + m1
else:
mtot = samples[self.system.basis.standard_basis_idx["mtot"], :]
m1 = None
min_epoch = self.epoch_idx[body_num - 1]
if min_epoch is None:
# Don't need to rotate and scale if no astrometric measurments for this body. Brute force rejection sampling
continue
period_prescale = np.sqrt(
4 * np.pi**2 * (sma * u.AU) ** 3 / (consts.G * (mtot * u.Msun))
)
period_prescale = period_prescale.to(u.day).value
meananno = self.epochs[min_epoch] / period_prescale - tau
# compute sep/PA of generated orbits
ra, dec, _ = orbitize.kepler.calc_orbit(
self.epochs[min_epoch],
sma,
ecc,
inc,
argp,
lan,
tau,
plx,
mtot,
tau_ref_epoch=0,
mass_for_Kamp=m1,
)
sep, pa = orbitize.system.radec2seppa(ra, dec) # sep[mas], PA[deg]
# generate Gaussian offsets from observational uncertainties
sep_offset = np.random.normal(0, self.sep_err[min_epoch], size=num_samples)
pa_offset = np.random.normal(0, self.pa_err[min_epoch], size=num_samples)
# calculate correction factors
sma_corr = (sep_offset + self.sep_observed[min_epoch]) / sep
lan_corr = pa_offset + self.pa_observed[min_epoch] - pa
# perform scale-and-rotate
sma *= sma_corr # [AU]
lan += np.radians(lan_corr) # [rad]
lan = (lan + 2 * np.pi) % (2 * np.pi)
if self.system.restrict_angle_ranges:
argp[lan >= np.pi] += np.pi
argp = argp % (2 * np.pi)
lan[lan >= np.pi] -= np.pi
period_new = np.sqrt(
4 * np.pi**2 * (sma * u.AU) ** 3 / (consts.G * (mtot * u.Msun))
)
period_new = period_new.to(u.day).value
tau = (self.epochs[min_epoch] / period_new - meananno) % 1
# updates samples with new values of sma, pan, tau
samples[
self.system.basis.standard_basis_idx["sma{}".format(body_num)], :
] = sma
samples[
self.system.basis.standard_basis_idx["aop{}".format(body_num)], :
] = argp
samples[
self.system.basis.standard_basis_idx["pan{}".format(body_num)], :
] = lan
samples[
self.system.basis.standard_basis_idx["tau{}".format(body_num)], :
] = tau
return samples
[docs] def reject(self, samples):
"""
Runs rejection sampling on some prepared samples.
Args:
samples (np.array): array of prepared samples. The first dimension \
has size ``num_samples``. This should be the output of \
``prepare_samples()``.
Return:
tuple:
np.array: a subset of ``samples`` that are accepted based on the
data.
np.array: the log likelihood values of the accepted orbits.
"""
lnp = self._logl(samples)
# we just want the chi2 term for rejection, so compute the Gaussian normalization term and remove it
errs = np.array(
[self.system.data_table["quant1_err"], self.system.data_table["quant2_err"]]
).T
if self.has_corr:
corrs = self.system.data_table["quant12_corr"]
else:
corrs = None
lnp_scaled = lnp - orbitize.lnlike.chi2_norm_term(errs, corrs)
# account for user-set priors on PAN that were destroyed by scale-and-rotate
for body_num in np.arange(self.system.num_secondary_bodies) + 1:
pan_idx = self.system.basis.standard_basis_idx["pan{}".format(body_num)]
pan_prior = self.system.sys_priors[pan_idx]
if pan_prior is not orbitize.priors.UniformPrior:
# apply PAN prior
lnp_scaled += pan_prior.compute_lnprob(samples[pan_idx, :])
# prior is uniform but with different bounds that OFTI expects
elif (pan_prior.minval != 0) or (
(pan_prior.maxval != np.pi) or (pan_prior.maxval != 2 * np.pi)
):
samples_outside_pan_prior = np.where(
(samples[pan_idx, :] < pan_prior.minval)
or (samples[pan_idx, :] > pan_prior.maxval)
)[0]
lnp_scaled[samples_outside_pan_prior] = -np.inf
# reject orbits with probability less than a uniform random number
random_samples = np.log(np.random.random(len(lnp)))
saved_orbit_idx = np.where(lnp_scaled > random_samples)[0]
saved_orbits = np.array([samples[:, i] for i in saved_orbit_idx])
lnlikes = np.array([lnp[i] for i in saved_orbit_idx])
return saved_orbits, lnlikes
def _sampler_process(
self, output, total_orbits, num_samples=10000, Value=0, lock=None
):
"""
Runs OFTI until it finds the number of total accepted orbits desired.
Meant to be called by run_sampler.
Args:
output (manager.Queue): manager.Queue object to store results
total_orbits (int): total number of accepted orbits desired by user
num_samples (int): number of orbits to prepare for OFTI to run
rejection sampling on
Value (mp.Value(int)): global counter for the orbits generated
lock: mp.lock object to prevent issues caused by access to shared
memory by multiple processes
Returns:
tuple:
output_orbits (np.array): array of accepted orbits,
size: total_orbits
output_lnlikes (np.array): array of log probabilities,
size: total_orbits
"""
np.random.seed()
n_orbits_saved = 0
output_orbits = np.empty((total_orbits, len(self.priors)))
output_lnlikes = np.empty(total_orbits)
# add orbits to `output_orbits` until `total_orbits` are saved
while n_orbits_saved < total_orbits:
samples = self.prepare_samples(num_samples)
accepted_orbits, lnlikes = self.reject(samples)
if len(accepted_orbits) == 0:
pass
else:
n_accepted = len(accepted_orbits)
maxindex2save = np.min([n_accepted, total_orbits - n_orbits_saved])
output_orbits[n_orbits_saved : n_orbits_saved + n_accepted] = (
accepted_orbits[0:maxindex2save]
)
output_lnlikes[n_orbits_saved : n_orbits_saved + n_accepted] = lnlikes[
0:maxindex2save
]
n_orbits_saved += maxindex2save
# add to the value of the global variable
with lock:
Value.value += maxindex2save
output.put((np.array(output_orbits), output_lnlikes))
return (np.array(output_orbits), output_lnlikes)
[docs] def run_sampler(
self, total_orbits, num_samples=10000, num_cores=None, OFTI_warning=60.0
):
"""
Runs OFTI in parallel on multiple cores until we get the number of total accepted orbits we want.
Args:
total_orbits (int): total number of accepted orbits desired by user
num_samples (int): number of orbits to prepare for OFTI to run
rejection sampling on. Defaults to 10000.
num_cores (int): the number of cores to run OFTI on. Defaults to
number of cores availabe.
OFTI_warning (float): if OFTI doesn't accept a single orbit before
this amount of time (in seconds), print a warning suggesting to
try MCMC. If None, don't print a warning.
Return:
np.array: array of accepted orbits. Size: total_orbits.
Written by: Vighnesh Nagpal(2019)
"""
start_time = time.time()
if num_cores != 1:
if num_cores is None:
num_cores = mp.cpu_count()
results = []
# orbits_saved is a global counter for the number of orbits generated
orbits_saved = mp.Value("i", 0)
manager = mp.Manager()
output = manager.Queue()
# setup the processes
lock = mp.Lock()
nrun_per_core = int(np.ceil(float(total_orbits) / float(num_cores)))
processes = [
mp.Process(
target=self._sampler_process,
args=(output, nrun_per_core, num_samples, orbits_saved, lock),
)
for x in range(num_cores)
]
# start the processes
for p in processes:
p.start()
# print out the number of orbits generated every second
while orbits_saved.value < total_orbits:
if OFTI_warning is not None and orbits_saved.value == 0:
check_time = time.time() - start_time
if check_time > OFTI_warning:
print(
"Warning! OFTI is taking a while, and you may want to consider MCMC, check out the MCMC vs OFTI tutorial.",
end="\n",
)
OFTI_warning = None
print(
str(orbits_saved.value) + "/" + str(total_orbits) + " orbits found",
end="\r",
)
time.sleep(0.1)
print(
str(total_orbits) + "/" + str(total_orbits) + " orbits found", end="\r"
)
# join the processes
for p in processes:
p.join()
# get the results of each process from the queue
for p in processes:
results.append(output.get())
# filling up the output_orbits array
output_orbits = np.zeros((total_orbits, len(self.priors)))
output_lnlikes = np.empty(total_orbits)
pos = 0
for p in results:
num_to_fill = np.min([len(p[0]), total_orbits - pos])
output_orbits[pos : pos + num_to_fill] = p[0][0:num_to_fill]
output_lnlikes[pos : pos + num_to_fill] = p[1][0:num_to_fill]
pos += num_to_fill
self.results.add_samples(np.array(output_orbits), output_lnlikes)
return output_orbits
else:
# this block is executed if num_cores=1
n_orbits_saved = 0
output_orbits = np.empty((total_orbits, len(self.priors)))
output_lnlikes = np.empty(total_orbits)
# add orbits to `output_orbits` until `total_orbits` are saved
while n_orbits_saved < total_orbits:
samples = self.prepare_samples(num_samples)
accepted_orbits, lnlikes = self.reject(samples)
if len(accepted_orbits) == 0:
check_time = time.time() - start_time
if OFTI_warning is not None:
if check_time > OFTI_warning:
print(
"Warning! OFTI is taking a while, and you may want to consider MCMC, check out the MCMC vs OFTI tutorial.",
end="\n",
)
OFTI_warning = None
pass
else:
n_accepted = len(accepted_orbits)
maxindex2save = np.min([n_accepted, total_orbits - n_orbits_saved])
output_orbits[n_orbits_saved : n_orbits_saved + n_accepted] = (
accepted_orbits[0:maxindex2save]
)
output_lnlikes[n_orbits_saved : n_orbits_saved + n_accepted] = (
lnlikes[0:maxindex2save]
)
n_orbits_saved += maxindex2save
# print progress statement
print(
str(n_orbits_saved) + "/" + str(total_orbits) + " orbits found",
end="\r",
)
self.results.add_samples(np.array(output_orbits), output_lnlikes)
return output_orbits
[docs]class MCMC(Sampler):
"""
MCMC sampler. Supports either parallel tempering or just regular MCMC. Parallel tempering will be run if ``num_temps`` > 1
Parallel-Tempered MCMC Sampler uses ptemcee, a fork of the emcee Affine-infariant sampler
Affine-Invariant Ensemble MCMC Sampler uses emcee.
.. Warning:: may not work well for multi-modal distributions
Args:
system (system.System): system.System object
num_temps (int): number of temperatures to run the sampler at.
Parallel tempering will be used if num_temps > 1 (default=20)
num_walkers (int): number of walkers at each temperature (default=1000)
num_threads (int): number of threads to use for parallelization (default=1)
chi2_type (str, optional): either "standard", or "log"
like (str): name of likelihood function in ``lnlike.py``
custom_lnlike (func): ability to include an addition custom likelihood
function in the fit. The function looks like
``clnlikes = custon_lnlike(params)`` where ``params`` is a RxM array
of fitting parameters, where R is the number of orbital paramters
(can be passed in system.compute_model()), and M is the number of
orbits we need model predictions for. It returns ``clnlikes``
which is an array of length M, or it can be a single float if M = 1.
prev_result_filename (str): if passed a filename to an HDF5 file
containing a orbitize.Result data, MCMC will restart from where it
left off.
Written: Jason Wang, Henry Ngo, 2018
"""
def __init__(
self,
system,
num_temps=20,
num_walkers=1000,
num_threads=1,
chi2_type="standard",
like="chi2_lnlike",
custom_lnlike=None,
prev_result_filename=None,
):
super(MCMC, self).__init__(
system, like=like, chi2_type=chi2_type, custom_lnlike=custom_lnlike
)
self.num_temps = num_temps
self.num_walkers = num_walkers
self.num_threads = num_threads
# create an empty results object
self.results = orbitize.results.Results(
self.system,
sampler_name=self.__class__.__name__,
post=None,
lnlike=None,
version_number=orbitize.__version__,
)
if self.num_temps > 1:
self.use_pt = True
else:
self.use_pt = False
self.num_temps = 1
# get priors from the system class. need to remove and record fixed priors
self.priors = []
self.fixed_params = []
self.sampled_param_idx = {}
sampled_param_counter = 0
for i, prior in enumerate(system.sys_priors):
# check for fixed parameters
if not hasattr(prior, "draw_samples"):
self.fixed_params.append((i, prior))
else:
self.priors.append(prior)
self.sampled_param_idx[self.system.labels[i]] = sampled_param_counter
sampled_param_counter += 1
# initialize walkers initial postions
self.num_params = len(self.priors)
if prev_result_filename is None:
# initialize walkers initial postions
init_positions = []
for prior in self.priors:
# draw them uniformly becase we don't know any better right now
# TODO: be smarter in the future
random_init = prior.draw_samples(num_walkers * self.num_temps)
if self.num_temps > 1:
random_init = random_init.reshape([self.num_temps, num_walkers])
init_positions.append(random_init)
# save this as the current position for the walkers
if self.use_pt:
# make this an numpy array, but combine the parameters into a shape of (ntemps, nwalkers, nparams)
# we currently have a list of [ntemps, nwalkers] with nparam arrays. We need to make nparams the third dimension
self.curr_pos = np.dstack(init_positions)
else:
# make this an numpy array, but combine the parameters into a shape of (nwalkers, nparams)
# we currently have a list of arrays where each entry is num_walkers prior draws for each parameter
# We need to make nparams the second dimension, so we have to transpose the stacked array
self.curr_pos = np.stack(init_positions).T
else:
# restart from previous walker positions
self.results.load_results(prev_result_filename, append=True)
prev_pos = self.results.curr_pos
# check previous positions has the correct dimensions as we need given how this sampler was created.
expected_shape = (self.num_walkers, len(self.priors))
if self.use_pt:
expected_shape = (self.num_temps,) + expected_shape
if prev_pos.shape != expected_shape:
raise ValueError(
"Unable to restart chain. Saved walker positions has shape {0}, while current sampler needs {1}".format(
prev_pos.shape, expected_shape
)
)
self.curr_pos = prev_pos
def _fill_in_fixed_params(self, sampled_params):
"""
Fills in the missing parameters from the chain that aren't being sampled
Args:
sampled_params (np.array): either 1-D array of size = number of
sampled params, or 2-D array of shape (num_models, num_params)
Returns:
np.array: same number of dimensions as sampled_params,
but with num_params including the fixed parameters
"""
if len(self.fixed_params) == 0:
# nothing to add
return sampled_params
# check if 1-D or 2-D
twodim = np.ndim(sampled_params) == 2
# insert in params
for index, value in self.fixed_params:
if twodim:
sampled_params = np.insert(sampled_params, index, value, axis=1)
else:
sampled_params = np.insert(sampled_params, index, value)
return sampled_params
def _logl(self, params, include_logp=False):
"""
log likelihood function that interfaces with the orbitize objects
Comptues the sum of the log likelihoods of the data given the input model
Args:
params (np.array of float): MxR array
of fitting parameters, where R is the number of
parameters being fit, and M is the number of orbits
we need model predictions for. Must be in the same order
documented in System() above. If M=1, this can be a 1d array.
include_logp (bool): if True, also include log prior in this function
Returns:
lnlikes (float): sum of all log likelihoods of the data given input model
"""
if include_logp:
if np.ndim(params) == 1:
logp = orbitize.priors.all_lnpriors(params, self.priors)
# escape if logp == -np.inf
if np.isinf(logp):
return -np.inf
else:
logp = np.array(
[orbitize.priors.all_lnpriors(pset, self.priors) for pset in params]
)
else:
logp = 0 # don't include prior
full_params = self._fill_in_fixed_params(params)
if np.ndim(full_params) == 2:
full_params = full_params.T
return super(MCMC, self)._logl(full_params) + logp
def _update_chains_from_sampler(self, sampler, num_steps=None):
"""
Updates self.post, self.chain, and self.lnlike from the MCMC sampler
Args:
sampler (emcee.EnsembleSampler or ptemcee.Sampler): sampler object.
num_steps (int): if not None, only stores the first num_steps number of steps
"""
if num_steps is None:
# use all the steps, grab total number of steps from dimension of chains
num_steps = sampler.chain.shape[-2]
self.chain = sampler.chain
num_params = self.chain.shape[-1]
if self.use_pt:
# chain is shape: Ntemp x Nwalkers x Nsteps x Nparams
self.post = sampler.chain[0, :, :num_steps].reshape(
-1, num_params
) # the reshaping flattens the chain
# should also be picking out the lowest temperature logps
self.lnlikes = sampler.loglikelihood[0, :, :num_steps].flatten()
self.lnlikes_alltemps = sampler.loglikelihood[:, :, :num_steps]
else:
# chain is shape: Nwalkers x Nsteps x Nparams
self.post = sampler.chain[:, :num_steps].reshape(-1, num_params)
self.lnlikes = sampler.lnprobability[:, :num_steps].flatten()
# convert posterior probability (returned by sampler objects) to likelihood (required by orbitize.results.Results)
for i, orb in enumerate(self.post):
self.lnlikes[i] -= orbitize.priors.all_lnpriors(orb, self.priors)
# include fixed parameters in posterior
self.post = self._fill_in_fixed_params(self.post)
[docs] def validate_xyz_positions(self):
"""
If using the XYZ basis, walkers might be initialized in an invalid
region of parameter space. This function fixes that by replacing invalid
positions by new randomly generated positions until all are valid.
"""
if self.system.fitting_basis == "XYZ":
if self.use_pt:
all_valid = False
while not all_valid:
total_invalids = 0
for temp in range(self.num_temps):
to_stand = self.system.basis.to_standard_basis(
self.curr_pos[temp, :, :].T.copy()
).T
# Get invalids by checking ecc values for each companion
indices = [
((i * 6) + 1)
for i in range(self.system.num_secondary_bodies)
]
invalids = np.where(
(to_stand[:, indices] < 0.0) | (to_stand[:, indices] >= 1.0)
)[0]
# Redraw samples for the invalid ones
if len(invalids) > 0:
newpos = []
for prior in self.priors:
randompos = prior.draw_samples(len(invalids))
newpos.append(randompos)
self.curr_pos[temp, invalids, :] = np.stack(newpos).T
total_invalids += len(invalids)
if total_invalids == 0:
all_valid = True
print("All walker positions validated.")
else:
all_valid = False
while not all_valid:
total_invalids = 0
to_stand = self.system.basis.to_standard_basis(
self.curr_pos[:, :].T.copy()
).T
# Get invalids by checking ecc values for each companion
indices = [
((i * 6) + 1) for i in range(self.system.num_secondary_bodies)
]
invalids = np.where(
(to_stand[:, indices] < 0.0) | (to_stand[:, indices] >= 1.0)
)[0]
# Redraw saples for the invalid ones
if len(invalids) > 0:
newpos = []
for prior in self.priors:
randompos = prior.draw_samples(len(invalids))
newpos.append(randompos)
self.curr_pos[invalids, :] = np.stack(newpos).T
total_invalids += len(invalids)
if total_invalids == 0:
all_valid = True
print("All walker positions validated.")
[docs] def run_sampler(
self,
total_orbits,
burn_steps=0,
thin=1,
examine_chains=False,
output_filename=None,
periodic_save_freq=None,
):
"""
Runs PT MCMC sampler. Results are stored in ``self.chain`` and ``self.lnlikes``.
Results also added to ``orbitize.results.Results`` object (``self.results``)
.. Note:: Can be run multiple times if you want to pause and inspect things.
Each call will continue from the end state of the last execution.
Args:
total_orbits (int): total number of accepted possible
orbits that are desired. This equals
``num_steps_per_walker`` x ``num_walkers``
burn_steps (int): optional paramter to tell sampler
to discard certain number of steps at the beginning
thin (int): factor to thin the steps of each walker
by to remove correlations in the walker steps
examine_chains (boolean): Displays plots of walkers at each step by
running `examine_chains` after `total_orbits` sampled.
output_filename (str): Optional filepath for where results file can be saved.
periodic_save_freq (int): Optionally, save the current results into ``output_filename``
every nth step while running, where n is value passed into this variable.
Returns:
``emcee.sampler`` object: the sampler used to run the MCMC
"""
if periodic_save_freq is not None and output_filename is None:
raise ValueError(
"output_filename must be defined for periodic saving of the chains"
)
if periodic_save_freq is not None and not isinstance(periodic_save_freq, int):
raise TypeError("periodic_save_freq must be an integer")
nsteps = int(np.ceil(total_orbits / self.num_walkers))
if nsteps <= 0:
raise ValueError("Total_orbits must be greater than num_walkers.")
with Pool(processes=self.num_threads) as pool:
if self.use_pt:
sampler = ptemcee.Sampler(
self.num_walkers,
self.num_params,
self._logl,
orbitize.priors.all_lnpriors,
ntemps=self.num_temps,
threads=self.num_threads,
logpargs=[
self.priors,
],
)
else:
sampler = emcee.EnsembleSampler(
self.num_walkers,
self.num_params,
self._logl,
pool=pool,
kwargs={"include_logp": True},
)
print("Starting Burn in")
for i, state in enumerate(
sampler.sample(self.curr_pos, iterations=burn_steps, thin=thin)
):
if self.use_pt:
self.curr_pos = state[0]
else:
self.curr_pos = state.coords
if (i + 1) % 5 == 0:
print(
str(i + 1)
+ "/"
+ str(burn_steps)
+ " steps of burn-in complete",
end="\r",
)
if periodic_save_freq is not None:
if (i + 1) % periodic_save_freq == 0: # we've completed i+1 steps
self.results.curr_pos = self.curr_pos
self.results.save_results(output_filename)
sampler.reset()
print("")
print("Burn in complete. Sampling posterior now.")
saved_upto = 0 # keep track of how many steps of this chain we've saved. this is the next index that needs to be saved
for i, state in enumerate(
sampler.sample(self.curr_pos, iterations=nsteps, thin=thin)
):
if self.use_pt:
self.curr_pos = state[0]
else:
self.curr_pos = state.coords
# print progress statement
if (i + 1) % 5 == 0:
print(str(i + 1) + "/" + str(nsteps) + " steps completed", end="\r")
if periodic_save_freq is not None:
if (i + 1) % periodic_save_freq == 0: # we've completed i+1 steps
self._update_chains_from_sampler(sampler, num_steps=i + 1)
# figure out what is the new chunk of the chain and corresponding lnlikes that have been computed before last save
# grab the current posterior and lnlikes and reshape them to have the Nwalkers x Nsteps dimension again
post_shape = self.post.shape
curr_chain_shape = (
self.num_walkers,
post_shape[0] // self.num_walkers,
post_shape[-1],
)
curr_chain = self.post.reshape(curr_chain_shape)
curr_lnlike_chain = self.lnlikes.reshape(curr_chain_shape[:2])
# use the reshaped arrays and find the new steps we computed
curr_chunk = curr_chain[:, saved_upto : i + 1]
curr_chunk = curr_chunk.reshape(
-1, curr_chunk.shape[-1]
) # flatten nwalkers x nsteps dim
curr_lnlike_chunk = curr_lnlike_chain[
:, saved_upto : i + 1
].flatten()
# add this current chunk to the results object (which already has all the previous chunks saved)
self.results.add_samples(
curr_chunk, curr_lnlike_chunk, curr_pos=self.curr_pos
)
self.results.save_results(output_filename)
saved_upto = i + 1
print("")
self._update_chains_from_sampler(sampler)
if periodic_save_freq is None:
# need to save everything
self.results.add_samples(
self.post, self.lnlikes, curr_pos=self.curr_pos
)
elif saved_upto < nsteps:
# just need to save the last few
# same code as above except we just need to grab the last few
post_shape = self.post.shape
curr_chain_shape = (
self.num_walkers,
post_shape[0] // self.num_walkers,
post_shape[-1],
)
curr_chain = self.post.reshape(curr_chain_shape)
curr_lnlike_chain = self.lnlikes.reshape(curr_chain_shape[:2])
curr_chunk = curr_chain[:, saved_upto:]
curr_chunk = curr_chunk.reshape(
-1, curr_chunk.shape[-1]
) # flatten nwalkers x nsteps dim
curr_lnlike_chunk = curr_lnlike_chain[:, saved_upto:].flatten()
self.results.add_samples(
curr_chunk, curr_lnlike_chunk, curr_pos=self.curr_pos
)
if output_filename is not None:
self.results.save_results(output_filename)
print("Run complete")
# Close pool
if examine_chains:
self.examine_chains()
return sampler
[docs] def examine_chains(
self,
param_list=None,
walker_list=None,
n_walkers=None,
step_range=None,
transparency=1,
):
"""
Plots position of walkers at each step from Results object. Returns list of figures, one per parameter
Args:
param_list: List of strings of parameters to plot (e.g. "sma1")
If None (default), all parameters are plotted
walker_list: List or array of walker numbers to plot
If None (default), all walkers are plotted
n_walkers (int): Randomly select `n_walkers` to plot
Overrides walker_list if this is set
If None (default), walkers selected as per `walker_list`
step_range (array or tuple): Start and end values of step numbers to plot
If None (default), all the steps are plotted
transparency (int or float): Determines visibility of the plotted function
If 1 (default) results plot at 100% opacity
Returns:
List of ``matplotlib.pyplot.Figure`` objects:
Walker position plot for each parameter selected
(written): Henry Ngo, 2019
"""
# Get the flattened chain from Results object (nwalkers*nsteps, nparams)
flatchain = np.copy(self.results.post)
total_samples, n_params = flatchain.shape
n_steps = int(total_samples / self.num_walkers)
# Reshape it to (nwalkers, nsteps, nparams)
chn = flatchain.reshape((self.num_walkers, n_steps, n_params))
# Get list of walkers to use
if (
n_walkers is not None
): # If n_walkers defined, randomly choose that many walkers
walkers_to_plot = np.random.choice(
self.num_walkers, size=n_walkers, replace=False
)
elif walker_list is not None: # if walker_list is given, use that list
walkers_to_plot = np.array(walker_list)
else: # both n_walkers and walker_list are none, so use all walkers
walkers_to_plot = np.arange(self.num_walkers)
# Get list of parameters to use
if param_list is None:
params_to_plot = np.arange(n_params)
else: # build list from user input strings
params_plot_list = []
for i in param_list:
if i in self.system.basis.param_idx:
params_plot_list.append(self.system.basis.param_idx[i])
else:
raise Exception(
"Invalid param name: {}. See system.basis.param_idx.".format(i)
)
params_to_plot = np.array(params_plot_list)
# Loop through each parameter and make plot
output_figs = []
for pp in params_to_plot:
fig, ax = plt.subplots()
for ww in walkers_to_plot:
ax.plot(chn[ww, :, pp], "k-", alpha=transparency)
ax.set_xlabel("Step")
if step_range is not None: # Limit range shown if step_range is set
ax.set_xlim(step_range)
output_figs.append(fig)
# Return
return output_figs
[docs] def chop_chains(self, burn, trim=0):
"""
Permanently removes steps from beginning (and/or end) of chains from the
Results object. Also updates `curr_pos` if steps are removed from the
end of the chain.
Args:
burn (int): The number of steps to remove from the beginning of the chains
trim (int): The number of steps to remove from the end of the chians (optional)
.. Warning:: Does not update bookkeeping arrays within `MCMC` sampler object.
(written): Henry Ngo, 2019
"""
# Retrieve information from results object
flatchain = np.copy(self.results.post)
total_samples, n_params = flatchain.shape
n_steps = int(total_samples / self.num_walkers)
flatlnlikes = np.copy(self.results.lnlike)
# Reshape chain to (nwalkers, nsteps, nparams)
chn = flatchain.reshape((self.num_walkers, n_steps, n_params))
# Reshape lnlike to (nwalkers, nsteps)
lnlikes = flatlnlikes.reshape((self.num_walkers, n_steps))
# Find beginning and end indices for steps to keep
keep_start = burn
keep_end = n_steps - trim
n_chopped_steps = n_steps - trim - burn
# Update arrays in `sampler`: chain, lnlikes, lnlikes_alltemps (if PT), post
chopped_chain = chn[:, keep_start:keep_end, :]
chopped_lnlikes = lnlikes[:, keep_start:keep_end]
# Update current position if trimmed from edge
if trim > 0:
self.curr_pos = chopped_chain[:, -1, :]
# Flatten likelihoods and samples
flat_chopped_chain = chopped_chain.reshape(
self.num_walkers * n_chopped_steps, n_params
)
flat_chopped_lnlikes = chopped_lnlikes.reshape(
self.num_walkers * n_chopped_steps
)
# Update results object associated with this sampler
self.results = orbitize.results.Results(
self.system,
sampler_name=self.__class__.__name__,
post=flat_chopped_chain,
lnlike=flat_chopped_lnlikes,
version_number=orbitize.__version__,
curr_pos=self.curr_pos,
)
# Print a confirmation
print("Chains successfully chopped. Results object updated.")
[docs] def check_prior_support(self):
"""
Review the positions of all MCMC walkers, to verify that they are supported by the prior space.
This function will raise a descriptive ValueError if any positions lie outside prior support.
Otherwise, it will return nothing.
(written): Adam Smith, 2021
"""
# Flatten the walker/temperature positions for ease of manipulation.
all_positions = self.curr_pos.reshape(
self.num_walkers * self.num_temps, self.num_params
)
# Placeholder list to track any bad parameters that come up.
bad_parameters = []
# If there are no covarient priors, loop on each variable to locate any out-of-place parameters. (this is why we transpose the walkers)
if not np.any([prior.is_correlated for prior in self.priors]):
for i, x in enumerate(all_positions.T):
# Any issues with this parameter?
lnprob = self.priors[i].compute_lnprob(np.array(x))
supported = np.isfinite(lnprob).all() == True
if supported == False:
# Problem detected. Take note and continue the loop - we want to catch all the problem parameters.
bad_parameters.append(str(i))
# Throw our ValueError if necessary,
if len(bad_parameters) > 0:
raise ValueError(
"Attempting to start with walkers outside of prior support: check parameter(s) "
+ ", ".join(bad_parameters)
)
# We're not done yet, however. There may be errors in covariant priors; run a check for that.
else:
for y in all_positions:
lnprob = orbitize.priors.all_lnpriors(y, self.priors)
if not np.isfinite(lnprob).all():
raise ValueError(
"Attempting to start with walkers outside of prior support: covariant prior failure."
)
# otherwise exit the function and continue.
return
[docs]class NestedSampler(Sampler):
"""
Implements nested sampling using Dynesty package.
Thea McKenna, Sarah Blunt, & Lea Hirsch 2024
"""
def __init__(self, system):
super(NestedSampler, self).__init__(system)
# create an empty results object
self.results = orbitize.results.Results(
self.system,
sampler_name=self.__class__.__name__,
post=None,
lnlike=None,
version_number=orbitize.__version__,
)
self.start = time.time()
self.dynesty_sampler = None
[docs] def run_sampler(
self,
static=False,
bound="multi",
pfrac=1.0,
num_threads=1,
start_method="fork",
run_nested_kwargs={},
):
"""Runs the nested sampler from the Dynesty package.
Args:
static (bool): True if using static nested sampling,
False if using dynamic.
bound (str): Method used to approximately bound the prior
using the current set of live points. Conditions the
sampling methods used to propose new live points. See
https://dynesty.readthedocs.io/en/latest/quickstart.html#bounding-options
for complete list of options.
pfrac (float): posterior weight, between 0 and 1. Can only be
altered for the Dynamic nested sampler, otherwise this
keyword is unused.
num_threads (int): number of threads to use for parallelization
(default=1)
start_method (str): multiprocessing start method. Default "fork," which
won't work on all OS. Change to "spawn" if you get an error,
and make sure you run your orbitize! script inside an
if __name__=='__main__' condition to protect entry points.
run_nested_kwargs (dict): dictionary of keywords to be passed into dynesty.Sampler.run_nested()
Returns:
tuple:
numpy.array of float: posterior samples
int: number of iterations it took to converge
"""
mp.set_start_method(start_method, force=True)
if static and pfrac != 1.0:
raise ValueError(
"""The static nested sampler does not take alternate values for pfrac."""
)
if num_threads > 1:
with dynesty.pool.Pool(num_threads, self._logl, self.ptform) as pool:
if static:
self.dynesty_sampler = dynesty.NestedSampler(
pool.loglike,
pool.prior_transform,
len(self.system.sys_priors),
pool=pool,
bound=bound,
bootstrap=False,
)
self.dynesty_sampler.run_nested(**run_nested_kwargs)
else:
self.dynesty_sampler = dynesty.DynamicNestedSampler(
pool.loglike,
pool.prior_transform,
len(self.system.sys_priors),
pool=pool,
bound=bound,
bootstrap=False,
)
self.dynesty_sampler.run_nested(
wt_kwargs={"pfrac": pfrac}, **run_nested_kwargs
)
else:
if static:
self.dynesty_sampler = dynesty.NestedSampler(
self._logl,
self.ptform,
len(self.system.sys_priors),
bound=bound,
)
self.dynesty_sampler.run_nested(**run_nested_kwargs)
else:
self.dynesty_sampler = dynesty.DynamicNestedSampler(
self._logl,
self.ptform,
len(self.system.sys_priors),
bound=bound,
)
self.dynesty_sampler.run_nested(
wt_kwargs={"pfrac": pfrac}, **run_nested_kwargs
)
self.results.add_samples(
self.dynesty_sampler.results["samples"],
self.dynesty_sampler.results["logl"],
)
num_iter = self.dynesty_sampler.results["niter"]
return self.dynesty_sampler.results["samples"], num_iter