test/core/sampler/dnest4_test.py

@@ -2,6 +2,7 @@ import unittest
 from unittest.mock import MagicMock
 
 import bilby
+import bilby.core.sampler.dnest4
 
 
 class TestDnest4(unittest.TestCase):
@@ -10,7 +11,7 @@ class TestDnest4(unittest.TestCase):
         self.priors = bilby.core.prior.PriorDict(
             dict(a=bilby.core.prior.Uniform(0, 1), b=bilby.core.prior.Uniform(0, 1))
         )
-        self.sampler = bilby.core.sampler.DNest4(
+        self.sampler = bilby.core.sampler.dnest4.DNest4(
             self.likelihood,
             self.priors,
             outdir="outdir",

test/core/sampler/dynamic_dynesty_test.py

@@ -2,6 +2,7 @@ import unittest
 from unittest.mock import MagicMock
 
 import bilby
+import bilby.core.sampler.dynamic_dynesty
 
 
 class TestDynamicDynesty(unittest.TestCase):
@@ -10,7 +11,7 @@ class TestDynamicDynesty(unittest.TestCase):
         self.priors = bilby.core.prior.PriorDict(
             dict(a=bilby.core.prior.Uniform(0, 1), b=bilby.core.prior.Uniform(0, 1))
         )
-        self.sampler = bilby.core.sampler.DynamicDynesty(
+        self.sampler = bilby.core.sampler.dynamic_dynesty.DynamicDynesty(
             self.likelihood,
             self.priors,
             outdir="outdir",

test/core/sampler/dynesty_test.py

@@ -5,6 +5,7 @@ from attr import define
 import bilby
 import numpy as np
 import parameterized
+import bilby.core.sampler.dynesty
 from bilby.core.sampler import dynesty_utils
 from scipy.stats import gamma, ks_1samp, uniform, powerlaw
 import shutil
@@ -41,7 +42,7 @@ class TestDynesty(unittest.TestCase):
         self.priors = bilby.core.prior.PriorDict(
             dict(a=bilby.core.prior.Uniform(0, 1), b=bilby.core.prior.Uniform(0, 1))
         )
-        self.sampler = bilby.core.sampler.Dynesty(
+        self.sampler = bilby.core.sampler.dynesty.Dynesty(
             self.likelihood,
             self.priors,
             outdir="outdir",
@@ -84,7 +85,7 @@ class TestDynesty(unittest.TestCase):
         self.priors["c"] = bilby.core.prior.Prior(boundary=None)
         self.priors["d"] = bilby.core.prior.Prior(boundary="reflective")
         self.priors["e"] = bilby.core.prior.Prior(boundary="periodic")
-        self.sampler = bilby.core.sampler.Dynesty(
+        self.sampler = bilby.core.sampler.dynesty.Dynesty(
             self.likelihood,
             self.priors,
             outdir="outdir",

test/core/sampler/emcee_test.py

@@ -2,6 +2,7 @@ import unittest
 from unittest.mock import MagicMock
 
 import bilby
+import bilby.core.sampler.emcee
 
 
 class TestEmcee(unittest.TestCase):
@@ -10,7 +11,7 @@ class TestEmcee(unittest.TestCase):
         self.priors = bilby.core.prior.PriorDict(
             dict(a=bilby.core.prior.Uniform(0, 1), b=bilby.core.prior.Uniform(0, 1))
         )
-        self.sampler = bilby.core.sampler.Emcee(
+        self.sampler = bilby.core.sampler.emcee.Emcee(
             self.likelihood,
             self.priors,
             outdir="outdir",

test/core/sampler/general_sampler_tests.py

+from bilby.core.sampler import (
+    get_implemented_samplers,
+    get_sampler_class,
+)
+import pytest
+
+
+def test_get_implemented_samplers():
+    """Assert the function returns a list of the correct length"""
+    from bilby.core.sampler import IMPLEMENTED_SAMPLERS
+
+    out = get_implemented_samplers()
+    assert isinstance(out, list)
+    assert len(out) == len(IMPLEMENTED_SAMPLERS)
+    assert "dynesty" in out
+
+
+def test_get_sampler_class():
+    """Assert the function returns the correct class"""
+    from bilby.core.sampler.dynesty import Dynesty
+
+    sampler_class = get_sampler_class("dynesty")
+    assert sampler_class is Dynesty
+
+
+def test_get_sampler_class_not_implemented():
+    """Assert an error is raised if the sampler is not recognized"""
+    with pytest.raises(
+        ValueError,
+        match=r"Sampler not_a_valid_sampler not yet implemented"
+    ):
+        get_sampler_class("not_a_valid_sampler")

test/core/sampler/kombine_test.py

@@ -2,6 +2,7 @@ import unittest
 from unittest.mock import MagicMock
 
 import bilby
+import bilby.core.sampler.kombine
 
 
 class TestKombine(unittest.TestCase):
@@ -10,7 +11,7 @@ class TestKombine(unittest.TestCase):
         self.priors = bilby.core.prior.PriorDict(
             dict(a=bilby.core.prior.Uniform(0, 1), b=bilby.core.prior.Uniform(0, 1))
         )
-        self.sampler = bilby.core.sampler.Kombine(
+        self.sampler = bilby.core.sampler.kombine.Kombine(
             self.likelihood,
             self.priors,
             outdir="outdir",

test/core/sampler/nessai_test.py

@@ -2,6 +2,7 @@ import unittest
 from unittest.mock import MagicMock, patch, mock_open
 
 import bilby
+import bilby.core.sampler.nessai
 
 
 class TestNessai(unittest.TestCase):
@@ -12,7 +13,7 @@ class TestNessai(unittest.TestCase):
         self.priors = bilby.core.prior.PriorDict(
             dict(a=bilby.core.prior.Uniform(0, 1), b=bilby.core.prior.Uniform(0, 1))
         )
-        self.sampler = bilby.core.sampler.Nessai(
+        self.sampler = bilby.core.sampler.nessai.Nessai(
             self.likelihood,
             self.priors,
             outdir="outdir",

test/core/sampler/nestle_test.py

@@ -2,6 +2,7 @@ import unittest
 from unittest.mock import MagicMock
 
 import bilby
+import bilby.core.sampler.nestle
 
 
 class TestNestle(unittest.TestCase):
@@ -10,7 +11,7 @@ class TestNestle(unittest.TestCase):
         self.priors = bilby.core.prior.PriorDict(
             dict(a=bilby.core.prior.Uniform(0, 1), b=bilby.core.prior.Uniform(0, 1))
         )
-        self.sampler = bilby.core.sampler.Nestle(
+        self.sampler = bilby.core.sampler.nestle.Nestle(
             self.likelihood,
             self.priors,
             outdir="outdir",

test/core/sampler/ptemcee_test.py

@@ -2,7 +2,7 @@ import unittest
 
 from bilby.core.likelihood import GaussianLikelihood
 from bilby.core.prior import Uniform, PriorDict
-from bilby.core.sampler import Ptemcee
+from bilby.core.sampler.ptemcee import Ptemcee
 from bilby.core.sampler.base_sampler import MCMCSampler
 import numpy as np
 

test/core/sampler/pymc_test.py

@@ -2,6 +2,7 @@ import unittest
 from unittest.mock import MagicMock
 
 import bilby
+import bilby.core.sampler.pymc
 
 
 class TestPyMC(unittest.TestCase):
@@ -10,7 +11,7 @@ class TestPyMC(unittest.TestCase):
         self.priors = bilby.core.prior.PriorDict(
             dict(a=bilby.core.prior.Uniform(0, 1), b=bilby.core.prior.Uniform(0, 1))
         )
-        self.sampler = bilby.core.sampler.Pymc(
+        self.sampler = bilby.core.sampler.pymc.Pymc(
             self.likelihood,
             self.priors,
             outdir="outdir",

test/core/sampler/pymultinest_test.py

@@ -2,6 +2,7 @@ import unittest
 from unittest.mock import MagicMock
 
 import bilby
+import bilby.core.sampler.pymultinest
 
 
 class TestPymultinest(unittest.TestCase):
@@ -12,7 +13,7 @@ class TestPymultinest(unittest.TestCase):
         )
         self.priors["a"] = bilby.core.prior.Prior(boundary="periodic")
         self.priors["b"] = bilby.core.prior.Prior(boundary="reflective")
-        self.sampler = bilby.core.sampler.Pymultinest(
+        self.sampler = bilby.core.sampler.pymultinest.Pymultinest(
             self.likelihood,
             self.priors,
             outdir="outdir",

test/core/sampler/ultranest_test.py

@@ -3,6 +3,7 @@ import unittest
 from unittest.mock import MagicMock
 
 import bilby
+import bilby.core.sampler.ultranest
 
 
 class TestUltranest(unittest.TestCase):
@@ -15,10 +16,15 @@ class TestUltranest(unittest.TestCase):
                  b=bilby.core.prior.Uniform(0, 1)))
         self.priors["a"] = bilby.core.prior.Prior(boundary="periodic")
         self.priors["b"] = bilby.core.prior.Prior(boundary="reflective")
-        self.sampler = bilby.core.sampler.Ultranest(self.likelihood, self.priors,
-                                                    outdir="outdir", label="label",
-                                                    use_ratio=False, plot=False,
-                                                    skip_import_verification=True)
+        self.sampler = bilby.core.sampler.ultranest.Ultranest(
+            self.likelihood,
+            self.priors,
+            outdir="outdir",
+            label="label",
+            use_ratio=False,
+            plot=False,
+            skip_import_verification=True,
+        )
 
     def tearDown(self):
         del self.likelihood

test/core/sampler/zeus_test.py

@@ -2,6 +2,7 @@ import unittest
 from unittest.mock import MagicMock
 
 import bilby
+import bilby.core.sampler.zeus
 
 
 class TestZeus(unittest.TestCase):
@@ -10,7 +11,7 @@ class TestZeus(unittest.TestCase):
         self.priors = bilby.core.prior.PriorDict(
             dict(a=bilby.core.prior.Uniform(0, 1), b=bilby.core.prior.Uniform(0, 1))
         )
-        self.sampler = bilby.core.sampler.Zeus(
+        self.sampler = bilby.core.sampler.zeus.Zeus(
             self.likelihood,
             self.priors,
             outdir="outdir",

test/integration/sampler_run_test.py

@@ -66,6 +66,8 @@ sampler_imports = dict(
 
 no_pool_test = ["dnest4", "pymultinest", "nestle", "ptmcmcsampler", "ultranest", "pymc"]
 
+loaded_samplers = {k: v.load() for k, v in bilby.core.sampler.IMPLEMENTED_SAMPLERS.items()}
+
 
 def slow_func(x, m, c):
     time.sleep(0.01)
@@ -154,7 +156,7 @@ class TestRunningSamplers(unittest.TestCase):
 
     def _run_with_signal_handling(self, sampler, pool_size=1):
         pytest.importorskip(sampler_imports.get(sampler, sampler))
-        if bilby.core.sampler.IMPLEMENTED_SAMPLERS[sampler.lower()].hard_exit:
+        if loaded_samplers[sampler.lower()].hard_exit:
             pytest.skip(f"{sampler} hard exits, can't test signal handling.")
         if pool_size > 1 and sampler.lower() in no_pool_test:
             pytest.skip(f"{sampler} cannot be parallelized")

test/test_samplers_import.py

-"""
-Tests that all of the implemented samplers can be initialized.
-
-The :code:`FakeSampler` is omitted as that doesn't require importing
-any package.
-"""
 import bilby
+import pytest
+
+
+@pytest.mark.parametrize(
+    "sampler_name", bilby.core.sampler.IMPLEMENTED_SAMPLERS.keys()
+)
+def test_sampler_import(sampler_name):
+    """
+    Tests that all of the implemented samplers can be initialized.
 
-bilby.core.utils.logger.setLevel("ERROR")
-IMPLEMENTED_SAMPLERS = bilby.core.sampler.IMPLEMENTED_SAMPLERS
-likelihood = bilby.core.likelihood.Likelihood(dict())
-priors = bilby.core.prior.PriorDict(dict(a=bilby.core.prior.Uniform(0, 1)))
-for sampler in IMPLEMENTED_SAMPLERS:
-    if sampler in ["fake_sampler", "pypolychord"]:
-        continue
-    sampler_class = IMPLEMENTED_SAMPLERS[sampler]
+    Do not test :code:`FakeSampler` since it requires an additional argument.
+    """
+    if sampler_name in ["fake_sampler", "pypolychord"]:
+        pytest.skip(f"Skipping import test for {sampler_name}")
+    bilby.core.utils.logger.setLevel("ERROR")
+    likelihood = bilby.core.likelihood.Likelihood(dict())
+    priors = bilby.core.prior.PriorDict(dict(a=bilby.core.prior.Uniform(0, 1)))
+    sampler_class = bilby.core.sampler.IMPLEMENTED_SAMPLERS[sampler_name].load()
     sampler = sampler_class(likelihood=likelihood, priors=priors)
+    assert sampler is not None