diff --git a/bilby/bilby_mcmc/sampler.py b/bilby/bilby_mcmc/sampler.py index fcd4c9c5aa82d02334b1ffe02e5fbe501d4e14b8..a3732add081bacf3c2174f8c8df6dc5dcc414e09 100644 --- a/bilby/bilby_mcmc/sampler.py +++ b/bilby/bilby_mcmc/sampler.py @@ -547,6 +547,29 @@ class Bilby_MCMC(MCMCSampler): all_samples=ptsampler.samples, ) + @classmethod + def get_expected_outputs(cls, outdir=None, label=None): + """Get lists of the expected outputs directories and files. + + These are used by :code:`bilby_pipe` when transferring files via HTCondor. + + Parameters + ---------- + outdir : str + The output directory. + label : str + The label for the run. + + Returns + ------- + list + List of file names. + list + List of directory names. Will always be empty for bilby_mcmc. + """ + filenames = [os.path.join(outdir, f"{label}_resume.pickle")] + return filenames, [] + class BilbyPTMCMCSampler(object): def __init__( diff --git a/bilby/core/sampler/base_sampler.py b/bilby/core/sampler/base_sampler.py index 3f73cb795840926e3140796656294d51b63f4cc3..3c573afe2b1ec3e5d73d8b035fab2a83767f4baa 100644 --- a/bilby/core/sampler/base_sampler.py +++ b/bilby/core/sampler/base_sampler.py @@ -173,6 +173,13 @@ class Sampler(object): Whether the implemented sampler exits hard (:code:`os._exit` rather than :code:`sys.exit`). The latter can be escaped as :code:`SystemExit`. The former cannot. + sampler_name : str + Name of the sampler. This is used when creating the output directory for + the sampler. + abbreviation : str + Abbreviated name of the sampler. Does not have to be specified in child + classes. If set to a value other than :code:`None`, this will be used + instead of :code:`sampler_name` when creating the output directory. Raises ====== @@ -187,6 +194,8 @@ class Sampler(object): """ + sampler_name = "sampler" + abbreviation = None default_kwargs = dict() npool_equiv_kwargs = [ "npool", @@ -779,8 +788,37 @@ class Sampler(object): def write_current_state(self): raise NotImplementedError() + @classmethod + def get_expected_outputs(cls, outdir=None, label=None): + """Get lists of the expected outputs directories and files. + + These are used by :code:`bilby_pipe` when transferring files via HTCondor. + Both can be empty. Defaults to a single directory: + :code:`"{outdir}/{name}_{label}/"`, where :code:`name` + is :code:`abbreviation` if it is defined for the sampler class, otherwise + it defaults to :code:`sampler_name`. + + Parameters + ---------- + outdir : str + The output directory. + label : str + The label for the run. + + Returns + ------- + list + List of file names. + list + List of directory names. + """ + name = cls.abbreviation or cls.sampler_name + dirname = os.path.join(outdir, f"{name}_{label}", "") + return [], [dirname] + class NestedSampler(Sampler): + sampler_name = "nested_sampler" npoints_equiv_kwargs = [ "nlive", "nlives", @@ -854,6 +892,7 @@ class NestedSampler(Sampler): class MCMCSampler(Sampler): + sampler_name = "mcmc_sampler" nwalkers_equiv_kwargs = ["nwalker", "nwalkers", "draws", "Niter"] nburn_equiv_kwargs = ["burn", "nburn"] diff --git a/bilby/core/sampler/cpnest.py b/bilby/core/sampler/cpnest.py index e1f3ae19e39f37854643aa86aaac9beece78465d..e777ebc67a5143fa8f41469b727df748f54735e1 100644 --- a/bilby/core/sampler/cpnest.py +++ b/bilby/core/sampler/cpnest.py @@ -40,6 +40,7 @@ class Cpnest(NestedSampler): """ + sampler_name = "cpnest" default_kwargs = dict( verbose=3, nthreads=1, diff --git a/bilby/core/sampler/dnest4.py b/bilby/core/sampler/dnest4.py index a767ef89d47a013b3cba9dc3360c903ee5552832..87717f6fd6346b7c89548a5b579b6e3c0f365141 100644 --- a/bilby/core/sampler/dnest4.py +++ b/bilby/core/sampler/dnest4.py @@ -99,6 +99,7 @@ class DNest4(_TemporaryFileSamplerMixin, NestedSampler): If True, prints information during run """ + sampler_name = "d4nest" default_kwargs = dict( max_num_levels=20, num_steps=500, diff --git a/bilby/core/sampler/dynamic_dynesty.py b/bilby/core/sampler/dynamic_dynesty.py index 294d8fd6d9ceee2e878ba0239cf64bb5fe8c94cf..8c7f2966ee6ea4e8934ade73b6ccd2582c019f1b 100644 --- a/bilby/core/sampler/dynamic_dynesty.py +++ b/bilby/core/sampler/dynamic_dynesty.py @@ -14,6 +14,7 @@ class DynamicDynesty(Dynesty): """ external_sampler_name = "dynesty" + sampler_name = "dynamic_dynesty" @property def nlive(self): diff --git a/bilby/core/sampler/dynesty.py b/bilby/core/sampler/dynesty.py index faebfc6bfad1c11a1f116c0d6780c3cfce60aa2a..852fb88c192195c7fb08d7bd9654dda9126e7ca5 100644 --- a/bilby/core/sampler/dynesty.py +++ b/bilby/core/sampler/dynesty.py @@ -151,6 +151,7 @@ class Dynesty(NestedSampler): specified. """ + sampler_name = "dynesty" sampling_seed_key = "seed" @property @@ -299,6 +300,32 @@ class Dynesty(NestedSampler): ) Sampler._verify_kwargs_against_default_kwargs(self) + @classmethod + def get_expected_outputs(cls, outdir=None, label=None): + """Get lists of the expected outputs directories and files. + + These are used by :code:`bilby_pipe` when transferring files via HTCondor. + + Parameters + ---------- + outdir : str + The output directory. + label : str + The label for the run. + + Returns + ------- + list + List of file names. + list + List of directory names. Will always be empty for dynesty. + """ + filenames = [] + for kind in ["resume", "dynesty"]: + filename = os.path.join(outdir, f"{label}_{kind}.pickle") + filenames.append(filename) + return filenames, [] + def _print_func( self, results, diff --git a/bilby/core/sampler/emcee.py b/bilby/core/sampler/emcee.py index 7253a0fa4c49affc236454e703310d55f45fc066..db88ee5a2207f85616f081aa3ee31585c5ad2da1 100644 --- a/bilby/core/sampler/emcee.py +++ b/bilby/core/sampler/emcee.py @@ -45,6 +45,7 @@ class Emcee(MCMCSampler): """ + sampler_name = "emcee" default_kwargs = dict( nwalkers=500, a=2, diff --git a/bilby/core/sampler/fake_sampler.py b/bilby/core/sampler/fake_sampler.py index 5f375fdbad8055e6a5bdaf7dd7e99caabe330f33..9795631fb4951398bfaf474d2d338ae55d94d997 100644 --- a/bilby/core/sampler/fake_sampler.py +++ b/bilby/core/sampler/fake_sampler.py @@ -17,6 +17,8 @@ class FakeSampler(Sampler): A string pointing to the posterior data file to be loaded. """ + sampler_name = "fake_sampler" + default_kwargs = dict( verbose=True, logl_args=None, logl_kwargs=None, print_progress=True ) diff --git a/bilby/core/sampler/kombine.py b/bilby/core/sampler/kombine.py index 1f09387cc33520a7c8408db7cd7af2a924fa85cf..bda7c6d4f06686bf48c142acf726a145da05c760 100644 --- a/bilby/core/sampler/kombine.py +++ b/bilby/core/sampler/kombine.py @@ -39,6 +39,7 @@ class Kombine(Emcee): """ + sampler_name = "kombine" default_kwargs = dict( nwalkers=500, args=[], diff --git a/bilby/core/sampler/nessai.py b/bilby/core/sampler/nessai.py index b6f40f8aaca0c9e0bfe8f3a361aa54b1e2be4e60..65a650efd33f99fbbff4211c96c2b4647e53a5ae 100644 --- a/bilby/core/sampler/nessai.py +++ b/bilby/core/sampler/nessai.py @@ -20,6 +20,7 @@ class Nessai(NestedSampler): Documentation: https://nessai.readthedocs.io/ """ + sampler_name = "nessai" _default_kwargs = None _run_kwargs_list = None sampling_seed_key = "seed" @@ -300,5 +301,30 @@ class Nessai(NestedSampler): self._log_interruption(signum=signum) sys.exit(self.exit_code) + @classmethod + def get_expected_outputs(cls, outdir=None, label=None): + """Get lists of the expected outputs directories and files. + + These are used by :code:`bilby_pipe` when transferring files via HTCondor. + + Parameters + ---------- + outdir : str + The output directory. + label : str + The label for the run. + + Returns + ------- + list + List of file names. This will be empty for nessai. + list + List of directory names. + """ + dirs = [os.path.join(outdir, f"{label}_{cls.sampler_name}", "")] + dirs += [os.path.join(dirs[0], d, "") for d in ["proposal", "diagnostics"]] + filenames = [] + return filenames, dirs + def _setup_pool(self): pass diff --git a/bilby/core/sampler/nestle.py b/bilby/core/sampler/nestle.py index ebd955376050c80bcb43277ff95688c6fb572be8..75d93bf69b242463fbfce9ce64011e5864796185 100644 --- a/bilby/core/sampler/nestle.py +++ b/bilby/core/sampler/nestle.py @@ -24,6 +24,7 @@ class Nestle(NestedSampler): """ + sampler_name = "nestle" default_kwargs = dict( verbose=True, method="multi", diff --git a/bilby/core/sampler/polychord.py b/bilby/core/sampler/polychord.py index e43c5d50b248ba0fb12cd8d5bca97b0fee726c45..9391dd2026d2256df2cb9b89efc870c6a40b557c 100644 --- a/bilby/core/sampler/polychord.py +++ b/bilby/core/sampler/polychord.py @@ -1,3 +1,5 @@ +import os + import numpy as np from .base_sampler import NestedSampler, signal_wrapper @@ -20,6 +22,7 @@ class PyPolyChord(NestedSampler): To see what the keyword arguments are for, see the docstring of PyPolyChordSettings """ + sampler_name = "pypolychord" default_kwargs = dict( use_polychord_defaults=False, nlive=None, @@ -130,6 +133,28 @@ class PyPolyChord(NestedSampler): physical_parameters = samples[:, -self.ndim :] return log_likelihoods, physical_parameters + @classmethod + def get_expected_outputs(cls, outdir=None, label=None): + """Get lists of the expected outputs directories and files. + + These are used by :code:`bilby_pipe` when transferring files via HTCondor. + + Parameters + ---------- + outdir : str + The output directory. + label : str + Ignored for pypolychord. + + Returns + ------- + list + List of file names. This will always be empty for pypolychord. + list + List of directory names. + """ + return [], [os.path.join(outdir, "chains", "")] + @property def _sample_file_directory(self): return self.outdir + "/chains" diff --git a/bilby/core/sampler/ptemcee.py b/bilby/core/sampler/ptemcee.py index 531fb102a0614015f4053e7326b87c5ec923dbfe..fd927235d370bbb0800ea805b64f8e9e495c9b78 100644 --- a/bilby/core/sampler/ptemcee.py +++ b/bilby/core/sampler/ptemcee.py @@ -128,6 +128,7 @@ class Ptemcee(MCMCSampler): """ + sampler_name = "ptemcee" # Arguments used by ptemcee default_kwargs = dict( ntemps=10, @@ -710,6 +711,29 @@ class Ptemcee(MCMCSampler): except Exception as e: logger.info(f"mean_logl plot failed with exception {e}") + @classmethod + def get_expected_outputs(cls, outdir=None, label=None): + """Get lists of the expected outputs directories and files. + + These are used by :code:`bilby_pipe` when transferring files via HTCondor. + + Parameters + ---------- + outdir : str + The output directory. + label : str + The label for the run. + + Returns + ------- + list + List of file names. + list + List of directory names. Will always be empty for ptemcee. + """ + filenames = [f"{outdir}/{label}_checkpoint_resume.pickle"] + return filenames, [] + def get_minimum_stable_itertion(mean_array, frac, nsteps_min=10): nsteps = mean_array.shape[1] diff --git a/bilby/core/sampler/ptmcmc.py b/bilby/core/sampler/ptmcmc.py index 42279e018ed124cd117118e75949b60d74e3a302..f2a771cb0e24758b977926760e820f096268ba43 100644 --- a/bilby/core/sampler/ptmcmc.py +++ b/bilby/core/sampler/ptmcmc.py @@ -41,6 +41,8 @@ class PTMCMCSampler(MCMCSampler): """ + sampler_name = "ptmcmcsampler" + abbreviation = "ptmcmc_temp" default_kwargs = { "p0": None, "Niter": 2 * 10**4 + 1, diff --git a/bilby/core/sampler/pymc.py b/bilby/core/sampler/pymc.py index e72aace4948c8b46e6b4d68e74940d378f42ff99..fd138b985e7a2e5628066426987704fca1625d6a 100644 --- a/bilby/core/sampler/pymc.py +++ b/bilby/core/sampler/pymc.py @@ -52,6 +52,7 @@ class Pymc(MCMCSampler): """ + sampler_name = "pymc" default_kwargs = dict( draws=500, step=None, diff --git a/bilby/core/sampler/pymultinest.py b/bilby/core/sampler/pymultinest.py index 303acb705d3ed3d912ded1a132ebb2f4e138e8fe..0a9bb0aaf7c00bb376f377be3d652b1e5154d4be 100644 --- a/bilby/core/sampler/pymultinest.py +++ b/bilby/core/sampler/pymultinest.py @@ -34,6 +34,8 @@ class Pymultinest(_TemporaryFileSamplerMixin, NestedSampler): """ + sampler_name = "pymultinest" + abbreviation = "pm" default_kwargs = dict( importance_nested_sampling=False, resume=True, diff --git a/bilby/core/sampler/ultranest.py b/bilby/core/sampler/ultranest.py index 542f862468c290d007f556ca6033bb30e78f2729..6aacfa99999b00e13dfafa4f69cd3657af1f6ec5 100644 --- a/bilby/core/sampler/ultranest.py +++ b/bilby/core/sampler/ultranest.py @@ -38,6 +38,8 @@ class Ultranest(_TemporaryFileSamplerMixin, NestedSampler): stepping behaviour is used. """ + sampler_name = "ultranest" + abbreviation = "ultra" default_kwargs = dict( resume=True, show_status=True, diff --git a/bilby/core/sampler/zeus.py b/bilby/core/sampler/zeus.py index c7ae40da222201e5b29c53635c16c3edc94744f0..ad6e7edb8b1c91c694ab0a3e06382c3864e5fb5e 100644 --- a/bilby/core/sampler/zeus.py +++ b/bilby/core/sampler/zeus.py @@ -38,6 +38,7 @@ class Zeus(Emcee): """ + sampler_name = "zeus" default_kwargs = dict( nwalkers=500, args=[], diff --git a/test/bilby_mcmc/test_sampler.py b/test/bilby_mcmc/test_sampler.py index 746eb1a9e1150732e93d5c31664751040e7b639c..7e636e1ab07c06449a254433ded0014a210f2e19 100644 --- a/test/bilby_mcmc/test_sampler.py +++ b/test/bilby_mcmc/test_sampler.py @@ -85,5 +85,16 @@ class TestBilbyMCMCSampler(unittest.TestCase): self.assertTrue(isinstance(sampler.samples, pd.DataFrame)) +def test_get_expected_outputs(): + label = "par0" + outdir = os.path.join("some", "bilby_pipe", "dir") + filenames, directories = Bilby_MCMC.get_expected_outputs( + outdir=outdir, label=label + ) + assert len(filenames) == 1 + assert len(directories) == 0 + assert os.path.join(outdir, f"{label}_resume.pickle") in filenames + + if __name__ == "__main__": unittest.main() diff --git a/test/core/sampler/base_sampler_test.py b/test/core/sampler/base_sampler_test.py index 1250fa0d682d126f8d9881c8ef51a254479fe534..d20ee978a3bb1fdbd88bc4656ae66690a2ed2ae4 100644 --- a/test/core/sampler/base_sampler_test.py +++ b/test/core/sampler/base_sampler_test.py @@ -148,6 +148,30 @@ class TestSampler(unittest.TestCase): ) +def test_get_expected_outputs(): + outdir = os.path.join("some", "bilby_pipe", "dir") + label = "par0" + filenames, directories = bilby.core.sampler.Sampler.get_expected_outputs( + outdir=outdir, label=label + ) + assert len(filenames) == 0 + assert len(directories) == 1 + assert directories[0] == os.path.join(outdir, f"sampler_{label}", "") + + +def test_get_expected_outputs_abbreviation(): + outdir = os.path.join("some", "bilby_pipe", "dir") + label = "par0" + bilby.core.sampler.Sampler.abbreviation = "abbr" + filenames, directories = bilby.core.sampler.Sampler.get_expected_outputs( + outdir=outdir, label=label + ) + assert len(filenames) == 0 + assert len(directories) == 1 + assert directories[0] == os.path.join(outdir, f"abbr_{label}", "") + bilby.core.sampler.Sampler.abbreviation = None + + samplers = [ "bilby_mcmc", "dynamic_dynesty", diff --git a/test/core/sampler/dynesty_test.py b/test/core/sampler/dynesty_test.py index 39ac6f2318fd3e4d7f7d2fa0ac2b433126cdfb26..d33cc2e23d9f4219f3cb114688e1d5eff9af3ad7 100644 --- a/test/core/sampler/dynesty_test.py +++ b/test/core/sampler/dynesty_test.py @@ -9,6 +9,7 @@ import bilby.core.sampler.dynesty from bilby.core.sampler import dynesty_utils from scipy.stats import gamma, ks_1samp, uniform, powerlaw import shutil +import os @define @@ -101,6 +102,18 @@ class TestDynesty(unittest.TestCase): self.sampler._run_test() +def test_get_expected_outputs(): + label = "par0" + outdir = os.path.join("some", "bilby_pipe", "dir") + filenames, directories = bilby.core.sampler.dynesty.Dynesty.get_expected_outputs( + outdir=outdir, label=label + ) + assert len(filenames) == 2 + assert len(directories) == 0 + assert os.path.join(outdir, f"{label}_resume.pickle") in filenames + assert os.path.join(outdir, f"{label}_dynesty.pickle") in filenames + + class ProposalsTest(unittest.TestCase): def test_boundaries(self): diff --git a/test/core/sampler/nessai_test.py b/test/core/sampler/nessai_test.py index cca5d22b035480e3b46eb7db7730ef86d82f0239..3246c74e752a060725b0a0accfcb08fc9ebe345d 100644 --- a/test/core/sampler/nessai_test.py +++ b/test/core/sampler/nessai_test.py @@ -3,6 +3,7 @@ from unittest.mock import MagicMock, patch, mock_open import bilby import bilby.core.sampler.nessai +import os class TestNessai(unittest.TestCase): @@ -84,5 +85,19 @@ class TestNessai(unittest.TestCase): self.assertDictEqual(expected, self.sampler.kwargs) +def test_get_expected_outputs(): + label = "par0" + outdir = os.path.join("some", "bilby_pipe", "dir") + filenames, directories = bilby.core.sampler.nessai.Nessai.get_expected_outputs( + outdir=outdir, label=label + ) + assert len(filenames) == 0 + assert len(directories) == 3 + base_dir = os.path.join(outdir, f"{label}_nessai", "") + assert base_dir in directories + assert os.path.join(base_dir, "proposal", "") in directories + assert os.path.join(base_dir, "diagnostics", "") in directories + + if __name__ == "__main__": unittest.main() diff --git a/test/core/sampler/ptemcee_test.py b/test/core/sampler/ptemcee_test.py index ec135eeef09bf5bdaa0fb3c7ff87b8fd7ea536da..4708a12b09bbb310a708727e0482b6b80d415633 100644 --- a/test/core/sampler/ptemcee_test.py +++ b/test/core/sampler/ptemcee_test.py @@ -5,6 +5,7 @@ from bilby.core.prior import Uniform, PriorDict from bilby.core.sampler.ptemcee import Ptemcee from bilby.core.sampler.base_sampler import MCMCSampler import numpy as np +import os class TestPTEmcee(unittest.TestCase): @@ -89,5 +90,16 @@ class TestPTEmcee(unittest.TestCase): self.assertEqual(old, new) +def test_get_expected_outputs(): + label = "par0" + outdir = os.path.join("some", "bilby_pipe", "dir") + filenames, directories = Ptemcee.get_expected_outputs( + outdir=outdir, label=label + ) + assert len(filenames) == 1 + assert len(directories) == 0 + assert os.path.join(outdir, f"{label}_checkpoint_resume.pickle") in filenames + + if __name__ == "__main__": unittest.main()