Implement nautilus sampler interface

cosmosis/main.py

@@ -263,17 +263,38 @@ def run_cosmosis(args, pool=None, ini=None, pipeline=None, values=None):
         # It's not fully rolled out to all the suitable samplers yet though.
         resume = ini.getboolean(RUNTIME_INI_SECTION, "resume", fallback=False)
 
+        # Polychord, multinest, and nautilus have their own internal
+        # mechanism for resuming chains.
+        if sampler_class.internal_resume:
+            resume2 = ini.getboolean(sampler_name, "resume", fallback=False)
+            resume = resume or resume2
+
+            if resume and is_root:
+                print(f"Resuming sampling using {sampler_name} internal mechanism, "
+                      "so starting a new output chain file.")
+
+            # Tell the sampler to resume directly
+            if not ini.has_section(sampler_name):
+                ini.add_section(sampler_name)
+            ini.set(sampler_name, "resume", str(resume))
+
+            # Switch off the main cosmosis resume mechanism
+            resume = False
+
         # Not all samplers can be resumed.
         if resume and not sampler_class.supports_resume:
             print("NOTE: You set resume=T in the [runtime] section but the sampler {} does not support resuming yet.  I will ignore this option.".format(sampler_name))
             resume=False
 
+
         if is_root:
             print("****************************")
             print("* Running sampler {}/{}: {}".format(sampler_number+1,number_samplers, sampler_name))
 
         output = setup_output(sampler_class, sampler_number, ini, pool, number_samplers, sample_method, resume)
-        print("****************************")
+
+        if is_root:
+            print("****************************")
 
         #Initialize our sampler, with the class we got above.
         #It needs an extra pool argument if it is a ParallelSampler.

cosmosis/postprocessing/postprocess.py

@@ -316,6 +316,8 @@ def reduced_col(self, name, stacked=True):
 		col = self.get_col(name)
 		return col
 
+class NautilusProcess(DynestyProcessor):
+	sampler = "nautilus"
 
 
 class PMCPostProcessor(WeightedMetropolisProcessor):

cosmosis/postprocessing/statistics.py

@@ -709,10 +709,13 @@ def run(self):
         # which is just read from the file
         files = super(MultinestStatistics,self).run()
         logz = self.source.final_metadata[0]["log_z"]
-        logz_sigma = self.source.final_metadata[0]["log_z_error"]
+        try:
+            logz_sigma = self.source.final_metadata[0]["log_z_error"]
+        except KeyError:
+            logz_sigma = "(error not computed)"
         #First print to screen
         print("Bayesian evidence:")
-        print("    log(Z) = %g ± %g" % (logz,logz_sigma))
+        print(f"    log(Z) = {logz:.2f} ± {logz_sigma}")
         print()
 
 
@@ -725,7 +728,7 @@ def run(self):
         #Now save to file
         header = '#logz    logz_sigma'
         f, filename, new_file  = self.get_text_output("evidence", header, self.source.name)
-        f.write('%e    %e\n'%(logz,logz_sigma))
+        f.write(f'{logz}    {logz_sigma}\n')
 
         #Include evidence in list of created files
         files.append(filename)

cosmosis/samplers/__init__.py

@@ -23,3 +23,4 @@
 from .dynesty.dynesty_sampler import DynestySampler
 from .zeus.zeus_sampler import ZeusSampler
 from .poco.poco_sampler import PocoSampler
+from .nautilus.nautilus_sampler import NautilusSampler

cosmosis/samplers/multinest/multinest_sampler.py

@@ -64,6 +64,7 @@ class MultinestSampler(ParallelSampler):
     parallel_output = False
     sampler_outputs = [("prior", float), ("like", float), ("post", float), ("weight", float)]
     supports_smp=False
+    internal_resume = True
 
     def config(self):
         if self.pool:

cosmosis/samplers/nautilus/nautilus_sampler.py

