Scale scaling constants by number of atoms

globus-labs · Dec 25, 2023 · b8e4beb · b8e4beb
1 parent 725db57
commit b8e4beb
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/jitterbug/model/dscribe/local.py b/jitterbug/model/dscribe/local.py
@@ -316,6 +316,6 @@ def train_calculator(self, data: list[Atoms]) -> Calculator:
             model=model,
             desc=self.descriptors,
             desc_scaling=(offset_x, scale_x),
-            energy_scaling=(offset_y, scale_y),
+            energy_scaling=(offset_y / len(data[0]), scale_y / len(data[0])),
             device=self.device
         )
diff --git a/tests/models/test_soap.py b/tests/models/test_soap.py
@@ -117,6 +117,10 @@ def test_model(soap, train_set):
     # Run the fitting
     calcs = model.train(train_set)
 
+    # Make sure the forces are reasonable
+    eng = calcs[0].get_potential_energy(train_set[0])
+    assert np.isclose(eng, train_set[0].get_potential_energy(), atol=1e-2)
+
     # Test the mean hessian function
     mean_hess = model.mean_hessian(calcs)
     assert mean_hess.shape == (9, 9), 'Wrong shape'