Add pycbc support for the snowline sampler

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]
 

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

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 <https://johannesbuchner.github.io/snowline/>`_
+============================================================
+
+.. literalinclude:: ../../../examples/inference/samplers/snowline_stub.ini
+   :language: ini
+
+
+============================================================
+`nessai <https://github.com/mj-will/nessai>`_
+============================================================
+
+.. literalinclude:: ../../../examples/inference/samplers/nessai_stub.ini
+   :language: ini
+
 If we run these samplers, we create the following plot:
 
 .. image:: ../../_include/sample.png

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 \

examples/inference/samplers/snowline_stub.ini

@@ -0,0 +1,6 @@
+[sampler]
+name = snowline
+
+num_gauss_samples = 500
+min_ess = 500
+max_improvement_loops = 3

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:

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'

pycbc/inference/sampler/__init__.py

@@ -25,13 +25,15 @@
 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 (
     MultinestSampler,
     UltranestSampler,
     DummySampler,
     RefineSampler,
+    SnowlineSampler,
 )}
 
 try:

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']

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
     """