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