feat(pt): add edge readout

Author: iProzd

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -148,6 +148,13 @@ def compute_input_stats(
             """Update mean and stddev for descriptor elements."""
             raise NotImplementedError
 
+        def get_norm_fact(self) -> list[float]:
+            """Returns the norm factor."""
+            raise NotImplementedError
+
+        def get_additional_output_for_fitting(self):
+            raise NotImplementedError
+
         def enable_compression(
             self,
             min_nbor_dist: float,

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/model/atomic_model/dp_atomic_model.py

@@ -260,15 +260,29 @@ def forward_atomic(
         if self.enable_eval_descriptor_hook:
             self.eval_descriptor_list.append(descriptor.detach())
         # 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,
+                sw=sw,
+                edge_index=add_input.get("edge_index", None),
+            )
         if self.enable_eval_fitting_last_layer_hook:
             assert "middle_output" in fit_ret, (
                 "eval_fitting_last_layer not supported for this fitting net!"

deepmd/pt/model/descriptor/dpa1.py

@@ -356,6 +356,13 @@ def get_dim_out(self) -> int:
     def get_dim_emb(self) -> int:
         return self.se_atten.dim_emb
 
+    def get_norm_fact(self) -> list[float]:
+        """Returns the norm factor."""
+        return [float(self.get_nnei())]
+
+    def get_additional_output_for_fitting(self) -> dict[str, Optional[torch.Tensor]]:
+        return {}
+
     def mixed_types(self) -> bool:
         """If true, the descriptor
         1. assumes total number of atoms aligned across frames;

deepmd/pt/model/descriptor/dpa3.py

@@ -248,6 +248,13 @@ def get_dim_emb(self) -> int:
         """Returns the embedding dimension of this descriptor."""
         return self.repflows.dim_emb
 
+    def get_norm_fact(self) -> list[float]:
+        """Returns the norm factor."""
+        return self.repflows.get_norm_fact()
+
+    def get_additional_output_for_fitting(self):
+        return self.repflows.get_additional_output_for_fitting()
+
     def mixed_types(self) -> bool:
         """If true, the descriptor
         1. assumes total number of atoms aligned across frames;

deepmd/pt/model/descriptor/repflows.py

@@ -257,6 +257,8 @@ def __init__(
         self.use_exp_switch = use_exp_switch
         self.use_dynamic_sel = use_dynamic_sel
         self.sel_reduce_factor = sel_reduce_factor
+        self.dynamic_e_sel = self.nnei / self.sel_reduce_factor
+        self.dynamic_a_sel = self.a_sel / self.sel_reduce_factor
         if self.use_dynamic_sel and not self.smooth_edge_update:
             raise NotImplementedError(
                 "smooth_edge_update must be True when use_dynamic_sel is True!"
@@ -334,6 +336,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)
@@ -344,6 +347,8 @@ def __init__(
         self.register_buffer("stddev", stddev)
         self.stats = None
 
+    additional_output_for_fitting: dict[str, Optional[torch.Tensor]]
+
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.e_rcut
@@ -376,6 +381,16 @@ def get_dim_emb(self) -> int:
         """Returns the embedding dimension e_dim."""
         return self.e_dim
 
+    def get_additional_output_for_fitting(self):
+        return self.additional_output_for_fitting
+
+    def get_norm_fact(self) -> list[float]:
+        """Returns the norm factor."""
+        return [
+            float(self.dynamic_e_sel if self.use_dynamic_sel else self.nnei),
+            # float(self.dynamic_a_sel if self.use_dynamic_sel else self.a_sel),
+        ]
+
     def __setitem__(self, key, value) -> None:
         if key in ("avg", "data_avg", "davg"):
             self.mean = value
@@ -548,10 +563,12 @@ def forward(
             angle_input = angle_input[a_nlist_mask]
             # n_angle x 1
             a_sw = (a_sw[:, :, :, None] * a_sw[:, :, None, :])[a_nlist_mask]
+            self.additional_output_for_fitting["edge_index"] = edge_index
         else:
             # avoid jit assertion
             edge_index = torch.zeros([2, 1], device=nlist.device, dtype=nlist.dtype)
             angle_index = torch.zeros([3, 1], device=nlist.device, dtype=nlist.dtype)
+            self.additional_output_for_fitting["edge_index"] = None
         # get edge and angle embedding
         # nb x nloc x nnei x e_dim [OR] n_edge x e_dim
         if not self.edge_init_use_dist:

deepmd/pt/model/model/__init__.py

@@ -90,8 +90,10 @@ 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_readout"]:
         fitting_net["embedding_width"] = descriptor.get_dim_emb()
+    if fitting_net["type"] in ["ener_readout"]:
+        fitting_net["norm_fact"] = descriptor.get_norm_fact()
     fitting_net["dim_descrpt"] = descriptor.get_dim_out()
     grad_force = "direct" not in fitting_net["type"]
     if not grad_force:
@@ -262,7 +264,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_readout"]:
         modelcls = EnergyModel
     elif fitting_net_type == "property":
         modelcls = PropertyModel

deepmd/pt/model/task/ener.py

@@ -13,9 +13,19 @@
     OutputVariableDef,
     fitting_check_output,
 )
+from deepmd.dpmodel.utils.seed import (
+    child_seed,
+)
+from deepmd.pt.model.network.mlp import (
+    FittingNet,
+    NetworkCollection,
+)
 from deepmd.pt.model.network.network import (
     ResidualDeep,
 )
+from deepmd.pt.model.network.utils import (
+    aggregate,
+)
 from deepmd.pt.model.task.fitting import (
     Fitting,
     GeneralFitting,
@@ -257,3 +267,155 @@ def forward(
             "energy": outs.to(env.GLOBAL_PT_FLOAT_PRECISION),
             "dforce": vec_out,
         }
+
+
+@Fitting.register("ener_readout")
+@fitting_check_output
+class EnergyFittingNetReadout(InvarFitting):
+    def __init__(
+        self,
+        ntypes: int,
+        dim_descrpt: int,
+        neuron: list[int] = [128, 128, 128],
+        bias_atom_e: Optional[torch.Tensor] = None,
+        resnet_dt: bool = True,
+        numb_fparam: int = 0,
+        numb_aparam: int = 0,
+        dim_case_embd: int = 0,
+        embedding_width: int = 128,
+        activation_function: str = "tanh",
+        precision: str = DEFAULT_PRECISION,
+        mixed_types: bool = True,
+        seed: Optional[Union[int, list[int]]] = None,
+        type_map: Optional[list[str]] = None,
+        norm_fact: list[float] = [120.0],
+        add_edge_readout: bool = True,
+        slim_edge_readout: bool = False,
+        **kwargs,
+    ) -> None:
+        """Construct a fitting net for energy.
+
+        Args:
+        - ntypes: Element count.
+        - embedding_width: Embedding width per atom.
+        - neuron: Number of neurons in each hidden layers of the fitting net.
+        - bias_atom_e: Average energy per atom for each element.
+        - resnet_dt: Using time-step in the ResNet construction.
+        """
+        self.add_edge_readout = add_edge_readout
+        super().__init__(
+            "energy",
+            ntypes,
+            dim_descrpt,
+            1,
+            neuron=neuron,
+            bias_atom_e=bias_atom_e,
+            resnet_dt=resnet_dt,
+            numb_fparam=numb_fparam,
+            numb_aparam=numb_aparam,
+            dim_case_embd=dim_case_embd,
+            activation_function=activation_function,
+            precision=precision,
+            mixed_types=mixed_types,
+            seed=seed,
+            type_map=type_map,
+            **kwargs,
+        )
+
+        # embedding for edge readout
+        self.embedding_width = embedding_width
+        self.slim_edge_readout = slim_edge_readout
+        self.norm_e_fact = norm_fact[0]
+
+        if self.add_edge_readout:
+            self.edge_embed = NetworkCollection(
+                1 if not self.mixed_types else 0,
+                self.ntypes,
+                network_type="fitting_network",
+                networks=[
+                    FittingNet(
+                        self.embedding_width,
+                        1,
+                        self.neuron if not self.slim_edge_readout else self.neuron[:1],
+                        self.activation_function,
+                        self.resnet_dt,
+                        self.precision,
+                        bias_out=True,
+                        seed=child_seed(self.seed + 100, ii),
+                    )
+                    for ii in range(self.ntypes if not self.mixed_types else 1)
+                ],
+            )
+        else:
+            self.edge_embed = None
+
+        # set trainable
+        for param in self.parameters():
+            param.requires_grad = self.trainable
+
+    # make jit happy with torch 2.0.0
+    exclude_types: list[int]
+
+    def need_additional_input(self) -> bool:
+        return True
+
+    def serialize(self) -> dict:
+        raise NotImplementedError
+
+    @classmethod
+    def deserialize(cls, data: dict) -> "EnergyFittingNetReadout":
+        raise NotImplementedError
+
+    def forward(
+        self,
+        descriptor: torch.Tensor,
+        atype: torch.Tensor,
+        gr: Optional[torch.Tensor] = None,
+        g2: Optional[torch.Tensor] = None,
+        h2: Optional[torch.Tensor] = None,
+        fparam: Optional[torch.Tensor] = None,
+        aparam: Optional[torch.Tensor] = None,
+        sw: Optional[torch.Tensor] = None,
+        edge_index: Optional[torch.Tensor] = None,
+    ):
+        """Based on embedding net output, alculate total energy.
+
+        Args:
+        - inputs: Embedding matrix. Its shape is [nframes, natoms[0], self.dim_descrpt].
+        - natoms: Tell atom count and element count. Its shape is [2+self.ntypes].
+
+        Returns
+        -------
+        - `torch.Tensor`: Total energy with shape [nframes, natoms[0]].
+        """
+        out = self._forward_common(descriptor, atype, gr, g2, h2, fparam, aparam)[
+            self.var_name
+        ]
+        nf, nloc, _ = descriptor.shape
+
+        if self.add_edge_readout:
+            assert g2 is not None
+            assert sw is not None
+            assert self.edge_embed is not None
+            # nf x nloc x nnei x d [OR] nedge x d
+            edge_feature = g2
+            # nf x nloc x nnei x 1 [OR] nedge x 1
+            edge_atomic_contrib = self.edge_embed.networks[0](edge_feature)
+            # nf x nloc x nnei x 1 [OR] nedge x 1
+            edge_atomic_contrib = edge_atomic_contrib * sw.unsqueeze(-1)
+            if edge_index is not None:
+                # use dynamic sel
+                n2e_index, n_ext2e_index = edge_index[0], edge_index[1]
+                # nf x nloc x 1
+                edge_energy = aggregate(
+                    edge_atomic_contrib,
+                    n2e_index,
+                    average=False,
+                    num_owner=nf * nloc,
+                ).reshape(nf, nloc, 1)
+            else:
+                # nf x nloc x 1
+                edge_energy = torch.sum(edge_atomic_contrib, dim=-2)
+            # energy
+            out = out + edge_energy / self.norm_e_fact
+        return {self.var_name: out.to(env.GLOBAL_PT_FLOAT_PRECISION)}

deepmd/pt/model/task/invar_fitting.py

@@ -170,6 +170,8 @@ def forward(
         h2: Optional[torch.Tensor] = None,
         fparam: Optional[torch.Tensor] = None,
         aparam: Optional[torch.Tensor] = None,
+        sw: Optional[torch.Tensor] = None,
+        edge_index: Optional[torch.Tensor] = None,
     ):
         """Based on embedding net output, alculate total energy.