From ba51ada203549c8de7beffea04cc198d32973451 Mon Sep 17 00:00:00 2001 From: Joe Zuntz Date: Wed, 20 Nov 2024 12:10:24 +0000 Subject: [PATCH] Add LikelihoodPipeline.from_chain_file --- bin/cosmosis-extract | 30 +++--------------- cosmosis/runtime/config.py | 8 ++++- cosmosis/runtime/pipeline.py | 22 ++++++++++++- cosmosis/test/test_pipeline.py | 56 ++++++++++++++++++++++++++++++++++ cosmosis/utils.py | 25 +++++++++++++++ 5 files changed, 113 insertions(+), 28 deletions(-) diff --git a/bin/cosmosis-extract b/bin/cosmosis-extract index 0fb36c6a..7caef01e 100755 --- a/bin/cosmosis-extract +++ b/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__': diff --git a/cosmosis/runtime/config.py b/cosmosis/runtime/config.py index 13855fe1..d6fc7e93 100644 --- a/cosmosis/runtime/config.py +++ b/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. diff --git a/cosmosis/runtime/pipeline.py b/cosmosis/runtime/pipeline.py index 6d5c4dc8..11af7ed9 100644 --- a/cosmosis/runtime/pipeline.py +++ b/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): diff --git a/cosmosis/test/test_pipeline.py b/cosmosis/test/test_pipeline.py index 20de5bf6..8f1f83d6 100644 --- a/cosmosis/test/test_pipeline.py +++ b/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() diff --git a/cosmosis/utils.py b/cosmosis/utils.py index bf3f4c16..5ab5bb6e 100644 --- a/cosmosis/utils.py +++ b/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 \ No newline at end of file