improve SO3Net model class coverage and simplify TensorNet implementa…

…tions

src/matgl/models/_tensornet.py

@@ -123,10 +123,7 @@ def __init__(
                 f"Invalid activation type, please try using one of {[af.name for af in ActivationFunction]}"
             ) from None
 
-        if element_types is None:
-            self.element_types = DEFAULT_ELEMENTS
-        else:
-            self.element_types = element_types  # type: ignore
+        self.element_types = element_types  # type: ignore
 
         self.bond_expansion = BondExpansion(
             cutoff=cutoff,

tests/models/test_so3net.py

@@ -37,16 +37,18 @@ def test_model_intensive(self, graph_MoS):
         output = model(g=graph)
         assert torch.numel(output) == 2
 
-    def test_model_intensive_with_weighted_atom(self, graph_MoS):
+    def test_model_intensive_reduce_atom_classification(self, graph_MoS):
         structure, graph, state = graph_MoS
         lat = torch.tensor(np.array([structure.lattice.matrix]), dtype=matgl.float_th)
         graph.edata["pbc_offshift"] = torch.matmul(graph.edata["pbc_offset"], lat[0])
         graph.ndata["pos"] = graph.ndata["frac_coords"] @ lat[0]
-        model = SO3Net(element_types=["Mo", "S"], is_intensive=True, readout_type="weighted_atom")
+        model = SO3Net(
+            element_types=["Mo", "S"], is_intensive=True, readout_type="reduce_atom", target_property="graph"
+        )
         output = model(g=graph)
-        assert torch.numel(output) == 2
+        assert torch.numel(output) == 1
 
-    def test_model_intensive_with_classification(self, graph_MoS):
+    def test_model_intensive_weighted_atom_classification(self, graph_MoS):
         structure, graph, state = graph_MoS
         lat = torch.tensor(np.array([structure.lattice.matrix]), dtype=matgl.float_th)
         graph.edata["pbc_offshift"] = torch.matmul(graph.edata["pbc_offset"], lat[0])

tests/models/test_tensornet.py

@@ -21,8 +21,6 @@ def test_model(self, graph_MoS):
         os.remove("model.pt")
         os.remove("model.json")
         os.remove("state.pt")
-        model = TensorNet(is_intensive=False, equivariance_invariance_group="SO(3)")
-        assert torch.numel(output) == 1
 
     def test_exceptions(self):
         with pytest.raises(ValueError, match="Invalid activation type"):