Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add a helper method to re-create a pipeline object from a chain output file. #151

Merged
merged 1 commit into from
Nov 20, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 4 additions & 26 deletions bin/cosmosis-extract
Original file line number Diff line number Diff line change
@@ -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__':
Expand Down
8 changes: 7 additions & 1 deletion cosmosis/runtime/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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.
Expand Down
22 changes: 21 additions & 1 deletion cosmosis/runtime/pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import collections
import warnings
import traceback
import io
from . import config
from . import parameter
from . import prior
Expand Down Expand Up @@ -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):
Expand Down
56 changes: 56 additions & 0 deletions cosmosis/test/test_pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
25 changes: 25 additions & 0 deletions cosmosis/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading