diff --git a/dpgen2/entrypoint/args.py b/dpgen2/entrypoint/args.py index 7acb4c33..45504127 100644 --- a/dpgen2/entrypoint/args.py +++ b/dpgen2/entrypoint/args.py @@ -355,10 +355,93 @@ def caly_args(): ] +def run_diffcsp_args(): + doc_gen_tasks = "Number of DiffCSP generation tasks" + doc_gen_command = "Command for DiffCSP generation" + doc_relax_group_size = "Group size for relaxation" + return [ + Argument( + "gen_tasks", + int, + optional=True, + default=1, + doc=doc_gen_tasks, + ), + Argument( + "gen_command", + str, + optional=False, + doc=doc_gen_command, + ), + Argument( + "relax_group_size", + int, + optional=True, + default=100, + doc=doc_relax_group_size, + ), + ] + + +def diffcsp_args(): + doc_config = "Configuration of DiffCSP exploration" + doc_max_numb_iter = "Maximum number of iterations per stage" + doc_fatal_at_max = ( + "Fatal when the number of iteration per stage reaches the `max_numb_iter`" + ) + doc_output_nopbc = "Remove pbc of the output configurations" + doc_convergence = "The method of convergence check." + doc_stages = ( + "The definition of exploration stages of type `List[List[ExplorationTaskGroup]`. " + "The outer list provides the enumeration of the exploration stages. " + "Then each stage is defined by a list of exploration task groups. " + "Each task group is described in :ref:`the task group definition` " + ) + doc_filters = "A list of configuration filters" + + return [ + Argument( + "config", + dict, + run_diffcsp_args(), + optional=False, + doc=doc_config, + ), + Argument( + "max_numb_iter", int, optional=True, default=10, doc=doc_max_numb_iter + ), + Argument( + "fatal_at_max", bool, optional=True, default=True, doc=doc_fatal_at_max + ), + Argument( + "output_nopbc", bool, optional=True, default=False, doc=doc_output_nopbc + ), + Argument( + "convergence", + dict, + [], + [variant_conv()], + optional=False, + doc=doc_convergence, + ), + Argument("stages", List[List[dict]], optional=False, doc=doc_stages), + Argument( + "filters", + list, + [], + [variant_filter()], + optional=True, + default=[], + doc=doc_filters, + ), + ] + + def variant_explore(): doc = "The type of the exploration" doc_lmp = "The exploration by LAMMPS simulations" doc_calypso = "The exploration by CALYPSO structure prediction" + doc_diffcsp = "The exploration by DiffCSP" return Variant( "type", [ @@ -366,6 +449,7 @@ def variant_explore(): Argument("calypso", dict, caly_args(), doc=doc_calypso), Argument("calypso:default", dict, caly_args(), doc=doc_calypso), Argument("calypso:merge", dict, caly_args(), doc=doc_calypso), + Argument("diffcsp", dict, diffcsp_args(), doc=doc_diffcsp), ], doc=doc, ) diff --git a/dpgen2/entrypoint/showkey.py b/dpgen2/entrypoint/showkey.py index 47180adc..36055426 100644 --- a/dpgen2/entrypoint/showkey.py +++ b/dpgen2/entrypoint/showkey.py @@ -35,6 +35,6 @@ def showkey( all_step_keys = sum(folded_keys.values(), []) prt_str = print_keys_in_nice_format( all_step_keys, - ["run-train", "run-lmp", "run-fp"], + ["run-train", "run-lmp", "run-fp", "diffcsp-gen", "run-relax"], ) print(prt_str) diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py index 7e07659a..35242dfe 100644 --- a/dpgen2/entrypoint/submit.py +++ b/dpgen2/entrypoint/submit.py @@ -84,7 +84,9 @@ LmpTemplateTaskGroup, NPTTaskGroup, caly_normalize, + diffcsp_normalize, make_calypso_task_group_from_config, + make_diffcsp_task_group_from_config, make_lmp_task_group_from_config, normalize_lmp_task_group_config, ) @@ -97,15 +99,18 @@ from dpgen2.op import ( CollectData, CollRunCaly, + DiffCSPGen, PrepCalyDPOptim, PrepCalyInput, PrepCalyModelDevi, PrepDPTrain, PrepLmp, + PrepRelax, RunCalyDPOptim, RunCalyModelDevi, RunDPTrain, RunLmp, + RunRelax, SelectConfs, ) from dpgen2.op.caly_evo_step_merge import ( @@ -114,6 +119,7 @@ from dpgen2.superop import ( ConcurrentLearningBlock, PrepRunCaly, + PrepRunDiffCSP, PrepRunDPTrain, PrepRunFp, PrepRunLmp, @@ -223,6 +229,16 @@ def make_concurrent_learning_op( run_config=run_explore_config, upload_python_packages=upload_python_packages, ) + elif explore_style == "diffcsp": + prep_run_explore_op = PrepRunDiffCSP( + "prep-run-diffcsp", + DiffCSPGen, + PrepRelax, + RunRelax, + prep_config=prep_explore_config, + run_config=run_explore_config, + upload_python_packages=upload_python_packages, + ) else: raise RuntimeError(f"unknown explore_style {explore_style}") @@ -269,12 +285,10 @@ def make_naive_exploration_scheduler( if explore_style == "lmp": return make_lmp_naive_exploration_scheduler(config) - elif "calypso" in explore_style: - return make_calypso_naive_exploration_scheduler(config) + elif "calypso" in explore_style or explore_style == "diffcsp": + return make_naive_exploration_scheduler_without_conf(config, explore_style) else: - raise KeyError( - f"Unknown key `{explore_style}`, Only support `lmp`, `calypso`, `calypso:merge` and `calypso:default`." - ) + raise KeyError(f"Unknown explore_style `{explore_style}`") def get_conf_filters(config): @@ -288,7 +302,7 @@ def get_conf_filters(config): return conf_filters -def make_calypso_naive_exploration_scheduler(config): +def make_naive_exploration_scheduler_without_conf(config, explore_style): model_devi_jobs = config["explore"]["stages"] fp_task_max = config["fp"]["task_max"] max_numb_iter = config["explore"]["max_numb_iter"] @@ -300,6 +314,7 @@ def make_calypso_naive_exploration_scheduler(config): # report conv_style = convergence.pop("type") report = conv_styles[conv_style](**convergence) + # trajectory render, the format of the output trajs are assumed to be lammps/dump render = TrajRenderLammps(nopbc=output_nopbc) # selector selector = ConfSelectorFrames( @@ -317,9 +332,16 @@ def make_calypso_naive_exploration_scheduler(config): # stage stage = ExplorationStage() for jj in job: - jconf = caly_normalize(jj) - # make task group - tgroup = make_calypso_task_group_from_config(jconf) + if "calypso" in explore_style: + jconf = caly_normalize(jj) + # make task group + tgroup = make_calypso_task_group_from_config(jconf) + elif explore_style == "diffcsp": + jconf = diffcsp_normalize(jj) + # make task group + tgroup = make_diffcsp_task_group_from_config(jconf) + else: + raise KeyError(f"Unknown explore_style `{explore_style}`") # add the list to task group tasks = tgroup.make_task() stage.add_task_group(tasks) @@ -800,6 +822,12 @@ def get_superop(key): return re.sub("run-caly-model-devi-[0-9]*", "prep-run-explore", key) elif "caly-evo-step" in key: return re.sub("caly-evo-step-[0-9]*", "prep-run-explore", key) + elif "diffcsp-gen-" in key: + return re.sub("diffcsp-gen-[0-9]*", "prep-run-explore", key) + elif "prep-relax" in key: + return re.sub("prep-relax", "prep-run-explore", key) + elif "run-relax-" in key: + return re.sub("run-relax-[0-9]*", "prep-run-explore", key) return None @@ -843,6 +871,9 @@ def get_resubmit_keys( "prep-run-explore", "prep-lmp", "run-lmp", + "diffcsp-gen", + "prep-relax", + "run-relax", "select-confs", "prep-run-fp", "prep-fp", @@ -880,7 +911,7 @@ def get_resubmit_keys( ) all_step_keys = sort_slice_ops( all_step_keys, - ["run-train", "run-lmp", "run-fp"], + ["run-train", "run-lmp", "run-fp", "diffcsp-gen", "run-relax"], ) folded_keys = fold_keys(all_step_keys) return folded_keys @@ -905,7 +936,7 @@ def resubmit_concurrent_learning( if list_steps: prt_str = print_keys_in_nice_format( all_step_keys, - ["run-train", "run-lmp", "run-fp"], + ["run-train", "run-lmp", "run-fp", "diffcsp-gen", "run-relax"], ) print(prt_str) diff --git a/dpgen2/exploration/task/__init__.py b/dpgen2/exploration/task/__init__.py index a062d511..534a8828 100644 --- a/dpgen2/exploration/task/__init__.py +++ b/dpgen2/exploration/task/__init__.py @@ -7,12 +7,16 @@ from .customized_lmp_template_task_group import ( CustomizedLmpTemplateTaskGroup, ) +from .diffcsp_task_group import ( + DiffCSPTaskGroup, +) from .lmp_template_task_group import ( LmpTemplateTaskGroup, ) from .make_task_group_from_config import ( caly_normalize, caly_task_group_args, + diffcsp_normalize, ) from .make_task_group_from_config import ( lmp_normalize as normalize_lmp_task_group_config, @@ -20,6 +24,7 @@ from .make_task_group_from_config import ( lmp_task_group_args, make_calypso_task_group_from_config, + make_diffcsp_task_group_from_config, make_lmp_task_group_from_config, variant_task_group, ) diff --git a/dpgen2/exploration/task/diffcsp_task_group.py b/dpgen2/exploration/task/diffcsp_task_group.py new file mode 100644 index 00000000..91a002b7 --- /dev/null +++ b/dpgen2/exploration/task/diffcsp_task_group.py @@ -0,0 +1,47 @@ +from typing import ( + Optional, +) + +from .task import ( + ExplorationTask, +) +from .task_group import ( + ExplorationTaskGroup, +) + + +class DiffCSPTaskGroup(ExplorationTaskGroup): + def __init__( + self, + trj_freq: int = 10, + fmax: float = 1e-4, + steps: int = 200, + timeout: Optional[int] = None, + ): + super().__init__() + self.trj_freq = trj_freq + self.fmax = fmax + self.steps = steps + self.timeout = timeout + + def make_task(self) -> "DiffCSPTaskGroup": + """ + Make the DiffCSP task group. + + Returns + ------- + task_grp: DiffCSPTaskGroup + Return one DiffCSP task group. + """ + # clear all existing tasks + self.clear() + self.add_task(self._make_diffcsp_task()) + return self + + def _make_diffcsp_task(self) -> ExplorationTask: + task = ExplorationTask() + task.trj_freq = self.trj_freq # type: ignore + task.fmax = self.fmax # type: ignore + task.steps = self.steps # type: ignore + task.timeout = self.timeout # type: ignore + return task diff --git a/dpgen2/exploration/task/make_task_group_from_config.py b/dpgen2/exploration/task/make_task_group_from_config.py index e6058030..d82f3463 100644 --- a/dpgen2/exploration/task/make_task_group_from_config.py +++ b/dpgen2/exploration/task/make_task_group_from_config.py @@ -10,6 +10,9 @@ model_name_pattern, plm_input_name, ) +from dpgen2.exploration.task import ( + DiffCSPTaskGroup, +) from dpgen2.exploration.task.caly_task_group import ( CalyTaskGroup, ) @@ -508,6 +511,63 @@ def make_calypso_task_group_from_config(config): return tgroup +def diffcsp_task_group_args(): + doc_diffcsp_task_grp = "DiffCSP exploration tasks" + doc_trj_freq = "The frequency of dumping configurations and model devis" + doc_fmax = "Force tolerence in relaxation" + doc_steps = "Maximum number of steps in relaxation" + doc_timeout = "Timeout (seconds) in relaxation" + return Argument( + "task_group", + dict, + [ + Argument( + "trj_freq", + int, + optional=True, + default=10, + doc=doc_trj_freq, + alias=["t_freq", "trj_freq", "traj_freq"], + ), + Argument( + "fmax", + float, + optional=True, + default=1e-4, + doc=doc_fmax, + ), + Argument( + "steps", + int, + optional=True, + default=200, + doc=doc_steps, + ), + Argument( + "timeout", + int, + optional=True, + default=None, + doc=doc_timeout, + ), + ], + doc=doc_diffcsp_task_grp, + ) + + +def diffcsp_normalize(config): + args = diffcsp_task_group_args() + config = args.normalize_value(config, trim_pattern="_*") + args.check_value(config, strict=False) + return config + + +def make_diffcsp_task_group_from_config(config): + config = diffcsp_normalize(config) + tgroup = DiffCSPTaskGroup(**config) + return tgroup + + def make_lmp_task_group_from_config( numb_models, mass_map, diff --git a/dpgen2/op/__init__.py b/dpgen2/op/__init__.py index ff4c8a8d..a0b0b1a8 100644 --- a/dpgen2/op/__init__.py +++ b/dpgen2/op/__init__.py @@ -4,6 +4,9 @@ from .collect_run_caly import ( CollRunCaly, ) +from .diffcsp_gen import ( + DiffCSPGen, +) from .prep_caly_dp_optim import ( PrepCalyDPOptim, ) @@ -19,6 +22,9 @@ from .prep_lmp import ( PrepLmp, ) +from .prep_relax import ( + PrepRelax, +) from .run_caly_dp_optim import ( RunCalyDPOptim, ) @@ -31,6 +37,9 @@ from .run_lmp import ( RunLmp, ) +from .run_relax import ( + RunRelax, +) from .select_confs import ( SelectConfs, ) diff --git a/dpgen2/op/diffcsp_gen.py b/dpgen2/op/diffcsp_gen.py new file mode 100644 index 00000000..a7ed5ef1 --- /dev/null +++ b/dpgen2/op/diffcsp_gen.py @@ -0,0 +1,102 @@ +import os +import subprocess +from pathlib import ( + Path, +) +from typing import ( + List, +) + +from dflow.python import ( + OP, + OPIO, + Artifact, + OPIOSign, +) + + +def convert_pt_to_cif(input_file, output_dir): + import numpy as np + import torch # type: ignore + from pymatgen.core.lattice import ( # type: ignore + Lattice, + ) + from pymatgen.core.structure import ( # type: ignore + Structure, + ) + + data = torch.load(input_file, map_location=torch.device("cpu")) + if not os.path.exists(output_dir): + os.makedirs(output_dir) + lengths = data["lengths"] + angles = data["angles"] + num_atoms = data["num_atoms"] + frac_coors = data["frac_coords"] + atom_types = data["atom_types"] + + lengths_list = lengths.numpy() + angles_list = angles.numpy() + num_atoms_list = num_atoms + frac_coors_list = frac_coors.numpy() + atom_types_list = atom_types + if len(atom_types_list.shape) > 1: + atom_types_list = np.argmax(atom_types_list, axis=-1) + 1 + + now_atom = 0 + for i in range(len(num_atoms_list)): + length = lengths_list[i] + angle = angles_list[i] + atom_num = num_atoms_list[i] + + atom_type = atom_types_list[now_atom : now_atom + atom_num] + frac_coord = frac_coors_list[now_atom : now_atom + atom_num] + lattice = Lattice.from_parameters(*(length.tolist() + angle.tolist())) + structure = Structure( + lattice, atom_type, frac_coord, coords_are_cartesian=False + ) + + filename = "%s.cif" % i + file_path = os.path.join(output_dir, filename) + structure.to(filename=file_path) + now_atom += atom_num + + +class DiffCSPGen(OP): + @classmethod + def get_input_sign(cls): + return OPIOSign( + { + "config": dict, + "task_id": str, + } + ) + + @classmethod + def get_output_sign(cls): + return OPIOSign( + { + "cifs": Artifact(List[Path]), + } + ) + + @OP.exec_sign_check + def execute( + self, + ip: OPIO, + ) -> OPIO: + cmd = ip["config"]["gen_command"] + args = cmd.split() + try: + i = args.index("--model_path") + except ValueError: + raise RuntimeError("Path of DiffCSP model not provided.") + model_path = args[i + 1] + subprocess.run(cmd, shell=True, check=True) + result_file = os.path.join(model_path, "eval_gen.pt") + task_dir = "diffcsp.%s" % ip["task_id"] + convert_pt_to_cif(result_file, task_dir) + return OPIO( + { + "cifs": list(Path(task_dir).glob("*.cif")), + } + ) diff --git a/dpgen2/op/prep_relax.py b/dpgen2/op/prep_relax.py new file mode 100644 index 00000000..1ee2869a --- /dev/null +++ b/dpgen2/op/prep_relax.py @@ -0,0 +1,57 @@ +import os +from pathlib import ( + Path, +) +from typing import ( + List, +) + +from dflow.python import ( + OP, + OPIO, + Artifact, + OPIOSign, +) + + +class PrepRelax(OP): + @classmethod + def get_input_sign(cls): + return OPIOSign( + { + "expl_config": dict, + "cifs": Artifact(List[Path]), + } + ) + + @classmethod + def get_output_sign(cls): + return OPIOSign( + { + "ntasks": int, + "task_paths": Artifact(List[Path]), + } + ) + + @OP.exec_sign_check + def execute( + self, + ip: OPIO, + ) -> OPIO: + ncifs = len(ip["cifs"]) + config = ip["expl_config"] + group_size = config["relax_group_size"] + ntasks = int(ncifs / group_size) + task_paths = [] + for i in range(ntasks): + task_dir = Path("task.%06d" % i) + task_dir.mkdir(exist_ok=True) + for j in range(group_size * i, min(group_size * (i + 1), ncifs)): + os.symlink(ip["cifs"][j], task_dir / ("%s.cif" % j)) + task_paths.append(task_dir) + return OPIO( + { + "ntasks": ntasks, + "task_paths": task_paths, + } + ) diff --git a/dpgen2/op/run_relax.py b/dpgen2/op/run_relax.py new file mode 100644 index 00000000..110ad4d8 --- /dev/null +++ b/dpgen2/op/run_relax.py @@ -0,0 +1,180 @@ +import os +from pathlib import ( + Path, +) +from typing import ( + List, +) + +from dflow.python import ( + OP, + OPIO, + Artifact, + BigParameter, + OPIOSign, +) + +from dpgen2.exploration.task import ( + DiffCSPTaskGroup, +) + +from .run_caly_model_devi import ( + atoms2lmpdump, +) + + +class RunRelax(OP): + @classmethod + def get_input_sign(cls): + return OPIOSign( + { + "diffcsp_task_grp": BigParameter(DiffCSPTaskGroup), + "task_path": Artifact(Path), + "models": Artifact(List[Path]), + } + ) + + @classmethod + def get_output_sign(cls): + return OPIOSign( + { + "trajs": Artifact(List[Path]), + "model_devis": Artifact(List[Path]), + } + ) + + @OP.exec_sign_check + def execute( + self, + ip: OPIO, + ) -> OPIO: + import pickle + import traceback + + import ase # type: ignore + import numpy as np + from ase.calculators.singlepoint import ( # type: ignore + SinglePointCalculator, + ) + from deepmd.infer import ( # type: ignore + DeepPot, + ) + from deepmd.infer.model_devi import ( # type: ignore + calc_model_devi_f, + calc_model_devi_v, + ) + from lam_optimize.main import ( # type: ignore + relax_run, + ) + from lam_optimize.relaxer import ( # type: ignore + Relaxer, + ) + + task_group = ip["diffcsp_task_grp"] + task = next(iter(task_group)) # Only support single task + models = ip["models"] + relaxer = Relaxer(models[0]) + type_map = relaxer.calculator.dp.get_type_map() + fmax = task.fmax + steps = task.steps + timeout = task.timeout + os.makedirs("relax_trajs", exist_ok=True) + relax_run( + ip["task_path"], + relaxer, + fmax=fmax, + steps=steps, + traj_file="relax_trajs", + timeout=timeout, + check_convergence=False, + check_duplicate=False, + ) + + trajs = [] + model_devis = [] + graphs = [None] + [DeepPot(model) for model in models[1:]] + trj_freq = task.trj_freq + for fname in os.listdir("relax_trajs"): + with open(os.path.join("relax_trajs", fname), "rb") as f: + try: + data = pickle.load(f) + except Exception: + traceback.print_exc() + continue + nsteps = len(data["energy"]) + if nsteps <= 0: + continue + + step_list = [] + forces_list = [[] for _ in range(len(models))] + virial_list = [[] for _ in range(len(models))] + dump_str = "" + coords_list = [] + cell_list = [] + for i in range(0, nsteps, trj_freq): + atoms = ase.Atoms( + numbers=data["atomic_number"], + positions=data["atom_positions"][i], + pbc=True, + cell=data["cell"][i], + ) + calc = SinglePointCalculator( + atoms, + energy=data["energy"][i], + forces=data["forces"][i], + stress=data["stresses"][i], + ) + atoms.calc = calc + dump_str += atoms2lmpdump(atoms, i, type_map) + coords_list.append(data["atom_positions"][i]) + cell_list.append(data["cell"][i]) + step_list.append(i) + # Use results of model 0 directly + forces_list[0].append(data["forces"][i]) + virial_list[0].append( + -atoms.get_volume() + * atoms.get_stress(False).reshape(9) + / len(atoms) + ) + for j in range(1, len(models)): + dp = graphs[j] + atype = [ + dp.get_type_map().index(ase.Atom(i).symbol) # type: ignore + for i in data["atomic_number"] + ] + _, forces, virial = dp.eval( # type: ignore + np.array(coords_list), np.array(cell_list), atype + ) + forces_list[j] = forces + virial_list[j] = virial / len(atype) + traj_file = ip["task_path"] / ("traj.%s.dump" % fname) + traj_file.write_text(dump_str) + trajs.append(traj_file) + devi = [np.array(step_list)] + devi += list(calc_model_devi_v(np.array(virial_list))) + devi += list(calc_model_devi_f(np.array(forces_list))) + devi = np.vstack(devi).T + header = "%10s%19s%19s%19s%19s%19s%19s" % ( + "step", + "max_devi_v", + "min_devi_v", + "avg_devi_v", + "max_devi_f", + "min_devi_f", + "avg_devi_f", + ) + model_devi_file = ip["task_path"] / ("model_devi.%s.out" % fname) + np.savetxt( + model_devi_file, + devi, + fmt=["%12d"] + ["%19.6e"] * 6, + delimiter="", + header=header, + ) + model_devis.append(model_devi_file) + return OPIO( + { + "trajs": trajs, + "model_devis": model_devis, + } + ) diff --git a/dpgen2/superop/__init__.py b/dpgen2/superop/__init__.py index f573f7d1..0223605f 100644 --- a/dpgen2/superop/__init__.py +++ b/dpgen2/superop/__init__.py @@ -4,6 +4,9 @@ from .prep_run_calypso import ( PrepRunCaly, ) +from .prep_run_diffcsp import ( + PrepRunDiffCSP, +) from .prep_run_dp_train import ( PrepRunDPTrain, ) diff --git a/dpgen2/superop/block.py b/dpgen2/superop/block.py index a43c17c8..eea566d2 100644 --- a/dpgen2/superop/block.py +++ b/dpgen2/superop/block.py @@ -52,6 +52,9 @@ from .prep_run_calypso import ( PrepRunCaly, ) +from .prep_run_diffcsp import ( + PrepRunDiffCSP, +) from .prep_run_dp_train import ( PrepRunDPTrain, ) @@ -86,7 +89,7 @@ def __init__( self, name: str, prep_run_dp_train_op: PrepRunDPTrain, - prep_run_explore_op: Union[PrepRunLmp, PrepRunCaly], + prep_run_explore_op: Union[PrepRunLmp, PrepRunCaly, PrepRunDiffCSP], select_confs_op: Type[OP], prep_run_fp_op: PrepRunFp, collect_data_op: Type[OP], diff --git a/dpgen2/superop/prep_run_diffcsp.py b/dpgen2/superop/prep_run_diffcsp.py new file mode 100644 index 00000000..c7978e10 --- /dev/null +++ b/dpgen2/superop/prep_run_diffcsp.py @@ -0,0 +1,216 @@ +import os +from copy import ( + deepcopy, +) +from pathlib import ( + Path, +) +from typing import ( + Any, + Dict, + List, + Optional, + Type, + Union, +) + +from dflow import ( + InputArtifact, + InputParameter, + Inputs, + OPTemplate, + OutputArtifact, + OutputParameter, + Outputs, + Step, + Steps, + Workflow, + argo_len, + argo_range, + argo_sequence, + download_artifact, + upload_artifact, +) +from dflow.python import ( + OP, + OPIO, + Artifact, + OPIOSign, + PythonOPTemplate, + Slices, +) + +from dpgen2.utils.step_config import ( + init_executor, +) +from dpgen2.utils.step_config import normalize as normalize_step_dict + + +class PrepRunDiffCSP(Steps): + def __init__( + self, + name: str, + diffcsp_gen_op: Type[OP], + prep_relax_op: Type[OP], + run_relax_op: Type[OP], + prep_config: dict = normalize_step_dict({}), + run_config: dict = normalize_step_dict({}), + upload_python_packages: Optional[List[os.PathLike]] = None, + ): + self._input_parameters = { + "block_id": InputParameter(type=str, value=""), + "expl_task_grp": InputParameter(), + "explore_config": InputParameter(), + "type_map": InputParameter(), + } + self._input_artifacts = { + "models": InputArtifact(), + } + self._output_parameters = {} + self._output_artifacts = { + "trajs": OutputArtifact(), + "model_devis": OutputArtifact(), + } + + super().__init__( + name=name, + inputs=Inputs( + parameters=self._input_parameters, + artifacts=self._input_artifacts, + ), + outputs=Outputs( + parameters=self._output_parameters, + artifacts=self._output_artifacts, + ), + ) + self._keys = ["diffcsp-gen", "prep-relax", "run-relax"] + + self = _prep_run_diffcsp( + self, + diffcsp_gen_op, + prep_relax_op, + run_relax_op, + prep_config=prep_config, + run_config=run_config, + upload_python_packages=upload_python_packages, + ) + + @property + def input_parameters(self): + return self._input_parameters + + @property + def input_artifacts(self): + return self._input_artifacts + + @property + def output_parameters(self): + return self._output_parameters + + @property + def output_artifacts(self): + return self._output_artifacts + + @property + def keys(self): + return self._keys + + +def _prep_run_diffcsp( + prep_run_diffcsp_steps: Steps, + diffcsp_gen_op: Type[OP], + prep_relax_op: Type[OP], + run_relax_op: Type[OP], + prep_config: dict = normalize_step_dict({}), + run_config: dict = normalize_step_dict({}), + upload_python_packages: Optional[List[os.PathLike]] = None, +): + prep_config = deepcopy(prep_config) + run_config = deepcopy(run_config) + prep_template_config = prep_config.pop("template_config") + run_template_config = run_config.pop("template_config") + prep_executor = init_executor(prep_config.pop("executor")) + run_executor = init_executor(run_config.pop("executor")) + template_slice_config = run_config.pop("template_slice_config", {}) + + block_id = prep_run_diffcsp_steps.inputs.parameters["block_id"] + expl_task_grp = prep_run_diffcsp_steps.inputs.parameters["expl_task_grp"] + expl_config = prep_run_diffcsp_steps.inputs.parameters["explore_config"] + type_map = prep_run_diffcsp_steps.inputs.parameters["type_map"] + models = prep_run_diffcsp_steps.inputs.artifacts["models"] + + diffcsp_gen = Step( + "diffcsp-gen", + template=PythonOPTemplate( + diffcsp_gen_op, + python_packages=upload_python_packages, + slices=Slices( + "int('{{item}}')", + output_artifact=["cifs"], + **template_slice_config, + ), + **prep_template_config, + ), + parameters={ + "task_id": "{{item}}", + "config": expl_config, + }, + key="%s--diffcsp-gen-{{item}}" % block_id, + executor=prep_executor, + with_sequence=argo_sequence(expl_config["gen_tasks"], format="%06d"), # type: ignore + ) + prep_run_diffcsp_steps.add(diffcsp_gen) + + prep_relax = Step( + "prep-relax", + template=PythonOPTemplate( + prep_relax_op, + python_packages=upload_python_packages, + **prep_template_config, + ), + parameters={ + "expl_config": expl_config, + }, + artifacts={ + "cifs": diffcsp_gen.outputs.artifacts["cifs"], + }, + key="%s--prep-relax" % block_id, + executor=prep_executor, + ) + prep_run_diffcsp_steps.add(prep_relax) + + run_relax = Step( + "run-relax", + template=PythonOPTemplate( + run_relax_op, + python_packages=upload_python_packages, + slices=Slices( + "int('{{item}}')", + input_artifact=["task_path"], + output_artifact=["trajs", "model_devis"], + **template_slice_config, + ), + **run_template_config, + ), + parameters={ + "diffcsp_task_grp": expl_task_grp, + }, + artifacts={ + "models": models, + "task_path": prep_relax.outputs.artifacts["task_paths"], + }, + key="%s--run-relax-{{item}}" % block_id, + executor=run_executor, + with_sequence=argo_sequence( + prep_relax.outputs.parameters["ntasks"], format="%06d" + ), + ) + prep_run_diffcsp_steps.add(run_relax) + + prep_run_diffcsp_steps.outputs.artifacts[ + "trajs" + ]._from = run_relax.outputs.artifacts["trajs"] + prep_run_diffcsp_steps.outputs.artifacts[ + "model_devis" + ]._from = run_relax.outputs.artifacts["model_devis"] + return prep_run_diffcsp_steps diff --git a/examples/diffcsp/dpgen.json b/examples/diffcsp/dpgen.json new file mode 100644 index 00000000..b20a59c3 --- /dev/null +++ b/examples/diffcsp/dpgen.json @@ -0,0 +1,534 @@ +{ + "name": "diffcsp", + "bohrium_config": { + "username": "", + "password": "", + "project_id": 1, + "_comment": "all" + }, + "default_step_config": { + "template_config": { + "image": "", + "_comment": "all" + }, + "_comment": "all" + }, + "step_configs": { + "run_train_config": { + "template_config": { + "image": "", + "_comment": "all" + }, + "executor": { + "type": "dispatcher", + "retry_on_submission_error": 10, + "image_pull_policy": "IfNotPresent", + "machine_dict": { + "batch_type": "Bohrium", + "context_type": "Bohrium", + "remote_profile": { + "input_data": { + "job_type": "container", + "platform": "ali", + "on_demand": 0, + "scass_type": "1 * NVIDIA V100_16g" + } + } + } + }, + "_comment": "all" + }, + "prep_explore_config": { + "template_config": { + "image": "", + "_comment": "all" + }, + "executor": { + "type": "dispatcher", + "retry_on_submission_error": 10, + "image_pull_policy": "IfNotPresent", + "machine_dict": { + "batch_type": "Bohrium", + "context_type": "Bohrium", + "remote_profile": { + "input_data": { + "job_type": "container", + "platform": "ali", + "on_demand": 1, + "scass_type": "c12_m92_1 * NVIDIA V100" + } + } + } + }, + "_comment": "all" + }, + "run_explore_config": { + "template_config": { + "image": "", + "_comment": "all" + }, + "executor": { + "type": "dispatcher", + "retry_on_submission_error": 10, + "image_pull_policy": "IfNotPresent", + "machine_dict": { + "batch_type": "Bohrium", + "context_type": "Bohrium", + "remote_profile": { + "input_data": { + "job_type": "container", + "platform": "ali", + "on_demand": 1, + "scass_type": "c12_m92_1 * NVIDIA V100" + } + } + } + }, + "_comment": "all" + }, + "run_fp_config": { + "template_config": { + "image": "", + "_comment": "all" + }, + "continue_on_success_ratio": 0.8, + "executor": { + "type": "dispatcher", + "retry_on_submission_error": 10, + "image_pull_policy": "IfNotPresent", + "machine_dict": { + "batch_type": "Bohrium", + "context_type": "Bohrium", + "remote_profile": { + "input_data": { + "job_type": "container", + "platform": "ali", + "scass_type": "c8_m32_cpu" + } + } + } + }, + "template_slice_config": { + "group_size": 1, + "pool_size": 1 + }, + "_comment": "all" + }, + "_comment": "all" + }, + "upload_python_packages": [ + "/path/to/dpgen2" + ], + "inputs": { + "type_map": [ + "H", + "He", + "Li", + "Be", + "B", + "C", + "N", + "O", + "F", + "Ne", + "Na", + "Mg", + "Al", + "Si", + "P", + "S", + "Cl", + "Ar", + "K", + "Ca", + "Sc", + "Ti", + "V", + "Cr", + "Mn", + "Fe", + "Co", + "Ni", + "Cu", + "Zn", + "Ga", + "Ge", + "As", + "Se", + "Br", + "Kr", + "Rb", + "Sr", + "Y", + "Zr", + "Nb", + "Mo", + "Tc", + "Ru", + "Rh", + "Pd", + "Ag", + "Cd", + "In", + "Sn", + "Sb", + "Te", + "I", + "Xe", + "Cs", + "Ba", + "La", + "Ce", + "Pr", + "Nd", + "Pm", + "Sm", + "Eu", + "Gd", + "Tb", + "Dy", + "Ho", + "Er", + "Tm", + "Yb", + "Lu", + "Hf", + "Ta", + "W", + "Re", + "Os", + "Ir", + "Pt", + "Au", + "Hg", + "Tl", + "Pb", + "Bi", + "Po", + "At", + "Rn", + "Fr", + "Ra", + "Ac", + "Th", + "Pa", + "U", + "Np", + "Pu", + "Am", + "Cm", + "Bk", + "Cf", + "Es", + "Fm", + "Md", + "No", + "Lr", + "Rf", + "Db", + "Sg", + "Bh", + "Hs", + "Mt", + "Ds", + "Rg", + "Cn", + "Nh", + "Fl", + "Mc", + "Lv", + "Ts", + "Og" + ], + "mixed_type": true, + "mass_map": [ + 4.0, + 4.0026, + 6.94, + 9.0122, + 10.81, + 12.011, + 14.007, + 15.999, + 18.998, + 20.18, + 22.99, + 24.305, + 26.982, + 28.0855, + 30.974, + 32.06, + 35.45, + 39.95, + 39.098, + 40.078, + 44.956, + 47.867, + 50.942, + 51.996, + 54.938, + 55.845, + 58.933, + 58.693, + 63.546, + 65.38, + 69.723, + 72.63, + 74.922, + 78.971, + 79.904, + 83.798, + 85.468, + 87.62, + 88.906, + 91.224, + 92.906, + 95.95, + 97, + 101.07, + 102.91, + 106.42, + 107.87, + 112.41, + 114.82, + 118.71, + 121.76, + 127.6, + 126.9, + 131.29, + 132.91, + 137.33, + 138.91, + 140.12, + 140.91, + 144.24, + 145, + 150.36, + 151.96, + 157.25, + 158.93, + 162.5, + 164.93, + 167.26, + 168.93, + 173.05, + 174.97, + 178.49, + 180.95, + 183.84, + 186.21, + 190.23, + 192.22, + 195.08, + 196.97, + 200.59, + 204.38, + 207.2, + 208.98, + 209, + 210, + 222, + 223, + 226, + 227, + 232.04, + 231.04, + 238.03, + 237, + 244, + 243, + 247, + 247, + 251, + 252, + 257, + 258, + 259, + 262, + 267, + 268, + 269, + 270, + 269, + 277, + 281, + 282, + 285, + 286, + 290, + 290, + 293, + 294, + 294 + ], + "init_data_prefix": null, + "init_data_sys": [ + "init/data_0", + "init/data_1", + "init/data_2" + ], + "_comment": "all" + }, + "train": { + "type": "dp", + "numb_models": 4, + "config": { + "backend": "pytorch", + "init_model_policy": "yes", + "init_model_old_ratio": 0.9, + "init_model_numb_steps": 200, + "init_model_start_lr": 2e-05, + "init_model_start_pref_e": 0.25, + "init_model_start_pref_f": 100, + "_comment": "all" + }, + "template_script": "train.json", + "_comment": "all" + }, + "explore": { + "type": "diffcsp", + "config": { + "gen_tasks": 2, + "gen_command": "python /root/DiffCSP/scripts/generation.py --model_path /root/DiffCSP/models/mp_gen --dataset mp_20", + "relax_group_size": 10 + }, + "convergence": { + "type": "adaptive-lower", + "conv_tolerance": 0.005, + "_numb_candi_f": 3000, + "rate_candi_f": 0.15, + "level_f_hi": 0.5, + "n_checked_steps": 8, + "_command": "all" + }, + "max_numb_iter": 16, + "fatal_at_max": false, + "filters": [ + { + "type": "box_skewness", + "theta": 60.0 + }, + { + "type": "box_length", + "length_ratio": 5.0 + }, + { + "type": "distance", + "custom_safe_dist": {}, + "safe_dist_ratio": 1.0 + } + ], + "stages": [ + [ + { + "trj_freq": 10, + "fmax": 1e-4, + "steps": 200, + "timeout": null + } + ] + ], + "_comment": "all" + }, + "fp": { + "type": "vasp", + "task_max": 300, + "inputs_config": { + "pp_files": { + "Pd": "PBE/Pd/POTCAR", + "Sb": "PBE/Sb/POTCAR", + "Pm": "PBE/Pm/POTCAR", + "Se": "PBE/Se/POTCAR", + "Yb": "PBE/Yb/POTCAR", + "Sc": "PBE/Sc/POTCAR", + "Pb": "PBE/Pb/POTCAR", + "Sm": "PBE/Sm/POTCAR", + "Ga": "PBE/Ga/POTCAR", + "He": "PBE/He/POTCAR", + "Ac": "PBE/Ac/POTCAR", + "Mo": "PBE/Mo/POTCAR", + "Ni": "PBE/Ni/POTCAR", + "Am": "PBE/Am/POTCAR", + "Be": "PBE/Be/POTCAR", + "Al": "PBE/Al/POTCAR", + "Mg": "PBE/Mg/POTCAR", + "Mn": "PBE/Mn/POTCAR", + "Na": "PBE/Na/POTCAR", + "Ir": "PBE/Ir/POTCAR", + "In": "PBE/In/POTCAR", + "Cu": "PBE/Cu/POTCAR", + "Cr": "PBE/Cr/POTCAR", + "Co": "PBE/Co/POTCAR", + "Cf": "PBE/Cf/POTCAR", + "Fe": "PBE/Fe/POTCAR", + "Tb": "PBE/Tb/POTCAR", + "Te": "PBE/Te/POTCAR", + "U": "PBE/U/POTCAR", + "Ru": "PBE/Ru/POTCAR", + "Tl": "PBE/Tl/POTCAR", + "I": "PBE/I/POTCAR", + "N": "PBE/N/POTCAR", + "Rn": "PBE/Rn/POTCAR", + "Tm": "PBE/Tm/POTCAR", + "Tc": "PBE/Tc/POTCAR", + "S": "PBE/S/POTCAR", + "Rh": "PBE/Rh/POTCAR", + "F": "PBE/F/POTCAR", + "O": "PBE/O/POTCAR", + "H": "PBE/H/POTCAR", + "Zn": "PBE/Zn/POTCAR", + "Pr": "PBE/Pr/POTCAR", + "Pu": "PBE/Pu/POTCAR", + "Pt": "PBE/Pt/POTCAR", + "Po": "PBE/Po/POTCAR", + "Pa": "PBE/Pa/POTCAR", + "Sn": "PBE/Sn/POTCAR", + "Si": "PBE/Si/POTCAR", + "Ag": "PBE/Ag/POTCAR", + "Kr": "PBE/Kr/POTCAR", + "Nd": "PBE/Nd/POTCAR", + "Ge": "PBE/Ge/POTCAR", + "Hf": "PBE/Hf/POTCAR", + "Ar": "PBE/Ar/POTCAR", + "Ho": "PBE/Ho/POTCAR", + "Au": "PBE/Au/POTCAR", + "Dy": "PBE/Dy/POTCAR", + "Ne": "PBE/Ne/POTCAR", + "Bi": "PBE/Bi/POTCAR", + "At": "PBE/At/POTCAR", + "Np": "PBE/Np/POTCAR", + "As": "PBE/As/POTCAR", + "Hg": "PBE/Hg/POTCAR", + "Gd": "PBE/Gd/POTCAR", + "Br": "PBE/Br/POTCAR", + "Os": "PBE/Os/POTCAR", + "Lu": "PBE/Lu/POTCAR", + "Cd": "PBE/Cd/POTCAR", + "Cm": "PBE/Cm/POTCAR", + "Li": "PBE/Li/POTCAR", + "Eu": "PBE/Eu/POTCAR", + "Cl": "PBE/Cl/POTCAR", + "Er": "PBE/Er/POTCAR", + "La": "PBE/La/POTCAR", + "Ce": "PBE/Ce/POTCAR", + "C": "PBE/C/POTCAR", + "V": "PBE/V/POTCAR", + "Ta": "PBE/Ta/POTCAR", + "Th": "PBE/Th/POTCAR", + "B": "PBE/B/POTCAR", + "Re": "PBE/Re/POTCAR", + "Xe": "PBE/Xe/POTCAR", + "Ti": "PBE/Ti/POTCAR", + "P": "PBE/P/POTCAR", + "W": "PBE/W/POTCAR" + }, + "incar": "vasp/INCAR", + "kspacing": 0.32, + "kgamma": true + }, + "run_config": { + "command": "source /opt/intel/oneapi/setvars.sh && mpirun -n 16 vasp_std" + }, + "_comment": "all" + } +} diff --git a/tests/op/test_diffcsp_gen.py b/tests/op/test_diffcsp_gen.py new file mode 100644 index 00000000..68359ec9 --- /dev/null +++ b/tests/op/test_diffcsp_gen.py @@ -0,0 +1,485 @@ +import os +import shutil +import sys +import unittest +from unittest.mock import ( + Mock, +) + +import numpy as np +from dflow.python import ( + OPIO, +) + +from dpgen2.op import ( + DiffCSPGen, +) + + +class MockedTorchTensor: + def __init__(self, array): + self.array = array + + def numpy(self): + return np.array(self.array) + + +class Lattice: + @classmethod + def from_parameters(cls, *args, **kwargs): + return cls() + + +class Structure: + def __init__(self, *args, **kwargs): + pass + + def to(self, filename): + with open(filename, "w") as f: + f.write("# generated using pymatgen") + + +class TestDiffCSPGen(unittest.TestCase): + def testDiffCSPGen(self): + mocked_torch = Mock() + mocked_torch.load = lambda *args, **kwargs: { + "num_atoms": [4], + "lengths": MockedTorchTensor([[3.9783, 3.9739, 7.8712]]), + "angles": MockedTorchTensor([[75.4395, 75.2368, 59.9560]]), + "frac_coords": MockedTorchTensor( + [ + [0.9364, 0.8685, 0.3059], + [0.6652, 0.5965, 0.1211], + [0.3940, 0.3267, 0.9381], + [0.1671, 0.0943, 0.6182], + ] + ), + "atom_types": np.array( + [ + [ + 3.3592e-04, + 3.0107e-04, + 3.9129e-04, + -6.3763e-04, + 8.0520e-04, + 4.4414e-04, + -2.8986e-04, + -2.1392e-03, + 4.8722e-05, + -3.3231e-04, + -6.3191e-04, + 5.9795e-04, + 1.4732e-04, + -3.7197e-04, + 7.1938e-04, + -6.7943e-04, + 3.8498e-04, + -1.4324e-04, + -2.6901e-04, + -5.0780e-04, + -1.0818e-04, + -3.5262e-04, + -8.5311e-05, + -5.2838e-04, + 2.2596e-04, + -7.2861e-04, + 9.3495e-05, + 5.2268e-04, + -1.3498e-03, + 4.2470e-06, + 5.4609e-04, + -1.0146e-04, + 1.2022e-04, + -6.8679e-04, + 1.2461e-04, + 2.1435e-04, + -7.1471e-04, + 1.1662e-03, + -6.6708e-04, + 8.2295e-04, + 6.2853e-04, + -2.4491e-04, + 6.4987e-04, + -5.1100e-04, + -5.1476e-04, + -7.7254e-04, + 8.8924e-04, + 1.1588e-03, + 7.2993e-04, + -1.0635e-03, + 1.1115e-03, + 3.4901e-04, + -2.2457e-04, + 1.4041e-04, + 5.2349e-04, + -7.5847e-05, + 1.0064e00, + -1.3512e-04, + -5.7845e-05, + -5.7440e-04, + -3.5642e-04, + -1.0602e-03, + -7.2654e-04, + -3.3495e-04, + 4.6810e-04, + 2.2099e-04, + 1.2765e-03, + -3.9296e-04, + 2.3515e-04, + -9.7137e-04, + -2.0593e-04, + 1.8036e-05, + 1.9761e-04, + -6.6997e-04, + -2.3799e-04, + -1.0367e-03, + -7.8580e-04, + 1.6282e-04, + -2.2094e-05, + -6.1782e-04, + -8.0335e-04, + -6.1013e-05, + 5.1792e-04, + 8.8457e-04, + -1.4201e-04, + 8.3348e-05, + -8.7103e-04, + -3.7644e-04, + -7.9451e-04, + 2.1580e-04, + 2.4297e-04, + -6.2706e-04, + 2.1187e-04, + 2.1576e-04, + -7.6110e-04, + -6.6156e-04, + -1.6728e-04, + 7.9632e-04, + -2.4411e-04, + 8.1214e-05, + ], + [ + -7.0909e-04, + -4.9160e-05, + 6.4183e-04, + -2.3159e-05, + -2.1111e-04, + 9.9265e-01, + -2.5487e-04, + 2.0234e-03, + 3.7245e-04, + -2.1519e-04, + -1.2896e-03, + 6.5524e-04, + 5.8468e-04, + 6.9007e-04, + -6.8366e-04, + -1.5442e-03, + 1.1736e-03, + -2.7028e-04, + -8.7887e-04, + -6.9466e-04, + -4.2686e-04, + -1.3817e-03, + -5.9043e-04, + 2.4127e-05, + -5.0269e-04, + 8.5697e-04, + 8.0366e-05, + -2.3199e-04, + 1.0801e-03, + -5.6569e-04, + 1.3053e-03, + -4.3600e-04, + 2.7265e-04, + -1.4574e-03, + 4.8319e-04, + 1.1370e-04, + -1.9870e-05, + 2.5677e-04, + -2.1308e-04, + 4.2040e-04, + -1.7702e-03, + 3.6868e-04, + -6.4718e-05, + -1.2740e-03, + 5.9798e-05, + 5.6298e-04, + 1.5739e-04, + -1.2427e-03, + -2.3625e-04, + 1.2725e-03, + 2.7244e-03, + 2.2982e-04, + -5.3311e-04, + 3.9148e-04, + 6.8299e-05, + 1.0015e-03, + -9.0848e-04, + -5.1714e-04, + -1.0943e-03, + 2.5606e-05, + 2.8485e-04, + -1.1111e-03, + -1.0351e-04, + 3.5585e-04, + 5.3167e-04, + -1.8233e-04, + -7.3155e-04, + 4.4785e-04, + 1.0446e-03, + 1.3169e-04, + -6.6539e-04, + -5.3372e-04, + 3.3434e-04, + -5.2778e-04, + 4.0154e-04, + 1.0038e-03, + 8.9614e-04, + 8.2089e-04, + -9.7924e-04, + -2.7403e-04, + 2.7191e-04, + -8.9014e-04, + 1.0588e-03, + -1.1651e-03, + -8.0853e-04, + 9.6170e-05, + -4.3855e-04, + 2.9800e-04, + -6.9158e-05, + -5.0128e-04, + 1.6640e-03, + 1.6139e-03, + -2.5935e-04, + -6.6282e-04, + 3.5163e-05, + 1.8553e-04, + -6.3214e-04, + 1.1661e-03, + -3.2857e-04, + -3.5929e-04, + ], + [ + 3.4069e-04, + -2.0518e-04, + 5.5739e-04, + 7.0733e-06, + 9.7434e-04, + 2.2107e-04, + -2.1307e-04, + -1.1901e-03, + -1.0868e-04, + -4.3491e-04, + -4.6747e-04, + 8.8033e-04, + -5.3345e-04, + -4.5509e-04, + 3.9525e-04, + 3.9781e-05, + -4.7723e-04, + -4.7385e-04, + -2.4597e-04, + -1.0039e-03, + 1.0478e-04, + 1.1463e-03, + -4.7850e-04, + -3.1328e-05, + -3.7674e-06, + -4.1706e-04, + -1.5533e-05, + 8.7933e-04, + -4.6539e-04, + -2.1604e-04, + -2.5042e-04, + -1.0196e-04, + -1.2501e-04, + -3.7640e-04, + 2.4961e-04, + -1.6852e-05, + -4.4518e-04, + 9.3337e-04, + -6.9345e-04, + 4.2971e-04, + -1.8360e-04, + -2.1426e-04, + 6.1983e-04, + -6.2752e-05, + 7.9834e-04, + 1.9826e-04, + 2.5389e-04, + 6.1837e-04, + -1.4870e-04, + -3.5978e-04, + 8.5050e-04, + 5.4130e-04, + 5.9721e-04, + -3.9270e-04, + 3.0747e-04, + -3.9029e-04, + 1.0000e00, + -6.1231e-05, + 1.9509e-04, + -9.3342e-04, + -9.5543e-05, + 1.6481e-04, + -3.4956e-04, + 1.8306e-04, + 1.6214e-04, + 6.7432e-04, + 3.2698e-04, + 4.6957e-04, + 6.6374e-04, + 2.4401e-04, + -2.8953e-04, + 4.5274e-04, + 7.7035e-04, + -1.3329e-04, + 2.1814e-04, + -4.4580e-04, + -1.6980e-04, + 4.3275e-05, + 4.4877e-04, + -3.5606e-04, + 3.7028e-05, + -3.2027e-06, + -5.4841e-04, + 3.4593e-04, + -6.1065e-05, + 1.6113e-04, + -4.5586e-04, + -9.2710e-05, + 7.3830e-05, + -8.2601e-05, + 3.6123e-05, + -6.0147e-04, + 6.8068e-04, + 5.1152e-04, + -1.6020e-04, + -4.6825e-04, + -4.4536e-04, + -1.2684e-04, + 3.4957e-05, + -3.4964e-04, + ], + [ + 6.9942e-04, + -2.3009e-04, + 3.9129e-04, + -4.5735e-05, + -1.5210e-04, + -9.5870e-05, + 3.2569e-04, + -1.1394e-03, + 7.7203e-04, + 1.9814e-04, + 1.1035e-04, + 7.0963e-04, + -2.9960e-04, + -3.1560e-04, + 6.6132e-04, + 2.2888e-04, + 5.6245e-04, + -5.1699e-05, + 3.0806e-04, + 7.6312e-04, + 3.6271e-04, + 5.5544e-04, + -5.3967e-04, + -1.0092e-04, + 1.7226e-04, + 2.1784e-04, + -1.9432e-04, + 9.4856e-04, + 5.7033e-04, + 6.6431e-05, + 3.8977e-05, + -9.3216e-04, + -1.3489e-04, + -3.3908e-04, + 9.9789e-01, + 5.3704e-04, + -7.5510e-05, + -5.2399e-04, + 7.2395e-04, + -7.5516e-04, + -2.5590e-04, + 2.5873e-04, + 7.3807e-04, + 2.2888e-07, + -1.4879e-04, + 5.1340e-04, + -1.3458e-04, + -3.0084e-04, + 2.7856e-04, + -2.3874e-04, + 3.5563e-04, + 7.3932e-04, + -7.4016e-05, + -3.3168e-04, + -6.0069e-05, + 5.1038e-04, + -8.4479e-05, + 2.1435e-04, + 1.1598e-04, + 6.8415e-05, + 2.7146e-05, + -2.2083e-04, + -5.2108e-04, + -3.8993e-04, + -9.7699e-04, + 5.9252e-04, + -1.6931e-04, + -9.4700e-04, + -2.9051e-04, + -5.3630e-04, + -4.2456e-04, + -3.9352e-05, + -1.3817e-04, + -2.2283e-04, + 3.3220e-05, + 2.5155e-04, + -7.0318e-07, + 4.8145e-04, + 3.5942e-04, + 5.4292e-04, + -1.6069e-04, + 1.8675e-04, + 1.6393e-04, + 5.2633e-05, + -5.0289e-05, + -1.9838e-04, + 1.0472e-04, + -2.1608e-05, + -5.2919e-05, + -1.9588e-04, + -2.3803e-05, + 1.1158e-05, + -6.2318e-06, + 7.9619e-05, + -3.4735e-06, + -1.6122e-05, + -3.8437e-04, + -5.6557e-04, + 5.6475e-05, + 3.4945e-04, + ], + ] + ), + } + sys.modules["torch"] = mocked_torch + sys.modules["pymatgen.core.lattice"] = sys.modules[__name__] + sys.modules["pymatgen.core.structure"] = sys.modules[__name__] + config = {"gen_command": "echo 'mocked generation' --model_path ."} + op = DiffCSPGen() + op_in = OPIO({"config": config, "task_id": "000000"}) + op_out = op.execute(op_in) + + cifs = op_out["cifs"] + self.assertEqual(len(cifs), 1) + self.assertEqual(str(cifs[0]), "diffcsp.000000/0.cif") + self.assertTrue(cifs[0].read_text().startswith("# generated using pymatgen")) + + def tearDown(self): + if os.path.isdir("diffcsp.000000"): + shutil.rmtree("diffcsp.000000") diff --git a/tests/op/test_prep_relax.py b/tests/op/test_prep_relax.py new file mode 100644 index 00000000..83dcc2b7 --- /dev/null +++ b/tests/op/test_prep_relax.py @@ -0,0 +1,46 @@ +import os +import shutil +import unittest +from pathlib import ( + Path, +) + +from dflow.python import ( + OPIO, +) + +from dpgen2.op import ( + PrepRelax, +) + + +class TestPrepRelax(unittest.TestCase): + def testPrepRelax(self): + cifs = [] + for i in range(4): + p = Path("%i.cif" % i) + p.write_text("Mocked cif.") + cifs.append(p) + op_in = OPIO( + { + "expl_config": { + "relax_group_size": 2, + }, + "cifs": cifs, + } + ) + op = PrepRelax() + op_out = op.execute(op_in) + self.assertEqual(op_out["ntasks"], 2) + self.assertEqual(len(op_out["task_paths"]), 2) + for i, task_path in enumerate(op_out["task_paths"]): + self.assertEqual(str(task_path), "task.%06d" % i) + self.assertEqual(len(list(task_path.iterdir())), 2) + + def tearDown(self): + for i in range(2): + if os.path.isdir("task.%06d" % i): + shutil.rmtree("task.%06d" % i) + for i in range(4): + if os.path.isfile("%s.cif" % i): + os.remove("%s.cif" % i) diff --git a/tests/op/test_run_relax.py b/tests/op/test_run_relax.py new file mode 100644 index 00000000..847056b6 --- /dev/null +++ b/tests/op/test_run_relax.py @@ -0,0 +1,299 @@ +import os +import pickle +import shutil +import sys +import unittest +from pathlib import ( + Path, +) +from unittest.mock import ( + patch, +) + +import numpy as np +from dflow.python import ( + OPIO, +) + +from dpgen2.exploration.task import ( + DiffCSPTaskGroup, +) +from dpgen2.op import ( + RunRelax, +) + +type_map = [ + "H", + "He", + "Li", + "Be", + "B", + "C", + "N", + "O", + "F", + "Ne", + "Na", + "Mg", + "Al", + "Si", + "P", + "S", + "Cl", + "Ar", + "K", + "Ca", + "Sc", + "Ti", + "V", + "Cr", + "Mn", + "Fe", + "Co", + "Ni", + "Cu", + "Zn", + "Ga", + "Ge", + "As", + "Se", + "Br", + "Kr", + "Rb", + "Sr", + "Y", + "Zr", + "Nb", + "Mo", + "Tc", + "Ru", + "Rh", + "Pd", + "Ag", + "Cd", + "In", + "Sn", + "Sb", + "Te", + "I", + "Xe", + "Cs", + "Ba", + "La", + "Ce", + "Pr", + "Nd", + "Pm", + "Sm", + "Eu", + "Gd", + "Tb", + "Dy", + "Ho", + "Er", + "Tm", + "Yb", + "Lu", + "Hf", + "Ta", + "W", + "Re", + "Os", + "Ir", + "Pt", + "Au", + "Hg", + "Tl", + "Pb", + "Bi", + "Po", + "At", + "Rn", + "Fr", + "Ra", + "Ac", + "Th", + "Pa", + "U", + "Np", + "Pu", + "Am", + "Cm", + "Bk", + "Cf", + "Es", + "Fm", + "Md", + "No", + "Lr", + "Rf", + "Db", + "Sg", + "Bh", + "Hs", + "Mt", + "Ds", + "Rg", + "Cn", + "Nh", + "Fl", + "Mc", + "Lv", + "Ts", + "Og", +] + + +class DeepPot: + def __init__(self, *args, **kwargs): + pass + + def get_type_map(self): + return type_map + + def eval(*args, **kwargs): + e = np.array([[-24.51994085]]) + f = np.array( + [ + [ + [0.09876919, 0.05018996, -0.38124716], + [-0.08633465, -0.05225886, 0.32088113], + [-0.00362617, -0.00345137, 0.011181], + [-0.00880837, 0.00552027, 0.04918503], + ] + ] + ) + v = np.array( + [ + [ + -2.35308634, + 0.02806883, + -0.31607685, + 0.02806883, + -2.29484913, + -0.15927731, + -0.31607685, + -0.15927731, + -1.31155533, + ] + ] + ) + return e, f, v + + +def calc_model_devi_f(fs): + fs_devi = np.linalg.norm(np.std(fs, axis=0), axis=-1) + max_devi_f = np.max(fs_devi, axis=-1) + min_devi_f = np.min(fs_devi, axis=-1) + avg_devi_f = np.mean(fs_devi, axis=-1) + return max_devi_f, min_devi_f, avg_devi_f + + +def calc_model_devi_v(vs): + vs_devi = np.std(vs, axis=0) + max_devi_v = np.max(vs_devi, axis=-1) + min_devi_v = np.min(vs_devi, axis=-1) + avg_devi_v = np.linalg.norm(vs_devi, axis=-1) / 3 + return max_devi_v, min_devi_v, avg_devi_v + + +def relax_run(*args, **kwargs): + os.makedirs("relax_trajs", exist_ok=True) + with open("relax_trajs/0", "wb") as f: + pickle.dump( + { + "energy": [-24.48511280170925], + "forces": [ + np.array( + [ + [0.07237937, 0.0548471, -0.25049795], + [-0.05103608, -0.04855114, 0.14486373], + [-0.00699195, -0.00202532, 0.05028293], + [-0.01435133, -0.00427063, 0.05535129], + ] + ) + ], + "stresses": [ + np.array( + [ + 0.02573293, + 0.02483516, + 0.02152565, + -0.00011215, + 0.0011876, + 0.00021871, + ] + ) + ], + "atom_positions": [ + np.array( + [ + [5.16042463, 2.95461897, 4.22479426], + [2.10195627, 1.11158311, 8.10963874], + [3.62923674, 2.02945063, 2.22314751], + [0.81206749, 0.32068803, 5.12969607], + ] + ) + ], + "cell": [ + np.array( + [ + [3.84691675, 0.0, 1.01375962], + [1.79422664, 3.40209201, 0.99903783], + [0.0, 0.0, 7.87117147], + ] + ) + ], + "atomic_number": np.array([57, 57, 6, 35]), + }, + f, + ) + + +class DPCalculator: + def __init__(self, dp): + self.dp = dp + + +class Relaxer: + def __init__(self, *args, **kwargs): + dp = DeepPot() + self.calculator = DPCalculator(dp) + + +class TestRunRelax(unittest.TestCase): + @patch("dpgen2.op.run_relax.atoms2lmpdump") + def testRunRelax(self, mocked_run): + mocked_run.side_effect = ["ITEM: TIMESTEP"] + sys.modules["deepmd.infer"] = sys.modules[__name__] + sys.modules["deepmd.infer.model_devi"] = sys.modules[__name__] + sys.modules["lam_optimize.main"] = sys.modules[__name__] + sys.modules["lam_optimize.relaxer"] = sys.modules[__name__] + + task_group = DiffCSPTaskGroup() + task_group.make_task() + os.makedirs("task.000000", exist_ok=True) + op_in = OPIO( + { + "diffcsp_task_grp": task_group, + "task_path": Path("task.000000"), + "models": [Path("model_0.pt"), Path("model_1.pt")], + } + ) + op = RunRelax() + op_out = op.execute(op_in) + self.assertEqual(len(op_out["trajs"]), 1) + self.assertTrue(op_out["trajs"][0].read_text().startswith("ITEM: TIMESTEP")) + self.assertEqual(len(op_out["model_devis"]), 1) + model_devi = np.array( + [0.0, 0.1132373, 0.00632493, 0.04404319, 0.0897801, 0.006415, 0.04564122] + ) + np.testing.assert_array_almost_equal( + np.loadtxt(op_out["model_devis"][0]), model_devi + ) + + def tearDown(self): + if os.path.isdir("task.000000"): + shutil.rmtree("task.000000") + if os.path.isdir("relax_trajs"): + shutil.rmtree("relax_trajs") diff --git a/tests/test_check_examples.py b/tests/test_check_examples.py index f7309093..1a4e40d9 100644 --- a/tests/test_check_examples.py +++ b/tests/test_check_examples.py @@ -23,6 +23,7 @@ p_examples / "chno" / "input.json", p_examples / "water" / "input_dpgen_abacus.json", p_examples / "water" / "input_dpgen_cp2k.json", + p_examples / "diffcsp" / "dpgen.json", ) diff --git a/tests/test_prep_run_diffcsp.py b/tests/test_prep_run_diffcsp.py new file mode 100644 index 00000000..d5569d17 --- /dev/null +++ b/tests/test_prep_run_diffcsp.py @@ -0,0 +1,124 @@ +import glob +import os +import shutil +import unittest +from pathlib import ( + Path, +) + +from dflow import ( + Step, + Workflow, + download_artifact, + upload_artifact, +) +from dflow.python import ( + OP, + OPIO, +) + +from dpgen2.exploration.task import ( + DiffCSPTaskGroup, +) +from dpgen2.op import ( + DiffCSPGen, + PrepRelax, + RunRelax, +) +from dpgen2.superop import ( + PrepRunDiffCSP, +) + + +class MockedDiffCSPGen(DiffCSPGen): + @OP.exec_sign_check + def execute( + self, + ip: OPIO, + ) -> OPIO: + task_dir = Path("diffcsp.%s" % ip["task_id"]) + task_dir.mkdir(exist_ok=True) + for i in range(2): + fpath = task_dir / ("%s.cif" % i) + fpath.write_text("Mocked cif.") + return OPIO( + { + "cifs": list(Path(task_dir).glob("*.cif")), + } + ) + + +class MockedRunRelax(RunRelax): + @OP.exec_sign_check + def execute( + self, + ip: OPIO, + ) -> OPIO: + cifs = os.listdir(ip["task_path"]) + assert len(cifs) == 2 + trajs = [] + model_devis = [] + for cif in cifs: + name = cif[:-4] + traj = ip["task_path"] / ("traj.%s.dump" % name) + traj.write_text("Mocked traj.") + trajs.append(traj) + model_devi = ip["task_path"] / ("model_devi.%s.out" % name) + model_devi.write_text("Mocked model_devi.") + model_devis.append(model_devi) + return OPIO( + { + "trajs": trajs, + "model_devis": model_devis, + } + ) + + +class TestPrepRunDiffCSP(unittest.TestCase): + def testPrepRunDiffCSP(self): + task_group = DiffCSPTaskGroup() + task_group.make_task() + + wf = Workflow("test-prep-run-diffcsp") + upload_packages = [] + if "__file__" in globals(): + upload_packages.append(__file__) + upload_packages.append(os.path.dirname(__file__)) + steps = PrepRunDiffCSP( + "prep-run-diffcsp", + MockedDiffCSPGen, + PrepRelax, + MockedRunRelax, + upload_python_packages=upload_packages, + ) + step = Step( + "main", + template=steps, + parameters={ + "block_id": "iter-000000", + "expl_task_grp": task_group, + "explore_config": { + "gen_tasks": 2, + "gen_command": "echo 'mocked generation' --model_path .", + "relax_group_size": 2, + }, + "type_map": [], + }, + artifacts={ + "models": upload_artifact([]), + }, + ) + wf.add(step) + wf.submit() + wf.wait() + self.assertEqual(wf.query_status(), "Succeeded") + step = wf.query_step("main")[0] + trajs = download_artifact(step.outputs.artifacts["trajs"]) + self.assertEqual(len(trajs), 4) + model_devis = download_artifact(step.outputs.artifacts["model_devis"]) + self.assertEqual(len(model_devis), 4) + + def tearDown(self): + for d in glob.glob("test-prep-run-diffcsp-*") + ["task.000000", "task.000001"]: + if os.path.isdir(d): + shutil.rmtree(d)