108 lines
3.1 KiB
Python
108 lines
3.1 KiB
Python
|
#+
|
||
|
|
# 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")
|