diff --git a/docs/tutorials.rst b/docs/tutorials.rst index 46798550..37988a5c 100644 --- a/docs/tutorials.rst +++ b/docs/tutorials.rst @@ -45,6 +45,8 @@ us if you are still confused). tutorials/Plotting_tutorial.ipynb tutorials/MCMC_vs_OFTI.ipynb tutorials/Modifying_MCMC_initial_positions.ipynb + tutorials/dynesty_tutorial.ipynb + tutorials/Nautilus_tutorial.ipynb .. toctree:: :maxdepth: 1 @@ -57,6 +59,7 @@ us if you are still confused). tutorials/Hipparcos_IAD.ipynb tutorials/HGCA_tutorial.ipynb tutorials/abs_astrometry.ipynb + tutorials/ONeil-ObsPriors.ipynb diff --git a/docs/tutorials/Nautilus_tutorial.ipynb b/docs/tutorials/Nautilus_tutorial.ipynb new file mode 100644 index 00000000..1ba4d876 --- /dev/null +++ b/docs/tutorials/Nautilus_tutorial.ipynb @@ -0,0 +1,287 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Nautilus Introduction #\n", + "\n", + "by Quinn Blackstone, Aniruddh Chalagulla, Eshel Dror and Niklas Naworal (2026)\n", + "\n", + "This is a tutorial for using the [Nautilus](https://arxiv.org/abs/2306.16923) sampler in `orbitize`!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import orbitize\n", + "import numpy as np\n", + "import multiprocessing as mp" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Basic Orbit Generating\n", + "\n", + "Before we get on with orbit generation, lets get some data ready. For this tutorial we generate some synthetic data with the `generate_synthetic_data()` function." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from orbitize.system import generate_synthetic_data\n", + "data_table, sma = generate_synthetic_data(\n", + " 95, # orbital fraction, how much of the orbit is covered (0-100%)\n", + " 1.2, # mass of the system\n", + " 60, # paralax\n", + " 0.1, # eccentricity (float)\n", + " np.pi/6, # inclination (float)\n", + " unc=2, # uncertainty (int)\n", + " num_obs=30, # number of observations to be made (int)\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from orbitize import system\n", + "mySys = system.System(1, data_table, 1.2, 60)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Initialize the sampler" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from orbitize import sampler\n", + "mySampler = sampler.NautilusSampler(mySys)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Nice, now it's time to run the sampler\n", + "\n", + "A bit more on the hyperparameters\n", + "\n", + "#### Sampler arguments\n", + "`n_networks`: The number of neural networks trained to determine the iso-likelihood shells at each iteration. The likelihood scores prediceted by the neural networks are averaged to determine if a point is included in a shell. A higher value of `n_networks` leads to greater sampling efficiency at the cost of higher overhead per iteration. Neural network training is not parallelized.\n", + "`n_batch`: The number of orbits evaluated in each iteration, must be a multiple of num_threads. Each thread evaluates a set of `n_batch/num_threads` orbits. Higher `n_batch` takes advantage of the vectorized prior transform to improve computation speed, at the cost of some extra likelihood evaluations being performed in each shell (i.e. if `n_update` is 2500 samples and n_batch is 1000 then 3000 samples would be taken).\n", + "\n", + "#### Run arguments\n", + "`f_live`: The maximum fraction of the evidence left in the live set before switching from the exploration to the sampling phase. This can be used to act as an ending point for the exploration phase (when shells are created), larger values lead to faster convergence but may not converge fully.\n", + "`n_eff`: The total effective size of the sample. Larger values produce higher quality posterior estimations but take longer to run. Acts as the ending point of the sampling phase.\n", + "\n", + "`n_live`: Number of live points used in the sampling process. Greater `n_live` yields better evidence estimation and shell accuracy but increased run time and computational cost.\n", + "num_threads: Number of threads used in parallel, greater number increases speed.\n", + "`n_update`: Number of additions to the live set before creating a new shell. If None defaults to `n_live`. Greater `n_update` leads to better shell accuracy but increased run time.\n", + "`verbose`: Lets you choose wether or not to display live updates of the sampler.\n", + "`savefile`: The name for a file that you want to save or resume nautilus progress to or from.\n", + "\n", + "See more on the possible hyperparameters in the [Nautilus docs](https://nautilus-sampler.readthedocs.io/en/latest/api_full.html)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sampler_arg = {\n", + " \"n_networks\": 4,\n", + " \"n_batch\": None, #(int) \n", + "}\n", + "run_arg = {\"f_live\": 0.01, \"n_eff\": 1000} \n", + "samples = mySampler.run_sampler(\n", + " n_live=500, #(int)\n", + " num_threads=mp.cpu_count(), #(int)\n", + " n_update=None, #(int)\n", + " verbose=True, #(bool) \n", + " savefile=None, #(str)\n", + " sampler_kwargs=sampler_arg,\n", + " run_kwargs=run_arg\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Plotting and results ##\n", + "\n", + "For a more detailed guide on data visualization capabilities within orbitize, see the [Orbitize plotting tutorial](https://orbitize.readthedocs.io/en/latest/tutorials/Plotting_tutorial.html)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "After generating the samples, the `run_sampler` method also creates a `Results` object that can be accessed\n", + "with `mySampler.results`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "myResults = mySampler.results" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Weighted results" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "What is the difference between a normal post and a weighted posterior?\n", + "Weighted results are all the points that have been sampled by Nautilus weighted by shell sampling density and likelihood. The main benefit of weighted results is that they contain all points instead of equal-weighted results which are downsampled from it with no duplicates. `plot_corner` is able to make use of the weighted results for better accuracy but `plot_orbits` is not and therefore defaults to the unweighted results.\n", + "\n", + "`myResults.weighted_post` returns weighted posteriors if it exists, and unweighted if it does not, so you can still get the posterior from Nautilus without weights.\n", + "`myResults.weighted_lnlike` functions similarly for getting the log-likelihoods\n", + "\n", + "`myResults.post` is always unweighted. The same goes for `myResults.lnlike`.\n", + "Weight can be acquired from `myResults.weight`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "myResults.post.shape" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "myResults.weighted_post.shape" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In case you want more points than that of the unweighted results but fewer points than the weighted, you can use the downsample function. Note that if you do not allow duplicates, you will not obtain a completely true posterior (especially if your `amount` approaches the weighted posterior size); if you want a true equal-weighted posterior without duplicates use myResults.post." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Samples from the posterior a number of times. If the posterior is weighted, samples proportionally to weight to get an equal-weighted posterior.\n", + "post, lnlike = myResults.downsample(\n", + " 150000, # amount\n", + " True # allow duplicates\n", + ")\n", + "post.shape, lnlike.shape" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You might be wondering what kind of cases you might want `duplicates` to be `True`, it would be when you want more points that still have equal weight.\n", + "For example this would return an error:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " post, lnlike = myResults.downsample(250000, duplicates=False)\n", + " print(post.shape, lnlike.shape)\n", + "\n", + "except ValueError as error:\n", + " print(f\"Downsampling error: {error}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In addition to our `orbits` array, Orbitize also creates a `Results` class that contains built-in plotting capabilities for two types of plots: corner plots and orbit plots. These cornerplots can also take in weighted results, in addition to downsizing them.\n", + "\n", + "### Corner Plot ###" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can now create a corner plot using the function `plot_corner` within the `Results` class. This function requires an input list of the parameters, in string format, that you wish to include in your corner plot. We can even plot all of the orbital parameters at once! You may wish to use the `downsample` keyword if you want a quick graph as the full weighted posteriors tend to take a while to plot. As shown below:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "corner_figure = myResults.plot_corner(downsample=100000, param_list=['sma1', 'ecc1', 'inc1','tau1'])" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.19" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/orbitize/plot.py b/orbitize/plot.py index afaf85c1..997d9683 100644 --- a/orbitize/plot.py +++ b/orbitize/plot.py @@ -28,8 +28,7 @@ cmap(np.linspace(0.0, 0.7, 1000)), ) - -def plot_corner(results, param_list=None, **corner_kwargs): +def plot_corner(results, param_list=None, downsample=None, **corner_kwargs): """ Make a corner plot of posterior on orbit fit from any sampler @@ -55,6 +54,9 @@ def plot_corner(results, param_list=None, **corner_kwargs): sigma: rv jitter mi: mass of individual body i, for i = 0, 1, 2, ... (only if fit_secondary_mass == True) mtot: total mass (only if fit_secondary_mass == False) + + downsample (int): + amount of samples to randomly draw from the posterior using ``results.downsample`` **corner_kwargs: any remaining keyword args are sent to ``corner.corner``. See `here `_. @@ -125,8 +127,15 @@ def plot_corner(results, param_list=None, **corner_kwargs): else: fixed_indices.append(i) + if downsample is not None: + post, _ = results.downsample(downsample) + weights = None + else: + post = results.weighted_post + weights = results.weights + samples = np.copy( - results.post[:, param_indices] + post[:, param_indices] ) # keep only chains for selected parameters samples[:, angle_indices] = np.degrees( samples[:, angle_indices] @@ -161,7 +170,7 @@ def plot_corner(results, param_list=None, **corner_kwargs): corner_kwargs["labels"] = reduced_labels_list - figure = corner.corner(samples, **corner_kwargs) + figure = corner.corner(samples, weights=weights, **corner_kwargs) return figure diff --git a/orbitize/results.py b/orbitize/results.py index fea92e13..0d7533b1 100644 --- a/orbitize/results.py +++ b/orbitize/results.py @@ -9,7 +9,6 @@ import orbitize.plot import orbitize.gaia, orbitize.hipparcos - class Results(object): """ A class to store accepted orbital configurations from the sampler @@ -28,6 +27,15 @@ class Results(object): data (astropy.table.Table): output from ``orbitize.read_input.read_file()`` curr_pos (np.array of float): for MCMC only. A multi-D array of the current walker positions that is used for restarting a MCMC sampler. + weighted_post (np.array of float): RxN array of orbital parameters + (posterior output from weighted orbit-fitting process), where R is the + number of orbits generated, and N is the number of varying orbital + parameters in the fit (default: None). + weighted_lnlike (np.array of float): R array of log-likelihoods corresponding to + the orbits described in ``weighted_post`` (default: None). + lnweights (np.array of float): R array of log-weights corresponding to the orbits + and log-likelihoods described in ``weighted_post`` and ``weighted_lnlike`` (default: None) + Written: Henry Ngo, Sarah Blunt, 2018 @@ -36,7 +44,7 @@ class Results(object): def __init__( self, system=None, sampler_name=None, post=None, lnlike=None, - version_number=None, curr_pos=None + version_number=None, curr_pos=None, weighted_post=None, weighted_lnlike=None, lnweight=None ): self.system = system @@ -45,6 +53,9 @@ def __init__( self.lnlike = lnlike self.curr_pos = curr_pos self.version_number = version_number + self._weighted_post = weighted_post + self._weighted_lnlike = weighted_lnlike + self.lnweight = lnweight self.ln_evidence = None self.ln_evidence_err = None @@ -58,7 +69,60 @@ def __init__( self.param_idx = self.system.param_idx self.standard_param_idx = self.system.basis.standard_basis_idx - def add_samples(self, orbital_params, lnlikes, curr_pos=None): + @property + def weighted_post(self): + """ + Returns the weighted posterior if it exists, + otherwise returns the unweighted posterior. + """ + if self._weighted_post is not None: + return self._weighted_post + return self.post + + @property + def weighted_lnlike(self): + """ + Returns the weighted log-likelihoods if it exists, + otherwise returns the unweighted log-likelihoods. + """ + if self._weighted_lnlike is not None: + return self._weighted_lnlike + return self.lnlike + + @property + def weights(self): + """ + Returns the weights of ``weighted_post`` and ``weighted_lnlike`` + if it exists, otherwise returns None. + """ + if self.lnweight is None: + return None + return np.exp(self.lnweight) + + def downsample(self, amount, duplicates=True): + """ + Samples from the posterior, or the weighted posterior if it exists. + + Args: + amount (int): number of samples to draw from the posetrior + duplicates (bool): whether to replace sampled orbits from the posterior, + which may cause duplicate samples (default: True) + + Returns: + tuple: + + numpy.array of float: orbital parameters of the samples (``amount``xN, where N is number of varying orbital parameters) + + numpy.array of float: log-likelihoods of the samples (length ``amount``) + + """ + + indexes = np.random.choice(len(self.weighted_post), amount, duplicates, self.weights) + if self.weighted_lnlike is None: + return self.weighted_post[indexes], None + return self.weighted_post[indexes], self.weighted_lnlike[indexes] + + def add_samples(self, orbital_params, lnlikes, curr_pos=None, weighted_post=None, weighted_lnlike=None, lnweight=None): """ Add accepted orbits, their likelihoods, and the orbitize version number to the results @@ -69,6 +133,12 @@ def add_samples(self, orbital_params, lnlikes, curr_pos=None): lnlike (np.array): add corresponding lnlike values to results curr_pos (np.array of float): for MCMC only. A multi-D array of the current walker positions + weighted_post (np.array): for Nautilus only. Sets of orbital params + associated with ``lnweight`` to add to results. + weighted_lnlike (np.array): for Nautilus only. Corresponding weighted lnlike + values for ``weighted_post`` to add to results. + weighted_post (np.array): for Nautilus only. Log of weights associated + with ``weighted_post`` and ``weighted_lnlike``. Written: Henry Ngo, 2018 @@ -88,6 +158,16 @@ def add_samples(self, orbital_params, lnlikes, curr_pos=None): else: self.post = np.vstack((self.post, orbital_params)) self.lnlike = np.append(self.lnlike, lnlikes) + + if weighted_post is not None: + if self._weighted_post is None: + self._weighted_post = weighted_post + self._weighted_lnlike = weighted_lnlike + self.lnweight = lnweight + else: + self._weighted_post = np.vstack((self._weighted_post, weighted_post)) + self._weighted_lnlike = np.append(self._weighted_lnlike, weighted_lnlike) + self.lnweight = np.append(self.lnweight, lnweight) if curr_pos is not None: self.curr_pos = curr_pos @@ -132,6 +212,12 @@ def save_results(self, filename): if self.curr_pos is not None: hf.create_dataset("curr_pos", data=self.curr_pos) + if self._weighted_post is not None and self._weighted_lnlike is not None and self.lnweight is not None: + hf.create_dataset("weighted_post", data=self._weighted_post) + hf.create_dataset("weighted_lnlike", data=self._weighted_lnlike) + hf.create_dataset("lnweight", data=self.lnweight) + + self.system.save(hf) hf.close() # Closes file object, which writes file to disk @@ -160,12 +246,12 @@ def load_results(self, filename, append=False): version_number = str(hf.attrs['version_number']) else: version_number = "<= 1.13" - post = hf.get('post') - if post is not None: - post = np.array(post) - lnlike = hf.get('lnlike') - if lnlike is not None: - lnlike = np.array(lnlike) + + post = array_not_none(hf.get('post')) + lnlike = array_not_none(hf.get('lnlike')) + weighted_post = array_not_none(hf.get('weighted_post')) + weighted_lnlike = array_not_none(hf.get('weighted_lnlike')) + lnweight = array_not_none(hf.get("lnweight")) if 'num_secondary_bodies' in hf.attrs: num_secondary_bodies = int(hf.attrs['num_secondary_bodies']) @@ -287,10 +373,7 @@ def load_results(self, filename, append=False): if 'ln_evidence_err' in hf.attrs: self.ln_evidence_err = hf.attrs['ln_evidence_err'] - try: - curr_pos = np.array(hf.get('curr_pos')) - except KeyError: - curr_pos = None + curr_pos = array_not_none(hf.get('curr_pos')) hf.close() # Closes file object @@ -316,14 +399,14 @@ def load_results(self, filename, append=False): 'Unable to append file {} to Results object. version_number of object and file do not match'.format(filename)) # Now append post and lnlike - self.add_samples(post, lnlike)#, self.labels) + self.add_samples(post, lnlike, weighted_post=weighted_post, weighted_lnlike=weighted_lnlike, lnweight=lnweight)#, self.labels) else: # Only proceed if object is completely empty - if self.sampler_name is None and self.post is None and self.lnlike is None and self.version_number is None:# and self.tau_ref_epoch is None : + if self.sampler_name is None and self.post is None and self.lnlike is None and self.version_number is None and self._weighted_lnlike is None and self._weighted_post is None and self.lnweight is None:# and self.tau_ref_epoch is None : self.sampler_name = sampler_name self.version_number = version_number - self.add_samples(post, lnlike)#, self.labels) + self.add_samples(post, lnlike, weighted_post=weighted_post, weighted_lnlike=weighted_lnlike, lnweight=lnweight)#, self.labels) else: raise Exception( @@ -390,11 +473,11 @@ def print_results(self): self.results_str += '-------------------\n' print(self.results_str) - def plot_corner(self, param_list=None, **corner_kwargs): + def plot_corner(self, param_list=None, downsample=None, **corner_kwargs): """ Wrapper for orbitize.plot.plot_corner """ - return orbitize.plot.plot_corner(self, param_list, **corner_kwargs) + return orbitize.plot.plot_corner(self, param_list, downsample, **corner_kwargs) def plot_orbits(self, object_to_plot=1, start_mjd=51544., num_orbits_to_plot=100, num_epochs_to_plot=100, @@ -444,4 +527,12 @@ def plot_propermotion(self, cmap=cmap, cbar_param=cbar_param, # fig=fig - ) \ No newline at end of file + ) + +def array_not_none(raw): + """ + Returns a numpy.array of the input if it is not None, else returns None + """ + if raw is not None: + return np.array(raw) + return raw diff --git a/orbitize/sampler.py b/orbitize/sampler.py index 09c23cc3..1b8c3a97 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -8,6 +8,7 @@ import dynesty import emcee import matplotlib.pyplot as plt +import nautilus import numpy as np import ptemcee @@ -18,6 +19,8 @@ import orbitize.priors import orbitize.results +import sys + class Sampler(abc.ABC): """ @@ -1309,7 +1312,32 @@ def check_prior_support(self): return -class NestedSampler(Sampler): +class BaseNestedSampler(Sampler): + def ptform(self, u): + """ + Prior transform function. + + Args: + u (np.array of floats): RxM array of uniform + samples with values 0 < u < 1, + where R is the number of parameters + and M is the number of orbits + + Returns: + numpy RxM array of floats: u samples transformed to + a chosen Prior Class distribution. + """ + utform = np.zeros(u.shape) + for i in range(u.shape[0]): + if hasattr(self.system.sys_priors[i], 'transform_samples'): + utform[i] = self.system.sys_priors[i].transform_samples(u[i]) + else: + # prior is a fixed number + utform[i] = self.system.sys_priors[i] + return utform + + +class NestedSampler(BaseNestedSampler): """ Implements nested sampling using the Dynesty package. @@ -1352,26 +1380,6 @@ def __init__(self, self.start = time.time() self.dynesty_sampler = None - def ptform(self, u): - """ - Prior transform function. - - Args: - u (array of floats): list of samples with values 0 < u < 1. - - Returns: - numpy array of floats: 1D u samples transformed to a chosen Prior - Class distribution. - """ - utform = np.zeros(len(u)) - for i in range(len(u)): - if hasattr(self.system.sys_priors[i], 'transform_samples'): - utform[i] = self.system.sys_priors[i].transform_samples(u[i]) - else: - # prior is a fixed number - utform[i] = self.system.sys_priors[i] - return utform - def run_sampler( self, nlive=500, @@ -1675,3 +1683,133 @@ def _loglike_multinest(param_cube, n_dim, n_param): self.results.save_results(filename=hdf5_file) return post_samples[:, :-1] + +class NautilusSampler(BaseNestedSampler): + """ + Implements nested sampling using the Nautilus-Sampler package. + + Args: + system (system.System): system.System object + 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. + + Eshel Dror, Quinton Blackston, Aniruddh Chalagulla, & Niklas Naworal 2026 + """ + def __init__(self, + system, + chi2_type="standard", + like="chi2_lnlike", + custom_lnlike=None, + ): + super(NautilusSampler, self).__init__( + system, + like=like, + chi2_type=chi2_type, + custom_lnlike=custom_lnlike, + ) + + # 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() + + def nautilus_ptform(self, u): + return self.ptform(u.T).T + + def nautilus_logl(self, u: np.ndarray): + return self._logl(u.T).T + + def run_sampler( + self, + n_live = 2000, + n_update = None, + verbose = False, + num_threads = 1, + savefile = None, + sampler_kwargs = {}, + run_kwargs = {} + ): + """Runs the nested sampler from the Nautilus package. + + Args: + n_live (int): Number of live points. A larger numbers results + in a more finely sampled posterior and a more accurate + evidence, but also a larger number of iterations is + required to converge (default: 2000). + n_update (int): Number of points added to the live set before + creating a new shell. When None defaults to `n_live` (default: None). + verbose (bool): Print progress when running sampler (default: False). + num_threads (int, tuple, pool): number of threads to use for parallelization. + If a tuple of two integers, the number of threads for + likelihood evaluations and sampler calculations respectively. If a thread + pool, the pool used for parallelization (default=1). + savefile (str): File used by Nautilus to save progress. + If file already exists, resumes from saved progress. + sampler_kwargs (dict): dictionary of keywords to be passed into nautilus.Sampler, + such as `periodic` + run_kwargs (dict): dictionary of keywords to be passed into nautilus.Sampler.run, + such as `n_eff` + + + Returns: + numpy.array of float: equal-weighted posterior samples + """ + if sys.version_info < (3,9,0) and isinstance(num_threads, int) and num_threads > 1: + with mp.Pool(processes=num_threads) as pool: + self.naut_sampler = nautilus.Sampler( + prior=self.nautilus_ptform, + likelihood=self.nautilus_logl, + n_dim=len(self.system.sys_priors), + vectorized=True, + n_live=n_live, + n_update=n_update, + pool=pool, + filepath=savefile, + **sampler_kwargs + ) + + success = self.naut_sampler.run( + verbose=verbose, + **run_kwargs + ) + else: + self.naut_sampler = nautilus.Sampler( + prior=self.nautilus_ptform, + likelihood=self.nautilus_logl, + n_dim=len(self.system.sys_priors), + vectorized=True, + n_live=n_live, + n_update=n_update, + pool=num_threads, + filepath=savefile, + **sampler_kwargs + ) + + success = self.naut_sampler.run( + verbose=verbose, + **run_kwargs + ) + points, _, log_l = self.naut_sampler.posterior(equal_weight = True) + weighted_points, log_w, weighted_log_l = self.naut_sampler.posterior() + + self.results.add_samples( + points, + log_l, + weighted_post=weighted_points, + weighted_lnlike=weighted_log_l, + lnweight=log_w + ) + + return points diff --git a/pyproject.toml b/pyproject.toml index 1e7ec9b6..2698bd55 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,8 @@ dependencies = [ "rebound", "dynesty", "pymultinest", - "tk" + "tk", + "nautilus-sampler" ] [tool.setuptools.dynamic] diff --git a/requirements.txt b/requirements.txt index 2c215d7e..fea4658c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,3 +16,4 @@ rebound dynesty pymultinest tk +nautilus-sampler diff --git a/tests/end-to-end-tests/GJ504_naut.py b/tests/end-to-end-tests/GJ504_naut.py new file mode 100644 index 00000000..58948f9f --- /dev/null +++ b/tests/end-to-end-tests/GJ504_naut.py @@ -0,0 +1,74 @@ +import matplotlib.pyplot as plt + +import orbitize + +# Based on driver.py + +import numpy as np +from orbitize import read_input, system, sampler, priors +import multiprocessing as mp +import orbitize +import time + +savedir = "" + +# System parameters +datafile = "GJ504.csv" +num_secondary_bodies = 1 +system_mass = 1.22 # Msol +plx = 56.95 # mas +mass_err = 0.08 # Msol +plx_err = 0.26 # mas + +# Sampler parameters +likelihood_func_name = "chi2_lnlike" +n_threads = mp.cpu_count() +n_live = 2000 +n_update = None +sampler_args = {"n_networks": 4} +run_args = {"f_live": 0.01, "n_eff": 10000} + +naut_file = f"{savedir}GJ504_naut.hdf5" + +results_file = f"{savedir}GJ504_results.hdf5" + +tau_ref_epoch = 50000 + + +# Read in data +data_table = read_input.read_file(orbitize.DATADIR + datafile) + +# Initialize System object which stores data & sets priors +my_system = system.System( + num_secondary_bodies, + data_table, + system_mass, + plx, + mass_err=mass_err, + plx_err=plx_err, + tau_ref_epoch=tau_ref_epoch, +) + +my_sampler = sampler.NautilusSampler(my_system, like=likelihood_func_name) + +print(f"Running {datafile} sampler with {n_threads} {n_live} {n_update} {sampler_args} {run_args}") +# Run the sampler to compute some orbits, yeah! +my_sampler.run_sampler(n_live, + n_update, + verbose=False, + num_threads=n_threads, + savefile=naut_file, + sampler_kwargs=sampler_args, + run_kwargs=run_args) + +end = time.time() +print(f"Processing time: {end} - {my_sampler.start} = {(end-my_sampler.start) / 60} minutes") + + +my_sampler.results.save_results(results_file) + +# make corner plot +fig = my_sampler.results.plot_corner(downsample=10000) +plt.savefig(f"{savedir}GJ504_naut_corner.png", dpi=250) + +my_sampler.results.print_results() \ No newline at end of file diff --git a/tests/test_mcmc.py b/tests/test_mcmc.py index 329434c2..3f85c3ed 100644 --- a/tests/test_mcmc.py +++ b/tests/test_mcmc.py @@ -96,6 +96,7 @@ def do_mcmc_runs(num_temps=0, num_threads=1, make_corner_plot=False): if make_corner_plot: assert myDriver.system.plx_err == 0 # (check that we're actually fixing at least one param: plx) plot_corner(new_sampler.results) + plot_corner(new_sampler.results, downsample=1000) # Test unweighted downsampling # clean up os.system(f'rm {output_filename} {output_filename_2}') diff --git a/tests/test_nautilus.py b/tests/test_nautilus.py new file mode 100644 index 00000000..0da3921a --- /dev/null +++ b/tests/test_nautilus.py @@ -0,0 +1,56 @@ +import pytest +from orbitize import sampler, system, results +import numpy as np +from orbitize.system import generate_synthetic_data +import matplotlib.pyplot as plt +""" This is pytest for Nautilus_Sampler, it assumes values for all the priors except + + eccentricity and trys to calculate the true eccentricity """ + +def test_nautilus_general(make_plot=False): + # generate synthetic data + mtot = 1.2 + plx = 60.0 + orbit_frac = 95 + ecc = 0.1 + inc = np.pi/4 + data_table, sma = generate_synthetic_data( + orbit_frac, + mtot, + plx, + num_obs=30, + ecc = ecc, + inc = inc + ) + + #initlialize the orbit + + mySys = system.System(1, data_table, mtot, plx) + lab = mySys.param_idx + + #set all paremeters except eccentricity + + #mySys.sys_priors[lab["inc1"]] = np.pi / 4, + mySys.sys_priors[lab["sma1"]] = sma + mySys.sys_priors[lab["aop1"]] = np.pi / 4 + mySys.sys_priors[lab["pan1"]] = np.pi / 4 + mySys.sys_priors[lab["tau1"]] = 0.8 + mySys.sys_priors[lab["plx"]] = plx + mySys.sys_priors[lab["mtot"]] = mtot + + my_sampler = sampler.NautilusSampler(mySys) + my_sampler.run_sampler(n_live=800, n_update=None, verbose=False) + + nautilus_eccentricities = my_sampler.results.post[:, lab["ecc1"]] + assert np.mean(nautilus_eccentricities) == pytest.approx(0.1, abs=0.1) + + nautilus_inclination = my_sampler.results.post[:, lab["inc1"]] + assert np.median(nautilus_inclination) == pytest.approx(inc, abs=0.1) + + if make_plot: + myResults = my_sampler.results + myResults.plot_corner(param_list=["ecc1","inc1"]) # No downsampling + myResults.plot_corner(param_list=["ecc1","inc1"], downsample = 1000) # With downsampling + +if __name__ == "__main__": + test_nautilus_general(make_plot = True) diff --git a/tests/test_results.py b/tests/test_results.py index 5dbd91a9..3f380f70 100644 --- a/tests/test_results.py +++ b/tests/test_results.py @@ -75,7 +75,7 @@ def simulate_orbit_sampling(n_sim_orbits): return sim_post -def test_init_and_add_samples(radec_input=False): +def test_init_and_add_samples(radec_input=False, weighted=False): """ Tests object creation and add_samples() with some simulated posterior samples, and returns results.Results object @@ -96,8 +96,18 @@ def test_init_and_add_samples(radec_input=False): n_orbit_draws1 = 1000 sim_post = simulate_orbit_sampling(n_orbit_draws1) sim_lnlike = np.random.uniform(size=n_orbit_draws1) - # Test adding samples - results_obj.add_samples(sim_post, sim_lnlike) # , labels=std_labels) + if weighted: + n_weighted_orbit_draws1 = 4000 + sim_weighted_post = simulate_orbit_sampling(n_weighted_orbit_draws1) + sim_weighted_lnlike = np.random.uniform(size=n_weighted_orbit_draws1) + sim_weight = np.random.uniform(size=n_weighted_orbit_draws1) + sim_lnweight = np.log(sim_weight / np.sum(sim_weight)) # Normalize + # Test adding samples + results_obj.add_samples(sim_post, sim_lnlike, weighted_post=sim_weighted_post, + weighted_lnlike=sim_weighted_lnlike, lnweight=sim_lnweight) + else: + # Test adding samples + results_obj.add_samples(sim_post, sim_lnlike) # , labels=std_labels) # Simulate some more sample draws n_orbit_draws2 = 2000 sim_post = simulate_orbit_sampling(n_orbit_draws2) @@ -108,6 +118,12 @@ def test_init_and_add_samples(radec_input=False): expected_length = n_orbit_draws1 + n_orbit_draws2 assert results_obj.post.shape == (expected_length, 8) assert results_obj.lnlike.shape == (expected_length,) + if weighted: + expected_length_weighted = n_weighted_orbit_draws1 + assert results_obj.weighted_post.shape == (expected_length_weighted, 8) + assert results_obj.weighted_lnlike.shape == (expected_length_weighted,) + assert results_obj.weights.shape == (expected_length_weighted,) + assert results_obj.tau_ref_epoch == 58849 assert results_obj.labels == std_labels @@ -132,8 +148,15 @@ def results_to_test(): n_orbit_draws2 = 2000 sim_post = simulate_orbit_sampling(n_orbit_draws2) sim_lnlike = np.random.uniform(size=n_orbit_draws2) - # Test adding more samples - results_obj.add_samples(sim_post, sim_lnlike) + # Simulate some weighted draws + n_weighted_orbit_draws = 4000 + sim_weighted_post = simulate_orbit_sampling(n_weighted_orbit_draws) + sim_weighted_lnlike = np.random.uniform(size=n_weighted_orbit_draws) + sim_weight = np.random.uniform(size=n_weighted_orbit_draws) + sim_lnweight = np.log(sim_weight / np.sum(sim_weight)) # Normalize + # Test adding weighted and more samples + results_obj.add_samples(sim_post, sim_lnlike, weighted_post=sim_weighted_post, + weighted_lnlike=sim_weighted_lnlike, lnweight=sim_lnweight) # Return object for testing return results_obj @@ -173,8 +196,11 @@ def test_save_and_load_results(results_to_test, has_lnlike=True): assert results_to_save.sampler_name == loaded_results.sampler_name assert results_to_save.version_number == loaded_results.version_number assert np.array_equal(results_to_save.post, loaded_results.post) + assert np.array_equal(results_to_save.weighted_post, loaded_results.weighted_post) + assert np.array_equal(results_to_save.lnweight, loaded_results.lnweight) if has_lnlike: assert np.array_equal(results_to_save.lnlike, loaded_results.lnlike) + assert np.array_equal(results_to_save.weighted_lnlike, loaded_results.weighted_lnlike) # Try to load the saved results again, this time appending loaded_results.load_results(save_filename, append=True) # Now check that the loaded results object has the expected size @@ -200,7 +226,8 @@ def test_save_and_load_results(results_to_test, has_lnlike=True): def test_plot_corner(results_to_test): """ Tests plot_corner() with plotting simulated posterior samples - for all 8 parameters and for just four selected parameters + for all 8 parameters, for just four selected parameters, + with fixed parameters, and downsampled """ Figure1 = results_to_test.plot_corner() @@ -213,10 +240,15 @@ def test_plot_corner(results_to_test): # test that fixing parameters doesn't crash corner plot code results_to_test.post[:, -1] = np.ones(len(results_to_test.post[:, -1])) Figure3 = results_to_test.plot_corner() + assert Figure3 is not None results_to_test.post[:, -1] = mass_vals - return Figure1, Figure2, Figure3 + Figure4 = results_to_test.plot_corner(downsample=1000) + assert Figure4 is not None + + + return Figure1, Figure2, Figure3, Figure4 def test_plot_orbits(results_to_test): @@ -245,6 +277,25 @@ def test_plot_orbits(results_to_test): assert Figure5 is not None return (Figure1, Figure2, Figure3, Figure4, Figure5) +def test_downsample(results_to_test): + """ + Test downsample() with simulated posterior samples + """ + size = results_to_test.weighted_post.shape[0] + post, lnlikes = results_to_test.downsample(size*2, duplicates=True) + assert post.shape[0] == size*2 + assert lnlikes.shape[0] == size*2 + post, lnlikes = results_to_test.downsample(size, duplicates=False) + assert post.shape[0] == size + assert lnlikes.shape[0] == size + try: + post, lnlikes = results_to_test.downsample(size+1, duplicates=False) + except ValueError: + pass # Expected error when taking too many samples without replacement + else: + assert False, "Expected ValueError for drawing more samples than possible without duplicates did not occur" + + def test_save_and_load_hipparcos_only(): """ @@ -339,29 +390,50 @@ def test_save_and_load_gaia_and_hipparcos(): test_save_and_load_hipparcos_only() test_save_and_load_gaia_and_hipparcos() - # test_results = test_init_and_add_samples() - - # test_results_printing(test_results) - # test_plot_long_periods(test_results) - # test_results_radec = test_init_and_add_samples(radec_input=True) - - # test_save_and_load_results(test_results, has_lnlike=True) - # test_save_and_load_results(test_results, has_lnlike=True) - # test_save_and_load_results(test_results, has_lnlike=False) - # test_save_and_load_results(test_results, has_lnlike=False) - # test_corner_fig1, test_corner_fig2, test_corner_fig3 = test_plot_corner( - # test_results - # ) - # test_orbit_figs = test_plot_orbits(test_results) - # test_orbit_figs = test_plot_orbits(test_results_radec) - # test_corner_fig1.savefig("test_corner1.png") - # test_corner_fig2.savefig("test_corner2.png") - # test_corner_fig3.savefig("test_corner3.png") - # test_orbit_figs[0].savefig("test_orbit1.png") - # test_orbit_figs[1].savefig("test_orbit2.png") - # test_orbit_figs[2].savefig("test_orbit3.png") - # test_orbit_figs[3].savefig("test_orbit4.png") - # test_orbit_figs[4].savefig("test_orbit5.png") - - # # clean up - # os.system("rm test_*.png") + # Not weighted + test_results = test_init_and_add_samples() + + test_results_printing(test_results) + test_plot_long_periods(test_results) + test_downsample(test_results) + # TODO: Update failing test with radec_input=True + test_results_radec = test_init_and_add_samples(radec_input=True) + + test_save_and_load_results(test_results, has_lnlike=True) + test_save_and_load_results(test_results, has_lnlike=True) + test_save_and_load_results(test_results, has_lnlike=False) + test_save_and_load_results(test_results, has_lnlike=False) + test_corner_fig1, test_corner_fig2, test_corner_fig3, test_corner_fig4 = test_plot_corner( + test_results + ) + test_orbit_figs = test_plot_orbits(test_results) + test_orbit_figs = test_plot_orbits(test_results_radec) + test_corner_fig1.savefig("test_corner1.png") + test_corner_fig2.savefig("test_corner2.png") + test_corner_fig3.savefig("test_corner3.png") + test_corner_fig4.savefig("test_corner4.png") + test_orbit_figs[0].savefig("test_orbit1.png") + test_orbit_figs[1].savefig("test_orbit2.png") + test_orbit_figs[2].savefig("test_orbit3.png") + test_orbit_figs[3].savefig("test_orbit4.png") + test_orbit_figs[4].savefig("test_orbit5.png") + + # Weighted + test_results_weighted = test_init_and_add_samples(weighted=True) + + test_results_printing(test_results_weighted) + test_downsample(test_results_weighted) + + test_save_and_load_results(test_results_weighted, has_lnlike=True) + test_save_and_load_results(test_results_weighted, has_lnlike=True) + test_save_and_load_results(test_results_weighted, has_lnlike=False) + test_save_and_load_results(test_results_weighted, has_lnlike=False) + test_corner_fig5, test_corner_fig6, test_corner_fig7, test_corner_fig8 = test_plot_corner( + test_results_weighted + ) + test_corner_fig5.savefig("test_corner5.png") + test_corner_fig6.savefig("test_corner6.png") + test_corner_fig7.savefig("test_corner7.png") + test_corner_fig8.savefig("test_corner8.png") + # clean up + os.system("rm test_*.png")