Add LikelihoodPipeline.from_chain_file

bin/cosmosis-extract

@@ -1,45 +1,23 @@
 #!/usr/bin/env python
 
 import argparse
+from cosmosis.utils import read_chain_header,  extract_inis_from_chain_header
 
 parser = argparse.ArgumentParser("Extract the input parameter files that were used to generate a cosmosis chain from the output.")
 parser.add_argument("chain", help="Name of the chain file to read")
 parser.add_argument("prefix", help="Prefix for output files {prefix}_params.ini, {prefix}_values.ini, {prefix}_priors.ini")
 
-def read_comment_section(filename):
-    lines = []
-    for line in open(filename):
-        if not line.startswith('#'):
-            break
-        lines.append(line)
-    return lines
-
-
-
-def extract_section(lines, section):
-    start = "## START_OF_{}_INI".format(section).upper()
-    end = "## END_OF_{}_INI".format(section).upper()
-    in_section = False
-    output_lines = []
-    for line in lines:
-        if line.startswith(start):
-            in_section = True
-            continue
-        elif line.startswith(end):
-            break
-        elif in_section:
-            output_lines.append(line[3:])
-    return output_lines
+
 
 def save(lines, section, prefix):
     filename = "{}_{}.ini".format(prefix, section)
     open(filename,'w').writelines(lines)
 
 def main(chain, prefix):
-    lines = read_comment_section(chain)
+    lines = read_chain_header(chain)
 
     for section in ['params', 'values', 'priors']:
-        section_lines = extract_section(lines, section)
+        section_lines = extract_inis_from_chain_header(lines, section)
         save(section_lines, section, prefix)
 
 if __name__ == '__main__':

cosmosis/runtime/config.py

@@ -128,6 +128,8 @@ def __init__(self, filename, defaults=None, override=None, print_include_message
             filename.write(s)
             s.seek(0)
             self.read_file(s)
+        elif hasattr(filename, "read"):
+            self.read_file(filename)
         # default read behaviour is to ignore unreadable files which
         # is probably not what we want here
         elif filename is not None:
@@ -145,7 +147,11 @@ def __init__(self, filename, defaults=None, override=None, print_include_message
                         self.add_section(section)
                     self.set(section, name, override[(section, name)])
 
-
+    @classmethod
+    def from_lines(cls, lines, *args, **kwargs):
+        u"""Create an Inifile from a list of lines."""
+        s = io.StringIO("\n".join(lines))
+        return cls(s, *args, **kwargs)
 
     def __iter__(self):
         u"""Iterate over all the parameters.

cosmosis/runtime/pipeline.py

@@ -10,6 +10,7 @@
 import collections
 import warnings
 import traceback
+import io
 from . import config
 from . import parameter
 from . import prior
@@ -814,7 +815,26 @@ def __init__(self, arg=None, id="", override=None, modules=None, load=True, valu
         else:
             self.likelihood_names = likelihood_names.split()
 
-
+    @classmethod
+    def from_chain_file(cls, filename, **kwargs):
+        from ..utils import extract_inis_from_chain_header, read_chain_header
+        # Pull out all of the comment bits in the header of the chain
+        # that start with a #
+        header = read_chain_header(filename)
+
+        # Parse he header to pull out the three chunks of INI files
+        # that we save there
+        param_lines = extract_inis_from_chain_header(header, "params")
+        value_lines = extract_inis_from_chain_header(header, "values")
+        prior_lines = extract_inis_from_chain_header(header, "priors")
+
+        # convert all these into Inifile objects
+        params = config.Inifile.from_lines(param_lines)
+        values = config.Inifile.from_lines(value_lines)
+        priors = config.Inifile.from_lines(prior_lines)
+
+        # Build the pipeline from these
+        return cls(arg=params, values=values, priors=priors, **kwargs)
 
 
     def print_priors(self):

cosmosis/test/test_pipeline.py

@@ -305,6 +305,62 @@ def log_like(p):
         assert lines[1].strip() == "-2.0 0.0 0.0 8.0"
 
 
+def test_recreate_pipeline():
+    with tempfile.TemporaryDirectory() as dirname:
+
+        values = os.path.join(dirname, "values.ini")
+        priors = os.path.join(dirname, "priors.ini")
+        output = os.path.join(dirname, "output.txt")
+        with open(values, "w") as f:
+            f.write(
+                "[parameters]\n"
+                "p1=-10.0  0.0  10.0\n"
+                "p2=-1000.0  0.0  1000.0\n")
+
+        with open(priors, "w") as f:
+            f.write(
+                "[parameters]\n"
+                "p1=uniform -5.0 5.0\n"
+                "p2=gaussian 0.0 1.0"
+            )
+
+        override = {
+            ('runtime', 'root'): root,
+            ('runtime', 'sampler'): "emcee",
+            ("pipeline", "debug"): "F",
+            ("pipeline", "modules"): "test2",
+            ("pipeline", "values"): values,
+            ("pipeline", "priors"): priors,
+            ("pipeline", "extra_output"): "parameters/p3",
+            ("output", "filename"): output,
+            ("output", "format"): "text",
+            ("test2", "file"): "example_module.py",
+            ("emcee", "walkers"): "8",
+            ("emcee", "samples"): "10"
+        }
+
+        ini = Inifile(None, override=override)
+
+        status = run_cosmosis(ini)
+
+        # Now we want to recreate the pipeline from the output
+        pipeline = LikelihoodPipeline.from_chain_file(output)
+
+        # check the basic pipeline configuration
+        assert pipeline.modules[0].name == "test2"
+        assert len(pipeline.modules) == 1
+
+        # check it produces the same results
+        r = pipeline.run_results([1.,2.])
+        assert np.isclose(r.like, -2.5)
+        assert np.isclose(r.extra[0], 3.0)
+
+        # check that the priors have been correctly passed through
+        # the recreated pipeline
+        assert np.isclose(r.prior, np.log(0.1) - 2.0 - 0.5*np.log(2*np.pi))
+
+
+
 
 if __name__ == '__main__':
     test_script_skip()

cosmosis/utils.py

@@ -402,3 +402,28 @@ def under_over_line(s, char='-'):
         The character to use for the line
     """
     return underline(overline(s, char), char)
+
+
+def read_chain_header(filename):
+    lines = []
+    for line in open(filename):
+        if not line.startswith('#'):
+            break
+        lines.append(line)
+    return lines
+
+
+def extract_inis_from_chain_header(lines, section):
+    start = "## START_OF_{}_INI".format(section).upper()
+    end = "## END_OF_{}_INI".format(section).upper()
+    in_section = False
+    output_lines = []
+    for line in lines:
+        if line.startswith(start):
+            in_section = True
+            continue
+        elif line.startswith(end):
+            break
+        elif in_section:
+            output_lines.append(line[3:])
+    return output_lines