feat(pt): add direct fitting

Author: iProzd

deepmd/dpmodel/fitting/make_base_fitting.py

@@ -67,6 +67,9 @@ def compute_output_stats(self, merged) -> NoReturn:
             """Update the output bias for fitting net."""
             raise NotImplementedError
 
+        def need_additional_input(self) -> bool:
+            return False
+
         @abstractmethod
         def get_type_map(self) -> list[str]:
             """Get the name to each type of atoms."""

deepmd/pt/infer/deep_eval.py

@@ -467,10 +467,19 @@ def _eval_model(
                 out = batch_output[pt_name].reshape(shape).detach().cpu().numpy()
                 results.append(out)
             else:
-                shape = self._get_output_shape(odef, nframes, natoms)
-                results.append(
-                    np.full(np.abs(shape), np.nan, dtype=prec)
-                )  # this is kinda hacky
+                if (
+                    self._OUTDEF_DP2BACKEND[odef.name] == "force"
+                    and "dforce" in batch_output
+                ):
+                    # if no force, use dforce if possible
+                    shape = self._get_output_shape(odef, nframes, natoms)
+                    out = batch_output["dforce"].reshape(shape).detach().cpu().numpy()
+                    results.append(out)
+                else:
+                    shape = self._get_output_shape(odef, nframes, natoms)
+                    results.append(
+                        np.full(np.abs(shape), np.nan, dtype=prec)
+                    )  # this is kinda hacky
         return tuple(results)
 
     def _eval_model_spin(

deepmd/pt/loss/ener.py

@@ -187,6 +187,9 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
             Other losses for display.
         """
         model_pred = model(**input_dict)
+
+        if "force" not in model_pred and "dforce" in model_pred:
+            model_pred["force"] = model_pred["dforce"]
         coef = learning_rate / self.starter_learning_rate
         pref_e = self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * coef
         pref_f = self.limit_pref_f + (self.start_pref_f - self.limit_pref_f) * coef

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -245,15 +245,30 @@ def forward_atomic(
         if self.enable_eval_descriptor_hook:
             self.eval_descriptor_list.append(descriptor)
         # energy, force
-        fit_ret = self.fitting_net(
-            descriptor,
-            atype,
-            gr=rot_mat,
-            g2=g2,
-            h2=h2,
-            fparam=fparam,
-            aparam=aparam,
-        )
+        if not self.fitting_net.need_additional_input():
+            fit_ret = self.fitting_net(
+                descriptor,
+                atype,
+                gr=rot_mat,
+                g2=g2,
+                h2=h2,
+                fparam=fparam,
+                aparam=aparam,
+            )
+        else:
+            add_input = self.descriptor.get_additional_output_for_fitting()
+            fit_ret = self.fitting_net(
+                descriptor,
+                atype,
+                gr=rot_mat,
+                g2=g2,
+                h2=h2,
+                fparam=fparam,
+                aparam=aparam,
+                diff=add_input["diff"],
+                edge_index=add_input["edge_index"],
+                sw=add_input["sw"],
+            )
         return fit_ret
 
     def get_out_bias(self) -> torch.Tensor:

deepmd/pt/model/descriptor/dpa3.py

@@ -252,6 +252,9 @@ def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.rcut
 
+    def get_additional_output_for_fitting(self):
+        return self.repflows.get_additional_output_for_fitting()
+
     def get_rcut_smth(self) -> float:
         """Returns the radius where the neighbor information starts to smoothly decay to 0."""
         return self.rcut_smth

deepmd/pt/model/descriptor/repflows.py

@@ -447,6 +447,7 @@ def __init__(
                 )
             )
         self.layers = torch.nn.ModuleList(layers)
+        self.additional_output_for_fitting: dict[str, Optional[torch.Tensor]] = {}
 
         wanted_shape = (self.ntypes, self.nnei, 4)
         mean = torch.zeros(wanted_shape, dtype=self.prec, device=env.DEVICE)
@@ -461,6 +462,8 @@ def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.e_rcut
 
+    additional_output_for_fitting: dict[str, Optional[torch.Tensor]]
+
     def get_rcut_smth(self) -> float:
         """Returns the radius where the neighbor information starts to smoothly decay to 0."""
         return self.e_rcut_smth
@@ -548,6 +551,9 @@ def reinit_exclude(
         self.exclude_types = exclude_types
         self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types)
 
+    def get_additional_output_for_fitting(self):
+        return self.additional_output_for_fitting
+
     def forward(
         self,
         nlist: torch.Tensor,
@@ -782,6 +788,8 @@ def forward(
             sw = sw[nlist_mask]
             # n_edge x 4
             dmatrix = dmatrix[nlist_mask]
+            # n_edge x 3
+            diff = diff[nlist_mask]
 
             if self.edge_use_esen_atom_ebd:
                 assert source_type is not None
@@ -809,12 +817,16 @@ def forward(
                     * d_sw[:, :, None, :, None]
                     * d_sw[:, :, None, None, :]
                 )[d_nlist_mask]
+            self.additional_output_for_fitting["edge_index"] = edge_index
         else:
             # avoid jit assertion
             edge_index = angle_index = torch.zeros(
                 [1, 3], device=nlist.device, dtype=nlist.dtype
             )
             dihedral_index = None
+            self.additional_output_for_fitting["edge_index"] = None
+        self.additional_output_for_fitting["diff"] = diff
+        self.additional_output_for_fitting["sw"] = sw
         # get edge and angle embedding
         # nb x nloc x nnei x e_dim [OR] n_edge x e_dim
         if self.edge_use_esen_rbf:

deepmd/pt/model/model/__init__.py

@@ -90,14 +90,14 @@ def _get_standard_model_components(model_params, ntypes):
     fitting_net["ntypes"] = descriptor.get_ntypes()
     fitting_net["type_map"] = copy.deepcopy(model_params["type_map"])
     fitting_net["mixed_types"] = descriptor.mixed_types()
-    if fitting_net["type"] in ["dipole", "polar"]:
+    if fitting_net["type"] in ["dipole", "polar", "ener_direct"]:
         fitting_net["embedding_width"] = descriptor.get_dim_emb()
     fitting_net["dim_descrpt"] = descriptor.get_dim_out()
-    grad_force = "direct" not in fitting_net["type"]
-    if not grad_force:
-        fitting_net["out_dim"] = descriptor.get_dim_emb()
-        if "ener" in fitting_net["type"]:
-            fitting_net["return_energy"] = True
+    # grad_force = "direct" not in fitting_net["type"]
+    # if not grad_force:
+    #     fitting_net["out_dim"] = descriptor.get_dim_emb()
+    #     if "ener" in fitting_net["type"]:
+    #         fitting_net["return_energy"] = True
     fitting = BaseFitting(**fitting_net)
     return descriptor, fitting, fitting_net["type"]
 
@@ -261,7 +261,7 @@ def get_standard_model(model_params):
         modelcls = PolarModel
     elif fitting_net_type == "dos":
         modelcls = DOSModel
-    elif fitting_net_type in ["ener", "direct_force_ener"]:
+    elif fitting_net_type in ["ener", "direct_force_ener", "ener_direct"]:
         modelcls = EnergyModel
     elif fitting_net_type == "property":
         modelcls = PropertyModel

deepmd/pt/model/model/ener_model.py

@@ -81,8 +81,8 @@ def forward(
                     model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(
                         -3
                     )
-            else:
-                model_predict["force"] = model_ret["dforce"]
+            if "dforce" in model_ret:
+                model_predict["dforce"] = model_ret["dforce"]
             if "mask" in model_ret:
                 model_predict["mask"] = model_ret["mask"]
         else:
@@ -119,15 +119,16 @@ def forward_lower(
             model_predict["energy"] = model_ret["energy_redu"]
             if self.do_grad_r("energy"):
                 model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2)
+            else:
+                assert model_ret["dforce"] is not None
+                model_predict["dforce"] = model_ret["dforce"]
+
             if self.do_grad_c("energy"):
                 model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2)
                 if do_atomic_virial:
                     model_predict["extended_virial"] = model_ret[
                         "energy_derv_c"
                     ].squeeze(-3)
-            else:
-                assert model_ret["dforce"] is not None
-                model_predict["dforce"] = model_ret["dforce"]
         else:
             model_predict = model_ret
         return model_predict