feat(pt): add fixed gaussian angle

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

@@ -86,6 +86,7 @@ def __init__(
         dropout_rate: float = 0.1,
         angle_use_sh_init: bool = False,
         angle_sh_init_lmax: int = 3,
+        angle_use_fixed_gaussian: bool = False,
     ) -> None:
         r"""The constructor for the RepFlowArgs class which defines the parameters of the repflow block in DPA3 descriptor.
 
@@ -220,6 +221,7 @@ def __init__(
         self.dropout_rate = dropout_rate
         self.angle_use_sh_init = angle_use_sh_init
         self.angle_sh_init_lmax = angle_sh_init_lmax
+        self.angle_use_fixed_gaussian = angle_use_fixed_gaussian
         assert (
             fix_stat_std == 0.3
         ), "fix_stat_std is not implemented in this version, please use skip_stat instead."

deepmd/pt/model/descriptor/dpa3.py

@@ -225,6 +225,7 @@ def init_subclass_params(sub_data, sub_class):
             dropout_rate=self.repflow_args.dropout_rate,
             angle_use_sh_init=self.repflow_args.angle_use_sh_init,
             angle_sh_init_lmax=self.repflow_args.angle_sh_init_lmax,
+            angle_use_fixed_gaussian=self.repflow_args.angle_use_fixed_gaussian,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,

deepmd/pt/model/descriptor/repflows.py

@@ -17,6 +17,7 @@
     prod_env_mat,
 )
 from deepmd.pt.model.network.mlp import (
+    AnglePriorEncoder,
     AngleSH,
     MLPLayer,
 )
@@ -168,6 +169,7 @@ def __init__(
         dropout_rate: float = 0.1,
         angle_use_sh_init: bool = False,
         angle_sh_init_lmax: int = 3,
+        angle_use_fixed_gaussian: bool = False,
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
         r"""