@@ -0,0 +1,102 @@
+from .. import ParallelSampler
+import numpy as np
+import sys
+import os
+
+def log_probability_function(p):
+    r = pipeline.run_results(p)
+    out = [r.post, r.prior]
+
+    # Flatten any vector outputs here
+    for e in r.extra:
+        if np.isscalar(e):
+            out.append(e)
+        else:
+            out.extend(e)
+    out = tuple(out)
+    return out
+
+
+def prior_transform(p):
+    return pipeline.denormalize_vector_from_prior(p)
+
+
+class NautilusSampler(ParallelSampler):
+    parallel_output = False
+    internal_resume = True
+    sampler_outputs = [ ('log_weight', float), ('prior', float), ("post", float)]
+
+    def config(self):
+        global pipeline
+        pipeline = self.pipeline
+
+        if self.is_master():
+            self.n_live = self.read_ini("n_live", int, 1500)
+            self.n_update = self.read_ini("n_update", int, self.n_live)
+            self.enlarge = self.read_ini(
+                "enlarge", float, 1.1**self.pipeline.nvaried)
+            self.n_batch = self.read_ini("n_batch", int, 100)
+            self.random_state = self.read_ini("random_state", int, -1)
+            if self.random_state < 0:
+                self.random_state = None
+            self.resume_ = self.read_ini("resume", bool, False)
+            self.f_live = self.read_ini("f_live", float, 0.01)
+            self.n_shell = self.read_ini("n_shell", int, self.n_batch)
+            self.n_eff = self.read_ini("n_eff", float, 10000.0)
+            self.discard_exploration = self.read_ini(
+                "discard_exploration", bool, False)
+            self.verbose = self.read_ini("verbose", bool, False)
+
+        self.converged = False
+
+    def execute(self):
+        from nautilus import Sampler
+
+        n_dim = self.pipeline.nvaried
+
+        try:
+            resume_filepath = self.output.name_for_sampler_resume_info()
+        except NotImplementedError:
+            resume_filepath = None
+
+        if resume_filepath is not None:
+            resume_filepath = resume_filepath + ".hdf5"
+            if self.resume_ and os.path.exists(resume_filepath):
+                print(f"Resuming Nautilus from file {resume_filepath}")
+
+        sampler = Sampler(
+            prior_transform,
+            log_probability_function,
+            n_dim,
+            n_live=self.n_live,
+            n_update=self.n_update,
+            enlarge=self.enlarge,
+            n_batch=self.n_batch,
+            random_state=self.random_state,
+            filepath=resume_filepath,
+            resume=self.resume_,
+            pool=self.pool,
+            blobs_dtype=float
+        )
+
+        sampler.run(f_live=self.f_live,
+                    n_shell=self.n_shell,
+                    n_eff=self.n_eff,
+                    discard_exploration=self.discard_exploration,
+                    verbose=self.verbose)
+
+        for sample, logwt, logl, blob in zip(*sampler.posterior(return_blobs=True)):
+            prior = blob[0]
+            extra = blob[1:]
+            logp = logl + prior
+            self.output.parameters(sample, extra, logwt, prior, logp)
+
+        self.output.final(
+            "efficiency", sampler.effective_sample_size() / sampler.n_like)
+        self.output.final("neff", sampler.effective_sample_size())
+        self.output.final("nsample", len(sampler.posterior()[0]))
+        self.output.final("log_z", sampler.evidence())
+        self.converged = True
+
+    def is_converged(self):
+        return self.converged

cosmosis/samplers/nautilus/sampler.yaml

@@ -0,0 +1,35 @@
+name: "nautilus"
+version: "0.4.1"
+parallel: parallel
+purpose: "Neural Network-Boosted Importance Nested Sampling"
+url: "https://github.com/johannesulf/nautilus"
+attribution: ["Johannes U. Lange"]
+
+explanation: >
+    Nautilus is an MIT-licensed pure-Python package for Bayesian posterior and
+    evidence estimation. It utilizes importance sampling and efficient space
+    exploration using neural networks. Compared to traditional MCMC and Nested
+    Sampling codes, it often needs fewer likelihood calls and produces much
+    larger posterior samples. Additionally, nautilus is highly accurate and
+    can produce Bayesian evidence estimates with percent precision.
+
+
+installation: >
+    pip install nautilus-sampler
+    conda install -c conda-forge nautilus-sampler
+
+
+# List of configuration options for this sampler
+params:
+    n_live: (integer; default=1500) number of live points
+    n_update: (integer; default=n_live) number of additions to the live set before a new bound is created
+    enlarge: (float; default=1.1*n_dim) factor by which the volume of ellipsoidal bounds is increased
+    n_batch: (integer; default=100) number of likelihood evaluations that are performed at each step
+    random_state: (int; default=-1) random seed, negative values give a random random seed
+    filepath: (string; default='None') file used for checkpointing, must have .hdf5 ending
+    resume: (bool; default=True) if True, resume from previous run stored in `filepath`
+    f_live: (float; default=0.01) live set evidence fraction when exploration phase terminates
+    n_shell: (int; default=n_batch) minimum number of points in each shell
+    n_eff: (float; default=10000.0) minimum effective sample size
+    discard_exploration: (bool; default=False) whether to discard points drawn in the exploration phase
+    verbose: (bool; default=False) If true, print information about sampler progress

cosmosis/samplers/polychord/polychord_sampler.py

@@ -79,6 +79,7 @@ class PolychordSampler(ParallelSampler):
     parallel_output = False
     sampler_outputs = [("prior", float), ("like", float), ("post", float), ("weight", float)]
     supports_smp=False
+    internal_resume = True
     understands_fast_subspaces = True
 
     def config(self):

cosmosis/samplers/sampler.py

@@ -36,6 +36,7 @@ class Sampler(metaclass=RegisteredSampler):
     parallel_output = False
     is_parallel_sampler = False
     supports_resume = False
+    internal_resume = False
 
 
     def __init__(self, ini, pipeline, output=None):

cosmosis/test/test_samplers.py

@@ -138,6 +138,9 @@ def test_polychord():
 def test_snake():
         run('snake', True)
 
+def test_nautilus():
+    run('nautilus', True)
+
 def test_star():
         run('star', False)
 

requirements.txt

@@ -13,3 +13,4 @@ mpi4py
 dulwich
 urllib3
 scikit-learn
+nautilus-sampler

setup.py

@@ -183,6 +183,7 @@ def run(self):
     "emcee",
     "dynesty",
     "zeus-mcmc",
+    "nautilus-sampler",
     "dulwich",
     "scikit-learn",
     "future",