From bac77c8f6f00a068e57e0e806470e3e303bd332d Mon Sep 17 00:00:00 2001
From: wiederm <marcus.wieder@gmail.com>
Date: Sat, 9 Nov 2024 23:36:52 +0100
Subject: [PATCH] fix tests

---
 .../data/potential_defaults/aimnet2.toml      |   2 +-
 modelforge/tests/test_aimnet2.py              | 128 +++++++++---------
 modelforge/tests/test_potentials.py           |  14 ++
 3 files changed, 79 insertions(+), 65 deletions(-)

diff --git a/modelforge/tests/data/potential_defaults/aimnet2.toml b/modelforge/tests/data/potential_defaults/aimnet2.toml
index 8edadb21..223df414 100644
--- a/modelforge/tests/data/potential_defaults/aimnet2.toml
+++ b/modelforge/tests/data/potential_defaults/aimnet2.toml
@@ -21,6 +21,6 @@ number_of_per_atom_features = 64
 [potential.postprocessing_parameter]
 properties_to_process = ['per_atom_energy']
 [potential.postprocessing_parameter.per_atom_energy]
-normalize = false
+normalize = true
 from_atom_to_system_reduction = true
 keep_per_atom_property = true
diff --git a/modelforge/tests/test_aimnet2.py b/modelforge/tests/test_aimnet2.py
index 27dc37a6..ce6385f8 100644
--- a/modelforge/tests/test_aimnet2.py
+++ b/modelforge/tests/test_aimnet2.py
@@ -152,70 +152,70 @@ def test_forward(single_batch_with_batchsize, prep_temp_dir):
 
     ref_per_system_energy = torch.tensor(
         [
-            [-1.6222e00],
-            [-1.7771e-01],
-            [1.5974e-01],
-            [-1.2089e-02],
-            [-1.8864e-01],
-            [-2.7185e-01],
-            [-4.3214e00],
-            [-1.3357e00],
-            [-1.1657e00],
-            [-1.4146e00],
-            [-1.8898e00],
-            [-1.1582e00],
-            [-9.1212e00],
-            [-4.8285e00],
-            [-5.0907e00],
-            [-5.4467e00],
-            [-1.8100e00],
-            [-4.9845e00],
-            [-3.7676e00],
-            [-2.5988e00],
-            [-1.5824e01],
-            [-1.0948e01],
-            [-2.8324e-01],
-            [-4.5179e-01],
-            [-6.8437e-01],
-            [-3.1547e-01],
-            [-5.7387e-01],
-            [-4.6788e-01],
-            [-1.9818e00],
-            [-3.8900e00],
-            [-4.2745e00],
-            [-2.8107e00],
-            [-1.2960e00],
-            [-1.5892e00],
-            [-5.7663e00],
-            [-4.2937e00],
-            [-3.0977e00],
-            [-2.2906e00],
-            [-1.4034e01],
-            [-9.6701e00],
-            [-7.9657e00],
-            [-6.4762e00],
-            [-9.7999e00],
-            [-5.6619e00],
-            [-9.1679e00],
-            [-6.8304e00],
-            [-1.0582e01],
-            [-6.0419e00],
-            [-7.2018e00],
-            [-5.0521e00],
-            [-4.0748e00],
-            [-3.5285e00],
-            [-2.5017e00],
-            [-2.5237e01],
-            [-1.9461e01],
-            [-1.7413e00],
-            [-2.1273e00],
-            [-2.5887e00],
-            [-1.1963e00],
-            [-2.4938e00],
-            [-3.1271e00],
-            [-1.7812e00],
-            [-8.0866e00],
-            [-8.7542e00],
+            [0.2630],
+            [-0.5150],
+            [-0.2999],
+            [-0.0297],
+            [-0.4382],
+            [-0.1805],
+            [0.5974],
+            [0.1769],
+            [0.0842],
+            [-0.2955],
+            [0.1295],
+            [-0.4067],
+            [0.4135],
+            [0.3202],
+            [0.2481],
+            [0.6696],
+            [0.0380],
+            [0.0834],
+            [-0.2613],
+            [-0.8373],
+            [0.2033],
+            [0.1554],
+            [0.0624],
+            [-0.3643],
+            [-0.7861],
+            [-0.0398],
+            [-0.4675],
+            [-0.1000],
+            [0.3265],
+            [0.2546],
+            [-0.1597],
+            [-0.9611],
+            [0.0653],
+            [-0.4411],
+            [0.2587],
+            [-0.1082],
+            [0.0461],
+            [0.0407],
+            [0.6725],
+            [0.3874],
+            [0.3393],
+            [0.1747],
+            [0.4048],
+            [0.1001],
+            [0.1496],
+            [0.2432],
+            [0.3578],
+            [0.2792],
+            [-0.3365],
+            [-0.3329],
+            [-0.8465],
+            [0.0463],
+            [-0.4385],
+            [0.1224],
+            [-0.0442],
+            [0.1029],
+            [-0.4559],
+            [-1.1701],
+            [-0.2714],
+            [0.0318],
+            [-0.8579],
+            [-0.3836],
+            [0.2487],
+            [-0.2728],
         ],
     )
 
diff --git a/modelforge/tests/test_potentials.py b/modelforge/tests/test_potentials.py
index f5ff183b..cb69e05a 100644
--- a/modelforge/tests/test_potentials.py
+++ b/modelforge/tests/test_potentials.py
@@ -287,6 +287,20 @@ def test_energy_scaling_and_offset(
     assert torch.allclose(scaled_output["per_system_energy"], compare_to.sum())
 
 
+"""
+tensor([[-406.9472],
+        [-397.2831],
+        [-397.2831],
+        [-397.2831],
+        [-397.2831]])
+tensor([[-402.9324],
+        [-401.2677],
+        [-401.2677],
+        [-401.2683],
+        [-401.2684]])
+"""
+
+
 @pytest.mark.parametrize(
     "potential_name", _Implemented_NNPs.get_all_neural_network_names()
 )