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)