Support spin in dpgen2.

dpgen2/entrypoint/submit.py

@@ -63,6 +63,7 @@
 )
 from dpgen2.exploration.render import (
     TrajRenderLammps,
+    TrajRenderLammpsSpin,
 )
 from dpgen2.exploration.report import (
     ExplorationReportTrustLevelsRandom,
@@ -81,6 +82,7 @@
     CustomizedLmpTemplateTaskGroup,
     ExplorationStage,
     ExplorationTask,
+    LmpSpinTaskGroup,
     LmpTemplateTaskGroup,
     NPTTaskGroup,
     caly_normalize,
@@ -377,7 +379,10 @@ def make_lmp_naive_exploration_scheduler(config):
     # report
     conv_style = convergence.pop("type")
     report = conv_styles[conv_style](**convergence)
-    render = TrajRenderLammps(nopbc=output_nopbc, use_ele_temp=use_ele_temp)
+    if "spin" in conv_style:
+        render = TrajRenderLammpsSpin(nopbc=output_nopbc, use_ele_temp=use_ele_temp)
+    else:
+        render = TrajRenderLammps(nopbc=output_nopbc, use_ele_temp=use_ele_temp)
     # selector
     selector = ConfSelectorFrames(
         render,

dpgen2/exploration/deviation/__init__.py

@@ -1,6 +1,9 @@
 from .deviation_manager import (
     DeviManager,
 )
+from .deviation_spin import (
+    DeviManagerSpin,
+)
 from .deviation_std import (
     DeviManagerStd,
 )

dpgen2/exploration/deviation/deviation_spin.py

@@ -0,0 +1,127 @@
+from collections import (
+    defaultdict,
+)
+from typing import (
+    Dict,
+    List,
+    Optional,
+)
+
+import numpy as np
+
+from .deviation_manager import (
+    DeviManager,
+)
+
+
+class DeviManagerSpin(DeviManager):
+    r"""The class which is responsible for DeepSPIN model deviation management.
+
+    This is the implementation of DeviManager for DeepSPIN model. Each deviation
+    (e.g. max_devi_af, max_devi_mf in file `model_devi.out`) is stored
+    as a List[Optional[np.ndarray]], where np.array is a one-dimensional
+    array.
+    A List[np.ndarray][ii][jj] is the force model deviation of the jj-th
+    frame of the ii-th trajectory.
+    The model deviation can be List[None], where len(List[None]) is
+    the number of trajectory files.
+
+    """
+
+    MAX_DEVI_AF = "max_devi_af"
+    MIN_DEVI_AF = "min_devi_af"
+    AVG_DEVI_AF = "avg_devi_af"
+    MAX_DEVI_MF = "max_devi_mf"
+    MIN_DEVI_MF = "min_devi_mf"
+    AVG_DEVI_MF = "avg_devi_mf"
+
+    def __init__(self):
+        super().__init__()
+        self._data = defaultdict(list)
+
+    def _check_name(self, name: str):
+        assert name in (
+            DeviManager.MAX_DEVI_V,
+            DeviManager.MIN_DEVI_V,
+            DeviManager.AVG_DEVI_V,
+            self.MAX_DEVI_AF,
+            self.MIN_DEVI_AF,
+            self.AVG_DEVI_AF,
+            self.MAX_DEVI_MF,
+            self.MIN_DEVI_MF,
+            self.AVG_DEVI_MF,
+        ), f"Error: unknown deviation name {name}"
+
+    def _add(self, name: str, deviation: np.ndarray) -> None:
+        assert isinstance(
+            deviation, np.ndarray
+        ), f"Error: deviation(type: {type(deviation)}) is not a np.ndarray"
+        assert len(deviation.shape) == 1, (
+            f"Error: deviation(shape: {deviation.shape}) is not a "
+            + f"one-dimensional array"
+        )
+        self._data[name].append(deviation)
+        self.ntraj = max(self.ntraj, len(self._data[name]))
+
+    def _get(self, name: str) -> List[Optional[np.ndarray]]:
+        if self.ntraj == 0:
+            return []
+        elif len(self._data[name]) == 0:
+            return [None for _ in range(self.ntraj)]
+        else:
+            return self._data[name]
+
+    def clear(self) -> None:
+        self.__init__()
+        return None
+
+    def _check_data(self) -> None:
+        r"""Check if data is valid"""
+        model_devi_names = (
+            DeviManager.MAX_DEVI_V,
+            DeviManager.MIN_DEVI_V,
+            DeviManager.AVG_DEVI_V,
+            self.MAX_DEVI_AF,
+            self.MIN_DEVI_AF,
+            self.AVG_DEVI_AF,
+            self.MAX_DEVI_MF,
+            self.MIN_DEVI_MF,
+            self.AVG_DEVI_MF,
+        )
+        # check the length of model deviations
+        frames = {}
+        for name in model_devi_names:
+            if len(self._data[name]) > 0:
+                assert len(self._data[name]) == self.ntraj, (
+                    f"Error: the number of model deviation {name} "
+                    + f"({len(self._data[name])}) and trajectory files ({self.ntraj}) "
+                    + f"are not equal."
+                )
+            for idx, ndarray in enumerate(self._data[name]):
+                assert isinstance(ndarray, np.ndarray), (
+                    f"Error: model deviation in {name} is not ndarray, "
+                    + f"index: {idx}, type: {type(ndarray)}"
+                )
+
+            frames[name] = [arr.shape[0] for arr in self._data[name]]
+            if len(frames[name]) == 0:
+                frames.pop(name)
+
+        # check if "max_devi_af" and "max_devi_mf" exist
+        assert (
+            len(self._data[self.MAX_DEVI_AF]) == self.ntraj
+        ), f"Error: cannot find model deviation {self.MAX_DEVI_AF}"
+        assert (
+            len(self._data[self.MAX_DEVI_MF]) == self.ntraj
+        ), f"Error: cannot find model deviation {self.MAX_DEVI_MF}"
+
+        # check if the length of the arrays corresponding to the same
+        # trajectory has the same number of frames
+        non_empty_deviations = list(frames.keys())
+        for name in non_empty_deviations[1:]:
+            assert frames[name] == frames[non_empty_deviations[0]], (
+                f"Error: the number of frames in {name} is different "
+                + f"with that in {non_empty_deviations[0]}.\n"
+                + f"{name}: {frames[name]}\n"
+                + f"{non_empty_deviations[0]}: {frames[non_empty_deviations[0]]}\n"
+            )

dpgen2/exploration/render/__init__.py

@@ -4,3 +4,6 @@
 from .traj_render_lammps import (
     TrajRenderLammps,
 )
+from .traj_render_lammps_spin import (
+    TrajRenderLammpsSpin,
+)

dpgen2/exploration/render/traj_render_lammps_spin.py

@@ -0,0 +1,76 @@
+from pathlib import (
+    Path,
+)
+from typing import (
+    TYPE_CHECKING,
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
+
+import dpdata
+import numpy as np
+
+from ..deviation import (
+    DeviManager,
+    DeviManagerSpin,
+)
+from .traj_render import (
+    TrajRender,
+)
+
+if TYPE_CHECKING:
+    from dpgen2.exploration.selector import (
+        ConfFilters,
+    )
+
+
+class TrajRenderLammpsSpin(TrajRender):
+    def __init__(
+        self,
+        nopbc: bool = False,
+        use_ele_temp: int = 0,
+    ):
+        self.nopbc = nopbc
+
+    def get_model_devi(
+        self,
+        files: List[Path],
+    ) -> DeviManagerSpin:
+        ntraj = len(files)
+
+        model_devi = DeviManagerSpin()
+        for ii in range(ntraj):
+            self._load_one_model_devi(files[ii], model_devi)
+
+        return model_devi
+
+    def _load_one_model_devi(self, fname, model_devi):
+        dd = np.loadtxt(fname)
+        model_devi.add(DeviManagerSpin.MAX_DEVI_AF, dd[:, 4])
+        model_devi.add(DeviManagerSpin.MIN_DEVI_AF, dd[:, 5])
+        model_devi.add(DeviManagerSpin.AVG_DEVI_AF, dd[:, 6])
+        model_devi.add(DeviManagerSpin.MAX_DEVI_MF, dd[:, 7])
+        model_devi.add(DeviManagerSpin.MIN_DEVI_MF, dd[:, 8])
+        model_devi.add(DeviManagerSpin.AVG_DEVI_MF, dd[:, 9])
+
+    def get_confs(
+        self,
+        trajs: List[Path],
+        id_selected: List[List[int]],
+        type_map: Optional[List[str]] = None,
+        conf_filters: Optional["ConfFilters"] = None,
+        optional_outputs: Optional[List[Path]] = None,
+    ) -> dpdata.MultiSystems:
+        del conf_filters  # by far does not support conf filters
+        ntraj = len(trajs)
+        traj_fmt = "lammps/dump"
+        ms = dpdata.MultiSystems(type_map=type_map)
+        for ii in range(ntraj):
+            if len(id_selected[ii]) > 0:
+                ss = dpdata.System(trajs[ii], fmt=traj_fmt, type_map=type_map)
+                ss.nopbc = self.nopbc
+                ss = ss.sub_system(id_selected[ii])
+                ms.append(ss)
+        return ms

dpgen2/exploration/report/__init__.py

@@ -10,9 +10,13 @@
 from .report_trust_levels_random import (
     ExplorationReportTrustLevelsRandom,
 )
+from .report_trust_levels_spin import (
+    ExplorationReportTrustLevelsSpin,
+)
 
 conv_styles = {
     "fixed-levels": ExplorationReportTrustLevelsRandom,
     "fixed-levels-max-select": ExplorationReportTrustLevelsMax,
+    "fixed-levels-max-select-spin": ExplorationReportTrustLevelsSpin,
     "adaptive-lower": ExplorationReportAdaptiveLower,
 }