diff --git a/orbitize/sampler.py b/orbitize/sampler.py index dcf0765d..dbf95910 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 @@ -1675,3 +1676,69 @@ def _loglike_multinest(param_cube, n_dim, n_param): self.results.save_results(filename=hdf5_file) return post_samples[:, :-1] + +class NautilusSampler(Sampler): + def __init__(self, system): + super(NautilusSampler, 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() + + 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)): + try: + utform[i] = self.system.sys_priors[i].transform_samples(u[i]) + except AttributeError: # prior is a fixed number + utform[i] = self.system.sys_priors[i] + return utform + + def run_sampler( + self, + n_live: int = 2000, + n_update: None|int = None, + verbose: bool = False, + sampler_kwargs: dict = {}, + run_kwargs: dict = {} + ): + + sampler = nautilus.Sampler( + prior=self.ptform, + likelihood=self._logl, + n_dim=len(self.system.sys_priors), + n_live=n_live, + n_update=n_update, + **sampler_kwargs + ) + + success = sampler.run( + verbose=verbose, + **run_kwargs + ) + + # TODO: use weighted posterior for higher accuracy + points, low_w, log_l = sampler.posterior(equal_weight = True) + + self.results.add_samples( + points, + log_l + ) + + return points 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