add Trinity model and debug the lem (deepmodeling#256)

* update input and e3tb hands on doc

* debug lem and add trinity

* change init2b to only2b, move scalarmlpfunction to base.py

* compute onsite shift mu from whole onsite block

* drop E3statistics when model does is valid and does not have a scale shift

* add edge component in onsite shift computation

* feat(transforms): add full_mask_to_diag initialization for diagonal elements

Initialize full_mask_to_diag tensor to track diagonal orbital pairs in the reduced matrix. This helps in identifying diagonal elements during further processing.

* refactor(loss): extract shift_mu function and simplify onsite shift calculation

The shift_mu function was extracted to avoid code duplication across multiple loss classes. The onsite shift calculation was simplified by using a more accurate formula that accounts for both node and edge features.

* refactor(loss): simplify shift_mu calculation and usage

remove redundant diagonal terms in shift_mu and update all dependent loss classes to use simplified return values

---------

Co-authored-by: QG-phy <guqq_phy@qq.com>
Co-authored-by: Qiangqiang Gu <98570179+QG-phy@users.noreply.github.com>

Author: floatingCatty

dptb/data/dataset/_default_dataset.py

@@ -401,6 +401,10 @@ def E3statistics(self, model: torch.nn.Module=None, decay=False):
 
         if self.data[AtomicDataDict.EDGE_FEATURES_KEY].abs().sum() < 1e-7:
             return None
+    
+        if model is not None:
+            if not isinstance(model.node_prediction_h, torch.nn.Module):
+                return None
 
         typed_dataset = idp(self.data.clone().to_dict())
         e3h = E3Hamiltonian(basis=idp.basis, decompose=True)

dptb/data/dataset/lmdb_dataset.py

@@ -198,6 +198,10 @@ def E3statistics(self, model: torch.nn.Module=None):
         if not self.get_Hamiltonian and not self.get_DM:
             return None
 
+        if model is not None:
+            if not isinstance(model.node_prediction_h, torch.nn.Module):
+                return None
+        
         assert self.transform is not None
         idp = self.transform
 

dptb/data/transforms.py

@@ -644,7 +644,12 @@ def __init__(
                         indices += sli.start
                         assert indices.max() < sli.stop
                         self.mask_to_ndiag[self.chemical_symbol_to_type[ib]][indices] = True
-
+            
+            self.full_mask_to_diag = torch.zeros(self.reduced_matrix_element, dtype=torch.bool, device=self.device)
+            for orbs, islice in self.orbpair_maps.items():
+                fio, fjo = orbs.split('-')
+                if fio == fjo:
+                    self.full_mask_to_diag[islice] = True
 
     def get_orbpairtype_maps(self):
         """

dptb/nn/base.py

@@ -5,7 +5,11 @@
 from torch import Tensor
 from dptb.utils.constants import dtype_dict
 from dptb.utils.tools import _get_activation_fn
+from e3nn.util.codegen import CodeGenMixin
+from e3nn.math import normalize2mom
 import torch.nn.functional as F
+import math
+from torch import fx
 import torch.nn as nn
 
 class AtomicLinear(torch.nn.Module):
@@ -466,4 +470,115 @@ def forward(self, x):
             x = layer(x)
             x = self.activation(x)
 
-        return self.out_layer(x)
+        return self.out_layer(x)
+
+class ScalarMLPFunction(CodeGenMixin, torch.nn.Module):
+    """Module implementing an MLP according to provided options."""
+
+    in_features: int
+    out_features: int
+
+    def __init__(
+        self,
+        mlp_input_dimension: Optional[int],
+        mlp_latent_dimensions: List[int],
+        mlp_output_dimension: Optional[int],
+        mlp_nonlinearity: Optional[str] = "silu",
+        mlp_initialization: str = "normal",
+        mlp_dropout_p: float = 0.0,
+        mlp_batchnorm: bool = False,
+    ):
+        super().__init__()
+        nonlinearity = {
+            None: None,
+            "silu": torch.nn.functional.silu,
+            "ssp": ShiftedSoftPlus,
+        }[mlp_nonlinearity]
+        if nonlinearity is not None:
+            nonlin_const = normalize2mom(nonlinearity).cst
+        else:
+            nonlin_const = 1.0
+
+        dimensions = (
+            ([mlp_input_dimension] if mlp_input_dimension is not None else [])
+            + mlp_latent_dimensions
+            + ([mlp_output_dimension] if mlp_output_dimension is not None else [])
+        )
+        assert len(dimensions) >= 2  # Must have input and output dim
+        num_layers = len(dimensions) - 1
+
+        self.in_features = dimensions[0]
+        self.out_features = dimensions[-1]
+
+        # Code
+        params = {}
+        graph = fx.Graph()
+        tracer = fx.proxy.GraphAppendingTracer(graph)
+
+        def Proxy(n):
+            return fx.Proxy(n, tracer=tracer)
+
+        features = Proxy(graph.placeholder("x"))
+        norm_from_last: float = 1.0
+
+        base = torch.nn.Module()
+
+        for layer, (h_in, h_out) in enumerate(zip(dimensions, dimensions[1:])):
+            # do dropout
+            if mlp_dropout_p > 0:
+                # only dropout if it will do something
+                # dropout before linear projection- https://stats.stackexchange.com/a/245137
+                features = Proxy(graph.call_module("_dropout", (features.node,)))
+
+            # make weights
+            w = torch.empty(h_in, h_out)
+
+            if mlp_initialization == "normal":
+                w.normal_()
+            elif mlp_initialization == "uniform":
+                # these values give < x^2 > = 1
+                w.uniform_(-math.sqrt(3), math.sqrt(3))
+            elif mlp_initialization == "orthogonal":
+                # this rescaling gives < x^2 > = 1
+                torch.nn.init.orthogonal_(w, gain=math.sqrt(max(w.shape)))
+            else:
+                raise NotImplementedError(
+                    f"Invalid mlp_initialization {mlp_initialization}"
+                )
+
+            # generate code
+            params[f"_weight_{layer}"] = w
+            w = Proxy(graph.get_attr(f"_weight_{layer}"))
+            w = w * (
+                norm_from_last / math.sqrt(float(h_in))
+            )  # include any nonlinearity normalization from previous layers
+            features = torch.matmul(features, w)
+
+            if mlp_batchnorm:
+                # if we call batchnorm, do it after the nonlinearity
+                features = Proxy(graph.call_module(f"_bn_{layer}", (features.node,)))
+                setattr(base, f"_bn_{layer}", torch.nn.BatchNorm1d(h_out))
+
+            # generate nonlinearity code
+            if nonlinearity is not None and layer < num_layers - 1:
+                features = nonlinearity(features)
+                # add the normalization const in next layer
+                norm_from_last = nonlin_const
+
+        graph.output(features.node)
+
+        for pname, p in params.items():
+            setattr(base, pname, torch.nn.Parameter(p))
+
+        if mlp_dropout_p > 0:
+            # with normal dropout everything blows up
+            base._dropout = torch.nn.AlphaDropout(p=mlp_dropout_p)
+
+        self._codegen_register({"_forward": fx.GraphModule(base, graph)})
+
+    def forward(self, x):
+        return self._forward(x)
+
+@torch.jit.script
+def ShiftedSoftPlus(x: torch.Tensor):
+    return torch.nn.functional.softplus(x) - math.log(2.0)

dptb/nn/deeptb.py

@@ -169,33 +169,54 @@ def __init__(
                     self.overlaponsite_param = overlaponsite_param
 
         elif prediction_copy.get("method") == "e3tb":
-            self.node_prediction_h = E3PerSpeciesScaleShift(
-                field=AtomicDataDict.NODE_FEATURES_KEY,
-                num_types=n_species,
-                irreps_in=self.embedding.out_node_irreps,
-                out_field = AtomicDataDict.NODE_FEATURES_KEY,
-                shifts=0.,
-                scales=1.,
-                dtype=self.dtype,
-                device=self.device,
-                **prediction_copy,
-            )
-            
-            self.edge_prediction_h = E3PerEdgeSpeciesScaleShift(
-                field=AtomicDataDict.EDGE_FEATURES_KEY,
-                num_types=n_species,
-                irreps_in=self.embedding.out_edge_irreps,
-                out_field = AtomicDataDict.EDGE_FEATURES_KEY,
-                shifts=0.,
-                scales=1.,
-                dtype=self.dtype,
-                device=self.device,
-                **prediction_copy,
-            )
+            if embedding.get("method") == "trinity":
+                # hack to pass the dataset operation
+                self.node_prediction_h = lambda x: x
+                self.edge_prediction_h = lambda x: x
+                self.node_prediction_h.set_scale_shift = lambda scales, shifts: 0
+                self.edge_prediction_h.set_scale_shift = lambda scales, shifts: 0
+            else:
+                self.node_prediction_h = E3PerSpeciesScaleShift(
+                    field=AtomicDataDict.NODE_FEATURES_KEY,
+                    num_types=n_species,
+                    irreps_in=self.embedding.out_node_irreps,
+                    out_field = AtomicDataDict.NODE_FEATURES_KEY,
+                    shifts=0.,
+                    scales=1.,
+                    dtype=self.dtype,
+                    device=self.device,
+                    **prediction_copy,
+                )
+                
+                self.edge_prediction_h = E3PerEdgeSpeciesScaleShift(
+                    field=AtomicDataDict.EDGE_FEATURES_KEY,
+                    num_types=n_species,
+                    irreps_in=self.embedding.out_edge_irreps,
+                    out_field = AtomicDataDict.EDGE_FEATURES_KEY,
+                    shifts=0.,
+                    scales=1.,
+                    dtype=self.dtype,
+                    device=self.device,
+                    **prediction_copy,
+                )
+
+            if embedding.get("method") == "trinity":
+                self.idp_sk = OrbitalMapper(self.idp.basis, method="sktb", device=self.device)
+                prediction_copy = prediction_copy.copy()
+                prediction_copy["neurons"] = [self.embedding.latent_dim] + prediction_copy["neurons"] + [self.idp_sk.reduced_matrix_element]
+                prediction_copy["config"] = get_neuron_config(prediction_copy["neurons"])
+                self.edge_prediction_h2 = AtomicResNet(
+                    **prediction_copy,
+                    in_field=AtomicDataDict.EDGE_ATTRS_KEY,
+                    out_field=AtomicDataDict.EDGE_ATTRS_KEY,
+                    device=device,
+                    dtype=dtype
+                )
 
             if overlap:
                 self.idp_sk = OrbitalMapper(self.idp.basis, method="sktb", device=self.device)
                 self.idp_sk.get_skonsite_maps()
+                prediction_copy = prediction.copy()
                 prediction_copy["neurons"] = [self.embedding.latent_dim] + prediction_copy["neurons"] + [self.idp_sk.reduced_matrix_element]
                 prediction_copy["config"] = get_neuron_config(prediction_copy["neurons"])
                 self.edge_prediction_s = AtomicResNet(
@@ -257,6 +278,17 @@ def __init__(
                     dtype=self.dtype, 
                     device=self.device,
                     )
+            if hasattr(self, "edge_prediction_h2"):
+                self.h2miltonian = SKHamiltonian(
+                    idp_sk=self.idp_sk, 
+                    edge_field=AtomicDataDict.EDGE_ATTRS_KEY,
+                    node_field=AtomicDataDict.NODE_ATTRS_KEY,
+                    onsite=True,
+                    strain=False,
+                    soc=False,
+                    dtype=self.dtype, 
+                    device=self.device,
+                    )
 
 
     def forward(self, data: AtomicDataDict.Type):
@@ -274,10 +306,18 @@ def forward(self, data: AtomicDataDict.Type):
             data[AtomicDataDict.NODE_OVERLAP_KEY] = self.overlaponsite_param[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()]
             data[AtomicDataDict.NODE_OVERLAP_KEY][:,self.idp_sk.mask_diag] = 1.
 
+        # prediction for two-body part of e3tb
+        if hasattr(self, "edge_prediction_h2"):
+            data = self.edge_prediction_h2(data)
+        
         if self.transform:
             data = self.hamiltonian(data)
             if hasattr(self, "overlap"):
                 data = self.overlap(data)
+            if hasattr(self, "edge_prediction_h2"):
+                data = self.h2miltonian(data)
+                data[AtomicDataDict.NODE_FEATURES_KEY] += data[AtomicDataDict.NODE_ATTRS_KEY]
+                data[AtomicDataDict.EDGE_FEATURES_KEY] += data[AtomicDataDict.EDGE_ATTRS_KEY]
 
         return data
 

dptb/nn/embedding/__init__.py

@@ -6,6 +6,7 @@
 from .e3baseline_local6 import E3BaseLineModel6
 from .slem import Slem
 from .lem import Lem
+from .trinity import Trinity
 from .e3baseline_nonlocal import E3BaseLineModelNonLocal
 
 __all__ = [
@@ -15,6 +16,7 @@
     "E3DeePH",
     "Lem",
     "Slem",
+    "Trinity",
     "E3BaseLineModel6",
     "E3BaseLineModelNonLocal",
 ]