ENH: add get_expected_outputs to sampler classes

bilby/bilby_mcmc/sampler.py

@@ -547,6 +547,29 @@ def plot_progress(ptsampler, label, outdir, priors, diagnostic=False):
                         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__(

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 @@ def _setup_pool(self):
     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 @@ def log_likelihood(self, theta):
 
 
 class MCMCSampler(Sampler):
+    sampler_name = "mcmc_sampler"
     nwalkers_equiv_kwargs = ["nwalker", "nwalkers", "draws", "Niter"]
     nburn_equiv_kwargs = ["burn", "nburn"]
 

bilby/core/sampler/cpnest.py

@@ -40,6 +40,7 @@ class Cpnest(NestedSampler):
 
     """
 
+    sampler_name = "cpnest"
     default_kwargs = dict(
         verbose=3,
         nthreads=1,

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,

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):

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 @@ def _verify_kwargs_against_default_kwargs(self):
                 )
         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,

bilby/core/sampler/emcee.py

@@ -45,6 +45,7 @@ class Emcee(MCMCSampler):
 
     """
 
+    sampler_name = "emcee"
     default_kwargs = dict(
         nwalkers=500,
         a=2,

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
     )

bilby/core/sampler/kombine.py

@@ -39,6 +39,7 @@ class Kombine(Emcee):
 
     """
 
+    sampler_name = "kombine"
     default_kwargs = dict(
         nwalkers=500,
         args=[],

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 @@ def write_current_state_and_exit(self, signum=None, frame=None):
         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

bilby/core/sampler/nestle.py

@@ -24,6 +24,7 @@ class Nestle(NestedSampler):
 
     """
 
+    sampler_name = "nestle"
     default_kwargs = dict(
         verbose=True,
         method="multi",

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 @@ def _read_sample_file(self):
         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"

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 @@ def write_current_state(self, plot=True):
             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]

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,

bilby/core/sampler/pymc.py

@@ -52,6 +52,7 @@ class for further help. Under Other Parameters, we list commonly used
 
     """
 
+    sampler_name = "pymc"
     default_kwargs = dict(
         draws=500,
         step=None,

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,

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,

bilby/core/sampler/zeus.py

@@ -38,6 +38,7 @@ class Zeus(Emcee):
 
     """
 
+    sampler_name = "zeus"
     default_kwargs = dict(
         nwalkers=500,
         args=[],

test/bilby_mcmc/test_sampler.py

@@ -85,5 +85,16 @@ def test_default_proposal_cycle(self):
         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()