Repository: mir-group/flare

Issue:
Structure.from_ase_atoms does not copy forces

Discussion:

Patch:
diff --git a/docs/source/tutorials/ase.ipynb b/docs/source/tutorials/ase.ipynb
index 033984b6..e12b1afc 100644
--- a/docs/source/tutorials/ase.ipynb
+++ b/docs/source/tutorials/ase.ipynb
@@ -261,8 +261,11 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "3. Set up parameters for On-The-Fly (OTF) training. The descriptions of the parameters are in [ASE OTF module](https://flare-yuuuu.readthedocs.io/en/development/flare/ase/otf.html).\n",
+    "3. Set up parameters for On-The-Fly (OTF) training. The descriptions of the parameters are in [ASE OTF module](https://flare.readthedocs.io/en/latest/flare/ase/otf.html).\n",
     "4. Set up the ASE_OTF training engine, and run\n",
+    "\n",
+    "    **Note**: the ASE Trajectory is supported, but NOT recommended. \n",
+    "\n",
     "5. Check `otf.out` after the training is done."
    ]
   },
@@ -278,7 +281,8 @@
     "              'output_name': 'otf',\n",
     "              'std_tolerance_factor': 2,\n",
     "              'max_atoms_added' : 4,\n",
-    "              'freeze_hyps': 10}\n",
+    "              'freeze_hyps': 10,\n",
+    "              'write_model': 3}\n",
     "\n",
     "test_otf = ASE_OTF(super_cell, \n",
     "                   timestep = 1 * units.fs,\n",
@@ -295,7 +299,31 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Then check the `*.out` file for the training log."
+    "Then check the `*.out` file for the training log.\n",
+    "\n",
+    "## Step 5 (Optional): Resume Interrupted Training\n",
+    "\n",
+    "At the end of each OTF training step, there will be several checkpoint files dumpped\n",
+    "\n",
+    "- `<output_name>_checkpt.json`: checkpoint of the current MD step of OTF. In the above example, `al_otf_qe_checkpt.json`. \n",
+    "\n",
+    "- `<output_name>_flare.json`: If you've set `write_model=3`, then there will be another file saving the trained FLARE calculator, which will be loaded when restarting OTF.\n",
+    "\n",
+    "- `<output_name>_atoms.json`: The ASE Atoms of the current MD step in the format of `json`\n",
+    "\n",
+    "- `<output_name>_dft.pickle`: The DFT calculator saved in the format of `.pickle`.\n",
+    "\n",
+    "Then, use `ASE_OTF.from_checkpoint(<output_name>_checkpt.json)` to load the OTF state, and resume the training by `run()`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "new_otf = ASE_OTF.from_checkpoint(\"<output_name>_checkpt.json\")\n",
+    "new_otf.run()"
    ]
   }
  ],
diff --git a/docs/source/tutorials/otf.ipynb b/docs/source/tutorials/otf.ipynb
index 18423f41..568f491e 100644
--- a/docs/source/tutorials/otf.ipynb
+++ b/docs/source/tutorials/otf.ipynb
@@ -107,10 +107,14 @@
     "freeze_hyps = 3\n",
     "\n",
     "otf = OTF(qe_input, dt, number_of_steps, gp, dft_loc,\n",
-    "          std_tolerance_factor, init_atoms=[0],\n",
-    "          calculate_energy=True, output_name='al_otf_qe',\n",
-    "          freeze_hyps=freeze_hyps, skip=5,\n",
-    "          max_atoms_added=max_atoms_added)"
+    "          std_tolerance_factor, \n",
+    "          init_atoms=[0],\n",
+    "          calculate_energy=True, \n",
+    "          output_name='al_otf_qe',\n",
+    "          freeze_hyps=freeze_hyps, \n",
+    "          skip=5,\n",
+    "          max_atoms_added=max_atoms_added,\n",
+    "          write_model=3)"
    ]
   },
   {
@@ -137,6 +141,16 @@
     "\n",
     "- `skip`: record/dump the information every skip steps.\n",
     "\n",
+    "- `write_model`: identify the frequency of dumpping GP model during the training. \n",
+    "\n",
+    "    - `1`: default, dump GP model only when the OTF is complete\n",
+    "\n",
+    "    - `2`: dump GP model every time after training\n",
+    "    \n",
+    "    - `3`: dump GP model at every step\n",
+    "\n",
+    "    (We recommend using `3` if you are possibly gonna resume the OTF training, as instructed by the next section.)\n",
+    "\n",
     "## Step 4: Launch the OTF Training\n",
     "Finally, let’s run it!"
    ]
