Merge branch 'devel' into add_seed

Signed-off-by: Duo <50307526+iProzd@users.noreply.github.com>

.github/workflows/test_python.yml

@@ -24,13 +24,18 @@ jobs:
       with:
         python-version: ${{ matrix.python }}
     - run: python -m pip install -U uv
-    - run: uv pip install --system --only-binary=horovod -e .[cpu,test,torch] horovod[tensorflow-cpu] mpi4py mpich
+    - run: |
+        uv pip install --system mpich
+        uv pip install --system "torch==2.3.0+cpu.cxx11.abi" -i https://download.pytorch.org/whl/
+        export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
+        uv pip install --system --only-binary=horovod -e .[cpu,test] horovod[tensorflow-cpu] mpi4py
       env:
         # Please note that uv has some issues with finding
         # existing TensorFlow package. Currently, it uses
         # TensorFlow in the build dependency, but if it
         # changes, setting `TENSORFLOW_ROOT`.
         TENSORFLOW_VERSION: ${{ matrix.python == '3.8' && '2.13.1' || '2.16.1' }}
+        DP_ENABLE_PYTORCH: 1
         DP_BUILD_TESTING: 1
         UV_EXTRA_INDEX_URL: "https://pypi.anaconda.org/njzjz/simple https://pypi.anaconda.org/mpi4py/simple"
     - run: dp --version

deepmd/dpmodel/descriptor/dpa1.py

@@ -193,6 +193,12 @@ class DescrptDPA1(NativeOP, BaseDescriptor):
             Setting this parameter to `True` is equivalent to setting `tebd_input_mode` to 'strip'.
             Setting it to `False` is equivalent to setting `tebd_input_mode` to 'concat'.
             The default value is `None`, which means the `tebd_input_mode` setting will be used instead.
+    use_econf_tebd: bool, Optional
+            Whether to use electronic configuration type embedding.
+    type_map: List[str], Optional
+            A list of strings. Give the name to each type of atoms.
+            Only used if `use_econf_tebd` is `True` in type embedding net.
+
     spin
             (Only support None to keep consistent with other backend references.)
             (Not used in this version. Not-none option is not implemented.)
@@ -242,6 +248,8 @@ def __init__(
         concat_output_tebd: bool = True,
         spin: Optional[Any] = None,
         stripped_type_embedding: Optional[bool] = None,
+        use_econf_tebd: bool = False,
+        type_map: Optional[List[str]] = None,
         # consistent with argcheck, not used though
         seed: Optional[int] = None,
     ) -> None:
@@ -287,12 +295,16 @@ def __init__(
             trainable_ln=trainable_ln,
             ln_eps=ln_eps,
         )
+        self.use_econf_tebd = use_econf_tebd
+        self.type_map = type_map
         self.type_embedding = TypeEmbedNet(
             ntypes=ntypes,
             neuron=[tebd_dim],
             padding=True,
             activation_function="Linear",
             precision=precision,
+            use_econf_tebd=use_econf_tebd,
+            type_map=type_map,
         )
         self.tebd_dim = tebd_dim
         self.concat_output_tebd = concat_output_tebd
@@ -457,6 +469,8 @@ def serialize(self) -> dict:
             "smooth_type_embedding": obj.smooth,
             "type_one_side": obj.type_one_side,
             "concat_output_tebd": self.concat_output_tebd,
+            "use_econf_tebd": self.use_econf_tebd,
+            "type_map": self.type_map,
             # make deterministic
             "precision": np.dtype(PRECISION_DICT[obj.precision]).name,
             "embeddings": obj.embeddings.serialize(),

deepmd/dpmodel/descriptor/dpa2.py

