add prob charge

Author: YuzhiLiu-ai

deepmd/entrypoints/test.py

@@ -75,6 +75,7 @@ def test(
     detail_file: str,
     atomic: bool,
     head: str | None = None,
+    output_latent_charge: bool = False,
     **kwargs: Any,
 ) -> None:
     """Test model predictions.
@@ -185,6 +186,7 @@ def test(
                 detail_file,
                 atomic,
                 append_detail=(cc != 0),
+                output_latent_charge=output_latent_charge,
             )
         elif isinstance(dp, DeepDOS):
             err = test_dos(
@@ -305,6 +307,7 @@ def test_ener(
     detail_file: str | None,
     has_atom_ener: bool,
     append_detail: bool = False,
+    output_latent_charge: bool = False,
 ) -> tuple[list[np.ndarray], list[int]]:
     """Test energy type model.
 
@@ -402,32 +405,46 @@ def test_ener(
         efield=efield,
         mixed_type=mixed_type,
         spin=spin,
+        output_latent_charge=output_latent_charge,
     )
     energy = ret[0]
     force = ret[1]
     virial = ret[2]
     energy = energy.reshape([numb_test, 1])
     force = force.reshape([numb_test, -1])
     virial = virial.reshape([numb_test, 9])
-    if dp.has_hessian:
-        hessian = ret[3]
-        hessian = hessian.reshape([numb_test, -1])
+    idx = 3
     if has_atom_ener:
-        ae = ret[3]
-        av = ret[4]
+        ae = ret[idx]
+        idx += 1
+        av = ret[idx]
+        idx += 1
         ae = ae.reshape([numb_test, -1])
         av = av.reshape([numb_test, -1])
         if dp.has_spin:
-            force_m = ret[5]
+            force_m = ret[idx]
+            idx += 1
+            mask_mag = ret[idx]
+            idx += 1
             force_m = force_m.reshape([numb_test, -1])
-            mask_mag = ret[6]
             mask_mag = mask_mag.reshape([numb_test, -1])
     else:
         if dp.has_spin:
-            force_m = ret[3]
+            force_m = ret[idx]
+            idx += 1
+            mask_mag = ret[idx]
+            idx += 1
             force_m = force_m.reshape([numb_test, -1])
-            mask_mag = ret[4]
             mask_mag = mask_mag.reshape([numb_test, -1])
+    if dp.has_hessian:
+        hessian = ret[idx]
+        idx += 1
+        hessian = hessian.reshape([numb_test, -1])
+    latent_charge = None
+    if output_latent_charge:
+        latent_charge = ret[idx]
+        idx += 1
+        latent_charge = latent_charge.reshape([numb_test, -1])
     out_put_spin = dp.get_ntypes_spin() != 0 or dp.has_spin
     if out_put_spin:
         if dp.get_ntypes_spin() != 0:  # old tf support for spin
@@ -659,6 +676,13 @@ def test_ener(
                 header=f"{system}: data_h pred_h (3Na*3Na matrix in row-major order)",
                 append=append_detail,
             )
+        if output_latent_charge and latent_charge is not None:
+            save_txt_file(
+                detail_path.with_suffix(".q.out"),
+                latent_charge,
+                header=f"{system}: pred_q (latent charge per atom)",
+                append=append_detail,
+            )
 
     return dict_to_return
 

deepmd/infer/deep_eval.py

@@ -78,6 +78,7 @@ class DeepEvalBackend(ABC):
         "global_polar": "global_polar",
         "wfc": "wfc",
         "energy_derv_r_derv_r": "hessian",
+        "latent_charge": "latent_charge",
     }
 
     @abstractmethod

deepmd/infer/deep_pot.py

@@ -203,13 +203,34 @@ def eval(
             nframes,
             natoms,
         ) = self._standard_input(coords, cells, atom_types, fparam, aparam, mixed_type)
+        output_latent_charge = kwargs.pop("output_latent_charge", False)
+        extra_request_defs = []
+        if output_latent_charge:
+            # Try to get dim_out_lr from the model fitting net
+            dim_out_lr = 1
+            model = self.deep_eval.get_model()
+            if hasattr(model, "atomic_model") and hasattr(
+                model.atomic_model, "fitting_net"
+            ):
+                dim_out_lr = getattr(model.atomic_model.fitting_net, "dim_out_lr", 1)
+            extra_request_defs.append(
+                OutputVariableDef(
+                    "latent_charge",
+                    shape=[dim_out_lr],
+                    reducible=False,
+                    r_differentiable=False,
+                    c_differentiable=False,
+                    atomic=True,
+                )
+            )
         results = self.deep_eval.eval(
             coords,
             cells,
             atom_types,
             atomic,
             fparam=fparam,
             aparam=aparam,
+            extra_request_defs=extra_request_defs,
             **kwargs,
         )
         energy = results["energy_redu"].reshape(nframes, 1)
@@ -251,6 +272,9 @@ def eval(
                 nframes, 3 * natoms, 3 * natoms
             )
             result = (*list(result), hessian)
+        if output_latent_charge and "latent_charge" in results:
+            latent_charge = results["latent_charge"].reshape(nframes, natoms, -1)
+            result = (*result, latent_charge)
         return result
 
 

deepmd/main.py

@@ -454,6 +454,12 @@ def main_parser() -> argparse.ArgumentParser:
         default=False,
         help="(Supported backend: PyTorch) Disable JIT compilation when loading the model.",
     )
+    parser_tst.add_argument(
+        "--output-latent-charge",
+        action="store_true",
+        default=False,
+        help="Output latent charge predicted by SOG models to the detail file.",
+    )
 
     # * eval_desc script ***************************************************************
     parser_eval_desc = subparsers.add_parser(

deepmd/pt/infer/deep_eval.py

@@ -384,6 +384,8 @@ def eval(
             coords, atom_types, len(atom_types.shape) > 1
         )
         request_defs = self._get_request_defs(atomic)
+        if "extra_request_defs" in kwargs:
+            request_defs = request_defs + kwargs["extra_request_defs"]
         if "spin" not in kwargs or kwargs["spin"] is None:
             out = self._eval_func(self._eval_model, numb_test, natoms)(
                 coords, cells, atom_types, fparam, aparam, request_defs

deepmd/pt/model/model/sog_model.py

@@ -485,6 +485,8 @@ def forward(
                 model_predict["mask"] = model_ret["mask"]
             if self._hessian_enabled:
                 model_predict["hessian"] = model_ret["energy_derv_r_derv_r"].squeeze(-2)
+            if "latent_charge" in model_ret:
+                model_predict["latent_charge"] = model_ret["latent_charge"]
         else:
             model_predict = model_ret
             model_predict["updated_coord"] += coord
@@ -528,6 +530,8 @@ def forward_lower(
             else:
                 assert model_ret["dforce"] is not None
                 model_predict["dforce"] = model_ret["dforce"]
+            if "latent_charge" in model_ret:
+                model_predict["latent_charge"] = model_ret["latent_charge"]
         else:
             model_predict = model_ret
         return model_predict