diff --git a/bin/inference/pycbc_inference_model_stats b/bin/inference/pycbc_inference_model_stats
index 263b6868cd7..9abf8c6dba8 100644
--- a/bin/inference/pycbc_inference_model_stats
+++ b/bin/inference/pycbc_inference_model_stats
@@ -27,6 +27,7 @@ import shutil
import logging
import time
import numpy
+import tqdm
import pycbc
from pycbc.pool import (use_mpi, choose_pool)
@@ -116,7 +117,8 @@ logging.info('Getting samples')
with loadfile(opts.input_file, 'r') as fp:
# we'll need the shape; all the arrays in the samples group should have the
# same shape
- shape = fp[fp.samples_group]['loglikelihood'].shape
+ pick = list(fp[fp.samples_group].keys())[0]
+ shape = fp[fp.samples_group][pick].shape
samples = {}
for p in model.variable_params:
samples[p] = fp[fp.samples_group][p][()].flatten()
@@ -127,7 +129,8 @@ logging.info("Loaded %i samples", samples.size)
# get the stats
logging.info("Calculating stats")
-data = list(pool.map(model_call, enumerate(samples)))
+data = list(tqdm.tqdm(pool.imap(model_call, enumerate(samples)),
+ total=len(samples)))
stats = [x[0] for x in data]
rec = [x[1] for x in data]
diff --git a/companion.txt b/companion.txt
index 19ea6b0bb73..d49f4121e9e 100644
--- a/companion.txt
+++ b/companion.txt
@@ -18,6 +18,7 @@ https://github.com/willvousden/ptemcee/archive/master.tar.gz
--extra-index-url https://download.pytorch.org/whl/cpu
torch
nessai>=0.11.0
+snowline
# useful to look at PyCBC Live with htop
setproctitle
diff --git a/docs/inference/examples/sampler_platter.rst b/docs/inference/examples/sampler_platter.rst
index 963232188f1..35f4158da99 100644
--- a/docs/inference/examples/sampler_platter.rst
+++ b/docs/inference/examples/sampler_platter.rst
@@ -79,6 +79,21 @@ or an external package (multinest).
.. literalinclude:: ../../../examples/inference/samplers/multinest_stub.ini
:language: ini
+============================================================
+`Snowline `_
+============================================================
+
+.. literalinclude:: ../../../examples/inference/samplers/snowline_stub.ini
+ :language: ini
+
+
+============================================================
+`nessai `_
+============================================================
+
+.. literalinclude:: ../../../examples/inference/samplers/nessai_stub.ini
+ :language: ini
+
If we run these samplers, we create the following plot:
.. image:: ../../_include/sample.png
diff --git a/examples/inference/samplers/run.sh b/examples/inference/samplers/run.sh
index 91e41e7fbb7..bdeaefc8d6a 100755
--- a/examples/inference/samplers/run.sh
+++ b/examples/inference/samplers/run.sh
@@ -1,5 +1,5 @@
#!/bin/sh
-for f in cpnest_stub.ini emcee_stub.ini emcee_pt_stub.ini dynesty_stub.ini ultranest_stub.ini epsie_stub.ini nessai_stub.ini; do
+for f in cpnest_stub.ini emcee_stub.ini emcee_pt_stub.ini dynesty_stub.ini ultranest_stub.ini epsie_stub.ini nessai_stub.ini snowline_stub.ini; do
echo $f
pycbc_inference \
--config-files `dirname $0`/simp.ini `dirname $0`/$f \
@@ -17,6 +17,7 @@ ultranest_stub.ini.hdf:ultranest \
epsie_stub.ini.hdf:espie \
cpnest_stub.ini.hdf:cpnest \
nessai_stub.ini.hdf:nessai \
+snowline_stub.ini.hdf:snowline \
--output-file sample.png \
--plot-contours \
--plot-marginal \
diff --git a/examples/inference/samplers/snowline_stub.ini b/examples/inference/samplers/snowline_stub.ini
new file mode 100644
index 00000000000..ac17879fe74
--- /dev/null
+++ b/examples/inference/samplers/snowline_stub.ini
@@ -0,0 +1,6 @@
+[sampler]
+name = snowline
+
+num_gauss_samples = 500
+min_ess = 500
+max_improvement_loops = 3
diff --git a/pycbc/inference/io/__init__.py b/pycbc/inference/io/__init__.py
index 4b3fd0ce909..3157f4408f6 100644
--- a/pycbc/inference/io/__init__.py
+++ b/pycbc/inference/io/__init__.py
@@ -37,6 +37,7 @@
from .multinest import MultinestFile
from .dynesty import DynestyFile
from .ultranest import UltranestFile
+from .snowline import SnowlineFile
from .nessai import NessaiFile
from .posterior import PosteriorFile
from .txt import InferenceTXTFile
@@ -51,6 +52,7 @@
PosteriorFile.name: PosteriorFile,
UltranestFile.name: UltranestFile,
NessaiFile.name: NessaiFile,
+ SnowlineFile.name: SnowlineFile,
}
try:
diff --git a/pycbc/inference/io/snowline.py b/pycbc/inference/io/snowline.py
new file mode 100644
index 00000000000..4aa9edcbf13
--- /dev/null
+++ b/pycbc/inference/io/snowline.py
@@ -0,0 +1,32 @@
+# Copyright (C) 2024 Alex Nitz
+# This program is free software; you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the
+# Free Software Foundation; either version 3 of the License, or (at your
+# self.option) any later version.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
+# Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+
+
+#
+# =============================================================================
+#
+# Preamble
+#
+# =============================================================================
+#
+"""Provides IO for the snowline sampler.
+"""
+from .posterior import PosteriorFile
+
+
+class SnowlineFile(PosteriorFile):
+ """Class to handle file IO for the ``snowline`` sampler."""
+
+ name = 'snowline_file'
diff --git a/pycbc/inference/sampler/__init__.py b/pycbc/inference/sampler/__init__.py
index 41da16f39c6..a25573b7ee2 100644
--- a/pycbc/inference/sampler/__init__.py
+++ b/pycbc/inference/sampler/__init__.py
@@ -25,6 +25,7 @@
from .ultranest import UltranestSampler
from .dummy import DummySampler
from .refine import RefineSampler
+from .snowline import SnowlineSampler
# list of available samplers
samplers = {cls.name: cls for cls in (
@@ -32,6 +33,7 @@
UltranestSampler,
DummySampler,
RefineSampler,
+ SnowlineSampler,
)}
try:
diff --git a/pycbc/inference/sampler/snowline.py b/pycbc/inference/sampler/snowline.py
new file mode 100644
index 00000000000..6a54825879a
--- /dev/null
+++ b/pycbc/inference/sampler/snowline.py
@@ -0,0 +1,168 @@
+# Copyright (C) 2023 Alex Nitz
+# This program is free software; you can redistribute it and/or modify it
+# under the terms of the GNU General Public License as published by the
+# Free Software Foundation; either version 3 of the License, or (at your
+# option) any later version.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
+# Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+
+
+#
+# =============================================================================
+#
+# Preamble
+#
+# =============================================================================
+#
+"""
+This modules provides classes and functions for using the snowline sampler
+packages for parameter estimation.
+"""
+
+import sys
+import logging
+
+from pycbc.inference.io.snowline import SnowlineFile
+from pycbc.io.hdf import dump_state
+from pycbc.pool import use_mpi
+from .base import (BaseSampler, setup_output)
+from .base_cube import setup_calls
+
+
+#
+# =============================================================================
+#
+# Samplers
+#
+# =============================================================================
+#
+
+class SnowlineSampler(BaseSampler):
+ """This class is used to construct an Snowline sampler from the snowline
+ package.
+
+ Parameters
+ ----------
+ model : model
+ A model from ``pycbc.inference.models``
+ """
+ name = "snowline"
+ _io = SnowlineFile
+
+ def __init__(self, model, **kwargs):
+ super().__init__(model)
+
+ import snowline
+ log_likelihood_call, prior_call = setup_calls(model, copy_prior=True)
+
+ self._sampler = snowline.ReactiveImportanceSampler(
+ list(self.model.variable_params),
+ log_likelihood_call,
+ transform=prior_call)
+
+ do_mpi, _, rank = use_mpi()
+ self.main = (not do_mpi) or (rank == 0)
+ self.nlive = 0
+ self.ndim = len(self.model.variable_params)
+ self.kwargs = kwargs
+
+ def run(self):
+ self.result = self._sampler.run(**self.kwargs)
+ if not self.main:
+ sys.exit(0)
+ self._sampler.print_results()
+
+ @property
+ def io(self):
+ return self._io
+
+ @property
+ def niterations(self):
+ return self.result['niter']
+
+ @classmethod
+ def from_config(cls, cp, model, output_file=None, **kwds):
+ """
+ Loads the sampler from the given config file.
+ """
+ skeys = {}
+ opts = {'num_global_samples': int,
+ 'num_gauss_samples': int,
+ 'max_ncalls': int,
+ 'min_ess': int,
+ 'max_improvement_loops': int
+ }
+ for opt_name in opts:
+ if cp.has_option('sampler', opt_name):
+ value = cp.get('sampler', opt_name)
+ skeys[opt_name] = opts[opt_name](value)
+ inst = cls(model, **skeys)
+
+ do_mpi, _, rank = use_mpi()
+ if not do_mpi or (rank == 0):
+ setup_output(inst, output_file)
+ return inst
+
+ def checkpoint(self):
+ """ There is currently no checkpointing implemented"""
+ pass
+
+ def resume_from_checkpoint(self):
+ """ There is currently no checkpointing implemented"""
+ pass
+
+ def finalize(self):
+ logging.info("Writing samples to files")
+ for fn in [self.checkpoint_file, self.backup_file]:
+ self.write_results(fn)
+
+ @property
+ def model_stats(self):
+ return {}
+
+ @property
+ def samples(self):
+ samples = self.result['samples']
+ params = list(self.model.variable_params)
+ samples_dict = {p: samples[:, i] for i, p in enumerate(params)}
+ return samples_dict
+
+ def write_results(self, filename):
+ """Writes samples, model stats, acceptance fraction, and random state
+ to the given file.
+
+ Parameters
+ ----------
+ filename : str
+ The file to write to. The file is opened using the ``io`` class
+ in an an append state.
+ """
+ with self.io(filename, 'a') as fp:
+ # write samples
+ fp.write_samples(self.samples, self.samples.keys())
+ # write log evidence
+ fp.write_logevidence(self.logz, self.logz_err)
+
+ # write full results
+ dump_state(self.result, fp,
+ path='sampler_info',
+ dsetname='presult')
+
+ @property
+ def logz(self):
+ """Return bayesian evidence estimated by snowline sampler.
+ """
+ return self.result['logz']
+
+ @property
+ def logz_err(self):
+ """Return error in bayesian evidence estimated by snowline sampler.
+ """
+ return self.result['logzerr']
diff --git a/pycbc/pool.py b/pycbc/pool.py
index 9750b4f4c91..8f2095181ee 100644
--- a/pycbc/pool.py
+++ b/pycbc/pool.py
@@ -121,6 +121,12 @@ def broadcast(self, fcn, args):
def map(self, f, items):
return [f(a) for a in items]
+ # This is single core, so imap and map
+ # would not behave differently. This is defined
+ # so that the general pool interfaces can use
+ # imap irrespective of the pool type.
+ imap = map
+
def use_mpi(require_mpi=False, log=True):
""" Get whether MPI is enabled and if so the current size and rank
"""