@@ -323,6 +323,8 @@ def __init__(
         trainable: bool = True,
         seed: Optional[int] = None,
         add_tebd_to_repinit_out: bool = False,
+        use_econf_tebd: bool = False,
+        type_map: Optional[List[str]] = None,
     ):
         r"""The DPA-2 descriptor. see https://arxiv.org/abs/2312.15492.
 
@@ -350,6 +352,11 @@ def __init__(
             (Unused yet) Random seed for parameter initialization.
         add_tebd_to_repinit_out : bool, optional
             Whether to add type embedding to the output representation from repinit before inputting it into repformer.
+        use_econf_tebd : bool, Optional
+            Whether to use electronic configuration type embedding.
+        type_map : List[str], Optional
+            A list of strings. Give the name to each type of atoms.
+            Only used if `use_econf_tebd` is `True` in type embedding net.
 
         Returns
         -------
@@ -433,12 +440,16 @@ def init_subclass_params(sub_data, sub_class):
             trainable_ln=self.repformer_args.trainable_ln,
             ln_eps=self.repformer_args.ln_eps,
         )
+        self.use_econf_tebd = use_econf_tebd
+        self.type_map = type_map
         self.type_embedding = TypeEmbedNet(
             ntypes=ntypes,
             neuron=[self.repinit_args.tebd_dim],
             padding=True,
             activation_function="Linear",
             precision=precision,
+            use_econf_tebd=use_econf_tebd,
+            type_map=type_map,
         )
         self.concat_output_tebd = concat_output_tebd
         self.precision = precision
@@ -641,6 +652,8 @@ def serialize(self) -> dict:
             "env_protection": self.env_protection,
             "trainable": self.trainable,
             "add_tebd_to_repinit_out": self.add_tebd_to_repinit_out,
+            "use_econf_tebd": self.use_econf_tebd,
+            "type_map": self.type_map,
             "type_embedding": self.type_embedding.serialize(),
             "g1_shape_tranform": self.g1_shape_tranform.serialize(),
         }

deepmd/dpmodel/infer/deep_eval.py

@@ -76,10 +76,10 @@ def __init__(
         self,
         model_file: str,
         output_def: ModelOutputDef,
-        *args: List[Any],
+        *args: Any,
         auto_batch_size: Union[bool, int, AutoBatchSize] = True,
         neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         self.output_def = output_def
         self.model_path = model_file
@@ -161,12 +161,12 @@ def get_ntypes_spin(self):
     def eval(
         self,
         coords: np.ndarray,
-        cells: np.ndarray,
+        cells: Optional[np.ndarray],
         atom_types: np.ndarray,
         atomic: bool = False,
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> Dict[str, np.ndarray]:
         """Evaluate the energy, force and virial by using this DP.
 

deepmd/dpmodel/utils/type_embed.py

@@ -39,6 +39,11 @@ class TypeEmbedNet(NativeOP):
         Random seed for initializing the network parameters.
     padding
         Concat the zero padding to the output, as the default embedding of empty type.
+    use_econf_tebd: bool, Optional
+        Whether to use electronic configuration type embedding.
+    type_map: List[str], Optional
+        A list of strings. Give the name to each type of atoms.
+        Only used if `use_econf_tebd` is `True` in type embedding net.
     """
 
     def __init__(
@@ -52,6 +57,8 @@ def __init__(
         trainable: bool = True,
         seed: Optional[int] = None,
         padding: bool = False,
+        use_econf_tebd: bool = False,
+        type_map: Optional[List[str]] = None,
     ) -> None:
         self.ntypes = ntypes
         self.neuron = neuron
@@ -61,8 +68,33 @@ def __init__(
         self.activation_function = str(activation_function)
         self.trainable = trainable
         self.padding = padding
+        self.use_econf_tebd = use_econf_tebd
+        self.type_map = type_map
+        embed_input_dim = ntypes
+        if self.use_econf_tebd:
+            from deepmd.utils.econf_embd import (
+                ECONF_DIM,
+                electronic_configuration_embedding,
+            )
+            from deepmd.utils.econf_embd import type_map as periodic_table
+
+            assert (
+                self.type_map is not None
+            ), "When using electronic configuration type embedding, type_map must be provided!"
+
+            missing_types = [t for t in self.type_map if t not in periodic_table]
+            assert not missing_types, (
+                "When using electronic configuration type embedding, "
+                "all element in type_map should be in periodic table! "
+                f"Found these invalid elements: {missing_types}"
+            )
+            self.econf_tebd = np.array(
+                [electronic_configuration_embedding[kk] for kk in self.type_map],
+                dtype=PRECISION_DICT[self.precision],
+            )
+            embed_input_dim = ECONF_DIM
         self.embedding_net = EmbeddingNet(
-            ntypes,
+            embed_input_dim,
             self.neuron,
             self.activation_function,
             self.resnet_dt,
@@ -71,9 +103,12 @@ def __init__(
 
     def call(self) -> np.ndarray:
         """Compute the type embedding network."""
-        embed = self.embedding_net(
-            np.eye(self.ntypes, dtype=PRECISION_DICT[self.precision])
-        )
+        if not self.use_econf_tebd:
+            embed = self.embedding_net(
+                np.eye(self.ntypes, dtype=PRECISION_DICT[self.precision])
+            )
+        else:
+            embed = self.embedding_net(self.econf_tebd)
         if self.padding:
             embed = np.pad(embed, ((0, 1), (0, 0)), mode="constant")
         return embed
@@ -120,5 +155,7 @@ def serialize(self) -> dict:
             "activation_function": self.activation_function,
             "trainable": self.trainable,
             "padding": self.padding,
+            "use_econf_tebd": self.use_econf_tebd,
+            "type_map": self.type_map,
             "embedding": self.embedding_net.serialize(),
         }

deepmd/entrypoints/test.py

@@ -46,14 +46,7 @@
 )
 
 if TYPE_CHECKING:
-    from deepmd.tf.infer import (
-        DeepDipole,
-        DeepDOS,
-        DeepPolar,
-        DeepPot,
-        DeepWFC,
-    )
-    from deepmd.tf.infer.deep_tensor import (
+    from deepmd.infer.deep_tensor import (
         DeepTensor,
     )
 

deepmd/infer/deep_dos.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from typing import (
     Any,
-    Dict,
     List,
     Optional,
     Tuple,
@@ -70,7 +69,7 @@ def eval(
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
         mixed_type: bool = False,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> Tuple[np.ndarray, ...]:
         """Evaluate energy, force, and virial. If atomic is True,
         also return atomic energy and atomic virial.

