Merge branch 'devel' into rf_finetune

Signed-off-by: Duo <[email protected]>

.pre-commit-config.yaml

@@ -29,7 +29,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.4.5
+    rev: v0.4.7
     hooks:
       - id: ruff
         args: ["--fix"]

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -80,6 +80,10 @@ def mixed_types(self) -> bool:
         """
         return self.descriptor.mixed_types()
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the atomic model has message passing."""
+        return self.descriptor.has_message_passing()
+
     def forward_atomic(
         self,
         extended_coord: np.ndarray,

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -92,6 +92,10 @@ def mixed_types(self) -> bool:
         """
         return True
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the atomic model has message passing."""
+        return any(model.has_message_passing() for model in self.models)
+
     def get_rcut(self) -> float:
         """Get the cut-off radius."""
         return max(self.get_model_rcuts())

deepmd/dpmodel/atomic_model/make_base_atomic_model.py

@@ -115,6 +115,10 @@ def mixed_types(self) -> bool:
             """
             pass
 
+        @abstractmethod
+        def has_message_passing(self) -> bool:
+            """Returns whether the descriptor has message passing."""
+
         @abstractmethod
         def fwd(
             self,

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -131,6 +131,10 @@ def mixed_types(self) -> bool:
         # to match DPA1 and DPA2.
         return True
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the atomic model has message passing."""
+        return False
+
     def serialize(self) -> dict:
         dd = BaseAtomicModel.serialize(self)
         dd.update(

deepmd/dpmodel/descriptor/descriptor.py

@@ -125,3 +125,7 @@ def call(
     ):
         """Calculate DescriptorBlock."""
         pass
+
+    @abstractmethod
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor block has message passing."""

deepmd/dpmodel/descriptor/dpa1.py

@@ -355,6 +355,10 @@ def mixed_types(self) -> bool:
         """
         return self.se_atten.mixed_types()
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return self.se_atten.has_message_passing()
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.se_atten.get_env_protection()
@@ -906,6 +910,10 @@ def call(
             sw,
         )
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor block has message passing."""
+        return False
+
 
 class NeighborGatedAttention(NativeOP):
     def __init__(

deepmd/dpmodel/descriptor/dpa2.py

@@ -531,6 +531,12 @@ def mixed_types(self) -> bool:
         """
         return True
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return any(
+            [self.repinit.has_message_passing(), self.repformers.has_message_passing()]
+        )
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.env_protection

deepmd/dpmodel/descriptor/hybrid.py

@@ -142,6 +142,10 @@ def mixed_types(self):
         """
         return any(descrpt.mixed_types() for descrpt in self.descrpt_list)
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return any(descrpt.has_message_passing() for descrpt in self.descrpt_list)
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix. All descriptors should be the same."""
         all_protection = [descrpt.get_env_protection() for descrpt in self.descrpt_list]

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -100,6 +100,10 @@ def mixed_types(self) -> bool:
             """
             pass
 
+        @abstractmethod
+        def has_message_passing(self) -> bool:
+            """Returns whether the descriptor has message passing."""
+
         @abstractmethod
         def get_env_protection(self) -> float:
             """Returns the protection of building environment matrix."""

deepmd/dpmodel/descriptor/repformers.py

@@ -387,6 +387,10 @@ def call(
         rot_mat = np.transpose(h2g2, (0, 1, 3, 2))
         return g1, g2, h2, rot_mat.reshape(-1, nloc, self.dim_emb, 3), sw
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor block has message passing."""
+        return True
+
 
 # translated by GPT and modified
 def get_residual(

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -256,6 +256,10 @@ def mixed_types(self):
         """
         return False
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return False
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.env_protection

deepmd/dpmodel/descriptor/se_r.py

@@ -214,6 +214,10 @@ def mixed_types(self):
         """
         return False
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return False
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.env_protection

deepmd/dpmodel/descriptor/se_t.py

@@ -200,6 +200,10 @@ def mixed_types(self):
         """
         return False
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return False
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.env_protection

deepmd/dpmodel/model/make_model.py

@@ -475,6 +475,10 @@ def mixed_types(self) -> bool:
             """
             return self.atomic_model.mixed_types()
 
+        def has_message_passing(self) -> bool:
+            """Returns whether the model has message passing."""
+            return self.atomic_model.has_message_passing()
+
         def atomic_output_def(self) -> FittingOutputDef:
             """Get the output def of the atomic model."""
             return self.atomic_model.atomic_output_def()

deepmd/entrypoints/doc.py

@@ -4,6 +4,7 @@
 from deepmd.utils.argcheck import (
     gen_doc,
     gen_json,
+    gen_json_schema,
 )
 
 __all__ = ["doc_train_input"]
@@ -15,6 +16,8 @@ def doc_train_input(*, out_type: str = "rst", **kwargs):
         doc_str = gen_doc(make_anchor=True)
     elif out_type == "json":
         doc_str = gen_json()
+    elif out_type == "json_schema":
+        doc_str = gen_json_schema()
     else:
         raise RuntimeError(f"Unsupported out type {out_type}")
     print(doc_str)  # noqa: T201

deepmd/main.py

@@ -505,7 +505,7 @@ def main_parser() -> argparse.ArgumentParser:
     parsers_doc.add_argument(
         "--out-type",
         default="rst",
-        choices=["rst", "json"],
+        choices=["rst", "json", "json_schema"],
         type=str,
         help="The output type",
     )

deepmd/pt/entrypoints/main.py

@@ -288,10 +288,7 @@ def train(FLAGS):
 
 def freeze(FLAGS):
     model = torch.jit.script(inference.Tester(FLAGS.model, head=FLAGS.head).model)
-    if '"type": "dpa2"' in model.get_model_def_script():
-        extra_files = {"type": "dpa2"}
-    else:
-        extra_files = {"type": "else"}
+    extra_files = {}
     torch.jit.save(
         model,
         FLAGS.output,

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -102,6 +102,10 @@ def slim_type_map(self, type_map: List[str]) -> None:
         self.descriptor.slim_type_map(type_map=type_map)
         self.fitting_net.slim_type_map(type_map=type_map)
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the atomic model has message passing."""
+        return self.descriptor.has_message_passing()
+
     def serialize(self) -> dict:
         dd = BaseAtomicModel.serialize(self)
         dd.update(

deepmd/pt/model/atomic_model/linear_atomic_model.py

@@ -104,6 +104,10 @@ def mixed_types(self) -> bool:
         """
         return True
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the atomic model has message passing."""
+        return any(model.has_message_passing() for model in self.models)
+
     def get_out_bias(self) -> torch.Tensor:
         return self.out_bias
 

deepmd/pt/model/atomic_model/pairtab_atomic_model.py

@@ -160,6 +160,10 @@ def mixed_types(self) -> bool:
         # to match DPA1 and DPA2.
         return True
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the atomic model has message passing."""
+        return False
+
     def serialize(self) -> dict:
         dd = BaseAtomicModel.serialize(self)
         dd.update(

deepmd/pt/model/descriptor/descriptor.py

@@ -169,6 +169,10 @@ def forward(
         """Calculate DescriptorBlock."""
         pass
 
+    @abstractmethod
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor block has message passing."""
+
 
 def make_default_type_embedding(
     ntypes,

deepmd/pt/model/descriptor/dpa1.py

@@ -349,6 +349,10 @@ def mixed_types(self) -> bool:
         """
         return self.se_atten.mixed_types()
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return self.se_atten.has_message_passing()
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.se_atten.get_env_protection()

deepmd/pt/model/descriptor/dpa2.py

@@ -301,6 +301,12 @@ def mixed_types(self) -> bool:
         """
         return True
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return any(
+            [self.repinit.has_message_passing(), self.repformers.has_message_passing()]
+        )
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         # the env_protection of repinit is the same as that of the repformer

deepmd/pt/model/descriptor/hybrid.py

@@ -147,6 +147,10 @@ def mixed_types(self):
         """
         return any(descrpt.mixed_types() for descrpt in self.descrpt_list)
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return any(descrpt.has_message_passing() for descrpt in self.descrpt_list)
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix. All descriptors should be the same."""
         all_protection = [descrpt.get_env_protection() for descrpt in self.descrpt_list]

deepmd/pt/model/descriptor/repformers.py

@@ -548,3 +548,7 @@ def get_stats(self) -> Dict[str, StatItem]:
                 "The statistics of the descriptor has not been computed."
             )
         return self.stats
+
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor block has message passing."""
+        return True

deepmd/pt/model/descriptor/se_a.py

@@ -153,6 +153,10 @@ def mixed_types(self):
         """
         return self.sea.mixed_types()
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return self.sea.has_message_passing()
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.sea.get_env_protection()
@@ -687,3 +691,7 @@ def forward(
             None,
             sw,
         )
+
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor block has message passing."""
+        return False

deepmd/pt/model/descriptor/se_atten.py

@@ -584,6 +584,10 @@ def forward(
             sw,
         )
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor block has message passing."""
+        return False
+
 
 class NeighborGatedAttention(nn.Module):
     def __init__(

deepmd/pt/model/descriptor/se_r.py

@@ -176,6 +176,10 @@ def mixed_types(self) -> bool:
         """
         return False
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return False
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.env_protection

deepmd/pt/model/descriptor/se_t.py

@@ -182,6 +182,10 @@ def mixed_types(self):
         """
         return self.seat.mixed_types()
 
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return self.seat.has_message_passing()
+
     def get_env_protection(self) -> float:
         """Returns the protection of building environment matrix."""
         return self.seat.get_env_protection()
@@ -702,3 +706,7 @@ def forward(
             None,
             sw,
         )
+
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor block has message passing."""
+        return False

deepmd/pt/model/model/frozen.py

@@ -106,6 +106,11 @@ def mixed_types(self) -> bool:
         """
         return self.model.mixed_types()
 
+    @torch.jit.export
+    def has_message_passing(self) -> bool:
+        """Returns whether the descriptor has message passing."""
+        return self.model.has_message_passing()
+
     @torch.jit.export
     def forward(
         self,

deepmd/pt/model/model/make_model.py

@@ -535,6 +535,11 @@ def mixed_types(self) -> bool:
             """
             return self.atomic_model.mixed_types()
 
+        @torch.jit.export
+        def has_message_passing(self) -> bool:
+            """Returns whether the model has message passing."""
+            return self.atomic_model.has_message_passing()
+
         def forward(
             self,
             coord,