made vector_representation in SO3net optional (#465)

Author: mgastegger

Diff for: src/schnetpack/configs/model/representation/so3net.yaml

@@ -7,6 +7,7 @@ n_atom_basis: 128
 n_interactions: 3
 lmax: 2
 shared_interactions: False
+return_vector_representation: False
 cutoff_fn:
   _target_: schnetpack.nn.cutoff.CosineCutoff
   cutoff: ${globals.cutoff}

Diff for: src/schnetpack/representation/so3net.py

@@ -27,6 +27,7 @@ def __init__(
         cutoff_fn: Optional[Callable] = None,
         shared_interactions: bool = False,
         max_z: int = 100,
+        return_vector_representation: bool = False,
     ):
         """
         Args:
@@ -39,6 +40,8 @@ def __init__(
             shared_interactions:
             max_z:
             conv_layer:
+            return_vector_representation: return l=1 features in Cartesian XYZ order
+                (e.g. for DipoleMoment output module)
         """
         super(SO3net, self).__init__()
 
@@ -48,6 +51,7 @@ def __init__(
         self.cutoff_fn = hydra.utils.instantiate(cutoff_fn)
         self.cutoff = cutoff_fn.cutoff
         self.radial_basis = hydra.utils.instantiate(radial_basis)
+        self.return_vector_representation = return_vector_representation
 
         self.embedding = nn.Embedding(max_z, n_atom_basis, padding_idx=0)
         self.sphharm = so3.RealSphericalHarmonics(lmax=lmax)
@@ -118,7 +122,10 @@ def forward(self, inputs: Dict[str, torch.Tensor]):
             x = x + dx
 
         inputs["scalar_representation"] = x[:, 0]
-        # extract cartesian vector from multipoles: [y, z, x] -> [x, y, z]
-        inputs["vector_representation"] = torch.roll(x[:, 1:4], 1, 1)
         inputs["multipole_representation"] = x
+
+        # extract cartesian vector from multipoles: [y, z, x] -> [x, y, z]
+        if self.return_vector_representation:
+            inputs["vector_representation"] = torch.roll(x[:, 1:4], 1, 1)
+
         return inputs