Repository: mir-group/flare_pp (language: python)

Issue:
Write_model for SGP_calculator class?
Dear all,

is it possible to save the trained SGP after an OTF run, so that I can continue the training with a different phase/structure? I noticed that the SGP_calculator class  does not have a write_model() function.

Many thanks for your consideration,
Veniero. 


Discussion / hints:
Hi @Renevio 

Thanks for this suggestion. Yes the `write_model` needs to be implemented for SGP_Calculator. We will update it asap

Fix patch:
diff --git a/flare_pp/sparse_gp_calculator.py b/flare_pp/sparse_gp_calculator.py
index ca1ce600..81beb810 100644
--- a/flare_pp/sparse_gp_calculator.py
+++ b/flare_pp/sparse_gp_calculator.py
@@ -1,7 +1,9 @@
 from ase.calculators.calculator import Calculator, all_changes
-from flare_pp._C_flare import SparseGP, Structure
+from flare.utils.element_coder import NumpyEncoder
+from flare_pp._C_flare import Structure
+from flare_pp.sparse_gp import SGP_Wrapper
 import numpy as np
-import time
+import time, json
 
 
 class SGP_Calculator(Calculator):
@@ -108,6 +110,31 @@ def get_uncertainties(self, atoms):
     def calculation_required(self, atoms, quantities):
         return True
 
+    def as_dict(self):
+        out_dict = dict(vars(self))
+        out_dict["gp_model"] = self.gp_model.as_dict()
+        out_dict.pop("get_spin_polarized")
+        return out_dict
+
+    @staticmethod
+    def from_dict(dct):
+        sgp, _ = SGP_Wrapper.from_dict(dct["gp_model"])
+        calc = SGP_Calculator(sgp)
+        calc.results = dct["results"]
+        return calc
+
+    def write_model(self, name):
+        if ".json" != name[-5:]:
+            name += ".json"
+        with open(name, "w") as f:
+            json.dump(self.as_dict(), f, cls=NumpyEncoder)
+
+    @staticmethod
+    def from_file(name):
+        with open(name, "r") as f:
+            calc = SGP_Calculator.from_dict(json.loads(f.readline()))
+        return calc
+
 
 def sort_variances(structure_descriptor, variances):
     # Check that the variance length matches the number of atoms.
diff --git a/setup.py b/setup.py
index 99c16c39..a5ef49cf 100644
--- a/setup.py
+++ b/setup.py
@@ -161,7 +161,7 @@ def _decode(x):
 setuptools.setup(
     name="flare_pp",
     packages=setuptools.find_packages(exclude=[]),
-    version="0.0.21",
+    version="0.0.22",
     author="Materials Intelligence Research",
     author_email="mir@g.harvard.edu",
     description="Many-body extension of the flare code",


Test patch:
diff --git a/flare_pp/tests/test_sparse_gp.py b/flare_pp/tests/test_sparse_gp.py
index 9907ffa1..f66d009d 100644
--- a/flare_pp/tests/test_sparse_gp.py
+++ b/flare_pp/tests/test_sparse_gp.py
@@ -8,6 +8,7 @@
 from flare.lammps import lammps_calculator
 
 from flare_pp.sparse_gp import SGP_Wrapper
+from flare_pp.sparse_gp_calculator import SGP_Calculator
 from flare_pp._C_flare import NormalizedDotProduct, B2, SparseGP, Structure
 
 
@@ -97,7 +98,38 @@ def test_train():
     assert hyps_init != hyps_post
     assert sgp_py.likelihood != 0.0
 
-
+def test_dict():
+    """
+    Check the method from_dict and as_dict
+    """
+
+    out_dict = sgp_py.as_dict()
+    assert len(sgp_py) == len(out_dict["training_structures"])
+    new_sgp, _ = SGP_Wrapper.from_dict(out_dict)
+    assert len(sgp_py) == len(new_sgp)
+    assert len(sgp_py.sparse_gp.kernels) == len(new_sgp.sparse_gp.kernels)
+    assert np.allclose(sgp_py.hyps, new_sgp.hyps)
+
+def test_dump():
+    """
+    Check the method from_file and write_model of SGP_Wrapper
+    """
+
+    sgp_py.write_model("sgp.json")
+    new_sgp, _ = SGP_Wrapper.from_file("sgp.json")
+    assert len(sgp_py) == len(new_sgp)
+    assert len(sgp_py.sparse_gp.kernels) == len(new_sgp.sparse_gp.kernels)
+    assert np.allclose(sgp_py.hyps, new_sgp.hyps)
+
+def test_calc():
+    """
+    Check the method from_file and write_model of SGP_Calculator
+    """
+
+    calc = SGP_Calculator(sgp_py)
+    calc.write_model("sgp_calc.json")
+    new_calc = SGP_Calculator.from_file("sgp_calc.json")
+    assert len(calc.gp_model) == len(new_calc.gp_model)
 
 @pytest.mark.skipif(
     not os.environ.get("lmp", False),
@@ -281,6 +313,7 @@ def test_lammps():
 #                    py_var[i, comp] += desc[i, comp, s1].dot(desc[i, comp, s2].dot(beta_matr))
 ##                    py_var[i, comp] += np.sum(desc[i, comp, s1]) #np.sum(beta_matr[:, comp]) #np.sum(desc[i, comp, s1])
     print(py_var)
+    os.system("rm -r tmp* *.txt lmp.flare *.json")
 
     #for i in range(4):
     #    print([sgp_py.sparse_gp.varmap_coeffs[0, n_descriptors * i + j] for j in range(3)])