def __init__(self, batch_tasks, group_size=1, ingroup_parallel=1,
dispatcher=None, resources=None,
outlog='log', errlog='err',
forward_files=None, backward_files=None,
**task_args):
super().__init__(**task_args)
self.batch_tasks = [deepcopy(task) for task in batch_tasks]
for task in self.batch_tasks:
assert isinstance(task, BatchTask), f'given task is instance of {task.__class__}'
assert not task.workdir.is_absolute()
task.prepend_workdir(self.workdir)
if task.prev_folder is None:
task.prev_folder = self.prev_folder
self.group_size = group_size
self.para_deg = ingroup_parallel
if dispatcher is None:
dispatcher = Dispatcher()
elif isinstance(dispatcher, dict):
dispatcher = Dispatcher(**dispatcher)
assert isinstance(dispatcher, Dispatcher)
self.dispatcher = dispatcher
self.resources = resources
self.outlog = outlog
self.errlog = errlog
self.forward_files = check_list(forward_files)
self.backward_files = check_list(backward_files)
def get_required_labels(fields=None, penalty_dicts=None):
field_labels = [check_list(f.required_labels) for f in check_list(fields)]
penalty_labels = [
check_list(
p.get("required_labels",
select_penalty(p["type"]).required_labels))
for p in check_list(penalty_dicts)
]
return set(sum(field_labels + penalty_labels, []))
def build_penalty(pnt_dict, label_dict={}):
pnt_dict = pnt_dict.copy()
pnt_type = pnt_dict.pop("type")
PenaltyClass = select_penalty(pnt_type)
label_names = pnt_dict.pop("required_labels", PenaltyClass.required_labels)
label_arrays = [label_dict[lb] for lb in check_list(label_names)]
return PenaltyClass(*label_arrays, **pnt_dict)
def load_stat_grouped(systems,
dump_dir=".",
with_conv=True,
with_e=True,
e_name="e_tot",
with_f=True,
f_name="f_tot"):
systems = check_list(systems)
lbases = [get_sys_name(fl) for fl in systems]
c_res = e_err = f_err = None
if with_conv:
c_res = load_array(get_with_prefix("conv", dump_dir, ".npy"))
if with_e:
e_res = load_array(get_with_prefix(e_name, dump_dir, ".npy"))
e_lbl = np.concatenate([
load_array(get_with_prefix("energy", lb, ".npy")) for lb in lbases
], 0).reshape(-1, 1)
e_err = e_lbl - e_res
if with_f:
f_res = load_array(get_with_prefix(f_name, dump_dir, ".npy"))
f_lbl = np.concatenate([
load_array(get_with_prefix("force", lb, ".npy")) for lb in lbases
], 0).reshape(f_res.shape)
f_err = f_lbl - f_res
return c_res, e_err, f_err
def __init__(self, cmds,
dispatcher=None, resources=None,
outlog='log', errlog='err',
forward_files=None, backward_files=None,
**task_args):
super().__init__(**task_args)
self.cmds = check_list(cmds)
if dispatcher is None:
dispatcher = Dispatcher()
elif isinstance(dispatcher, dict):
dispatcher = Dispatcher(**dispatcher)
assert isinstance(dispatcher, Dispatcher)
self.dispatcher = dispatcher
self.resources = resources
self.outlog = outlog
self.errlog = errlog
self.forward_files = check_list(forward_files)
self.backward_files = check_list(backward_files)
def main(data_paths,
model_file="model.pth",
output_prefix='test',
group=False,
e_name='l_e_delta',
d_name=['dm_eig']):
data_paths = load_dirs(data_paths)
g_reader = GroupReader(data_paths, e_name=e_name, d_name=d_name)
model_file = check_list(model_file)
for f in model_file:
print(f)
p = os.path.dirname(f)
model = CorrNet.load(f).double().to(DEVICE)
dump = os.path.join(p, output_prefix)
dir_name = os.path.dirname(dump)
if dir_name:
os.makedirs(dir_name, exist_ok=True)
test(model, g_reader, dump_prefix=dump, group=group)
def load_stat(systems,
dump_dir,
with_conv=True,
with_e=True,
e_name="e_tot",
with_f=True,
f_name="f_tot"):
systems = check_list(systems)
c_res = []
e_err = []
f_err = []
for fl in systems:
lbase = get_sys_name(fl)
rbase = os.path.join(dump_dir, os.path.basename(lbase))
if with_conv:
try:
c_res.append(load_array(get_with_prefix("conv", rbase,
".npy")))
except FileNotFoundError as e:
print("Warning! conv.npy not found:", e, file=sys.stderr)
if with_e:
try:
re = load_array(get_with_prefix(e_name, rbase,
".npy")).reshape(-1, 1)
le = load_array(get_with_prefix("energy", lbase,
".npy")).reshape(-1, 1)
e_err.append(le - re)
except FileNotFoundError as e:
print("Warning! energy file not found:", e, file=sys.stderr)
if with_f:
try:
rf = load_array(get_with_prefix(f_name, rbase, ".npy"))
lf = load_array(get_with_prefix("force", lbase,
".npy")).reshape(rf.shape)
f_err.append(np.abs(lf - rf).mean((-1, -2)))
except FileNotFoundError as e:
print("Warning! force file not found:", e, file=sys.stderr)
return np.concatenate(c_res, 0) if c_res else None, \
np.concatenate(e_err, 0) if e_err else None, \
np.concatenate(f_err, 0) if f_err else None
def __init__(self, workdir='.', backup=False, prev_task=None,
prev_folder=None, link_prev_files=None, copy_prev_files=None,
share_folder=None, link_share_files=None, copy_share_files=None,
link_abs_files=None, copy_abs_files=None):
# workdir has to be relative in order to be chained
# prev_task is dereferenced to folder dynamically.
# folders are absolute.
super().__init__(workdir)
self.backup = backup
assert prev_task is None or prev_folder is None
self.prev_task = prev_task
self.prev_folder = get_abs_path(prev_folder)
self.share_folder = get_abs_path(share_folder)
self.link_prev_files = check_list(link_prev_files)
self.copy_prev_files = check_list(copy_prev_files)
self.link_share_files = check_list(link_share_files)
self.copy_share_files = check_list(copy_share_files)
self.link_abs_files = check_list(link_abs_files)
self.copy_abs_files = check_list(copy_abs_files)
def make_cleanup(pattern="slurm-*.out", workdir=".", **task_args):
pattern = check_list(pattern)
pattern = " ".join(pattern)
assert pattern
return ShellTask(f"rm -r {pattern}", workdir=workdir, **task_args)
def main(systems,
model_file="model.pth",
basis='ccpvdz',
proj_basis=None,
penalty_terms=None,
device=None,
dump_dir=".",
dump_fields=DEFAULT_FNAMES,
group=False,
mol_args=None,
scf_args=None,
verbose=0):
if model_file is None or model_file.upper() == "NONE":
model = None
default_scf_args = DEFAULT_HF_ARGS
else:
model = CorrNet.load(model_file).double()
default_scf_args = DEFAULT_SCF_ARGS
# check arguments
penalty_terms = check_list(penalty_terms)
if mol_args is None: mol_args = {}
if scf_args is None: scf_args = {}
scf_args = {**default_scf_args, **scf_args}
fields = select_fields(dump_fields)
# check label names from label fields and penalties
label_names = get_required_labels(fields["scf"] + fields["grad"],
penalty_terms)
if verbose:
print(f"starting calculation with OMP threads: {lib.num_threads()}",
f"and max memory: {lib.param.MAX_MEMORY}")
if verbose > 1:
print(f"basis: {basis}")
print(f"specified scf args:\n {scf_args}")
meta = old_meta = None
res_list = []
systems = load_sys_paths(systems)
for fl in systems:
fl = fl.rstrip(os.path.sep)
for atom, attrs, labels in system_iter(fl, label_names):
mol_input = {
**mol_args, "verbose": verbose,
"atom": atom,
"basis": basis,
**attrs
}
mol = build_mol(**mol_input)
penalties = [build_penalty(pd, labels) for pd in penalty_terms]
try:
meta, result = solve_mol(mol,
model,
fields,
labels,
proj_basis=proj_basis,
penalties=penalties,
device=device,
verbose=verbose,
**scf_args)
except Exception as e:
print(fl, 'failed! error:', e, file=sys.stderr)
# continue
raise
if group and old_meta is not None and np.any(meta != old_meta):
break
res_list.append(result)
if not group:
sub_dir = os.path.join(dump_dir,
get_sys_name(os.path.basename(fl)))
dump_meta(sub_dir, meta)
dump_data(sub_dir, **collect_fields(fields, meta, res_list))
res_list = []
elif old_meta is not None and np.any(meta != old_meta):
print(fl,
'meta does not match! saving previous results only.',
file=sys.stderr)
break
old_meta = meta
if verbose:
print(fl, 'finished')
if group:
dump_meta(dump_dir, meta)
dump_data(dump_dir, **collect_fields(fields, meta, res_list))
if verbose:
print('group finished')
def __init__(self, penalties=None):
self.penalties = check_list(penalties)
for pnt in self.penalties:
pnt.init_hook(self)