@@ -157,6 +171,29 @@
    "source": [
     "After OTF training is finished, we can check log file `al_otf_qe.out` for all the information dumped. This output file can be parsed using our `otf_parser.py` module, which will be introduced in the [after-training-tutorial](https://flare.readthedocs.io/en/latest/tutorials/after_training.html)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 5 (Optional): Resume Interrupted Training\n",
+    "\n",
+    "At the end of each OTF training step, there will be a checkpoint file dumpped of name `<output_name>_checkpt.json`, in the above example, `al_otf_qe_checkpt.json`. This file saves the state at the current OTF MD step. \n",
+    "\n",
+    "If you've set `write_model=3`, then there will be another file saving the trained GP model, `<output_name>_gp.json`, which will be loaded when restarting OTF.\n",
+    "\n",
+    "Then, use `OTF.from_checkpoint(<output_name>_checkpt.json)` to load the OTF state, and resume the training by `run()`.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "new_otf = OTF.from_checkpoint(log_name + \"_checkpt.json\")\n",
+    "new_otf.run()"
+   ]
   }
  ],
  "metadata": {
@@ -180,4 +217,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file
diff --git a/flare/ase/atoms.py b/flare/ase/atoms.py
index dc336b1a..c9650649 100644
--- a/flare/ase/atoms.py
+++ b/flare/ase/atoms.py
@@ -6,6 +6,7 @@ from copy import deepcopy
 
 import numpy as np
 from ase import Atoms
+from ase.io import read, write
 from flare.utils.learner import get_max_cutoff
 
 
@@ -112,3 +113,11 @@ class FLARE_Atoms(Atoms):
     @property
     def max_cutoff(self):
         return get_max_cutoff(self.cell)
+
+    def as_dict(self):
+        return self.todict()
+
+    @staticmethod
+    def from_dict(dct):
+        atoms = Atoms.fromdict(dct)
+        return FLARE_Atoms.from_ase_atoms(atoms)
diff --git a/flare/ase/calculator.py b/flare/ase/calculator.py
index 84a37065..a924855e 100644
--- a/flare/ase/calculator.py
+++ b/flare/ase/calculator.py
@@ -8,8 +8,10 @@ and `ASE Calculator tutorial <https://wiki.fysik.dtu.dk/ase/ase/atoms.html#addin
 import warnings
 import numpy as np
 import multiprocessing as mp
+import json
+
 from flare.env import AtomicEnvironment
-from flare.struc import Structure
+from flare.gp import GaussianProcess
 from flare.mgp import MappedGaussianProcess
 from flare.predict import (
     predict_on_structure_par_en,
@@ -17,6 +19,8 @@ from flare.predict import (
     predict_on_structure_efs,
     predict_on_structure_efs_par,
 )
+from flare.utils.element_coder import NumpyEncoder
+
 from ase.calculators.calculator import Calculator
 
 
@@ -54,7 +58,9 @@ class FLARE_Calculator(Calculator):
             by default.
     """
 
-    def __init__(self, gp_model, mgp_model=None, par=False, use_mapping=False):
+    def __init__(
+        self, gp_model, mgp_model=None, par=False, use_mapping=False, **kwargs
+    ):
         super().__init__()  # all set to default values, TODO: change
         self.mgp_model = mgp_model
         self.gp_model = gp_model
@@ -185,3 +191,54 @@ class FLARE_Calculator(Calculator):
 
     def calculation_required(self, atoms, quantities):
         return True
+
+    def as_dict(self):
+        outdict = {}
+
+        gp_dict = self.gp_model.as_dict()
+        outdict["gp_model"] = gp_dict
+
+        outdict["use_mapping"] = self.use_mapping
+        if self.use_mapping:
+            mgp_dict = self.mgp_model.as_dict()
+            outdict["mgp_model"] = mgp_dict
+        else:
+            outdict["mgp_model"] = None
+
+        outdict["par"] = self.par
+        outdict["results"] = self.results
+        return outdict
+
+    @staticmethod
+    def from_dict(dct):
+        dct["gp_model"] = GaussianProcess.from_dict(dct["gp_model"])
+        if dct["use_mapping"]:
+            dct["mgp_model"] = MappedGaussianProcess.from_dict(dct["mgp_model"])
+
+        calc = FLARE_Calculator(**dct)
+        res = dct["results"]
+        for key in res:
+            if isinstance(res[key], float):
+                calc.results[key] = res[key]
+            if isinstance(res[key], list):
+                calc.results[key] = np.array(res[key])
+
+        if dct["use_mapping"]:
+            for xb in calc.mgp_model.maps:
+                xb_map = calc.mgp_model.maps[xb]
+                xb_map.hyps_mask = calc.gp_model.hyps_mask
+
+        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 = FLARE_Calculator.from_dict(json.loads(f.readline()))
+
+        return calc
diff --git a/flare/ase/otf.py b/flare/ase/otf.py
index 32995704..0f1a565c 100644
--- a/flare/ase/otf.py
+++ b/flare/ase/otf.py
@@ -1,12 +1,14 @@
 """
-:class:`OTF` is the on-the-fly training module for ASE, WITHOUT molecular dynamics engine. 
+:class:`ASE_OTF` is the on-the-fly training module for ASE, WITHOUT molecular dynamics engine. 
 It needs to be used adjointly with ASE MD engine. 
 """
 import os
 import sys
 import inspect
+import pickle
 from time import time
 from copy import deepcopy
+import logging
 
 import numpy as np
 from ase.md.npt import NPT
@@ -15,7 +17,9 @@ from ase.md.nptberendsen import NPTBerendsen
 from ase.md.verlet import VelocityVerlet
 from ase.md.langevin import Langevin
 from ase import units
+from ase.io import read, write
 
+import flare
 from flare.otf import OTF
 from flare.utils.learner import is_std_in_bound
 
@@ -24,6 +28,17 @@ from flare.ase.calculator import FLARE_Calculator
 import flare.ase.dft as dft_source
 
 
+def reset_npt_momenta(npt_engine, force):
+    # in the last step, the momenta was set by flare forces, change to dft forces
+    npt_engine._calculate_q_future(force)
+    npt_engine.atoms.set_momenta(
+        np.dot(
+            npt_engine.q_future - npt_engine.q_past, npt_engine.h / (2 * npt_engine.dt)
+        )
+        * npt_engine._getmasses()
+    )
+
+
 class ASE_OTF(OTF):
 
     """
@@ -45,7 +60,7 @@ class ASE_OTF(OTF):
             currently in experiment.
 
     The following arguments are for on-the-fly training, the user can also
-    refer to :class:`OTF`
+    refer to :class:`flare.otf.OTF`
 
     Args:
         prev_pos_init ([type], optional): Previous positions. Defaults
@@ -97,7 +112,9 @@ class ASE_OTF(OTF):
     ):
 
         self.atoms = FLARE_Atoms.from_ase_atoms(atoms)
+        self.timestep = timestep
         self.md_engine = md_engine
+        self.md_kwargs = md_kwargs
 
         if md_engine == "VelocityVerlet":
             MD = VelocityVerlet
@@ -107,6 +124,10 @@ class ASE_OTF(OTF):
             MD = NPTBerendsen
         elif md_engine == "NPT":
             MD = NPT
+            # TODO: solve the md step
+            assert (
+                md_kwargs["pfactor"] is None
+            ), "Current MD OTF only supports pfactor=None"
         elif md_engine == "Langevin":
             MD = Langevin
         else:
@@ -132,6 +153,10 @@ class ASE_OTF(OTF):
             **otf_kwargs
         )
 
+        self.flare_name = self.output_name + "_flare.json"
+        self.dft_name = self.output_name + "_dft.pickle"
+        self.atoms_name = self.output_name + "_atoms.json"
+
     def get_structure_from_input(self, prev_pos_init):
         self.structure = self.atoms
         if prev_pos_init is None:
@@ -163,9 +188,10 @@ class ASE_OTF(OTF):
 
         # Reset FLARE calculator if necessary.
         if not isinstance(self.atoms.calc, FLARE_Calculator):
-            self.atoms.set_calculator(self.flare_calc)
+            self.atoms.calc = self.flare_calc
 
-        self.atoms.calc.calculate(self.atoms)
+        if not self.flare_calc.results:
+            self.atoms.calc.calculate(self.atoms)
 
     def md_step(self):
         """
@@ -176,12 +202,30 @@ class ASE_OTF(OTF):
         self.structure.prev_positions = np.copy(self.structure.positions)
 
         # Take MD step.
-        self.md.step()
+        f = self.atoms.get_forces()
+
+        if self.md_engine == "NPT":
+            self.flare_calc.results = {}  # init flare calculator
+
+            if self.dft_step:
+                reset_npt_momenta(self.md, f)
+                self.atoms.calc = self.flare_calc
+
+            self.md.step()  # use flare to get force for next step
+        else:
+            self.flare_calc.results = {}  # init flare calculator
+            if self.dft_step:
+                self.atoms.calc = self.flare_calc
+
+            self.md.step(f)
 
         # Update the positions and cell of the structure object.
         self.structure.cell = np.copy(self.atoms.cell)
         self.structure.positions = np.copy(self.atoms.positions)
 
+    def write_gp(self):
+        self.flare_calc.write_model(self.flare_name)
+
     def update_positions(self, new_pos):
         # call OTF method
         super().update_positions(new_pos)
@@ -201,10 +245,84 @@ class ASE_OTF(OTF):
         # Convert velocities to Angstrom / ps.
         self.velocities = self.atoms.get_velocities() * units.fs * 1e3
 
-    def update_gp(self, train_atoms, dft_frcs, dft_energy=None,
-                  dft_stress=None):
-        super().update_gp(train_atoms, dft_frcs, dft_energy=dft_energy,
-                          dft_stress=dft_stress)
+    def update_gp(self, train_atoms, dft_frcs, dft_energy=None, dft_stress=None):
+        self.output.add_atom_info(train_atoms, self.structure.stds)
+
+        # update gp model
+        self.gp.update_db(
+            self.structure,
+            dft_frcs,
+            custom_range=train_atoms,
+            energy=dft_energy,
+            stress=dft_stress,
+        )
+
+        self.gp.set_L_alpha()
 
+        # train model
+        if (self.dft_count - 1) < self.freeze_hyps:
+            self.train_gp()
+
+        # update mgp model
         if self.flare_calc.use_mapping:
             self.flare_calc.mgp_model.build_map(self.flare_calc.gp_model)
+
+        # write model
+        if (self.dft_count - 1) < self.freeze_hyps:
+            if self.write_model == 2:
+                self.write_gp()
+        if self.write_model == 3:
+            self.write_gp()
+
+    def as_dict(self):
+
+        # DFT module and Trajectory will cause issue in deepcopy
+        self.dft_module = self.dft_module.__name__
+        md = self.md
+        self.md = None
+
+        dct = deepcopy(dict(vars(self)))
+        self.dft_module = eval(self.dft_module)
+        self.md = md
+
+        # write atoms and flare calculator to separate files
+        write(self.atoms_name, self.atoms)
+        dct["atoms"] = self.atoms_name
+
+        self.flare_calc.write_model(self.flare_name)
+        dct["flare_calc"] = self.flare_name
+
+        # dump dft calculator as pickle
+        with open(self.dft_name, "wb") as f:
+            pickle.dump(self.dft_loc, f)  # dft_loc is the dft calculator
+        dct["dft_loc"] = self.dft_name
+
+        dct["gp"] = self.gp_name
+
+        for key in ["output", "pred_func", "structure", "dft_input", "md"]:
+            dct.pop(key)
+
+        return dct
+
+    @staticmethod
+    def from_dict(dct):
+        flare_calc = FLARE_Calculator.from_file(dct["flare_calc"])
+        dct["atoms"] = read(dct["atoms"])
+        dct["atoms"].calc = flare_calc
+        dct.pop("gp")
+
+        with open(dct["dft_loc"], "rb") as f:
+            dct["dft_calc"] = pickle.load(f)
+
+        for key in ["dt", "dft_loc"]:
+            dct.pop(key)
+
+        new_otf = ASE_OTF(**dct)
+        new_otf.dft_count = dct["dft_count"]
+        new_otf.curr_step = dct["curr_step"]
+
+        if new_otf.md_engine == "NPT":
+            if not new_otf.md.initialized:
+                new_otf.md.initialize()
+
+        return new_otf
diff --git a/flare/gp.py b/flare/gp.py
index 12697216..12c4a005 100644
--- a/flare/gp.py
+++ b/flare/gp.py
@@ -1,5 +1,4 @@
 import json
-import json
 import logging
 import multiprocessing as mp
 import pickle
@@ -735,9 +734,8 @@ class GaussianProcess:
         # Remove the callables
         for key in ['kernel', 'kernel_grad', 'energy_kernel',
                     'energy_force_kernel', 'efs_energy_kernel',
-                    'efs_force_kernel', 'efs_self_kernel']:
-            if out_dict.get(key) is not None:
-                del out_dict[key]
+                    'efs_force_kernel', 'efs_self_kernel', 'output']:
+            out_dict.pop(key)
 
         return out_dict
 
diff --git a/flare/gp_algebra.py b/flare/gp_algebra.py
index aa7849e3..be4a3912 100644
--- a/flare/gp_algebra.py
+++ b/flare/gp_algebra.py
@@ -844,8 +844,10 @@ def force_force_vector_unit(name, s, e, x, kernel, hyps, cutoffs, hyps_mask,
 
 def efs_force_vector_unit(name, s, e, x, efs_force_kernel, hyps, cutoffs,
                           hyps_mask):
-    size = e - s
 
+    training_data = _global_training_data[name]
+
+    size = e - s
     k_ef = np.zeros((1, size * 3))
     k_ff = np.zeros((3, size * 3))
     k_sf = np.zeros((6, size * 3))
@@ -853,7 +855,7 @@ def efs_force_vector_unit(name, s, e, x, efs_force_kernel, hyps, cutoffs,
     args = from_mask_to_args(hyps, cutoffs, hyps_mask)
 
     for m_index in range(size):
-        x_2 = _global_training_data[name][m_index + s]
+        x_2 = training_data[m_index + s]
         ef, ff, sf = efs_force_kernel(x, x_2, *args)
 
         ind1 = m_index * 3
diff --git a/flare/mgp/mapxb.py b/flare/mgp/mapxb.py
index 40440e9c..8dfaf65a 100644
--- a/flare/mgp/mapxb.py
+++ b/flare/mgp/mapxb.py
@@ -106,11 +106,6 @@ class MapXbody:
             m.build_map(GP)
 
     def predict(self, atom_env):
-
-        assert Parameters.compare_dict(
-            self.hyps_mask, atom_env.cutoffs_mask
-        ), "GP.hyps_mask is not the same as atom_env.cutoffs_mask"
-
         f_spcs = np.zeros(3)
         vir_spcs = np.zeros(6)
         v_spcs = 0
@@ -612,13 +607,22 @@ class SingleMapXbody:
         """
 
         min_dist = np.min(lengths)
-        if min_dist < np.max(self.bounds[0][0]):
+        if min_dist < np.max(self.bounds[0]):
             raise ValueError(
                 self.species,
                 min_dist,
                 f"The minimal distance {min_dist:.3f}"
                 f" is below the mgp lower bound {self.bounds[0]}",
             )
+        
+        max_dist = np.max(lengths)
+        if max_dist > np.min(self.bounds[1]):
+            raise Exception(
+                self.species,
+                max_dist,
+                f"The atomic environment should have cutoff smaller"
+                f" than the GP cutoff"
+            )
 
         lengths = np.array(lengths)
         xyzs = np.array(xyzs)
diff --git a/flare/mgp/mgp.py b/flare/mgp/mgp.py
index 623ec125..8681dfe5 100644
--- a/flare/mgp/mgp.py
+++ b/flare/mgp/mgp.py
@@ -54,9 +54,9 @@ class MappedGaussianProcess:
     For `grid_params`, the following keys and values are allowed
 
     Args:
-        'two_body' (dict, optional): if 2-body is present, set as a dictionary
+        'twobody' (dict, optional): if 2-body is present, set as a dictionary
             of parameters for 2-body mapping. Parameters see below.
-        'three_body' (dict, optional): if 3-body is present, set as a dictionary
+        'threebody' (dict, optional): if 3-body is present, set as a dictionary
             of parameters for 3-body mapping. Parameters see below.
         'load_grid' (str, optional): Default None. the path to the directory
             where the previously generated grids (``grid_*.npy``) are stored.
@@ -97,7 +97,7 @@ class MappedGaussianProcess:
         GP: GaussianProcess = None,
         var_map: str = None,
         container_only: bool = True,
-        lmp_file_name: str = "lmp.mgp",
+        lmp_file_name: str = "lmp",
         n_cpus: int = None,
         n_sample: int = 10,
     ):
diff --git a/flare/otf.py b/flare/otf.py
index 2bdd229b..d47476d8 100644
--- a/flare/otf.py
+++ b/flare/otf.py
@@ -1,18 +1,21 @@
 import logging
+import json
 import numpy as np
 import time
+import warnings
 
 from copy import deepcopy
 from datetime import datetime
 from shutil import copyfile
 from typing import List, Tuple, Union
 
+import flare
 import flare.predict as predict
 from flare import struc, gp, env, md
 from flare.dft_interface import dft_software
 from flare.output import Output
 from flare.utils.learner import is_std_in_bound
-
+from flare.utils.element_coder import NumpyEncoder
 
 class OTF:
     """Trains a Gaussian process force field on the fly during
@@ -106,47 +109,57 @@ class OTF:
      dft_loc: str = None, dft_input: str = None, dft_output='dft.out',
      dft_kwargs=None,
      store_dft_output: Tuple[Union[str, List[str]], str] = None,
-     # par args
-     n_cpus: int = 1):
+     # other args
+     n_cpus: int = 1, **kwargs):
+
 
+        # set DFT
+        self.dft_loc = dft_loc
         self.dft_input = dft_input
         self.dft_output = dft_output
-        self.dt = dt
-        self.number_of_steps = number_of_steps
-        self.gp = gp
-        self.dft_loc = dft_loc
-        self.std_tolerance = std_tolerance_factor
-        self.skip = skip
         self.dft_step = True
-        self.freeze_hyps = freeze_hyps
-
+        self.dft_count = 0
         if isinstance(force_source, str):
             self.dft_module = dft_software[force_source]
         else:
             self.dft_module = force_source
 
-        # parse input file
-        self.get_structure_from_input(prev_pos_init)
-
+        # set md
+        self.dt = dt
+        self.number_of_steps = number_of_steps
+        self.get_structure_from_input(prev_pos_init) # parse input file
         self.noa = self.structure.positions.shape[0]
-        self.atom_list = list(range(self.noa))
-        self.curr_step = 0
-
-        self.max_atoms_added = max_atoms_added
+        self.rescale_steps = rescale_steps
+        self.rescale_temps = rescale_temps
 
+        # set flare
+        self.gp = gp
         # initialize local energies
         if calculate_energy:
             self.local_energies = np.zeros(self.noa)
         else:
             self.local_energies = None
 
-        # set atom list for initial dft run
-        if init_atoms is None:
+        # set otf
+        self.std_tolerance = std_tolerance_factor
+        self.skip = skip
+        self.max_atoms_added = max_atoms_added
+        self.freeze_hyps = freeze_hyps
+        if init_atoms is None: # set atom list for initial dft run
             self.init_atoms = [int(n) for n in range(self.noa)]
         else:
             self.init_atoms = init_atoms
 
-        self.dft_count = 0
+        self.n_cpus = n_cpus # set number of cpus and npool for DFT runs
+        self.npool = npool
+        self.mpi = mpi
+
+        self.dft_kwargs = dft_kwargs
+        self.store_dft_output = store_dft_output
+
+        # other args
+        self.atom_list = list(range(self.noa))
+        self.curr_step = 0
 
         # Set the prediction function based on user inputs.
         # Force only prediction.
@@ -167,21 +180,12 @@ class OTF:
         else:
             self.pred_func = predict.predict_on_structure_efs
 
-        # set rescale attributes
-        self.rescale_steps = rescale_steps
-        self.rescale_temps = rescale_temps
-
         # set logger
         self.output = Output(output_name, always_flush=True)
         self.output_name = output_name
+        self.gp_name = self.output_name + '_gp.json'
+        self.checkpt_name = self.output_name + '_checkpt.json'
 
-        # set number of cpus and npool for DFT runs
-        self.n_cpus = n_cpus
-        self.npool = npool
-        self.mpi = mpi
-
-        self.dft_kwargs = dft_kwargs
-        self.store_dft_output = store_dft_output
         self.write_model = write_model
 
     def run(self):
@@ -193,9 +197,10 @@ class OTF:
         'Year.Month.Day:Hour:Minute:Second:'.
         """
 
+        optional_dict = {"Restart": self.curr_step}
         self.output.write_header(
             str(self.gp), self.dt, self.number_of_steps, self.structure,
-            self.std_tolerance)
+            self.std_tolerance, optional_dict)
 
         counter = 0
         self.start_time = time.time()
@@ -209,8 +214,6 @@ class OTF:
                 # When DFT is called, ASE energy, forces, and stresses should
                 # get updated.
                 self.initialize_train()
-                self.update_temperature()
-                self.record_state()
 
             # after step 1, try predicting with GP model
             else:
@@ -253,11 +256,13 @@ class OTF:
             # TODO: Reinstate velocity rescaling.
             self.md_step()
             self.curr_step += 1
+            self.checkpoint()
 
         self.output.conclude_run()
 
         if self.write_model >= 1:
-            self.gp.write_model(self.output_name+"_model")
+            self.write_gp()
+            self.checkpoint()
 
     def get_structure_from_input(self, prev_pos_init):
         positions, species, cell, masses = \
@@ -272,6 +277,9 @@ class OTF:
         self.run_dft()
         dft_frcs = deepcopy(self.structure.forces)
 
+        self.update_temperature()
+        self.record_state()
+
         # make initial gp model and predict forces
         self.update_gp(self.init_atoms, dft_frcs)
 
@@ -288,6 +296,9 @@ class OTF:
         '''
         md.update_positions(self.dt, self.noa, self.structure)
 
+    def write_gp(self):
+        self.gp.write_model(self.gp_name)
+
     def run_dft(self):
         """Calculates DFT forces on atoms in the current structure.
 
@@ -352,9 +363,9 @@ class OTF:
         if (self.dft_count-1) < self.freeze_hyps:
             self.train_gp()
             if self.write_model == 2:
-                self.gp.write_model(self.output_name+"_model")
+                self.write_gp()
         if self.write_model == 3:
-            self.gp.write_model(self.output_name+'_model')
+            self.write_gp()
 
     def train_gp(self):
         """Optimizes the hyperparameters of the current GP model."""
@@ -411,3 +422,53 @@ class OTF:
         self.output.write_md_config(
             self.dt, self.curr_step, self.structure, self.temperature,
             self.KE, self.start_time, self.dft_step, self.velocities)
+
+    def as_dict(self):
+        self.dft_module = self.dft_module.__name__
+        out_dict = deepcopy(dict(vars(self)))
+        self.dft_module = eval(self.dft_module)
+
+        out_dict['gp'] = self.gp_name
+        out_dict['structure'] = self.structure.as_dict()
+
+        for key in ['output', 'pred_func']:
+            out_dict.pop(key)
+
+        return out_dict
+
+    @staticmethod
+    def from_dict(in_dict):
+        if in_dict['write_model'] <= 1: # TODO: detect GP version
+            warnings.warn("The GP model might not be the latest")
+       
+        gp_model = gp.GaussianProcess.from_file(in_dict['gp'])
+        in_dict['gp'] = gp_model
+        in_dict['structure'] = struc.Structure.from_dict(in_dict['structure'])
+
+        if 'flare.dft_interface' in in_dict['dft_module']:
+            for dft_name in ['qe', 'cp2k', 'vasp']:
+                if dft_name in in_dict['dft_module']:
+                    in_dict['force_source'] = dft_name
+                    break
+        else: # if force source is a module
+            in_dict['force_source'] = eval(in_dict['dft_module'])
+
+        new_otf = OTF(**in_dict)
+        new_otf.structure = in_dict['structure']
+        new_otf.dft_count = in_dict['dft_count']
+        new_otf.curr_step = in_dict['curr_step']
+        return new_otf
+
+    def checkpoint(self):
+        name = self.checkpt_name
+        if '.json' != name[-5:]:
+            name += '.json'
+        with open(name, 'w') as f:
+            json.dump(self.as_dict(), f, cls=NumpyEncoder)
+
+    @classmethod
+    def from_checkpoint(cls, filename):
+        with open(filename, 'r') as f:
+            otf_model = cls.from_dict(json.loads(f.readline()))
+
+        return otf_model
diff --git a/flare/otf_parser.py b/flare/otf_parser.py
index 9e96cb56..30602d1d 100644
--- a/flare/otf_parser.py
+++ b/flare/otf_parser.py
@@ -1,4 +1,4 @@
-import sys
+import sys, subprocess
 import numpy as np
 
 from copy import deepcopy
@@ -13,76 +13,86 @@ class OtfAnalysis:
 
         self.calculate_energy = calculate_energy
 
-        self.header = parse_header_information(filename)
-
-        position_list, force_list, uncertainty_list, velocity_list,\
-            dft_frames, temperatures, times, msds, dft_times, energies = \
-            self.parse_pos_otf(filename)
-
-        self.position_list = position_list
-        self.force_list = force_list
-        self.uncertainty_list = uncertainty_list
-        self.velocity_list = velocity_list
-        self.dft_frames = dft_frames
-        self.temperatures = temperatures
-        self.times = times
-        self.msds = msds
-        self.dft_times = dft_times
+        blocks = split_blocks(filename)
+
+        self.header = parse_header_information(blocks[0])
+        self.noa = self.header["atoms"]
+        self.noh = self.header["n_hyps"]
+
+        self.position_list = []
+        self.cell_list = []
+        self.force_list = []
+        self.stress_list = []
+        self.uncertainty_list = []
+        self.velocity_list = []
+        self.temperatures = []
+        self.dft_frames = []
+        self.dft_times = []
+        self.times = []
+        self.msds = []
+        self.energies = []
+        self.thermostat = {}
+
+        self.gp_position_list = []
+        self.gp_cell_list = []
+        self.gp_force_list = []
+        self.gp_stress_list = []
+        self.gp_uncertainty_list = []
+        self.gp_velocity_list = []
+        self.gp_atom_list = []
+        self.gp_species_list = []
+        self.gp_atom_count = []
+        self.gp_thermostat = {}
+
+        self.gp_hyp_list = [self.header["hyps"]]
+
+        self.mae_list = []
+        self.maf_list = []
+
+        self.parse_pos_otf(blocks[1:])
 
         if self.calculate_energy:
             self.energies = energies
 
-        gp_position_list, gp_force_list, gp_uncertainty_list,\
-            gp_velocity_list, gp_atom_list, gp_hyp_list, \
-            gp_species_list, gp_atom_count = self.extract_gp_info(filename)
+    def make_gp(
+        self, cell=None, call_no=None, hyps=None, init_gp=None, hyp_no=None, **kwargs,
+    ):
 
-        self.gp_position_list = gp_position_list
-        self.gp_force_list = gp_force_list
-        self.gp_uncertainty_list = gp_uncertainty_list
-        self.gp_velocity_list = gp_velocity_list
-        self.gp_atom_list = gp_atom_list
-        self.gp_hyp_list = gp_hyp_list
-        self.gp_species_list = gp_species_list
-        self.gp_atom_count = gp_atom_count
-
-    def make_gp(self, cell=None, call_no=None, hyps=None, init_gp=None,
-                hyp_no=None, **kwargs,):
+        if "restart" in self.header and self.header["restart"] > 0:
+            assert init_gp is not None, "Please input the init_gp as the gp model dumpped"\
+                    "before restarting otf."
 
         if call_no is None:
             call_no = len(self.gp_position_list)
         if hyp_no is None:
-            hyp_no = call_no
+            hyp_no = len(self.gp_hyp_list)  # use the last hyps by default
         if hyps is None:
             # check out the last non-empty element from the list
-            for icall in reversed(range(hyp_no)):
-                if len(self.gp_hyp_list[icall]) > 0:
-                    hyps = self.gp_hyp_list[icall][-1]
-                    break
+            hyps = self.gp_hyp_list[hyp_no - 1]
         if cell is None:
-            cell = self.header['cell']
-
+            cell = self.header["cell"]
 
         if init_gp is None:
             # Use run's values as extracted from header
             # TODO Allow for kernel gradient in header
 
             dictionary = deepcopy(self.header)
-            dictionary['hyps'] = hyps
+            dictionary["hyps"] = hyps
             for k in kwargs:
                 if kwargs[k] is not None:
                     dictionary[k] = kwargs[k]
 
-            gp_model = \
-                GaussianProcess.from_dict(dictionary)
+            gp_model = GaussianProcess.from_dict(dictionary)
         else:
             gp_model = init_gp
             gp_model.hyps = hyps
 
-        for (positions, forces, atoms, _, species) in \
-            zip(self.gp_position_list[:call_no],
-                self.gp_force_list[:call_no],
-                self.gp_atom_list[:call_no], self.gp_hyp_list[:call_no],
-                self.gp_species_list[:call_no]):
+        for (positions, forces, atoms, species) in zip(
+            self.gp_position_list[:call_no],
+            self.gp_force_list[:call_no],
+            self.gp_atom_list[:call_no],
+            self.gp_species_list[:call_no],
+        ):
 
             struc_curr = struc.Structure(cell, species, positions)
 
@@ -96,149 +106,83 @@ class OtfAnalysis:
     def get_gp_activation(gp_model):
         pass
 
-    def parse_pos_otf(self, filename):
+    def parse_pos_otf(self, blocks):
         """
         Exclusively parses MD run information
         :param filename:
         :return:
         """
-        position_list = []
-        force_list = []
-        uncertainty_list = []
-        velocity_list = []
-        temperatures = []
-        dft_frames = []
-        dft_times = []
-        times = []
-        msds = []
-        energies = []
-
-        with open(filename, 'r') as f:
-            lines = f.readlines()
-
-        n_steps = 0
-
-        for index, line in enumerate(lines):
-            if line.startswith("number of atoms"):
-                at_line = line.split()
-                noa = int(at_line[3])
-
-            # number of hyperparameters
-            if line.startswith("number of hyperparameters"):
-                line_curr = line.split(':')
-                noh = int(line_curr[-1])
-
-            # DFT frame
-            if line.startswith("*-Frame"):
-                dft_frame_line = line.split()
-                dft_frames.append(int(dft_frame_line[1]))
-                dft_time_line = lines[index+1].split()
-                dft_times.append(float(dft_time_line[-2]))
-
-            # MD frame
-            if line.startswith("-Frame"):
-                n_steps += 1
-                time_line = lines[index+1].split()
-                sim_time = float(time_line[2])
-                times.append(sim_time)
-
-                _, positions, forces, uncertainties, velocities = \
-                    parse_snapshot(lines, index, noa, False, noh)
-
-                position_list.append(positions)
-                force_list.append(forces)
-                uncertainty_list.append(uncertainties)
-                velocity_list.append(velocities)
-
-                temp_line = lines[index+4+noa].split()
-                temperatures.append(float(temp_line[-2]))
-
-                if self.calculate_energy:
-                    en_line = lines[index+5+noa].split()
-                    energies.append(float(en_line[2]))
-
-                msds.append(np.mean((positions - position_list[0])**2))
-
-        return position_list, force_list, uncertainty_list, velocity_list,\
-            dft_frames, temperatures, times, msds, dft_times, energies
-
-    def extract_gp_info(self, filename):
-        """
-        Exclusively parses DFT run information
-        :param filename:
-        :return:
-        """
-        species_list = []
-        position_list = []
-        force_list = []
-        uncertainty_list = []
-        velocity_list = []
-        atom_list = []
-        atom_count = []
-        hyp_list = []
-
-        with open(filename, 'r') as f:
-            lines = f.readlines()
-
-        for index, line in enumerate(lines):
-
-            # number of atoms
-            if line.startswith("number of atoms"):
-                line_curr = line.split()
-                noa = int(line_curr[3])
-
-            # number of hyperparameters
-            if line.startswith("number of hyperparameters"):
-                line_curr = line.split()
-                noh = int(line_curr[3])
-
-            if line.startswith("*-"):
-                # keep track of atoms added to training set
-                ind_count = index+1
-                line_check = lines[ind_count]
-                atoms_added = []
-                hyps_added = []
-                while not line_check.startswith("-"):
-
-                    # keep track of atom number
-                    if line_check.startswith("Adding atom"):
-                        line_split = line_check.split()
-                        atom_strings = line_split[2:-4]
-                        for n, atom_string in enumerate(atom_strings):
-                            if n == 0:
-                                atoms_added.append(int(atom_string[1:-1]))
-                            else:
-                                atoms_added.append(int(atom_string[0:-1]))
-
-                    # keep track of hyperparameters
-                    if line_check.startswith("GP hyperparameters:"):
-                        hyps = []
-                        for hyp_line in lines[(ind_count+1):(ind_count+1+noh)]:
-                            hyp_line = hyp_line.split()
-                            hyps.append(float(hyp_line[-1]))
-                        hyps = np.array(hyps)
-                        hyps_added.append(hyps)
-
-                    ind_count += 1
-                    line_check = lines[ind_count]
-                    # catch case where final frame is a DFT call
-                    if line_check.startswith('Run complete'):
-                        break
-
-                hyp_list.append(hyps_added)
-                atom_list.append(atoms_added)
-                atom_count.append(len(atoms_added))
-
-                # add DFT positions and forces
-                line_curr = line.split()
-
-                # TODO: generalize this to account for arbitrary starting list
-                append_atom_lists(species_list, position_list, force_list,
-                                  uncertainty_list, velocity_list,
-                                  lines, index, noa, True, noh)
-
-        return position_list, force_list, uncertainty_list, velocity_list,\
-            atom_list, hyp_list, species_list, atom_count
+        n_steps = len(blocks) - 1
+
+        for block in blocks:
+            for index, line in enumerate(block):
+                # DFT frame
+                if line.startswith("*-Frame"):
+                    dft_frame_line = line.split()
+                    self.dft_frames.append(int(dft_frame_line[1]))
+                    dft_time_line = block[index + 1].split()
+                    self.dft_times.append(float(dft_time_line[-2]))
+
+                    # TODO: generalize this to account for arbitrary starting list
+                    append_atom_lists(
+                        self.gp_species_list,
+                        self.gp_position_list,
+                        self.gp_force_list,
+                        self.gp_uncertainty_list,
+                        self.gp_velocity_list,
+                        block,
+                        index,
+                        self.noa,
+                        True,
+                        self.noh,
+                    )
+
+                    post_frame = block[index + 3 + self.noa :]
+                    extract_global_info(
+                        self.gp_cell_list,
+                        self.gp_stress_list,
+                        self.gp_thermostat,
+                        post_frame,
+                    )
+
+                    extract_gp_info(
+                        post_frame,
+                        self.mae_list,
+                        self.maf_list,
+                        self.gp_atom_list,
+                        self.gp_hyp_list,
+                        self.noh,
+                    )
+
+                # MD frame
+                if line.startswith("-Frame"):
+                    n_steps += 1
+                    time_line = block[index + 1].split()
+                    sim_time = float(time_line[2])
+                    self.times.append(sim_time)
+
+                    # TODO: generalize this to account for arbitrary starting list
+                    append_atom_lists(
+                        [],
+                        self.position_list,
+                        self.force_list,
+                        self.uncertainty_list,
+                        self.velocity_list,
+                        block,
+                        index,
+                        self.noa,
+                        False,
+                        self.noh,
+                    )
+
+                    post_frame = block[index + 3 + self.noa :]
+                    extract_global_info(
+                        self.cell_list, self.stress_list, self.thermostat, post_frame,
+                    )
+
+                    self.msds.append(
+                        np.mean((self.position_list[-1] - self.position_list[0]) ** 2)
+                    )
 
     def output_md_structures(self):
         """
@@ -248,31 +192,163 @@ class OtfAnalysis:
 
         positions = self.position_list
         structures = []
-        cell = self.header['cell']
-        species = self.header['species']
+        cell = self.header["cell"]
+        species = self.header["species"]
         forces = self.force_list
         stds = self.uncertainty_list
         for i in range(len(positions)):
-            cur_struc = struc.Structure(cell=cell, species=species,
-                                        positions=positions[i])
+            cur_struc = struc.Structure(
+                cell=cell, species=species, positions=positions[i]
+            )
             cur_struc.forces = forces[i]
             cur_struc.stds = stds[i]
             structures.append(cur_struc)
         return structures
 
 
-def append_atom_lists(species_list: List[str],
-                      position_list: List[np.ndarray],
-                      force_list: List[np.ndarray],
-                      uncertainty_list: List[np.ndarray],
-                      velocity_list: List[np.ndarray],
-                      lines: List[str], index: int, noa: int,
-                      dft_call: bool, noh: int) -> None:
+def split_blocks(filename):
+    with open(filename, "r") as f:
+        lines = f.readlines()
+        head = 0
+        blocks = []
+        for index, line in enumerate(lines):
+            if line.startswith("---"):
+                blocks.append(lines[head:index])
+                head = index + 1
+    return blocks
+
+
+def parse_header_information(lines) -> dict:
+    """
+    Get information about the run from the header of the file
+    :param outfile:
+    :return:
+    """
+    header_info = {}
+
+    cutoffs_dict = {}
+    for i, line in enumerate(lines):
+        line_lower = line.lower()
+
+        # gp related
+        if "cutoffs" in line_lower:
+            if "{" in line_lower:
+                line = line[len("Cutoffs: ") :]
+                cutoffs = eval(line)
+            else:
+                line = line.split(":")[1].strip()
+                line = line.strip("[").strip("]")
+                line = line.split()
+                cutoffs = []
+                for val in line:
+                    try:
+                        cutoffs.append(float(val))
+                    except:
+                        cutoffs.append(float(val[:-1]))
+            header_info["cutoffs"] = cutoffs
+
+        if "number of hyperparameters" in line_lower:
+            n_hyps = int(line.split(":")[1].strip())
+            header_info["n_hyps"] = n_hyps 
+            
+            # parse hyps
+            new_line = lines[i+1].replace("[", "")
+            new_line = new_line.replace("]", "")
+            assert "hyperparameter" in new_line.lower()
+
+            hyps_array = new_line[new_line.find(":")+1:].split()
+            if len(hyps_array) < n_hyps:
+                next_new_line = lines[i+2].replace("[", "")
+                next_new_line = next_new_line.replace("]", "")
+                hyps_array += next_new_line.split()
+            assert len(hyps_array) == n_hyps
+
+            hyps_array = [float(h.strip()) for h in hyps_array]
+            header_info["hyps"] = np.array(hyps_array)
+
+        if "kernel_name" in line_lower:
+            header_info["kernel_name"] = line.split(":")[1].strip()
+        elif "kernels" in line_lower:
+            line = line.split(":")[1].strip()
+            line = line.strip("[").strip("]")
+            line = line.split()
+            header_info["kernels"] = line
+        elif "kernel" in line_lower:
+            header_info["kernel_name"] = line.split(":")[1].strip()
+
+        for kw in header_dict:
+            get_header_item(line, header_info, kw)
+
+        if "optimization algorithm" in line:
+            header_info["algo"] = str(line.split(":")[1].strip()).upper()
+
+        if "system species" in line_lower:
+            line = line.split(":")[1]
+            line = line.split("'")
+            species = [item for item in line if item.isalpha()]
+            header_info["species_set"] = set(species)
+        if "periodic cell" in line_lower:
+            vectors = []
+            for cell_line in lines[i + 1 : i + 4]:
+                cell_line = cell_line.strip().replace("[", "").replace("]", "")
+                vec = cell_line.split()
+                vector = [float(vec[0]), float(vec[1]), float(vec[2])]
+                vectors.append(vector)
+            header_info["cell"] = np.array(vectors)
+        if "previous positions" in line_lower:
+            struc_spec = []
+            prev_positions = []
+            for pos_line in lines[i + 1 : i + 1 + header_info.get("atoms", 0)]:
+                pos = pos_line.split()
+                struc_spec.append(pos[0])
+                prev_positions.append((float(pos[1]), float(pos[2]), float(pos[3])))
+            header_info["species"] = struc_spec
+            header_info["prev_positions"] = np.array(prev_positions)
+
+    return header_info
+
+
+def get_header_item(line, header_info, kw):
+    if not isinstance(line, str):
+        return
+
+    pattern = header_dict[kw][0]
+    value_type = header_dict[kw][1]
+
+    if header_dict[kw][2]:
+        pattern = pattern.lower()
+        line = line.lower()
+
+    if pattern in line:
+        header_info[kw] = value_type(line.split(":")[1].strip())
+
+
+header_dict = {
+    "restart": ["Restart", int, False],
+    "frames": ["Frames", int, True],
+    "atoms": ["Number of atoms", int, True],
+    "dt": ["Timestep", float, True],
+}
+
+
+def append_atom_lists(
+    species_list: List[str],
+    position_list: List[np.ndarray],
+    force_list: List[np.ndarray],
+    uncertainty_list: List[np.ndarray],
+    velocity_list: List[np.ndarray],
+    lines: List[str],
+    index: int,
+    noa: int,
+    dft_call: bool,
+    noh: int,
+) -> None:
 
     """Update lists containing atom information at each snapshot."""
 
-    species, positions, forces, uncertainties, velocities = \
-        parse_snapshot(lines, index, noa, dft_call, noh)
+    species, positions, forces, uncertainties, velocities = parse_snapshot(
+        lines, index, noa, dft_call, noh
+    )
 
     species_list.append(species)
     position_list.append(positions)
@@ -294,10 +370,9 @@ def parse_snapshot(lines, index, noa, dft_call, noh):
     # Current setting for # of lines to skip after Frame marker
     skip = 3
 
-    for count, frame_line in enumerate(lines[(index+skip):(index+skip+noa)]):
+    for count, frame_line in enumerate(lines[(index + skip) : (index + skip + noa)]):
         # parse frame line
-        spec, position, force, uncertainty, velocity = \
-            parse_frame_line(frame_line)
+        spec, position, force, uncertainty, velocity = parse_frame_line(frame_line)
 
         # update values
         species.append(spec)
@@ -322,102 +397,6 @@ def strip_and_split(line):
     return stripped_line
 
 
-def parse_header_information(outfile: str = 'otf_run.out') -> dict:
-    """
-    Get information about the run from the header of the file
-    :param outfile:
-    :return:
-    """
-    with open(outfile, 'r') as f:
-        lines = f.readlines()
-
-    header_info = {}
-
-    stopreading = None
-
-    for line in lines:
-        if '---' in line or '=' in line:
-            stopreading = lines.index(line)
-            break
-
-    if stopreading is None:
-        raise Exception("OTF File is malformed")
-
-    cutoffs_dict = {}
-    for i, line in enumerate(lines[:stopreading]):
-        line_lower = line.lower()
-
-        # gp related
-        if 'cutoffs' in line_lower:
-            line = line.split(':')[1].strip()
-            line = line.strip('[').strip(']')
-            line = line.split()
-            cutoffs = []
-            for val in line:
-                try:
-                    cutoffs.append(float(val))
-                except:
-                    cutoffs.append(float(val[:-1]))
-            header_info['cutoffs'] = cutoffs
-        elif 'cutoff' in line_lower:
-            line = line.split(':')
-            name = line[0][7:]
-            value = float(line[1])
-            cutoffs_dict[name] = value
-            header_info['cutoffs'] = cutoffs_dict
-
-        if 'kernel_name' in line_lower:
-            header_info['kernel_name'] = line.split(':')[1].strip()
-        elif 'kernels' in line_lower:
-            line = line.split(':')[1].strip()
-            line = line.strip('[').strip(']')
-            line = line.split()
-            header_info['kernels'] = line
-        elif 'kernel' in line_lower:
-            header_info['kernel_name'] = line.split(':')[1].strip()
-
-        if 'number of hyperparameters:' in line_lower:
-            header_info['n_hyps'] = int(line.split(':')[1])
-        if 'optimization algorithm' in line_lower:
-            header_info['algo'] = line.split(':')[1].strip()
-
-        # otf related
-        if 'frames' in line_lower:
-            header_info['frames'] = int(line.split(':')[1])
-        if 'number of atoms' in line_lower:
-            header_info['atoms'] = int(line.split(':')[1])
-        if 'timestep' in line_lower:
-            header_info['dt'] = float(line.split(':')[1])
-        if 'system species' in line_lower:
-            line = line.split(':')[1]
-            line = line.split("'")
-
-            species = [item for item in line if item.isalpha()]
-
-            header_info['species_set'] = set(species)
-        if 'periodic cell' in line_lower:
-            vectors = []
-            for cell_line in lines[i+1:i+4]:
-                cell_line = \
-                    cell_line.strip().replace('[', '').replace(']', '')
-                vec = cell_line.split()
-                vector = [float(vec[0]), float(vec[1]), float(vec[2])]
-                vectors.append(vector)
-            header_info['cell'] = np.array(vectors)
-        if 'previous positions' in line_lower:
-            struc_spec = []
-            prev_positions = []
-            for pos_line in lines[i+1:i+1+header_info.get('atoms', 0)]:
-                pos = pos_line.split()
-                struc_spec.append(pos[0])
-                prev_positions.append((float(pos[1]), float(pos[2]),
-                                       float(pos[3])))
-            header_info['species'] = struc_spec
-            header_info['prev_positions'] = np.array(prev_positions)
-
-    return header_info
-
-
 def parse_frame_line(frame_line):
     """parse a line in otf output.
     :param frame_line: frame line to be parsed
@@ -435,3 +414,78 @@ def parse_frame_line(frame_line):
     velocity = np.array([float(n) for n in frame_line[10:13]])
 
     return spec, position, force, uncertainty, velocity
+
+
+def extract_global_info(
+    cell_list, stress_list, thermostat, block,
+):
+
+    for ind, line in enumerate(block):
+        if "cell" in line:
+            vectors = []
+            for cell_line in block[ind + 1 : ind + 4]:
+                cell_line = cell_line.strip().replace("[", "").replace("]", "")
+                vec = cell_line.split()
+                vector = [float(vec[0]), float(vec[1]), float(vec[2])]
+                vectors.append(vector)
+            cell_list.append(vectors)
+        if "Stress" in line:
+            vectors = []
+            stress_line = block[ind + 2].split()
+            vectors = [float(s) for s in stress_line]
+            stress_list.append(vectors)
+
+        for t in [
+            "Pressure",
+            "Temperature",
+            "Kinetic energy",
+            "Potential energy",
+            "Total energy",
+        ]:
+            get_thermostat(thermostat, t, line)
+
+
+def get_thermostat(thermostat, kw, line):
+    if kw in line:
+        value = float(line.split()[-1])
+        kw = kw.lower()
+        if kw in thermostat:
+            thermostat[kw].append(value)
+        else:
+            thermostat[kw] = [value]
+
+
+def extract_gp_info(block, mae_list, maf_list, atoms_list, hyps_list, noh):
+    """
+    Exclusively parses DFT run information
+    :param filename:
+    :return:
+    """
+    for ind, line in enumerate(block):
+        if "mean absolute error" in line:
+            value = float(line.split()[3])
+            mae_list.append(value)
+        if "mean absolute dft component" in line:
+            value = float(line.split()[4])
+            maf_list.append(value)
+
+        # keep track of atom number
+        if line.startswith("Adding atom"):
+            atoms_added = []
+            line_split = line.split()
+            atom_strings = line_split[2:-4]
+            for n, atom_string in enumerate(atom_strings):
+                if n == 0:
+                    atoms_added.append(int(atom_string[1:-1]))
+                else:
+                    atoms_added.append(int(atom_string[0:-1]))
+            atoms_list.append(atoms_added)
+
+        # keep track of hyperparameters
+        if line.startswith("GP hyperparameters:"):
+            hyps = []
+            for hyp_line in block[(ind + 1) : (ind + 1 + noh)]:
+                hyp_line = hyp_line.split()
+                hyps.append(float(hyp_line[-1]))
+            hyps = np.array(hyps)
+            hyps_list.append(hyps)
diff --git a/flare/output.py b/flare/output.py
index 6614ba3a..67e165be 100644
--- a/flare/output.py
+++ b/flare/output.py
@@ -149,7 +149,7 @@ class Output:
             headerstring += \
                 f'System species: {set(structure.species_labels)}\n'
             headerstring += 'Periodic cell (A): \n'
-            headerstring += str(structure.cell)+'\n'
+            headerstring += str(np.array(structure.cell))+'\n'
 
         if optional:
             for key, value in optional.items():
@@ -170,6 +170,20 @@ class Output:
         if self.always_flush:
             f.handlers[0].flush()
 
+
+    def write_md_header(self, dt, curr_step, dft_step):
+        string = ''
+        # Mark if a frame had DFT forces with an asterisk
+        if not dft_step:
+            string += '-' * 80 + '\n'
+            string += f"-Frame: {curr_step} "
+        else:
+            string += f"\n*-Frame: {curr_step} "
+
+        string += f'\nSimulation Time: {(dt * curr_step):.3} ps \n'
+        return string
+
+
     def write_md_config(
             self, dt, curr_step, structure, temperature, KE, start_time, dft_step,
             velocities):
@@ -188,16 +202,7 @@ class Output:
         :return:
         """
 
-        string = ''
-
-        # Mark if a frame had DFT forces with an asterisk
-        if not dft_step:
-            string += '-' * 80 + '\n'
-            string += f"-Frame: {curr_step} "
-        else:
-            string += f"\n*-Frame: {curr_step} "
-
-        string += f'\nSimulation Time: {(dt * curr_step):.3} ps \n'
+        string = self.write_md_header(dt, curr_step, dft_step)
 
         # Construct Header line
         n_space = 30
@@ -235,7 +240,7 @@ class Output:
         # Report cell if stress attribute is present.
         if structure.stress is not None:
             string += 'Periodic cell (A): \n'
-            string += str(structure.cell)+'\n\n'
+            string += str(np.array(structure.cell))+'\n\n'
 
         # Report stress tensor.
         pressure = None
@@ -263,23 +268,23 @@ class Output:
         if pressure is not None:
             string += f'Pressure (GPa): {pressure:.6f} \n'
 
-        string += f'Temperature: {temperature:.2f} K \n'
-        string += f'Kinetic energy: {KE:.6f} eV \n'
+        string += f'Temperature (K): {temperature:.2f} \n'
+        string += f'Kinetic energy (eV): {KE:.6f} \n'
 
         # Report potential energy.
         if structure.potential_energy is not None:
             string += \
-                f'Potential energy: {structure.potential_energy:.6f} eV \n'
+                f'Potential energy (eV): {structure.potential_energy:.6f} \n'
 
         # Report potential energy uncertainty.
         if structure.local_energy_stds is not None:
             pot_en_std = np.sqrt(np.sum(structure.local_energy_stds**2))
-            string += f'Uncertainty: {pot_en_std:.6f} eV \n'
+            string += f'Uncertainty (eV): {pot_en_std:.6f} \n'
 
         # Report total energy.
         if structure.potential_energy is not None:
             tot_en = KE + structure.potential_energy
-            string += f'Total energy: {tot_en:.6f} eV \n'
+            string += f'Total energy (eV): {tot_en:.6f} \n'
 
         logger = logging.getLogger(self.basename+'log')
         logger.info(string)
diff --git a/flare/predict.py b/flare/predict.py
index db996907..02a50d23 100644
--- a/flare/predict.py
+++ b/flare/predict.py
@@ -179,7 +179,8 @@ def predict_on_structure_par(
     :rtype: (np.ndarray, np.ndarray)
     """
     # Just work in serial in the number of cpus is 1
-    if n_cpus == 1:
+    # or the gp is not parallelized by atoms
+    if (n_cpus is 1) or (not gp.per_atom_par):
         return predict_on_structure(
             structure=structure, gp=gp, n_cpus=n_cpus,
             write_to_structure=write_to_structure,
@@ -274,8 +275,9 @@ def predict_on_structure_efs_par(
  write_to_structure: bool = True, selective_atoms: List[int] = None,
  skipped_atom_value=0):
 
-    # Just work in serial in the number of cpus is 1
-    if n_cpus is 1:
+    # Just work in serial in the number of cpus is 1, 
+    # or the gp is not parallelized by atoms
+    if (n_cpus is 1) or (not gp.per_atom_par):
         return predict_on_structure_efs(
             structure=structure, gp=gp, n_cpus=n_cpus,
             write_to_structure=write_to_structure,
@@ -384,8 +386,6 @@ def predict_on_structure_en(structure: Structure, gp: GaussianProcess,
     forces = np.zeros((structure.nat, 3))
     stds = np.zeros((structure.nat, 3))
     local_energies = np.zeros(structure.nat)
-    forces = np.zeros(shape=(structure.nat, 3))
-    stds = np.zeros(shape=(structure.nat, 3))
 
     if selective_atoms:
         forces.fill(skipped_atom_value)
@@ -437,7 +437,8 @@ def predict_on_structure_par_en(structure: Structure, gp: GaussianProcess,
     :rtype: (np.ndarray, np.ndarray, np.ndarray)
     """
     # Work in serial if the number of cpus is 1
-    if n_cpus == 1:
+    # or the gp is not parallelized by atoms
+    if (n_cpus is 1) or (not gp.per_atom_par):
         predict_on_structure_en(structure, gp,
                                 n_cpus, write_to_structure,
                                 selective_atoms,
@@ -446,8 +447,6 @@ def predict_on_structure_par_en(structure: Structure, gp: GaussianProcess,
     forces = np.zeros((structure.nat, 3))
     stds = np.zeros((structure.nat, 3))
     local_energies = np.zeros(structure.nat)
-    forces = np.zeros(shape=(structure.nat, 3))
-    stds = np.zeros(shape=(structure.nat, 3))
 
     if selective_atoms:
         forces.fill(skipped_atom_value)
diff --git a/flare/struc.py b/flare/struc.py
index 73431af5..0f9c202e 100644
--- a/flare/struc.py
+++ b/flare/struc.py
@@ -99,7 +99,7 @@ class Structure:
         self.partial_stress_stds = None
         self.stress = None
         self.stress_stds = None
-        self.potential_energy = None
+        self.potential_energy = None # duplicated with self.energy?
 
         self.mass_dict = mass_dict
 
@@ -331,14 +331,53 @@ class Structure:
 
         if cell is None:
             cell = np.array(atoms.cell)
-        struc = Structure(cell=cell, positions=atoms.positions,
-                          species=atoms.get_chemical_symbols())
+
+        try:
+            forces = atoms.get_forces()
+        except:
+            forces = None 
+
+        try:
+            stds = atoms.get_uncertainties()
+        except:
+            stds = None
+
+        try:
+            energy = atoms.get_potential_energy()
+        except:
+            energy = None
+
+        try:
+            stress = atoms.get_stress()
+        except:
+            stress = None 
+
+        struc = Structure(
+            cell = cell, 
+            positions = atoms.positions,
+            species = atoms.get_chemical_symbols(),
+            forces = forces,
+            stds = stds,
+            energy = energy,
+        )
+        struc.stress = stress
         return struc
 
     def to_ase_atoms(self) -> 'ase.Atoms':
         from ase import Atoms
-        return Atoms(self.species_labels, positions=self.positions,
-                     cell=self.cell, pbc=True)
+        from ase.calculators.singlepoint import SinglePointCalculator
+
+        atoms =  Atoms(self.species_labels, positions=self.positions,
+                       cell=self.cell, pbc=True)
+
+        results = {}
+        properties = ["forces", "energy", "stress"]
+        for p in properties:
+            results[p] = getattr(self, p)
+        calculator = SinglePointCalculator(atoms, **results)
+        atoms.set_calculator(calculator)
+
+        return atoms
 
     def to_pmg_structure(self):
         """


Test patch:
diff --git a/tests/test_ase_otf.py b/tests/test_ase_otf.py
index bbb8108c..e31f1425 100644
--- a/tests/test_ase_otf.py
+++ b/tests/test_ase_otf.py
@@ -4,83 +4,100 @@ import pytest
 import numpy as np
 
 from flare import otf, kernels
+from flare.otf_parser import OtfAnalysis
 from flare.gp import GaussianProcess
 from flare.mgp import MappedGaussianProcess
 from flare.ase.calculator import FLARE_Calculator
 from flare.ase.otf import ASE_OTF
 from flare.utils.parameter_helper import ParameterHelper
-# from flare.ase.logger import OTFLogger
 
+from ase.constraints import FixAtoms
 from ase import units
-from ase.md.velocitydistribution import (MaxwellBoltzmannDistribution,
-                                         Stationary, ZeroRotation)
+from ase.md.velocitydistribution import (
+    MaxwellBoltzmannDistribution,
+    Stationary,
+    ZeroRotation,
+)
 from ase.spacegroup import crystal
 from ase.calculators.espresso import Espresso
 from ase import io
 
+
 def read_qe_results(self):
 
-    out_file = self.label + '.pwo'
+    out_file = self.label + ".pwo"
 
     # find out slurm job id
-    qe_slurm_dat = open('qe_slurm.dat').readlines()[0].split()
+    qe_slurm_dat = open("qe_slurm.dat").readlines()[0].split()
     qe_slurm_id = qe_slurm_dat[3]
 
     # mv scf.pwo to scp+nsteps.pwo
-    if ('forces' not in self.results.keys()) and (out_file in os.listdir()):
-        subprocess.call(['mv', out_file, f'{self.label}{self.nsteps}.pwo'])
-    
+    if ("forces" not in self.results.keys()) and (out_file in os.listdir()):
+        subprocess.call(["mv", out_file, f"{self.label}{self.nsteps}.pwo"])
+
     # sleep until the job is finished
-    job_list = subprocess.check_output(['showq', '-p', 'kozinsky']).decode('utf-8')
+    job_list = subprocess.check_output(["showq", "-p", "kozinsky"]).decode("utf-8")
     while qe_slurm_id in job_list:
         time.sleep(10)
-        job_list = subprocess.check_output(['showq', '-p', 'kozinsky']).decode('utf-8')
+        job_list = subprocess.check_output(["showq", "-p", "kozinsky"]).decode("utf-8")
 
     output = io.read(out_file)
     self.calc = output.calc
     self.results = output.calc.results
 
-md_list = ['VelocityVerlet', 'NVTBerendsen', 'NPTBerendsen', 'NPT', 'Langevin']
 
-@pytest.fixture(scope='module')
+md_list = ["VelocityVerlet", "NVTBerendsen", "NPTBerendsen", "NPT", "Langevin"]
+number_of_steps = 3
+
+@pytest.fixture(scope="module")
 def md_params():
 
-    md_dict = {'temperature': 500}
+    md_dict = {"temperature": 500}
+    print(md_list)
+
     for md_engine in md_list:
-        if md_engine == 'VelocityVerlet':
+        for f in glob.glob(md_engine + "*"):
+            os.remove(f)
+
+        if md_engine == "VelocityVerlet":
             md_dict[md_engine] = {}
         else:
-            md_dict[md_engine] = {'temperature': md_dict['temperature']}
+            md_dict[md_engine] = {"temperature": md_dict["temperature"]}
 
-    md_dict['NVTBerendsen'].update({'taut': 0.5e3 * units.fs})
-    md_dict['NPT'].update({'externalstress': 0, 'ttime': 25, 'pfactor': 3375})
-    md_dict['Langevin'].update({'friction': 0.02})
+    md_dict["NVTBerendsen"].update({"taut": 0.5e3 * units.fs})
+    md_dict["NPT"].update({"externalstress": 0, "ttime": 25, "pfactor": None})
+    md_dict["Langevin"].update({"friction": 0.02})
 
     yield md_dict
     del md_dict
 
 
-@pytest.fixture(scope='module')
+@pytest.fixture(scope="module")
 def super_cell():
 
     from ase.spacegroup import crystal
-    a = 3.855
+
+    a = 10.0
     alpha = 90
-    atoms = crystal(['H', 'He'],
-                    basis=[(0, 0, 0), (0.5, 0.5, 0.5)],
-                    size=(2, 1, 1),
-                    cellpar=[a, a, a, alpha, alpha, alpha])
+    atoms = crystal(
+        ["H", "He"],
+        basis=[(0, 0, 0), (0.5, 0.5, 0.5)],
+        size=(2, 1, 1),
+        cellpar=[a, a, a, alpha, alpha, alpha],
+    )
 
     # jitter positions to give nonzero force on first frame
     for atom_pos in atoms.positions:
         for coord in range(3):
-            atom_pos[coord] += (2*np.random.random()-1) * 0.5
+            atom_pos[coord] += (2 * np.random.random() - 1) * 0.5
+
+    atoms.set_constraint(FixAtoms(indices=[0]))
 
     yield atoms
     del atoms
 
 
-@pytest.fixture(scope='module')
+@pytest.fixture(scope="module")
 def flare_calc():
     flare_calc_dict = {}
     for md_engine in md_list:
@@ -88,42 +105,39 @@ def flare_calc():
         # ---------- create gaussian process model -------------------
 
         # set up GP hyperparameters
-        kernels = ['twobody', 'threebody'] # use 2+3 body kernel
-        parameters = {'cutoff_twobody': 5.0,
-                      'cutoff_threebody': 3.5}
-        pm = ParameterHelper(
-            kernels = kernels,
-            random = True,
-            parameters=parameters
-        )
+        kernels = ["twobody", "threebody"]  # use 2+3 body kernel
+        parameters = {"cutoff_twobody": 5.0, "cutoff_threebody": 3.5}
+        pm = ParameterHelper(kernels=kernels, random=True, parameters=parameters)
 
         hm = pm.as_dict()
-        hyps = hm['hyps']
-        cut = hm['cutoffs']
-        print('hyps', hyps)
+        hyps = hm["hyps"]
+        cut = hm["cutoffs"]
+        print("hyps", hyps)
 
         gp_model = GaussianProcess(
-            kernels = kernels,
-            component = 'sc', # single-component. For multi-comp, use 'mc'
-            hyps = hyps,
-            cutoffs = cut,
-            hyp_labels = ['sig2','ls2','sig3','ls3','noise'],
-            opt_algorithm = 'L-BFGS-B',
-            n_cpus = 1
+            kernels=kernels,
+            component="sc",  # single-component. For multi-comp, use 'mc'
+            hyps=hyps,
+            cutoffs=cut,
+            hyp_labels=["sig2", "ls2", "sig3", "ls3", "noise"],
+            opt_algorithm="L-BFGS-B",
+            n_cpus=1,
         )
 
         # ----------- create mapped gaussian process ------------------
-        grid_params = {'twobody':   {'grid_num': [64]},
-                       'threebody': {'grid_num': [16, 16, 16]}}
+        grid_params = {
+            "twobody": {"grid_num": [64]},
+            "threebody": {"grid_num": [16, 16, 16]},
+        }
 
-        mgp_model = MappedGaussianProcess(grid_params = grid_params,
-                                          unique_species = [1, 2],
-                                          n_cpus = 1,
-                                          var_map = 'pca')
+        mgp_model = MappedGaussianProcess(
+            grid_params=grid_params, unique_species=[1, 2], n_cpus=1, var_map="pca"
+        )
 
         # ------------ create ASE's flare calculator -----------------------
-        flare_calculator = FLARE_Calculator(gp_model, mgp_model=mgp_model,
-                                            par=True, use_mapping=True)
+        flare_calculator = FLARE_Calculator(
+            gp_model, mgp_model=mgp_model, par=True, use_mapping=True
+        )
 
         flare_calc_dict[md_engine] = flare_calculator
         print(md_engine)
@@ -131,53 +145,54 @@ def flare_calc():
     del flare_calc_dict
 
 
-@pytest.fixture(scope='module')
+@pytest.fixture(scope="module")
 def qe_calc():
     from ase.calculators.lj import LennardJones
+
     dft_calculator = LennardJones()
 
     yield dft_calculator
     del dft_calculator
 
 
-@pytest.mark.parametrize('md_engine', md_list)
+@pytest.mark.parametrize("md_engine", md_list)
 def test_otf_md(md_engine, md_params, super_cell, flare_calc, qe_calc):
     np.random.seed(12345)
 
     flare_calculator = flare_calc[md_engine]
     # set up OTF MD engine
-    otf_params = {'init_atoms': [0, 1, 2, 3],
-                  'output_name': md_engine,
-                  'std_tolerance_factor': 2,
-                  'max_atoms_added' : len(super_cell.positions),
-                  'freeze_hyps': 10}
-#                  'use_mapping': flare_calculator.use_mapping}
+    otf_params = {
+        "init_atoms": [0, 1, 2, 3],
+        "output_name": md_engine,
+        "std_tolerance_factor": 2,
+        "max_atoms_added": len(super_cell.positions),
+        "freeze_hyps": 10,
+    }
+    #                  'use_mapping': flare_calculator.use_mapping}
 
     md_kwargs = md_params[md_engine]
 
     # intialize velocity
-    temperature = md_params['temperature']
+    temperature = md_params["temperature"]
     MaxwellBoltzmannDistribution(super_cell, temperature * units.kB)
     Stationary(super_cell)  # zero linear momentum
     ZeroRotation(super_cell)  # zero angular momentum
 
     super_cell.set_calculator(flare_calculator)
-    test_otf = ASE_OTF(super_cell,
-                       timestep = 1 * units.fs,
-                       number_of_steps = 3,
-                       dft_calc = qe_calc,
-                       md_engine = md_engine,
-                       md_kwargs = md_kwargs,
-                       **otf_params)
+    test_otf = ASE_OTF(
+        super_cell,
+        timestep=1 * units.fs,
+        number_of_steps=number_of_steps,
+        dft_calc=qe_calc,
+        md_engine=md_engine,
+        md_kwargs=md_kwargs,
+        trajectory=md_engine+"_otf.traj",
+        **otf_params,
+    )
 
     # TODO: test if mgp matches gp
     # TODO: see if there's difference between MD timestep & OTF timestep
 
-    # set up logger
-#    otf_logger = OTFLogger(test_otf, super_cell,
-#        logfile=md_engine+'.log', mode="w", data_in_logfile=True)
-#    test_otf.attach(otf_logger, interval=1)
-
     test_otf.run()
 
     for f in glob.glob("scf*.pw*"):
@@ -190,9 +205,30 @@ def test_otf_md(md_engine, md_params, super_cell, flare_calc, qe_calc):
         shutil.rmtree(f, ignore_errors=True)
     for f in glob.glob("out"):
         shutil.rmtree(f, ignore_errors=True)
-
     for f in os.listdir("./"):
-        if md_engine in f or 'lmp.mgp' in f:
+        if ".mgp" in f or ".var" in f:
             os.remove(f)
-        if 'slurm' in f:
+        if "slurm" in f:
             os.remove(f)
+
+@pytest.mark.parametrize("md_engine", md_list)
+def test_load_checkpoint(md_engine):
+    new_otf = ASE_OTF.from_checkpoint(md_engine + "_checkpt.json")
+    assert new_otf.curr_step == number_of_steps
+    new_otf.number_of_steps = new_otf.number_of_steps + 2
+    new_otf.run()
+
+@pytest.mark.parametrize("md_engine", md_list)
+def test_otf_parser(md_engine):
+    output_name = f"{md_engine}.out"
+    otf_traj = OtfAnalysis(output_name)
+    try:
+        replicated_gp = otf_traj.make_gp()
+    except:
+        init_flare = FLARE_Calculator.from_file(md_engine + "_flare.json")
+        replicated_gp = otf_traj.make_gp(init_gp=init_flare.gp_model)
+
+    print("ase otf traj parsed")
+
+    for f in glob.glob(md_engine + "*"):
+        os.remove(f)
diff --git a/tests/test_mgp.py b/tests/test_mgp.py
index e91b2675..f3d69e69 100644
--- a/tests/test_mgp.py
+++ b/tests/test_mgp.py
@@ -104,7 +104,7 @@ def test_init(bodies, multihyps, all_mgp, all_gp):
     # grid parameters
     grid_params = {}
     if ('2' in bodies):
-        grid_params['twobody'] = {'grid_num': [128], 'lower_bound': [0.02]}
+        grid_params['twobody'] = {'grid_num': [160], 'lower_bound': [0.02]}
     if ('3' in bodies):
         grid_params['threebody'] = {'grid_num': [31, 32, 33], 'lower_bound':[0.02]*3}
 
diff --git a/tests/test_otf.py b/tests/test_otf.py
index c0f8ff3d..fecfae28 100644
--- a/tests/test_otf.py
+++ b/tests/test_otf.py
@@ -5,29 +5,45 @@ import glob, os, re, shutil
 import numpy as np
 
 from flare.otf import OTF
+from flare.otf_parser import OtfAnalysis
 from flare.gp import GaussianProcess
 from flare.struc import Structure
 
 
-cmd = {'cp2k':'CP2K_COMMAND', 'qe':'PWSCF_COMMAND'}
-software_list = ['cp2k', 'qe']
+cmd = {"cp2k": "CP2K_COMMAND", "qe": "PWSCF_COMMAND"}
+software_list = ["cp2k", "qe"]
 example_list = [1, 2]
-name_list = {1:'h2', 2:'al'}
+name_list = {1: "h2", 2: "al"}
 
-print('running test_otf.py')
-print('current working directory:')
+dt = 0.0001
+number_of_steps = 3
+
+print("running test_otf.py")
+print("current working directory:")
 print(os.getcwd())
 
+
 def get_gp(par=False, per_atom_par=False, n_cpus=1):
     hyps = np.array([1, 1, 1, 1, 1])
-    hyp_labels = ['Signal Std 2b', 'Length Scale 2b',
-                  'Signal Std 3b', 'Length Scale 3b',
-                  'Noise Std']
-    cutoffs = {'twobody':4, 'threebody':4}
-    return GaussianProcess(\
-            kernel_name='23mc', hyps=hyps, cutoffs=cutoffs,
-            hyp_labels=hyp_labels, maxiter=50, par=par,
-            per_atom_par=per_atom_par, n_cpus=n_cpus)
+    hyp_labels = [
+        "Signal Std 2b",
+        "Length Scale 2b",
+        "Signal Std 3b",
+        "Length Scale 3b",
+        "Noise Std",
+    ]
+    cutoffs = {"twobody": 4, "threebody": 4}
+    return GaussianProcess(
+        kernel_name="23mc",
+        hyps=hyps,
+        cutoffs=cutoffs,
+        hyp_labels=hyp_labels,
+        maxiter=50,
+        par=par,
+        per_atom_par=per_atom_par,
+        n_cpus=n_cpus,
+    )
+
 
 def cleanup(software="qe", casename="h2_otf_cp2k"):
 
@@ -38,13 +54,13 @@ def cleanup(software="qe", casename="h2_otf_cp2k"):
     for f in glob.glob(f"*{software}.out"):
         os.remove(f)
 
-    if (software == 'qe'):
+    if software == "qe":
         for f in glob.glob(f"pwscf.wfc*"):
             os.remove(f)
         for f in glob.glob(f"*pwscf.out"):
             os.remove(f)
         for f in os.listdir("./"):
-            if f in ['pwscf.save']:
+            if f in ["pwscf.save"]:
                 shutil.rmtree(f)
     else:
         for f in os.listdir("./"):
@@ -52,65 +68,66 @@ def cleanup(software="qe", casename="h2_otf_cp2k"):
                 os.remove(f)
 
 
-@pytest.mark.parametrize('software', software_list)
-@pytest.mark.parametrize('example', example_list)
+@pytest.mark.parametrize("software", software_list)
+@pytest.mark.parametrize("example", example_list)
 def test_otf(software, example):
     """
     Test that an otf run can survive going for more steps
     :return:
     """
 
-    print('running test_otf.py')
-    print('current working directory:')
+    print("running test_otf.py")
+    print("current working directory:")
     print(os.getcwd())
 
-    outdir = f'test_outputs_{software}'
+    outdir = f"test_outputs_{software}"
     if os.path.isdir(outdir):
         shutil.rmtree(outdir)
 
-    if (not os.environ.get(cmd[software], False)):
-        pytest.skip(f'{cmd[software]} not found in environment:'
-                    f' Please install the code '
-                    f' and set the {cmd[software]} env. '
-                    'variable to point to the executable.')
+    if not os.environ.get(cmd[software], False):
+        pytest.skip(
+            f"{cmd[software]} not found in environment:"
+            f" Please install the code "
+            f" and set the {cmd[software]} env. "
+            "variable to point to the executable."
+        )
 
-    dft_input = f'{software}.in'
-    dft_output = f'{software}.out'
-    shutil.copy(f'./test_files/{software}_input_{example}.in', dft_input)
+    dft_input = f"{software}.in"
+    dft_output = f"{software}.out"
+    shutil.copy(f"./test_files/{software}_input_{example}.in", dft_input)
 
-    dt = 0.0001
-    number_of_steps = 5
     dft_loc = os.environ.get(cmd[software])
     std_tolerance_factor = -0.1
 
-    casename=name_list[example]
+    casename = name_list[example]
 
     gp = get_gp()
 
-    otf = OTF(dt=dt, number_of_steps=number_of_steps,
-              gp=gp, write_model=3,
-              std_tolerance_factor=std_tolerance_factor,
-              init_atoms=[0],
-              calculate_energy=True, max_atoms_added=1,
-              freeze_hyps=1, skip=1,
-              force_source=software,
-              dft_input=dft_input, dft_loc=dft_loc,
-              dft_output=dft_output,
-              output_name=f'{casename}_otf_{software}',
-              store_dft_output=([dft_output, dft_input], '.'))
+    otf = OTF(
+        dt=dt,
+        number_of_steps=number_of_steps,
+        gp=gp,
+        write_model=3,
+        std_tolerance_factor=std_tolerance_factor,
+        init_atoms=[0],
+        calculate_energy=True,
+        max_atoms_added=1,
+        freeze_hyps=1,
+        skip=1,
+        force_source=software,
+        dft_input=dft_input,
+        dft_loc=dft_loc,
+        dft_output=dft_output,
+        output_name=f"{casename}_otf_{software}",
+        store_dft_output=([dft_output, dft_input], "."),
+    )
 
     otf.run()
 
-    if not os.path.isdir(outdir):
-         os.mkdir(outdir)
-    for f in os.listdir("./"):
-        if f'{casename}_otf_{software}' in f:
-            shutil.move(f, outdir)
-    cleanup(software, f'{casename}_otf_{software}')
 
-@pytest.mark.parametrize('software', software_list)
-@pytest.mark.parametrize('per_atom_par', [True, False])
-@pytest.mark.parametrize('n_cpus', [2])
+@pytest.mark.parametrize("software", software_list)
+@pytest.mark.parametrize("per_atom_par", [True, False])
+@pytest.mark.parametrize("n_cpus", [2])
 def test_otf_par(software, per_atom_par, n_cpus):
     """
     Test that an otf run can survive going for more steps
@@ -118,30 +135,29 @@ def test_otf_par(software, per_atom_par, n_cpus):
     """
 
     example = 1
-    outdir = f'test_outputs_{software}'
+    outdir = f"test_outputs_{software}"
     if os.path.isdir(outdir):
         shutil.rmtree(outdir)
 
-    print('running test_otf.py')
-    print('current working directory:')
+    print("running test_otf.py")
+    print("current working directory:")
     print(os.getcwd())
 
-    if (not os.environ.get(cmd[software], False)):
-        pytest.skip(f'{cmd[software]} not found in environment:'
-                    f' Please install the code '
-                    f' and set the {cmd[software]} env. '
-                    'variable to point to the executable.')
-    if (software == 'cp2k'):
-        if (not "popt" in os.environ.get(cmd[software])):
-            pytest.skip(f'cp2k is serial version'
-                        f' skipping the parallel test')
-
-    dft_input = f'{software}.in'
-    dft_output = f'{software}.out'
-    shutil.copy(f'./test_files/{software}_input_{example}.in', dft_input)
-
-    dt = 0.0001
-    number_of_steps = 5
+    if not os.environ.get(cmd[software], False):
+        pytest.skip(
+            f"{cmd[software]} not found in environment:"
+            f" Please install the code "
+            f" and set the {cmd[software]} env. "
+            "variable to point to the executable."
+        )
+    if software == "cp2k":
+        if not "popt" in os.environ.get(cmd[software]):
+            pytest.skip(f"cp2k is serial version" f" skipping the parallel test")
+
+    dft_input = f"{software}.in"
+    dft_output = f"{software}.out"
+    shutil.copy(f"./test_files/{software}_input_{example}.in", dft_input)
+
     dft_loc = os.environ.get(cmd[software])
     std_tolerance_factor = -0.1
 
@@ -149,22 +165,100 @@ def test_otf_par(software, per_atom_par, n_cpus):
 
     gp = get_gp(par=True, n_cpus=n_cpus, per_atom_par=per_atom_par)
 
-    otf = OTF(dt=dt, number_of_steps=number_of_steps, gp=gp,
-              std_tolerance_factor=std_tolerance_factor, init_atoms=[0],
-              calculate_energy=True, max_atoms_added=1,
-              n_cpus=n_cpus,
-              freeze_hyps=1, skip=1,
-              mpi="mpi", force_source=software,
-              dft_input=dft_input, dft_loc=dft_loc,
-              dft_output=dft_output,
-              output_name=f'{casename}_otf_{software}',
-              store_dft_output=([dft_output, dft_input], '.'))
+    otf = OTF(
+        dt=dt,
+        number_of_steps=number_of_steps,
+        gp=gp,
+        write_model=3,
+        std_tolerance_factor=std_tolerance_factor,
+        init_atoms=[0],
+        calculate_energy=True,
+        max_atoms_added=1,
+        n_cpus=n_cpus,
+        freeze_hyps=1,
+        skip=1,
+        mpi="mpi",
+        force_source=software,
+        dft_input=dft_input,
+        dft_loc=dft_loc,
+        dft_output=dft_output,
+        output_name=f"{casename}_otf_{software}",
+        store_dft_output=([dft_output, dft_input], "."),
+    )
 
     otf.run()
 
+@pytest.mark.parametrize("software", software_list)
+def test_otf_parser(software):
+
+    if not os.environ.get(cmd[software], False):
+        pytest.skip(
+            f"{cmd[software]} not found in environment:"
+            f" Please install the code "
+            f" and set the {cmd[software]} env. "
+            "variable to point to the executable."
+        )
+
+    example = 1
+    casename = name_list[example]
+    output_name = f"{casename}_otf_{software}.out"
+    otf_traj = OtfAnalysis(output_name)
+    replicated_gp = otf_traj.make_gp()
+
+
+@pytest.mark.parametrize("software", software_list)
+def test_load_checkpoint(software):
+
+    if not os.environ.get(cmd[software], False):
+        pytest.skip(
+            f"{cmd[software]} not found in environment:"
+            f" Please install the code "
+            f" and set the {cmd[software]} env. "
+            "variable to point to the executable."
+        )
+
+    if software == "cp2k":
+        pytest.skip()
+
+    example = 1
+    casename = name_list[example]
+    log_name = f"{casename}_otf_{software}"
+    new_otf = OTF.from_checkpoint(log_name + "_checkpt.json")
+    print("loaded from checkpoint", log_name)
+    assert new_otf.curr_step == number_of_steps
+    new_otf.number_of_steps = new_otf.number_of_steps + 2
+    new_otf.run()
+
+
+@pytest.mark.parametrize("software", software_list)
+def test_otf_parser_from_checkpt(software):
+
+    if not os.environ.get(cmd[software], False):
+        pytest.skip(
+            f"{cmd[software]} not found in environment:"
+            f" Please install the code "
+            f" and set the {cmd[software]} env. "
+            "variable to point to the executable."
+        )
+
+    if software == "cp2k":
+        pytest.skip()
+
+    example = 1
+    casename = name_list[example]
+    log_name = f"{casename}_otf_{software}"
+    output_name = f"{log_name}.out"
+    otf_traj = OtfAnalysis(output_name)
+    try:
+        replicated_gp = otf_traj.make_gp()
+    except:
+        init_gp = GaussianProcess.from_file(log_name + "_gp.json")
+        replicated_gp = otf_traj.make_gp(init_gp=init_gp)
+
+    outdir = f"test_outputs_{software}"
     if not os.path.isdir(outdir):
-         os.mkdir(outdir)
+        os.mkdir(outdir)
     for f in os.listdir("./"):
-        if f'{casename}_otf_{software}' in f:
+        if f"{casename}_otf_{software}" in f:
             shutil.move(f, outdir)
-    cleanup(software, f'{casename}_otf_{software}')
+    cleanup(software, f"{casename}_otf_{software}")
diff --git a/tests/test_parse_otf.py b/tests/test_parse_otf.py
index cd495b70..61c165aa 100644
--- a/tests/test_parse_otf.py
+++ b/tests/test_parse_otf.py
@@ -8,47 +8,48 @@ from flare.predict import predict_on_structure
 
 def test_parse_header():
 
-    os.system('cp test_files/sample_slab_otf.out .')
+    os.system("cp test_files/sample_slab_otf.out .")
 
-    header_dict = OtfAnalysis('sample_slab_otf.out').header
+    header_dict = OtfAnalysis("sample_slab_otf.out").header
 
-    assert header_dict['frames'] == 5000
-    assert header_dict['atoms'] == 28
-    assert header_dict['species_set'] == {'Al'}
-    assert header_dict['dt'] == .001
-    assert header_dict['kernel_name'] == 'two_plus_three_body'
-    assert header_dict['n_hyps'] == 5
-    assert header_dict['algo'] == 'BFGS'
-    assert np.equal(header_dict['cell'],
-                    np.array([[8.59135,  0.,       0.],
-                              [4.29567,  7.44033,  0.],
-                              [0., 0., 26.67654]])).all()
+    assert header_dict["frames"] == 5000
+    assert header_dict["atoms"] == 28
+    assert header_dict["species_set"] == {"Al"}
+    assert header_dict["dt"] == 0.001
+    assert header_dict["kernel_name"] == "two_plus_three_body"
+    assert header_dict["n_hyps"] == 5
+    assert header_dict["algo"] == "BFGS"
+    assert np.equal(
+        header_dict["cell"],
+        np.array([[8.59135, 0.0, 0.0], [4.29567, 7.44033, 0.0], [0.0, 0.0, 26.67654]]),
+    ).all()
+
+    os.system("rm sample_slab_otf.out")
 
-    os.system('rm sample_slab_otf.out')
 
 def test_gp_parser():
     """
     Test the capability of otf parser to read GP/DFT info
     :return:
     """
-    os.system('cp test_files/sample_slab_otf.out .')
-    parsed = OtfAnalysis('sample_slab_otf.out')
-    assert (parsed.gp_species_list == [['Al']*28] * 4)
+    os.system("cp test_files/sample_slab_otf.out .")
+    parsed = OtfAnalysis("sample_slab_otf.out")
+    assert parsed.gp_species_list == [["Al"] * 28] * 4
 
     gp_positions = parsed.gp_position_list
     assert len(gp_positions) == 4
 
     pos1 = 1.50245891
     pos2 = 10.06179079
-    assert(pos1 == gp_positions[0][-1][1])
-    assert(pos2 == gp_positions[-1][0][2])
+    assert pos1 == gp_positions[0][-1][1]
+    assert pos2 == gp_positions[-1][0][2]
 
     force1 = 0.29430943
     force2 = -0.02350709
-    assert(force1 == parsed.gp_force_list[0][-1][1])
-    assert(force2 == parsed.gp_force_list[-1][0][2])
+    assert force1 == parsed.gp_force_list[0][-1][1]
+    assert force2 == parsed.gp_force_list[-1][0][2]
 
-    os.system('rm sample_slab_otf.out')
+    os.system("rm sample_slab_otf.out")
 
 
 def test_md_parser():
@@ -56,19 +57,20 @@ def test_md_parser():
     Test the capability of otf parser to read MD info
     :return:
     """
-    os.system('cp test_files/sample_slab_otf.out .')
-    parsed = OtfAnalysis('sample_slab_otf.out')
+    os.system("cp test_files/sample_slab_otf.out .")
+    parsed = OtfAnalysis("sample_slab_otf.out")
 
     pos1 = 10.09769665
-    assert(pos1 == parsed.position_list[0][0][2])
-    assert(len(parsed.position_list[0]) == 28)
+    assert pos1 == parsed.position_list[0][0][2]
+    assert len(parsed.position_list[0]) == 28
+
+    os.system("rm sample_slab_otf.out")
 
-    os.system('rm sample_slab_otf.out')
 
 def test_output_md_structures():
 
-    os.system('cp test_files/sample_slab_otf.out .')
-    parsed = OtfAnalysis('sample_slab_otf.out')
+    os.system("cp test_files/sample_slab_otf.out .")
+    parsed = OtfAnalysis("sample_slab_otf.out")
 
     positions = parsed.position_list
     forces = parsed.force_list
@@ -78,7 +80,7 @@ def test_output_md_structures():
     assert np.isclose(structures[-1].positions, positions[-1]).all()
     assert np.isclose(structures[-1].forces, forces[-1]).all()
 
-    os.system('rm sample_slab_otf.out')
+    os.system("rm sample_slab_otf.out")
 
 
 def test_replicate_gp():
@@ -89,8 +91,8 @@ def test_replicate_gp():
     :return:
     """
 
-    os.system('cp test_files/sample_h2_otf.out .')
-    parsed = OtfAnalysis('sample_h2_otf.out')
+    os.system("cp test_files/sample_h2_otf.out .")
+    parsed = OtfAnalysis("sample_h2_otf.out")
 
     positions = parsed.position_list
     forces = parsed.force_list
@@ -114,4 +116,4 @@ def test_replicate_gp():
         pred_for, pred_stds = predict_on_structure(structure, gp_model)
         assert np.isclose(structure.forces, pred_for, atol=1e-6).all()
         assert np.isclose(structure.stds, pred_stds, atol=1e-6).all()
-    os.system('rm sample_slab_otf.out')
+    os.system("rm sample_slab_otf.out")
diff --git a/tests/test_struc.py b/tests/test_struc.py
index 93827dfa..76f668e8 100644
--- a/tests/test_struc.py
+++ b/tests/test_struc.py
@@ -6,8 +6,10 @@ from tests.test_gp import get_random_structure
 from flare.struc import Structure
 from json import loads, dumps
 from flare.utils.element_coder import Z_to_element, NumpyEncoder
+
 try:
     import pymatgen.core.structure as pmgstruc
+
     _test_pmg = True
 except ImportError:
     _test_pmg = False
@@ -16,9 +18,7 @@ from .test_gp import dumpcompare
 
 
 def test_random_structure_setup():
-    struct, forces = get_random_structure(cell=np.eye(3),
-                                          unique_species=[1, 2],
-                                          noa=2)
+    struct, forces = get_random_structure(cell=np.eye(3), unique_species=[1, 2], noa=2)
 
     assert np.equal(struct.cell, np.eye(3)).all()
     assert len(struct.positions) == 2
@@ -35,18 +35,15 @@ def test_prev_positions_arg():
         prev_positions.append(np.random.uniform(-1, 1, 3))
 
     test_structure1 = Structure(cell, species, positions)
-    test_structure2 = Structure(cell, species, positions,
-                                prev_positions=positions)
-    test_structure3 = Structure(cell, species, positions,
-                                prev_positions=prev_positions)
+    test_structure2 = Structure(cell, species, positions, prev_positions=positions)
+    test_structure3 = Structure(cell, species, positions, prev_positions=prev_positions)
 
     assert np.equal(test_structure1.positions, test_structure2.positions).all()
-    assert np.equal(test_structure1.prev_positions,
-                    test_structure2.prev_positions).all()
-    assert np.equal(test_structure2.positions,
-                    test_structure2.prev_positions).all()
-    assert not np.equal(test_structure3.positions,
-                        test_structure3.prev_positions).all()
+    assert np.equal(
+        test_structure1.prev_positions, test_structure2.prev_positions
+    ).all()
+    assert np.equal(test_structure2.positions, test_structure2.prev_positions).all()
+    assert not np.equal(test_structure3.positions, test_structure3.prev_positions).all()
 
 
 def test_raw_to_relative():
@@ -55,17 +52,20 @@ def test_raw_to_relative():
     cell = np.random.rand(3, 3)
     noa = 10
     positions = np.random.rand(noa, 3)
-    species = ['Al'] * len(positions)
+    species = ["Al"] * len(positions)
 
     test_struc = Structure(cell, species, positions)
-    rel_vals = test_struc.raw_to_relative(test_struc.positions,
-                                          test_struc.cell_transpose,
-                                          test_struc.cell_dot_inverse)
+    rel_vals = test_struc.raw_to_relative(
+        test_struc.positions, test_struc.cell_transpose, test_struc.cell_dot_inverse
+    )
 
     ind = np.random.randint(0, noa)
-    assert (np.isclose(positions[ind], rel_vals[ind, 0] * test_struc.vec1 +
-                       rel_vals[ind, 1] * test_struc.vec2 +
-                       rel_vals[ind, 2] * test_struc.vec3).all())
+    assert np.isclose(
+        positions[ind],
+        rel_vals[ind, 0] * test_struc.vec1
+        + rel_vals[ind, 1] * test_struc.vec2
+        + rel_vals[ind, 2] * test_struc.vec3,
+    ).all()
 
 
 def test_wrapped_coordinates():
@@ -74,27 +74,26 @@ def test_wrapped_coordinates():
 
     cell = np.random.rand(3, 3)
     positions = np.random.rand(10, 3)
-    species = ['Al'] * len(positions)
+    species = ["Al"] * len(positions)
 
     test_struc = Structure(cell, species, positions)
 
     wrap_diff = test_struc.positions - test_struc.wrapped_positions
-    wrap_rel = test_struc.raw_to_relative(wrap_diff,
-                                          test_struc.cell_transpose,
-                                          test_struc.cell_dot_inverse)
+    wrap_rel = test_struc.raw_to_relative(
+        wrap_diff, test_struc.cell_transpose, test_struc.cell_dot_inverse
+    )
 
-    assert (np.isclose(np.round(wrap_rel) - wrap_rel,
-                       np.zeros(positions.shape)).all())
+    assert np.isclose(np.round(wrap_rel) - wrap_rel, np.zeros(positions.shape)).all()
 
 
 @pytest.fixture
 def varied_test_struc():
-    struc = Structure(np.eye(3), species=[1, 2, 2, 3, 3, 4, 4, 4,
-                                          4, 3],
-                      positions=np.array([np.random.uniform(-1, 1, 3) for i
-                                          in range(10)]))
-    struc.forces = np.array([np.random.uniform(-1, 1, 3) for _ in
-                             range(10)])
+    struc = Structure(
+        np.eye(3),
+        species=[1, 2, 2, 3, 3, 4, 4, 4, 4, 3],
+        positions=np.array([np.random.uniform(-1, 1, 3) for i in range(10)]),
+    )
+    struc.forces = np.array([np.random.uniform(-1, 1, 3) for _ in range(10)])
     return struc
 
 
@@ -109,7 +108,7 @@ def test_to_from_methods(varied_test_struc):
     test_dict = varied_test_struc.as_dict()
 
     assert isinstance(test_dict, dict)
-    assert (test_dict['forces'] == varied_test_struc.forces).all()
+    assert (test_dict["forces"] == varied_test_struc.forces).all()
 
     new_struc_1 = Structure.from_dict(test_dict)
     new_struc_2 = Structure.from_dict(loads(varied_test_struc.as_str()))
@@ -125,69 +124,107 @@ def test_rep_methods(varied_test_struc):
     assert isinstance(str(varied_test_struc), str)
 
     thestr = str(varied_test_struc)
-    assert 'Structure with 10 atoms of types {' in thestr
+    assert "Structure with 10 atoms of types {" in thestr
+
 
 def test_struc_from_ase():
     from ase import Atoms
-    uc = Atoms(['Pd' for i in range(10)]+['Ag' for i in range(10)],
-               positions=np.random.rand(20, 3),
-               cell=np.random.rand(3, 3))
+    from ase.calculators.singlepoint import SinglePointCalculator
+
+    results = {
+        "forces": np.random.randn(20, 3),
+        "energy": np.random.rand(),
+        "stress": np.random.randn(6),
+    }
+
+    uc = Atoms(
+        ["Pd" for i in range(10)] + ["Ag" for i in range(10)],
+        positions=np.random.rand(20, 3),
+        cell=np.random.rand(3, 3),
+    )
+
+    calculator = SinglePointCalculator(uc, **results)
+    uc.set_calculator(calculator)
+
     new_struc = Structure.from_ase_atoms(uc)
+
     assert np.all(new_struc.species_labels == uc.get_chemical_symbols())
     assert np.all(new_struc.positions == uc.get_positions())
     assert np.all(new_struc.cell == uc.get_cell())
+    assert np.all(new_struc.forces == results["forces"])
+    assert np.all(new_struc.energy == results["energy"])
+    assert np.all(new_struc.stress == results["stress"])
+
 
 def test_struc_to_ase():
-    from ase import Atoms
-    uc = Structure(species=['Pd' for i in range(10)]+['Ag' for i in range(10)],
-                   positions=np.random.rand(20, 3),
-                   cell=np.random.rand(3, 3))
+    uc = Structure(
+        species=["Pd" for i in range(10)] + ["Ag" for i in range(10)],
+        positions=np.random.rand(20, 3),
+        cell=np.random.rand(3, 3),
+    )
+
+    uc.forces = np.random.randn(20, 3)
+    uc.energy = np.random.rand()
+    uc.stress = np.random.randn(6)
+
     new_atoms = Structure.to_ase_atoms(uc)
     assert np.all(uc.species_labels == new_atoms.get_chemical_symbols())
     assert np.all(uc.positions == new_atoms.get_positions())
     assert np.all(uc.cell == new_atoms.get_cell())
+    assert np.all(new_atoms.get_forces() == uc.forces)
+    assert np.all(new_atoms.get_potential_energy() == uc.energy)
+    assert np.all(new_atoms.get_stress() == uc.stress)
 
 
-@pytest.mark.skipif(not _test_pmg,reason='Pymatgen not present in available '
-                                        'packages.')
+@pytest.mark.skipif(
+    not _test_pmg, reason="Pymatgen not present in available " "packages."
+)
 def test_from_pmg_structure():
 
-    pmg_struc = pmgstruc.Structure(lattice= np.eye(3),
-                                   species=['H'],
-                                   coords=[[.25, .5, 0]],
-                                   site_properties={
-                                       'force': [np.array((1., 1.,  1.))],
-                                        'std':[np.array((1., 1., 1.))]},
-                                   coords_are_cartesian=True)
+    pmg_struc = pmgstruc.Structure(
+        lattice=np.eye(3),
+        species=["H"],
+        coords=[[0.25, 0.5, 0]],
+        site_properties={
+            "force": [np.array((1.0, 1.0, 1.0))],
+            "std": [np.array((1.0, 1.0, 1.0))],
+        },
+        coords_are_cartesian=True,
+    )
 
     new_struc = Structure.from_pmg_structure(pmg_struc)
 
     assert len(new_struc) == 1
 
-    assert np.equal(new_struc.positions, np.array([.25, .5, 0])).all()
+    assert np.equal(new_struc.positions, np.array([0.25, 0.5, 0])).all()
     assert new_struc.coded_species == [1]
-    assert new_struc.species_labels[0] == 'H'
-    assert np.equal(new_struc.forces, np.array([1., 1., 1.])).all()
-
-    pmg_struc = pmgstruc.Structure(lattice= np.diag([1.2,0.8,1.5]),
-                                   species=['H'],
-                                   coords=[[.25, .5, 0]],
-                                   site_properties={
-                                       'force': [np.array((1., 1.,  1.))],
-                                        'std':[np.array((1., 1., 1.))]},
-                                   coords_are_cartesian=True)
+    assert new_struc.species_labels[0] == "H"
+    assert np.equal(new_struc.forces, np.array([1.0, 1.0, 1.0])).all()
+
+    pmg_struc = pmgstruc.Structure(
+        lattice=np.diag([1.2, 0.8, 1.5]),
+        species=["H"],
+        coords=[[0.25, 0.5, 0]],
+        site_properties={
+            "force": [np.array((1.0, 1.0, 1.0))],
+            "std": [np.array((1.0, 1.0, 1.0))],
+        },
+        coords_are_cartesian=True,
+    )
 
     new_struc = Structure.from_pmg_structure(pmg_struc)
 
     assert len(new_struc) == 1
 
-    assert np.equal(new_struc.positions, np.array([.25, .5, 0])).all()
+    assert np.equal(new_struc.positions, np.array([0.25, 0.5, 0])).all()
     assert new_struc.coded_species == [1]
-    assert new_struc.species_labels[0] == 'H'
-    assert np.equal(new_struc.forces, np.array([1., 1., 1.])).all()
+    assert new_struc.species_labels[0] == "H"
+    assert np.equal(new_struc.forces, np.array([1.0, 1.0, 1.0])).all()
+
 
-@pytest.mark.skipif(not _test_pmg,reason='Pymatgen not present in available '
-                                        'packages.')
+@pytest.mark.skipif(
+    not _test_pmg, reason="Pymatgen not present in available " "packages."
+)
 def test_to_pmg_structure(varied_test_struc):
 
     new_struc = Structure.to_pmg_structure(varied_test_struc)
@@ -195,91 +232,89 @@ def test_to_pmg_structure(varied_test_struc):
     assert np.equal(new_struc.cart_coords, varied_test_struc.positions).all()
     assert (new_struc.atomic_numbers == varied_test_struc.coded_species).all()
 
+
 def test_to_xyz(varied_test_struc):
 
-    simple_str = varied_test_struc.to_xyz(extended_xyz=False,
-                print_stds=False, print_forces=False, print_max_stds=False)
+    simple_str = varied_test_struc.to_xyz(
+        extended_xyz=False, print_stds=False, print_forces=False, print_max_stds=False
+    )
 
-    simple_str_by_line = simple_str.split('\n')
+    simple_str_by_line = simple_str.split("\n")
 
-    assert len(simple_str_by_line)-2 == len(varied_test_struc)
+    assert len(simple_str_by_line) - 2 == len(varied_test_struc)
 
     for i, atom_line in enumerate(simple_str_by_line[2:-1]):
         split_line = atom_line.split()
-        assert split_line[0] == \
-               Z_to_element(int(varied_test_struc.species_labels[i]))
+        assert split_line[0] == Z_to_element(int(varied_test_struc.species_labels[i]))
         for j in range(3):
-            assert float(split_line[1+j]) == varied_test_struc.positions[i][j]
-
+            assert float(split_line[1 + j]) == varied_test_struc.positions[i][j]
 
-    complex_str = varied_test_struc.to_xyz(True,True,True,True)
-    complex_str_by_line = complex_str.split('\n')
+    complex_str = varied_test_struc.to_xyz(True, True, True, True)
+    complex_str_by_line = complex_str.split("\n")
 
-    assert len(complex_str_by_line)-2 == len(varied_test_struc)
+    assert len(complex_str_by_line) - 2 == len(varied_test_struc)
 
     for i, atom_line in enumerate(complex_str_by_line[2:-1]):
         split_line = atom_line.split()
-        assert split_line[0] == \
-               Z_to_element(int(varied_test_struc.species_labels[i]))
-        for j in range(1,4):
-            assert float(split_line[j]) == varied_test_struc.positions[i][j-1]
-        for j in range(4,7):
-            assert float(split_line[j]) == varied_test_struc.stds[i][j-4]
-        for j in range(7,10):
-            assert float(split_line[j]) == varied_test_struc.forces[i][j-7]
+        assert split_line[0] == Z_to_element(int(varied_test_struc.species_labels[i]))
+        for j in range(1, 4):
+            assert float(split_line[j]) == varied_test_struc.positions[i][j - 1]
+        for j in range(4, 7):
+            assert float(split_line[j]) == varied_test_struc.stds[i][j - 4]
+        for j in range(7, 10):
+            assert float(split_line[j]) == varied_test_struc.forces[i][j - 7]
         assert float(split_line[10]) == np.max(varied_test_struc.stds[i])
 
 
 def test_file_load():
-    struct1, forces = get_random_structure(cell=np.eye(3),
-                                           unique_species=[1, 2],
-                                           noa=2)
-    struct2, forces = get_random_structure(cell=np.eye(3),
-                                           unique_species=[1, 2],
-                                           noa=2)
-
-    with open("test_write.json",'w') as f:
+    struct1, forces = get_random_structure(cell=np.eye(3), unique_species=[1, 2], noa=2)
+    struct2, forces = get_random_structure(cell=np.eye(3), unique_species=[1, 2], noa=2)
+
+    with open("test_write.json", "w") as f:
         f.write(dumps(struct1.as_dict(), cls=NumpyEncoder))
 
     with pytest.raises(NotImplementedError):
-        Structure.from_file(file_name='test_write.json', format='xyz')
+        Structure.from_file(file_name="test_write.json", format="xyz")
 
-    struct1a = Structure.from_file('test_write.json')
+    struct1a = Structure.from_file("test_write.json")
     assert dumpcompare(struct1.as_dict(), struct1a.as_dict())
-    os.remove('test_write.json')
+    os.remove("test_write.json")
 
-    with open("test_write_set.json", 'w') as f:
-        f.write(dumps(struct1.as_dict(), cls=NumpyEncoder)+'\n')
-        f.write(dumps(struct2.as_dict(), cls=NumpyEncoder)+'\n')
+    with open("test_write_set.json", "w") as f:
+        f.write(dumps(struct1.as_dict(), cls=NumpyEncoder) + "\n")
+        f.write(dumps(struct2.as_dict(), cls=NumpyEncoder) + "\n")
 
-    structs = Structure.from_file('test_write_set.json')
+    structs = Structure.from_file("test_write_set.json")
     for struc in structs:
         assert isinstance(struc, Structure)
     assert len(structs) == 2
     assert dumpcompare(structs[0].as_dict(), struct1.as_dict())
     assert dumpcompare(structs[1].as_dict(), struct2.as_dict())
 
-    os.remove('test_write_set.json')
+    os.remove("test_write_set.json")
 
     with pytest.raises(FileNotFoundError):
-        Structure.from_file('fhqwhgads')
+        Structure.from_file("fhqwhgads")
 
-    vasp_struct = Structure.from_file('./test_files/test_POSCAR')
+    vasp_struct = Structure.from_file("./test_files/test_POSCAR")
     assert isinstance(vasp_struct, Structure)
     assert len(vasp_struct) == 6
 
+
 def test_is_valid():
     """
     Try a trivial case of 1-len structure and then one above and below
     tolerance
     :return:
     """
-    test_struc = Structure(cell=np.eye(3), species=['H'],
-                           positions=np.array([[0, 0, 0]]))
+    test_struc = Structure(
+        cell=np.eye(3), species=["H"], positions=np.array([[0, 0, 0]])
+    )
 
     assert test_struc.is_valid()
 
-    test_struc = Structure(cell=np.eye(3), species=['H', 'H'],
-                           positions=[[0, 0, 0], [.3, 0, 0]])
+    test_struc = Structure(
+        cell=np.eye(3), species=["H", "H"], positions=[[0, 0, 0], [0.3, 0, 0]]
+    )
     assert not test_struc.is_valid()
-    assert test_struc.is_valid(tolerance=.25)
+    assert test_struc.is_valid(tolerance=0.25)