borg_public/scripts/regraft_run.py

108 lines
3.1 KiB
Python
Raw Normal View History

2023-05-29 10:41:03 +02:00
#+
# ARES/HADES/BORG Package -- ./scripts/regraft_run.py
# Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
# Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
#
# Additional contributions from:
# Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
#
#+
import warnings
warnings.simplefilter("ignore", FutureWarning)
import h5py as h5
import argparse
import os
class NotCompatible(Exception):
def __init__(msg_):
self.msg = msg_
def __repr__():
return "NotCompatible: " + self.msg
def restart_name(prefix, cpu):
return os.path.join(prefix, "restart.h5_%d" % cpu)
def detect_ncpu(prefix):
ncpu = 0
while True:
try:
os.stat(os.path.join(prefix, "restart.h5_%d" % ncpu))
except OSError:
break
ncpu += 1
return ncpu
def bias_name(bid):
return "scalars/galaxy_bias_%d" % bid
def nmean_name(nid):
return "scalars/galaxy_nmean_%d" % nid
def check_compat(restart_prefix, mcmc):
print("Checking compatibility of MCMC and restart")
with h5.File(restart_name(restart_prefix, 0), mode="r") as f_r, \
h5.File(mcmc, mode="r") as f_m:
bname = bias_name(0)
while bname in f_m:
if not bname in f_r:
raise NotCompatible("Not enough catalogs")
if f_m[bname].size != f_r[bname].size:
raise NotCompatible("Incompatible bias model")
bias_num += 1
bname = bias_name(bias_num)
def transfer(src, dest, name, start):
if (not name in src) or (not name in dest):
return False
sz = dest[name].shape[0]
dest[name][...] = src[name][start:(start+sz),...]
return True
def checked_transfer(*args):
if not transfer(*args):
raise NotCompatible("Problems in grafting")
def graft(restart_prefix, ncpus, mcmc):
with h5.File(args.mcmc, mode="r") as f_m:
plane_start = 0
plane_len = 0
for cpu in range(ncpus):
print("Transplanting to restart CPU %d" % cpu)
with h5.File(restart_name(restart_prefix, cpu), mode="r") as f_r:
plane_len = f_r['scalars/BORG_final_density'].shape[0]
checked_transfer(f_m, f_r, 'scalars/BORG_final_density', plane_start)
checked_transfer(f_m, f_r, 'scalars/s_field', plane_start)
checked_transfer(f_m, f_r, 'scalars/s_hat_field', plane_start)
bias_num = 0
while True:
if not transfer(f_m, f_r, bias_name(bias_num), 0):
break
if not transfer(f_m, f_r, nmean_name(bias_num), 0):
break
bias_num += 1
plane_start += plane_len
p = argparse.ArgumentParser(description="Graft the state of a previous run on a given restart.")
p.add_argument('mcmc', type=str, help="MCMC state to import")
p.add_argument('restart', type=str, help="restart prefix directory")
args = p.parse_args()
ncpu = detect_ncpu(args.restart)
print("Found %d CPU restart file")
check_compat(args.restart, args.mcmc)
graft(args.restart, ncpu, args.mcmc)
print("Done")