deepmd/infer/deep_eval.py

@@ -11,6 +11,7 @@
     List,
     Optional,
     Tuple,
+    Type,
     Union,
 )
 
@@ -82,10 +83,10 @@ def __init__(
         self,
         model_file: str,
         output_def: ModelOutputDef,
-        *args: List[Any],
+        *args: Any,
         auto_batch_size: Union[bool, int, AutoBatchSize] = True,
         neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         pass
 
@@ -99,12 +100,12 @@ def __new__(cls, model_file: str, *args, **kwargs):
     def eval(
         self,
         coords: np.ndarray,
-        cells: np.ndarray,
+        cells: Optional[np.ndarray],
         atom_types: np.ndarray,
         atomic: bool = False,
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> Dict[str, np.ndarray]:
         """Evaluate the energy, force and virial by using this DP.
 
@@ -166,13 +167,13 @@ def get_dim_aparam(self) -> int:
     def eval_descriptor(
         self,
         coords: np.ndarray,
-        cells: np.ndarray,
+        cells: Optional[np.ndarray],
         atom_types: np.ndarray,
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
         efield: Optional[np.ndarray] = None,
         mixed_type: bool = False,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> np.ndarray:
         """Evaluate descriptors by using this DP.
 
@@ -246,11 +247,11 @@ def _check_mixed_types(self, atom_types: np.ndarray) -> bool:
             # assume mixed_types if there are virtual types, even when
             # the atom types of all frames are the same
             return False
-        return np.all(np.equal(atom_types, atom_types[0]))
+        return np.all(np.equal(atom_types, atom_types[0])).item()
 
     @property
     @abstractmethod
-    def model_type(self) -> "DeepEval":
+    def model_type(self) -> Type["DeepEval"]:
         """The the evaluator of the model type."""
 
     @abstractmethod
@@ -316,10 +317,10 @@ def __new__(cls, model_file: str, *args, **kwargs):
     def __init__(
         self,
         model_file: str,
-        *args: List[Any],
+        *args: Any,
         auto_batch_size: Union[bool, int, AutoBatchSize] = True,
         neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         self.deep_eval = DeepEvalBackend(
             model_file,
@@ -387,7 +388,7 @@ def eval_descriptor(
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
         mixed_type: bool = False,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> np.ndarray:
         """Evaluate descriptors by using this DP.
 

deepmd/infer/deep_polar.py

@@ -50,7 +50,7 @@ def eval(
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
         mixed_type: bool = False,
-        **kwargs: dict,
+        **kwargs,
     ) -> np.ndarray:
         """Evaluate the model.