Added vector output from skew-symmetric tensor

tests/test_equivariance.py

@@ -2,6 +2,7 @@
 # Distributed under the MIT License.
 # (See accompanying file README.md file or copy at http://opensource.org/licenses/MIT)
 
+import pytest
 import torch
 from torchmdnet.models.model import create_model
 from utils import load_example_args
@@ -27,7 +28,8 @@ def test_scalar_invariance():
     torch.testing.assert_allclose(y, y_rot)
 
 
-def test_vector_equivariance():
+@pytest.mark.parametrize("model_name", ["equivariant-transformer", "equivariant-tensornet"])
+def test_vector_equivariance(model_name):
     torch.manual_seed(1234)
     rotate = torch.tensor(
         [
@@ -39,7 +41,7 @@ def test_vector_equivariance():
 
     model = create_model(
         load_example_args(
-            "equivariant-transformer",
+            model_name,
             prior_model=None,
             output_model="VectorOutput",
         )

tests/utils.py

@@ -10,7 +10,7 @@
 
 def load_example_args(model_name, remove_prior=False, config_file=None, **kwargs):
     if config_file is None:
-        if model_name == "tensornet":
+        if model_name == "tensornet" or model_name == "equivariant-tensornet":
             config_file = join(dirname(dirname(__file__)), "examples", "TensorNet-QM9.yaml")
         else:
             config_file = join(dirname(dirname(__file__)), "examples", "ET-QM9.yaml")

torchmdnet/models/model.py

@@ -103,6 +103,17 @@ def create_model(args, prior_model=None, mean=None, std=None):
             static_shapes=args["static_shapes"],
             **shared_args,
         )
+    elif args["model"] == "equivariant-tensornet":
+        from torchmdnet.models.tensornet import TensorNet
+
+        # returns an equivariant vector
+        is_equivariant = True
+        representation_model = TensorNet(
+            equivariance_invariance_group=args["equivariance_invariance_group"],
+            static_shapes=args["static_shapes"],
+            vector_output=True,
+            **shared_args,
+        )
     else:
         raise ValueError(f'Unknown architecture: {args["model"]}')
 

torchmdnet/models/tensornet.py

@@ -48,6 +48,10 @@ def vector_to_symtensor(vector):
     S = 0.5 * (tensor + tensor.transpose(-2, -1)) - I
     return S
 
+def skewtensor_to_vector(tensor):
+    '''Converts a skew-symmetric tensor to a vector.'''
+    return torch.stack((tensor[:, :, 1, 2], tensor[:, :, 2, 0], tensor[:, :, 0, 1]), dim=-1)
+
 
 def decompose_tensor(tensor):
     """Full tensor decomposition into irreducible components."""
@@ -120,6 +124,7 @@ class TensorNet(nn.Module):
             (default: :obj:`True`)
         check_errors (bool, optional): Whether to check for errors in the distance module.
             (default: :obj:`True`)
+        vector_output (bool, optional): Whether to return 
     """
 
     def __init__(
@@ -139,6 +144,7 @@ def __init__(
         check_errors=True,
         dtype=torch.float32,
         box_vecs=None,
+        vector_output=False
     ):
         super(TensorNet, self).__init__()
 
@@ -210,6 +216,7 @@ def __init__(
             box=box_vecs,
             long_edge_index=True,
         )
+        self.vector_output = vector_output
 
         self.reset_parameters()
 
@@ -265,10 +272,22 @@ def forward(
         x = torch.cat((tensor_norm(I), tensor_norm(A), tensor_norm(S)), dim=-1)
         x = self.out_norm(x)
         x = self.act(self.linear((x)))
-        # # Remove the extra atom
+        # Remove the extra atom
         if self.static_shapes:
             x = x[:-1]
-        return x, None, z, pos, batch
+
+        # calculate vector_output if needed
+        v = None
+        if self.vector_output:
+            # (n_atoms, hidden_channels, 3, 3) -> (n_atoms, hidden_channels, 3)
+            v = skewtensor_to_vector(A)
+            # (n_atoms, hidden_channels, 3) -> (n_atoms, 3, hidden_channels)
+            v = v.transpose(1, 2)
+
+            if self.static_shapes:
+                v = v[:-1]
+
+        return x, v, z, pos, batch
 
 
 class TensorEmbedding(nn.Module):