Fix tests

Author: frostedoyster

src/metatrain/cli/train.py

@@ -421,10 +421,7 @@ def train_model(
     # TRAIN MODEL #############
     ###########################
 
-    # logger.info("Calling trainer")
-    # from torch.profiler import profile, ProfilerActivity
-    # with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
-    # if True:
+    logger.info("Calling trainer")
     trainer.train(
         model=model,
         dtype=dtype,
@@ -433,8 +430,6 @@ def train_model(
         val_datasets=val_datasets,
         checkpoint_dir=str(checkpoint_dir),
     )
-    # print(prof.key_averages().table(sort_by="self_cpu_time_total", row_limit=20))
-    # exit()
 
     if not is_main_process():
         return  # only save and evaluate on the main process

src/metatrain/experimental/phace/tests/test_equivariance.py

@@ -64,7 +64,7 @@ def test_rotational_invariance():
     torch.set_default_dtype(torch.float32)  # change back
 
 
-@pytest.mark.parametrize("o3_lambda", [0, 1, 2, 3, 4])
+@pytest.mark.parametrize("o3_lambda", [0, 1, 2, 3])
 @pytest.mark.parametrize("o3_sigma", [1])
 def test_equivariance_rotations(o3_lambda, o3_sigma):
     """Tests that the model is rotationally equivariant when predicting
@@ -128,7 +128,7 @@ def test_equivariance_rotations(o3_lambda, o3_sigma):
 
 
 @pytest.mark.parametrize("dataset_path", [DATASET_PATH, DATASET_PATH_PERIODIC])
-@pytest.mark.parametrize("o3_lambda", [0, 1, 2, 3, 4])
+@pytest.mark.parametrize("o3_lambda", [0, 1, 2, 3])
 @pytest.mark.parametrize("o3_sigma", [1])
 def test_equivariance_inversion(dataset_path, o3_lambda, o3_sigma):
     """Tests that the model is equivariant with respect to inversions."""

src/metatrain/experimental/phace/tests/test_functionality.py

@@ -164,7 +164,7 @@ def test_output_last_layer_features():
     ]
     assert features.values.shape == (
         4,
-        32,
+        192,
     )
     assert features.properties.names == [
         "properties",
@@ -178,7 +178,7 @@ def test_output_last_layer_features():
     assert last_layer_features.values.shape == (
         4,
         1,
-        32,
+        192,
     )
     assert last_layer_features.properties.names == [
         "properties",
@@ -208,7 +208,7 @@ def test_output_last_layer_features():
     ]
     assert features.values.shape == (
         1,
-        32,
+        192,
     )
     assert features.properties.names == [
         "properties",
@@ -220,7 +220,7 @@ def test_output_last_layer_features():
     assert outputs["mtt::aux::energy_last_layer_features"].block().values.shape == (
         1,
         1,
-        32,
+        192,
     )
     assert outputs["mtt::aux::energy_last_layer_features"].block().properties.names == [
         "properties",

src/metatrain/experimental/phace/tests/test_regression.py

@@ -45,17 +45,17 @@ def test_regression_init():
 
     expected_output = torch.tensor(
         [
-            [-3.736138916016e01],
-            [6.381711363792e-02],
-            [1.800283813477e01],
-            [-2.970426513672e03],
-            [1.789148712158e01],
+            [0.002085668733],
+            [-0.003157143714],
+            [0.000328244379],
+            [0.004316798877],
+            [0.001980246045],
         ]
     )
 
     # if you need to change the hardcoded values:
-    torch.set_printoptions(precision=12)
-    print(output["mtt::U0"].block().values)
+    # torch.set_printoptions(precision=12)
+    # print(output["mtt::U0"].block().values)
 
     torch.testing.assert_close(output["mtt::U0"].block().values, expected_output)
 
@@ -115,16 +115,16 @@ def test_regression_train():
 
     expected_output = torch.tensor(
         [
-            [2.120110988617],
-            [0.246357604861],
-            [0.113200485706],
-            [0.136439576745],
-            [0.023953542113],
+            [0.101170130074],
+            [0.038209509104],
+            [0.012803453952],
+            [0.151425197721],
+            [0.050753910094],
         ]
     )
 
     # if you need to change the hardcoded values:
-    torch.set_printoptions(precision=12)
-    print(output["mtt::U0"].block().values)
+    # torch.set_printoptions(precision=12)
+    # print(output["mtt::U0"].block().values)
 
     torch.testing.assert_close(output["mtt::U0"].block().values, expected_output)