add DiffCSP exploration engine

Signed-off-by: zjgemi <liuxin_zijian@163.com>

dpgen2/entrypoint/args.py

@@ -317,17 +317,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<task_group_sec>` "
+ )
+ 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[dict], 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",
  [
  Argument("lmp", dict, lmp_args(), doc=doc_lmp),
  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,
  )

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)

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}")
 
@@ -271,6 +287,8 @@ def make_naive_exploration_scheduler(
  return make_lmp_naive_exploration_scheduler(config)
  elif "calypso" in explore_style:
  return make_calypso_naive_exploration_scheduler(config)
+ elif explore_style == "diffcsp":
+ return make_diffcsp_naive_exploration_scheduler(config)
  else:
  raise KeyError(
  f"Unknown key `{explore_style}`, Only support `lmp`, `calypso`, `calypso:merge` and `calypso:default`."
@@ -406,6 +424,53 @@ def make_lmp_naive_exploration_scheduler(config):
  return scheduler
 
 
+def make_diffcsp_naive_exploration_scheduler(config):
+ model_devi_jobs = config["explore"]["stages"]
+ fp_task_max = config["fp"]["task_max"]
+ max_numb_iter = config["explore"]["max_numb_iter"]
+ fatal_at_max = config["explore"]["fatal_at_max"]
+ convergence = config["explore"]["convergence"]
+ output_nopbc = config["explore"]["output_nopbc"]
+ scheduler = ExplorationScheduler()
+ # 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(
+ render,
+ report,
+ fp_task_max,
+ )
+
+ for job_ in model_devi_jobs:
+ if not isinstance(job_, list):
+ job = [job_]
+ else:
+ job = job_
+ # stage
+ stage = ExplorationStage()
+ for jj in job:
+ jconf = diffcsp_normalize(jj)
+ # make task group
+ tgroup = make_diffcsp_task_group_from_config(jconf)
+ # add the list to task group
+ tasks = tgroup.make_task()
+ stage.add_task_group(tasks)
+ # stage_scheduler
+ stage_scheduler = ConvergenceCheckStageScheduler(
+ stage,
+ selector,
+ max_numb_iter=max_numb_iter,
+ fatal_at_max=fatal_at_max,
+ )
+ # scheduler
+ scheduler.add_stage_scheduler(stage_scheduler)
+
+ return scheduler
+
+
 def get_kspacing_kgamma_from_incar(
  fname,
 ):
@@ -800,6 +865,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 +914,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 +954,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 +979,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)
 

dpgen2/exploration/task/__init__.py

@@ -7,19 +7,24 @@
 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,
 )
 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,
 )

dpgen2/exploration/task/diffcsp_task_group.py

@@ -0,0 +1,53 @@
+import copy
+import logging
+import random
+from typing import (
+ List,
+ Optional,
+)
+
+import numpy as np
+
+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
+ task.fmax = self.fmax
+ task.steps = self.steps
+ task.timeout = self.timeout
+ return task