test inner cutoff to the density block

ACEsuit · Nov 22, 2024 · dfa93e9 · dfa93e9
1 parent 49fb8b2
commit dfa93e9
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/mace/modules/blocks.py b/mace/modules/blocks.py
@@ -329,6 +329,7 @@ def forward(
         edge_attrs: torch.Tensor,
         edge_feats: torch.Tensor,
         edge_index: torch.Tensor,
+        edge_lengths: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         raise NotImplementedError
 
@@ -659,6 +660,7 @@ def forward(
         edge_attrs: torch.Tensor,
         edge_feats: torch.Tensor,
         edge_index: torch.Tensor,
+        edge_lengths: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, None]:
         sender = edge_index[0]
         receiver = edge_index[1]
@@ -741,6 +743,7 @@ def forward(
         edge_attrs: torch.Tensor,
         edge_feats: torch.Tensor,
         edge_index: torch.Tensor,
+        edge_lengths: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         sender = edge_index[0]
         receiver = edge_index[1]
@@ -766,6 +769,9 @@ class RealAgnosticDensityInteractionBlock(InteractionBlock):
     def _setup(self) -> None:
         if not hasattr(self, "cueq_config"):
             self.cueq_config = None
+
+        # Inner radial basis
+        self.inner_cutoff = PolynomialCutoff(r_max=2.0, p=5)
         # First linear
         self.linear_up = Linear(
             self.node_feats_irreps,
@@ -833,13 +839,15 @@ def forward(
         edge_attrs: torch.Tensor,
         edge_feats: torch.Tensor,
         edge_index: torch.Tensor,
+        edge_lengths: torch.Tensor,
     ) -> Tuple[torch.Tensor, None]:
         sender = edge_index[0]
         receiver = edge_index[1]
         num_nodes = node_feats.shape[0]
         node_feats = self.linear_up(node_feats)
         tp_weights = self.conv_tp_weights(edge_feats)
-        edge_density = torch.tanh(self.density_fn(edge_feats) ** 2)
+        inner_cutoff = self.inner_cutoff(edge_lengths)
+        edge_density = torch.tanh(self.density_fn(edge_feats * inner_cutoff) ** 2)
         mji = self.conv_tp(
             node_feats[sender], edge_attrs, tp_weights
         )  # [n_edges, irreps]
@@ -862,7 +870,8 @@ class RealAgnosticDensityResidualInteractionBlock(InteractionBlock):
     def _setup(self) -> None:
         if not hasattr(self, "cueq_config"):
             self.cueq_config = None
-
+        # Cutoff
+        self.inner_cutoff = PolynomialCutoff(r_max=2.0, p=5)
         # First linear
         self.linear_up = Linear(
             self.node_feats_irreps,
@@ -931,14 +940,16 @@ def forward(
         edge_attrs: torch.Tensor,
         edge_feats: torch.Tensor,
         edge_index: torch.Tensor,
+        edge_lengths: torch.Tensor,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         sender = edge_index[0]
         receiver = edge_index[1]
         num_nodes = node_feats.shape[0]
         sc = self.skip_tp(node_feats, node_attrs)
         node_feats = self.linear_up(node_feats)
         tp_weights = self.conv_tp_weights(edge_feats)
-        edge_density = torch.tanh(self.density_fn(edge_feats) ** 2)
+        inner_cutoff = self.inner_cutoff(edge_lengths)
+        edge_density = torch.tanh(self.density_fn(edge_feats * inner_cutoff) ** 2)
         mji = self.conv_tp(
             node_feats[sender], edge_attrs, tp_weights
         )  # [n_edges, irreps]

diff --git a/mace/modules/models.py b/mace/modules/models.py
@@ -273,6 +273,7 @@ def forward(
                 edge_attrs=edge_attrs,
                 edge_feats=edge_feats,
                 edge_index=data["edge_index"],
+                edge_lengths=lengths,
             )
             node_feats = product(
                 node_feats=node_feats,
@@ -415,6 +416,7 @@ def forward(
                 edge_attrs=edge_attrs,
                 edge_feats=edge_feats,
                 edge_index=data["edge_index"],
+                edge_lengths=lengths,
             )
             node_feats = product(
                 node_feats=node_feats, sc=sc, node_attrs=data["node_attrs"]