add a test that new feature runs

joezuntz · Oct 30, 2024 · 9ad6a66 · 9ad6a66
1 parent 089ebe3
commit 9ad6a66
Showing 1 changed file with 28 additions and 12 deletions.
diff --git a/cosmosis/test/test_polychord.py b/cosmosis/test/test_polychord.py
@@ -2,7 +2,7 @@
 from cosmosis.samplers.sampler import Sampler
 import cosmosis.samplers.minuit.minuit_sampler
 from cosmosis.runtime.pipeline import LikelihoodPipeline
-from cosmosis.output.in_memory_output import InMemoryOutput
+from cosmosis.output import InMemoryOutput, TextColumnOutput
 from cosmosis.runtime.handler import activate_segfault_handling
 import tempfile
 import os
@@ -15,7 +15,7 @@
 
 minuit_compiled = os.path.exists(cosmosis.samplers.minuit.minuit_sampler.libname)
 
-def run(sampler, check_prior, check_extra=True, fast_slow=False, **options):
+def run(sampler, check_prior, inmem_output=True, check_extra=True, fast_slow=False, **options):
 
     sampler_class = Sampler.registry[sampler]
 
@@ -52,24 +52,35 @@ def run(sampler, check_prior, check_extra=True, fast_slow=False, **options):
     if pipeline.do_fast_slow:
         pipeline.setup_fast_subspaces()
 
-    output = InMemoryOutput()
+    if inmem_output:
+        output = InMemoryOutput()
+    else:
+        output = TextColumnOutput(options["base_dir"] + "/output.txt")
+
     sampler = sampler_class(ini, pipeline, output)
     sampler.config()
+    print("output = ", inmem_output, sampler.output, output)
 
 
     while not sampler.is_converged():
         sampler.execute()
 
-    if check_prior:
-        pr = np.array(output['prior'])
-        # but not all of them
-        assert not np.all(pr==-np.inf)
 
-    if check_extra:
-        p1 = output['parameters--p1']
-        p2 = output['parameters--p2']
-        p3 = output['PARAMETERS--P3']
-        assert np.all((p1+p2==p3)|(np.isnan(p3)))
+    if inmem_output:
+        if check_prior:
+            pr = np.array(output['prior'])
+            # but not all of them
+            assert not np.all(pr==-np.inf)
+
+        if check_extra:
+            p1 = output['parameters--p1']
+            p2 = output['parameters--p2']
+            p3 = output['PARAMETERS--P3']
+            assert np.all((p1+p2==p3)|(np.isnan(p3)))
+    else:
+        print(os.listdir(options["base_dir"]))
+        assert os.path.exists(options["base_dir"] + "/output.txt")
+        assert os.path.exists(options["base_dir"] + "/output_boosted.txt")
 
     return output
 
@@ -83,6 +94,11 @@ def test_polychord_hang():
         run('polychord', True, live_points=20, feedback=0, base_dir=base_dir, polychord_outfile_root='pc', fast_slow=True, fast_fraction=0.5)
 
 
+def test_polychord_boosting():
+    with tempfile.TemporaryDirectory() as base_dir:
+        run('polychord', True, inmem_output=False, live_points=20, boost_posteriors=2.0, feedback=0, base_dir=base_dir, polychord_outfile_root='pc', fast_slow=True, fast_fraction=0.5)
+
+
 if __name__ == '__main__':
     import sys
     locals()[sys.argv[1]]()