add extra_fact

Author: iProzd

deepmd/pt/model/task/ener.py

@@ -536,6 +536,8 @@ def __init__(
         add_angle_readout: bool = False,
         slim_edge_readout: bool = False,
         slim_angle_readout: bool = False,
+        edge_extra_fact: float = 1.0,
+        angle_extra_fact: float = 1.0,
         **kwargs,
     ) -> None:
         """Construct a fitting net for energy.
@@ -549,6 +551,8 @@ def __init__(
         """
         self.add_edge_readout = add_edge_readout
         self.add_angle_readout = add_angle_readout
+        self.edge_extra_fact = edge_extra_fact
+        self.angle_extra_fact = angle_extra_fact
         super().__init__(
             "energy",
             ntypes,
@@ -714,7 +718,7 @@ def forward(
                 # nf x nloc x 1
                 edge_energy = torch.sum(edge_atomic_contrib, dim=-2)
             # energy
-            out = out + edge_energy / self.norm_e_fact
+            out = out + (edge_energy * self.edge_extra_fact) / self.norm_e_fact
 
         if self.add_angle_readout:
             assert angle_embd is not None
@@ -747,5 +751,5 @@ def forward(
                 )
             # energy
             # self.norm_a_fact ** 2
-            out = out + angle_energy / (self.norm_a_fact**2)
+            out = out + (angle_energy * self.angle_extra_fact) / (self.norm_a_fact**2)
         return {self.var_name: out.to(env.GLOBAL_PT_FLOAT_PRECISION)}

deepmd/utils/argcheck.py

@@ -2220,6 +2220,18 @@ def fitting_ener_readout():
             optional=True,
             default=False,
         ),
+        Argument(
+            "edge_extra_fact",
+            float,
+            optional=True,
+            default=1.0,
+        ),
+        Argument(
+            "angle_extra_fact",
+            float,
+            optional=True,
+            default=1.0,
+        ),
     ]