borg_public/scripts/migrate_restart.py
2023-05-29 10:41:03 +02:00

86 lines
2.0 KiB
Python

#+
# ARES/HADES/BORG Package -- ./scripts/migrate_restart.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 os
import h5py as h5
import errno
from ares_tools import rebuild_spliced_h5
def detect_ncpus(path):
ncpu = 0
try:
while True:
with open("%s_%d" % (path,ncpu), mode= "rb") as f:
ncpu += 1
except IOError as e:
if e.errno != errno.ENOENT:
raise e
return ncpu
def detect_job(path):
obj_list=[]
group_list=[]
def _handle_item(name, obj):
if isinstance(obj, h5.Group):
group_list.append(name)
if not isinstance(obj, h5.Dataset):
return
obj_list.append(name)
with h5.File("%s_0" % path, mode="r") as f:
f.visititems(_handle_item)
return obj_list,group_list
def load_nonarray(path, outpath, objlist, group_list):
arr = {}
with h5.File(path, mode="r") as f, h5.File(outpath, mode="w") as of:
for g in group_list:
if g[:4] == 'info':
newname = g[4:]
elif g[:6] == 'markov':
newname = g[6:]
if len(newname)==0:
continue
if newname in of:
continue
print("Create group %s" % newname)
of.create_group(newname)
for oname in objlist:
print("Loading %s..." % oname)
if oname[:4] == 'info':
newname = oname[4:]
elif oname[:6] == 'markov':
newname = oname[6:]
else:
print("Beuh ! " + oname)
abort
if oname == '/info/scalars/BORG_version':
of[newname] = np.array([f[oname][0]], dtype='S')
else:
of[newname] = f[oname][:]
def migrate(path, newpath):
ncpu = detect_ncpus(path)
elem_list,group_list = detect_job(path)
print("Loading spliced arrays")
for n in range(ncpu):
load_nonarray("%s_%d" % (path,n), "%s_%d" % (newpath,n), elem_list, group_list)
if __name__=="__main__":
migrate("./restart.h5", "./new/restart.h5")