feat pt : Support property fitting

deepmd/pt/loss/__init__.py

@@ -14,6 +14,9 @@
 from .loss import (
     TaskLoss,
 )
+from .property import (
+    PropertyLoss,
+)
 from .tensor import (
     TensorLoss,
 )
@@ -25,4 +28,5 @@
     "TensorLoss",
     "TaskLoss",
     "DOSLoss",
+    "PropertyLoss",
 ]

deepmd/pt/loss/property.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import logging
+from typing import (
+    List,
+)
+
+import torch
+import torch.nn.functional as F
+
+from deepmd.pt.loss.loss import (
+    TaskLoss,
+)
+from deepmd.pt.utils import (
+    env,
+)
+from deepmd.utils.data import (
+    DataRequirementItem,
+)
+
+log = logging.getLogger(__name__)
+
+
+class PropertyLoss(TaskLoss):
+    def __init__(
+        self,
+        task_dim,
+        loss_func: str = "smooth_mae",
+        metric: list = ["mae"],
+        beta: float = 1.00,
+        intensive: bool = False,
+        **kwargs,
+    ):
+        r"""Construct a layer to compute loss on property.
+
+        Parameters
+        ----------
+        task_dim : float
+            The output dimension of property fitting net.
+        loss_func : str
+            The loss function, such as "smooth_mae", "mae", "rmse".
+        metric : list
+            The metric such as mae, rmse which will be printed.
+        beta:
+            TODO
+        """
+        super().__init__()
+        self.task_dim = task_dim
+        self.loss_func = loss_func
+        self.metric = metric
+        self.beta = beta
+        self.intensive = intensive
+
+    def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False):
+        """Return loss on properties .
+
+        Parameters
+        ----------
+        input_dict : dict[str, torch.Tensor]
+            Model inputs.
+        model : torch.nn.Module
+            Model to be used to output the predictions.
+        label : dict[str, torch.Tensor]
+            Labels.
+        natoms : int
+            The local atom number.
+
+        Returns
+        -------
+        model_pred: dict[str, torch.Tensor]
+            Model predictions.
+        loss: torch.Tensor
+            Loss for model to minimize.
+        more_loss: dict[str, torch.Tensor]
+            Other losses for display.
+        """
+        model_pred = model(**input_dict)
+        assert label["property"].shape[-1] == self.task_dim
+        assert model_pred["property"].shape[-1] == self.task_dim
+        loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0]
+        more_loss = {}
+
+        if not self.intensive:
+            label["property"] = label["property"] / natoms
+            model_pred["property"] = model_pred["property"] / natoms
+
+        # loss
+        if self.loss_func == "smooth_mae":
+            loss += F.smooth_l1_loss(
+                label["property"],
+                model_pred["property"],
+                reduction="sum",
+                beta=self.beta,
+            )
+        elif self.loss_func == "mae":
+            loss += F.l1_loss(
+                label["property"], model_pred["property"], reduction="sum"
+            )
+        elif self.loss_func == "mse":
+            loss += F.mse_loss(
+                label["property"],
+                model_pred["property"],
+                reduction="sum",
+            )
+        elif self.loss_func == "rmse":
+            loss += torch.sqrt(
+                F.mse_loss(
+                    label["property"],
+                    model_pred["property"],
+                    reduction="mean",
+                )
+            )
+        else:
+            raise RuntimeError(f"Unknown loss function : {self.loss_func}")
+
+        # more loss
+        if "smooth_mae" in self.metric:
+            more_loss["smooth_mae"] = F.smooth_l1_loss(
+                label["property"],
+                model_pred["property"],
+                reduction="mean",
+                beta=self.beta,
+            ).detach()
+        if "mae" in self.metric:
+            more_loss["mae"] = F.l1_loss(
+                label["property"],
+                model_pred["property"],
+                reduction="mean",
+            ).detach()
+        if "mse" in self.metric:
+            more_loss["mse"] = F.mse_loss(
+                label["property"],
+                model_pred["property"],
+                reduction="mean",
+            ).detach()
+        if "rmse" in self.metric:
+            more_loss["rmse"] = torch.sqrt(
+                F.mse_loss(
+                    label["property"],
+                    model_pred["property"],
+                    reduction="mean",
+                )
+            ).detach()
+
+        return model_pred, loss, more_loss
+
+    @property
+    def label_requirement(self) -> List[DataRequirementItem]:
+        """Return data label requirements needed for this loss calculation."""
+        label_requirement = []
+        label_requirement.append(
+            DataRequirementItem(
+                "property",
+                ndof=self.task_dim,
+                atomic=False,
+                must=False,
+                high_prec=True,
+            )
+        )
+        return label_requirement