@@ -302,6 +304,13 @@ def __init__(
         else:
             self.angle_sh = None
 
+        self.angle_use_fixed_gaussian = angle_use_fixed_gaussian
+        if self.angle_use_fixed_gaussian:
+            self.angle_gaussian_encoder = AnglePriorEncoder(
+                sigma_deg=6.0, learn_sigma=False, normalize=None
+            )
+        else:
+            self.angle_gaussian_encoder = None
         self.use_env_envelope = use_env_envelope
         self.use_new_sw = use_new_sw
         self.use_force_embedding = use_force_embedding
@@ -468,14 +477,17 @@ def __init__(
                 self.e_dim,
             ]
             self.edge_embd = RadialMLP(edge_channels_list)
-        if not self.angle_use_sh_init:
+
+        if self.angle_use_sh_init:
+            angle_input_dim = self.angle_sh_init_lmax + 1
+        elif self.angle_use_fixed_gaussian:
+            angle_input_dim = 10 + 1
+        else:
             angle_input_dim = (
                 len(self.angle_multi_freq_list_float) + 1
                 if not self.angle_init_use_sin
                 else 2 * (len(self.angle_multi_freq_list_float) + 1)
             )
-        else:
-            angle_input_dim = self.angle_sh_init_lmax + 1
 
         self.angle_embd = MLPLayer(
             angle_input_dim,
@@ -999,6 +1011,11 @@ def forward(
             assert self.angle_sh is not None
             # nf x nloc x a_nnei x a_nnei x sh_sim [OR] n_angle x sh_sim
             angle_input = self.angle_sh(angle_input * (torch.pi**0.5))
+        elif self.angle_use_fixed_gaussian:
+            assert not self.angle_init_use_sin and not self.angle_use_multi_freq
+            assert self.angle_gaussian_encoder is not None
+            # nf x nloc x a_nnei x a_nnei x 11 [OR] n_angle x 11
+            angle_input = self.angle_gaussian_encoder(angle_input)
 
         # nf x nloc x a_nnei x a_nnei x a_dim [OR] n_angle x a_dim
         angle_ebd = self.angle_embd(angle_input)

deepmd/pt/model/network/mlp.py

@@ -9,6 +9,7 @@
 import numpy as np
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 from deepmd.pt.utils import (
     env,
@@ -524,6 +525,115 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return P * self.norm.type(x.dtype)
 
 
+class AnglePriorEncoder(nn.Module):
+    """
+    Smooth delta encoder for bond angles a ∈ [0, π] (radians).
+    - Fixed 10 prior centers (in radians) taken from common molecular/material geometries.
+    - Kernel: Gaussian RBF on linear difference (no periodic wrapping).
+    - Output: softmax-normalized similarity vector of length 10 (smooth one-hot).
+    - Optional: learnable global width (sigma).
+
+    Centers (degrees → radians):
+      1) 180.0  : linear (sp), also octahedral/trans positions
+      2) 120.0  : trigonal planar (sp2), graphene etc.
+      3) 109.47 : ideal tetrahedral (sp3)
+      4) 104.5  : water H-O-H
+      5) 106.7  : ammonia H-N-H (trigonal pyramidal)
+      6) 90.0   : square planar / octahedral adjacent
+      7) 180.0  : (duplicate center kept intentionally for prior emphasis)
+      8) 120.0  : trigonal bipyramidal equatorial-equatorial
+      9) 90.0   : trigonal bipyramidal axial-equatorial
+     10) 60.0   : cyclopropane strained angle
+    """
+
+    def __init__(
+        self,
+        sigma_deg: float = 6.0,  # initial Gaussian width in degrees
+        learn_sigma: bool = True,  # make sigma trainable if desired
+        normalize: Optional[str] = "softmax",
+        eps: float = 1e-9,
+    ):
+        super().__init__()
+        assert normalize in ("softmax", "l1", None)
+        self.normalize = normalize
+        self.eps = eps
+
+        # --- Fixed prior centers (degrees) ---
+        centers_deg = torch.tensor(
+            [180.0, 120.0, 109.47, 104.5, 106.7, 90.0, 180.0, 120.0, 90.0, 60.0],
+            dtype=env.GLOBAL_PT_FLOAT_PRECISION,
+            device=device,
+        )
+
+        # Convert to radians and store as buffer: shape (K,)
+        centers_rad = centers_deg * (torch.pi / 180.0)
+        self.register_buffer("centers", centers_rad)  # (10,)
+
+        # --- Width parameter (global sigma, radians) ---
+        sigma_rad = float(sigma_deg) * math.pi / 180.0
+
+        # Softplus parameterization to keep sigma > 0
+        def inv_softplus(x: float) -> float:
+            x = max(x, 1e-12)
+            return float(math.log(math.exp(x) - 1.0))
+
+        raw = torch.tensor(
+            inv_softplus(sigma_rad), dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=device
+        )
+        if learn_sigma:
+            self.raw_sigma = nn.Parameter(data=raw)
+        else:
+            self.register_buffer("raw_sigma", raw)
+
+    @property
+    def sigma(self) -> torch.Tensor:
+        """Current positive width (radians)."""
+        return F.softplus(self.raw_sigma) + 1e-12
+
+    # @torch.no_grad()
+    # def auto_sigma_from_centers(self, factor: float = 0.6, min_sigma_deg: float = 1.0):
+    #     """
+    #     Set a reasonable global sigma from center spacing on [0, π].
+    #     Uses median nearest-neighbor distance x factor, with a lower bound.
+    #     """
+    #     c = self.centers  # (K,)
+    #     # Pairwise |c_i - c_j|
+    #     dmat = torch.abs(c[:, None] - c[None, :])
+    #     # Ignore self-distance
+    #     dmat = dmat + torch.eye(c.numel(), dtype=c.dtype, device=c.device) * 1e6
+    #     dmin = dmat.min(dim=1).values  # nearest neighbor distance per center
+    #     # Use median spacing to get a single global sigma
+    #     sigma = torch.clamp(torch.median(dmin) * factor,
+    #                         min=min_sigma_deg * math.pi / 180.0)
+    #     # Write into raw_sigma (inverse softplus)
+    #     with torch.no_grad():
+    #         self.raw_sigma.copy_(torch.log(torch.exp(sigma) - 1.0))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        a: (...,) tensor of angles in radians, expected in [0, π].
+        returns: (..., 10) similarity/weight vector.
+        """
+        theta = torch.acos(x)
+        centers = self.centers.type(x.dtype)
+        s = self.sigma.type(x.dtype)
+        # Linear difference (no periodicity)
+        diff = theta - centers  # (..., K)
+        # Gaussian kernel
+        sims = torch.exp(-0.5 * (diff / s).pow(2))  # (..., K)
+
+        # Normalization
+        if self.normalize is None:
+            codes = sims
+        elif self.normalize == "softmax":
+            codes = F.softmax(torch.log(sims + self.eps), dim=-1)
+        elif self.normalize == "l1":
+            codes = sims / (sims.sum(dim=-1, keepdim=True) + self.eps)
+        else:
+            raise ValueError(f"Unknown normalization: {self.normalize}")
+        return torch.cat([x, codes], dim=-1)
+
+
 def find_normalization(name: str, dim: int | None = None) -> nn.Module | None:
     """Return an normalization function using name."""
     if name is None:

deepmd/utils/argcheck.py

@@ -1623,6 +1623,12 @@ def dpa3_repflow_args():
             optional=True,
             default=3,
         ),
+        Argument(
+            "angle_use_fixed_gaussian",
+            bool,
+            optional=True,
+            default=False,
+        ),
         Argument(
             "use_dynamic_sel",
             bool,