update for average charge

Author: YuzhiLiu-ai

deepmd/pt/model/model/sog_model.py

@@ -180,13 +180,19 @@ def _corr_redu(positions: torch.Tensor, charges: torch.Tensor) -> torch.Tensor:
             )
             corr_redu = _corr_redu(coord_for_grad, latent_charge)
 
-            force_local = -torch.autograd.grad(
+            grad_result = torch.autograd.grad(
                 [corr_redu],
                 [coord_for_grad],
                 grad_outputs=[torch.ones_like(corr_redu)],
                 create_graph=self.training,
                 retain_graph=True,
+                allow_unused=True,
             )[0]
+            force_local = (
+                -grad_result
+                if grad_result is not None
+                else torch.zeros_like(coord_for_grad)
+            )
             out: dict[str, torch.Tensor] = {"corr_redu": corr_redu}
             out["force_local"] = force_local
             if need_virial:

deepmd/pt/model/task/lr_fitting.py

@@ -129,6 +129,7 @@ def __init__(
         type_map: list[str] | None = None,
         use_aparam_as_mask: bool = False,
         default_fparam: list[float] | None = None,
+        use_charge_constraint: bool = False,
         **kwargs: Any,
     ) -> None:
         super().__init__()
@@ -159,6 +160,7 @@ def __init__(
             all(self.trainable) if isinstance(self.trainable, list) else self.trainable
         )
         self.remove_vaccum_contribution = remove_vaccum_contribution
+        self.use_charge_constraint = bool(use_charge_constraint)
         self.bias_atom_q_bound = 3.0
 
         self.sr_net_dim_out = self._sr_net_out_dim()
@@ -393,6 +395,7 @@ def serialize(self) -> dict:
             "trainable_lr": [self.trainable] * (len(self.neuron_lr) + 1),
             "layer_name": None,
             "use_aparam_as_mask": self.use_aparam_as_mask,
+            "use_charge_constraint": self.use_charge_constraint,
             "spin": None,
         }
 
@@ -402,6 +405,8 @@ def deserialize(cls, data: dict) -> "LRFittingNet":
         # Compatibility with old checkpoints.
         data.pop("use_type_embed_for_bias_q", None)
         data.pop("bias_atom_q_type_embed", None)
+        if "use_charge_constraint" not in data:
+            data["use_charge_constraint"] = False
         variables = data.pop("@variables")
         nets_sr = data.pop("nets_sr")
         nets_lr = data.pop("nets_lr")
@@ -657,16 +662,17 @@ def _forward_common(
         lr_out = lr_out + self._get_lr_bias(atype)
 
         # Hard charge constraint: enforce sum of latent charges equals target total charge.
-        q_mean = lr_out.mean(dim=1)  # [nf, lr_net_dim_out]
-        q_target_per_atom = target_total_charge / float(nloc)  # [nf]
-        if lr_out.shape[-1] > 1:
-            correction = torch.zeros_like(lr_out)
-            correction[:, :, 0] = q_mean[:, 0] - q_target_per_atom
-            lr_out = lr_out - correction
-        else:
-            q_mean = q_mean.unsqueeze(1)  # [nf, 1, 1]
-            target = q_target_per_atom.view(nf, 1, 1)
-            lr_out = lr_out - (q_mean - target)
+        if self.use_charge_constraint:
+            q_mean = lr_out.mean(dim=1)  # [nf, lr_net_dim_out]
+            q_target_per_atom = target_total_charge / float(nloc)  # [nf]
+            if lr_out.shape[-1] > 1:
+                correction = torch.zeros_like(lr_out)
+                correction[:, :, 0] = q_mean[:, 0] - q_target_per_atom
+                lr_out = lr_out - correction
+            else:
+                q_mean = q_mean.unsqueeze(1)  # [nf, 1, 1]
+                target = q_target_per_atom.view(nf, 1, 1)
+                lr_out = lr_out - (q_mean - target)
 
         mask = self.emask(atype).to(torch.bool)
         sr_out = torch.where(mask[:, :, None], sr_out, 0.0)

deepmd/pt/model/task/sog_energy_fitting.py

@@ -140,6 +140,7 @@ def __init__(
         M: int | None = None,
         n_dl: float | int = 1.0,
         remove_self_interaction: bool = False,
+        use_charge_constraint: bool = False,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -167,6 +168,7 @@ def __init__(
             type_map=type_map,
             use_aparam_as_mask=use_aparam_as_mask,
             default_fparam=default_fparam,
+            use_charge_constraint=use_charge_constraint,
             **kwargs,
         )
         if b is None:

deepmd/utils/argcheck.py

@@ -2131,6 +2131,9 @@ def fitting_sog_energy() -> list[Argument]:
     doc_remove_self_interaction = (
         "Whether to remove self interaction term in long-range correction."
     )
+    doc_use_charge_constraint = (
+        "Whether to enforce the sum of latent charges equals the target total charge (from fparam column 0)."
+    )
 
     return [
         Argument(
@@ -2289,6 +2292,13 @@ def fitting_sog_energy() -> list[Argument]:
             default=False,
             doc=doc_only_pt_supported + doc_remove_self_interaction,
         ),
+        Argument(
+            "use_charge_constraint",
+            bool,
+            optional=True,
+            default=False,
+            doc=doc_only_pt_supported + doc_use_charge_constraint,
+        ),
     ]