deepmd/pt/model/atomic_model/__init__.py

@@ -39,12 +39,16 @@
 from .polar_atomic_model import (
     DPPolarAtomicModel,
 )
+from .property_atomic_model import (
+    DPPropertyAtomicModel,
+)
 
 __all__ = [
     "BaseAtomicModel",
     "DPAtomicModel",
     "DPDOSAtomicModel",
     "DPEnergyAtomicModel",
+    "DPPropertyAtomicModel",
     "PairTabAtomicModel",
     "LinearEnergyAtomicModel",
     "DPPolarAtomicModel",

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -462,6 +462,7 @@ def change_out_bias(
                 stat_file_path=stat_file_path,
                 rcond=self.rcond,
                 preset_bias=self.preset_out_bias,
+                intensive=self.fitting_net.intensive,
             )
             self._store_out_stat(bias_out, std_out)
         else:

deepmd/pt/model/atomic_model/property_atomic_model.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.pt.model.task.property import (
+    PropertyFittingNet,
+)
+
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+
+class DPPropertyAtomicModel(DPAtomicModel):
+    def __init__(self, descriptor, fitting, type_map, **kwargs):
+        assert isinstance(fitting, PropertyFittingNet)
+        super().__init__(descriptor, fitting, type_map, **kwargs)

deepmd/pt/model/model/__init__.py

@@ -60,6 +60,9 @@
 from .polar_model import (
     PolarModel,
 )
+from .property_model import (
+    PropertyModel,
+)
 from .spin_model import (
     SpinEnergyModel,
     SpinModel,
@@ -182,6 +185,8 @@
         modelcls = DOSModel
     elif fitting_net["type"] in ["ener", "direct_force_ener"]:
         modelcls = EnergyModel
+    elif fitting_net["type"] in ["prop", "property"]:
+        modelcls = PropertyModel
     else:
         raise RuntimeError(f"Unknown fitting type: {fitting_net['type']}")
 

deepmd/pt/model/model/property_model.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Dict,
+    Optional,
+)
+
+import torch
+
+from deepmd.pt.model.atomic_model import (
+    DPPropertyAtomicModel,
+)
+from deepmd.pt.model.model.model import (
+    BaseModel,
+)
+
+from .dp_model import (
+    DPModelCommon,
+)
+from .make_model import (
+    make_model,
+)
+
+DPPropertyModel_ = make_model(DPPropertyAtomicModel)
+
+
+@BaseModel.register("property")
+class PropertyModel(DPModelCommon, DPPropertyModel_):
+    model_type = "property"
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        DPModelCommon.__init__(self)
+        DPPropertyModel_.__init__(self, *args, **kwargs)
+
+    def forward(
+        self,
+        coord,
+        atype,
+        box: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+        do_atomic_virial: bool = False,
+    ) -> Dict[str, torch.Tensor]:
+        model_ret = self.forward_common(
+            coord,
+            atype,
+            box,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+        )
+        model_predict = {}
+        model_predict["atom_property"] = model_ret["property"]
+        model_predict["property"] = model_ret["property_redu"]
+        if self.atomic_model.fitting_net.intensive:
+            natoms = model_predict["atom_property"].shape[1]
+            model_predict["property"] = model_ret["property_redu"] / natoms
+        if "mask" in model_ret:
+            model_predict["mask"] = model_ret["mask"]
+        return model_predict
+
+    @torch.jit.export
+    def forward_lower(
+        self,
+        extended_coord,
+        extended_atype,
+        nlist,
+        mapping: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+        do_atomic_virial: bool = False,
+        comm_dict: Optional[Dict[str, torch.Tensor]] = None,
+    ):
+        model_ret = self.forward_common_lower(
+            extended_coord,
+            extended_atype,
+            nlist,
+            mapping,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+            comm_dict=comm_dict,
+        )
+        model_predict = {}
+        model_predict["atom_property"] = model_ret["property"]
+        natoms = model_predict["atom_property"].shape[1]
+        model_predict["property"] = model_ret["property_redu"]
+        if "mask" in model_ret:
+            model_predict["mask"] = model_ret["mask"]
+        return model_predict

deepmd/pt/model/task/__init__.py

@@ -24,6 +24,9 @@
 from .polarizability import (
     PolarFittingNet,
 )
+from .property import (
+    PropertyFittingNet,
+)
 from .type_predict import (
     TypePredictNet,
 )
@@ -39,4 +42,5 @@
     "TypePredictNet",
     "PolarFittingNet",
     "DOSFittingNet",
+    "PropertyFittingNet",
 ]