Skip to content

Commit

Permalink
Merge pull request #148 from joezuntz/polychord-use-boosted
Browse files Browse the repository at this point in the history
Add a version of the boosted polychord file with a full header.
  • Loading branch information
joezuntz authored Oct 31, 2024
2 parents ab828be + 9ad6a66 commit 5c5d1b5
Show file tree
Hide file tree
Showing 3 changed files with 112 additions and 24 deletions.
68 changes: 68 additions & 0 deletions cosmosis/samplers/polychord/polychord_sampler.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import numpy as np
import sys
from ...runtime.utils import mkdir
from ...output import TextColumnOutput
import warnings

prior_type = ct.CFUNCTYPE(None,
Expand Down Expand Up @@ -302,6 +303,73 @@ def sample(self):

self.converged = True

if self.is_master() and self.boost_posteriors:
self.make_boosted_posterior_file()

def make_boosted_posterior_file(self):
# two cases where there is nothing to boost
if not self.polychord_outfile_root:
return
if not self.boost_posteriors:
return

old = self.output

boosted_filename = os.path.join(self.base_dir, self.polychord_outfile_root + ".txt")

# This is a horribly awkward hack
if isinstance(self.output, TextColumnOutput):
# make sure everything is written out already
old.flush()

new_output_filename = old.filename_base + "_boosted.txt"
main_output_info = TextColumnOutput.load_from_options({"filename":old._filename, "delimiter":old.delimiter})

_, _, metadata, comments, final_metadata = main_output_info
metadata = metadata[0]
comments = comments[0]
final_metadata = final_metadata[0]

# make a new output object
new = TextColumnOutput.from_options({"filename":new_output_filename})

# write the main bits of metadata from the output file
new.metadata("sampler", "polychord")
new.comment("NOTE: Boosted posterior file")
self.write_header(new)
for md in metadata:
new.metadata(md, metadata[md])
for c in comments:
new.comment(c)


# write the parameter rows from the old file
with open(boosted_filename, "r") as f:
for line in f:
values = [float(x) for x in line.split()]
weight = values[0]
like = values[1]
params = values[2:]
prior = params[-1]
post = like + prior
new.parameters(params, like, post, weight)


# This hasn't been written yet because the original output file has not been closed.
# so we have to use the saved one in the old object. The * is because a comment is also included
for md in old._final_metadata:
new.final(md, *old._final_metadata[md])
new.final("log_z", self.log_z)
new.final("log_z_error", self.log_z_err)

# fake these first two, and add the completion marker.
new.final("evaluations", 0)
new.final("successes", 0)
new.final("complete", 1)
new.close()



def output_params(self, ndead, nlive, npars, live, dead, logweights, log_z, log_z_err):
# Polychord repeats output, but with changed weights, so reset to the start
# of the chain to overwrite them.
Expand Down
28 changes: 16 additions & 12 deletions cosmosis/samplers/sampler.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,19 +54,23 @@ def __init__(self, ini, pipeline, output=None):
self.distribution_hints = Hints()
self.write_header()

def write_header(self):
if self.output:
for p in self.pipeline.output_names():
self.output.add_column(p, float)
for p,ptype in self.sampler_outputs:
self.output.add_column(p, ptype)
self.output.metadata("n_varied", len(self.pipeline.varied_params))
self.attribution.write_output(self.output)
for key, value in self.collect_run_metadata().items():
self.output.metadata(key, value)
def write_header(self, output=None):
if output is None:
output = self.output
if output is None:
return

for p in self.pipeline.output_names():
output.add_column(p, float)
for p,ptype in self.sampler_outputs:
output.add_column(p, ptype)
output.metadata("n_varied", len(self.pipeline.varied_params))
self.attribution.write_output(output)
for key, value in self.collect_run_metadata().items():
output.metadata(key, value)
blinding_header = self.ini.getboolean("output","blinding-header", fallback=False)
if blinding_header and self.output:
self.output.blinding_header()
if blinding_header:
output.blinding_header()

def collect_run_metadata(self):
info = {
Expand Down
40 changes: 28 additions & 12 deletions cosmosis/test/test_polychord.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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]

Expand Down Expand Up @@ -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

Expand All @@ -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]]()

0 comments on commit 5c5d1b5

Please sign in to comment.