Merge pull request #347 from mir-group/development

1.4.2 release - Bug fixes, new features

.github/dependabot.yml

@@ -0,0 +1,6 @@
+version: 2
+updates:
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "monthly"

CMakeLists.txt

@@ -104,6 +104,7 @@ set(FLARE_SOURCES
     src/flare_pp/bffs/sparse_gp.cpp
     src/flare_pp/bffs/gp.cpp
     src/flare_pp/descriptors/descriptor.cpp
+    src/flare_pp/descriptors/b1.cpp
     src/flare_pp/descriptors/b2.cpp
     src/flare_pp/descriptors/b2_norm.cpp
     src/flare_pp/descriptors/b2_simple.cpp

ctests/test_descriptor.cpp

@@ -1,4 +1,3 @@
-#include "b3.h"
 #include "descriptor.h"
 #include "test_structure.h"
 #include "gtest/gtest.h"
@@ -43,7 +42,7 @@ class DescTest : public StructureTest {
   }
 };
 
-using DescTypes = ::testing::Types<B2, B3>;
+using DescTypes = ::testing::Types<B1, B2, B3>;
 TYPED_TEST_SUITE(DescTest, DescTypes);
 
 //TEST_P(DescRotTest, RotationTest) {

ctests/test_structure.h

@@ -1,3 +1,4 @@
+#include "b1.h"
 #include "b2.h"
 #include "b2_norm.h"
 #include "b2_simple.h"

examples/test_SGP_Fake_fresh.yaml

@@ -27,8 +27,8 @@ flare_calc:
         - 14
         - 6
     single_atom_energies:
-        - 0
-        - 0
+        - -7.4
+        - -8.9
     cutoff: 4.0
     variance_type: local
     max_iterations: 1

examples/test_SGP_LMP_fresh.yaml

@@ -26,8 +26,8 @@ flare_calc:
         - 14
         - 6
     single_atom_energies:
-        - 0
-        - 0
+        - -7.4
+        - -8.9
     cutoff: 4.0
     variance_type: local
     max_iterations: 1

flare/_version.py

@@ -1 +1 @@
-__version__ = "1.4.1"
+__version__ = "1.4.2"

flare/bffs/sgp/sparse_gp.py

@@ -288,12 +288,18 @@ def from_dict(in_dict):
             rel_efs_noise = train_struc.info.get("rel_efs_noise", [1, 1, 1])
             rel_e_noise, rel_f_noise, rel_s_noise = rel_efs_noise
 
+            # check for stress training since ASE atoms assert stress arrays of length 6
+            if not in_dict["stress_training"]:
+                struc_stress = np.array(None)
+            else:
+                struc_stress = train_struc.stress
+
             gp.update_db(
                 train_struc,
                 train_struc.forces,
                 custom_range=custom_range,
                 energy=energy,
-                stress=train_struc.stress,
+                stress=struc_stress,
                 mode="specific",
                 sgp=None,
                 update_qr=False,
@@ -516,6 +522,7 @@ def duplicate(self, new_hyps=None, new_kernels=None, new_powers=None):
 
         n_kern = len(kernels)
         new_gp = SparseGP(kernels, hyps[n_kern], hyps[n_kern + 1], hyps[n_kern + 2])
+        new_gp.Kuu_jitter = self.sparse_gp.Kuu_jitter
 
         # add training data
         sparse_indices = self.sparse_gp.sparse_indices

flare/io/output.py

@@ -214,7 +214,7 @@ def write_md_header(self, dt, curr_step, dft_step):
         else:
             string += f"\n*-Frame: {curr_step} "
 
-        string += f"\nSimulation Time: {(dt * curr_step):10.3f} ps \n"
+        string += f"\nSimulation Time: {(dt * curr_step):10.4f} ps \n"
         return string
 
     def write_md_config(

flare/learners/otf.py

@@ -156,6 +156,7 @@ def __init__(
     ):
 
         self.atoms = FLARE_Atoms.from_ase_atoms(atoms)
+
         if flare_calc is not None:
             self.atoms.calc = flare_calc
         self.md_engine = md_engine
@@ -184,9 +185,10 @@ def __init__(
         if self.md_engine == "PyLAMMPS":
             md_kwargs["output_name"] = output_name
             if isinstance(self.atoms.calc, SGP_Calculator):
-                assert self.atoms.calc.gp_model.variance_type == "local", \
-                        "LAMMPS training only supports variance_type='local'"
-
+                assert (
+                    self.atoms.calc.gp_model.variance_type == "local"
+                ), "LAMMPS training only supports variance_type='local'"
+
         self.md = MD(
             atoms=self.atoms, timestep=timestep, trajectory=trajectory, **md_kwargs
         )
@@ -227,7 +229,9 @@ def __init__(
         if init_atoms is None:  # set atom list for initial dft run
             self.init_atoms = [int(n) for n in range(self.noa)]
         elif isinstance(init_atoms, int):
-            self.init_atoms = np.random.choice(len(self.atoms), size=init_atoms, replace=False)
+            self.init_atoms = np.random.choice(
+                len(self.atoms), size=init_atoms, replace=False
+            )
         else:
             # detect if there are duplicated atoms
             assert len(set(init_atoms)) == len(
@@ -423,7 +427,7 @@ def run(self):
                                 "dft_s_mae": s_mae,
                                 "dft_s_mav": s_mav,
                             },
-                            step = self.curr_step,
+                            step=self.curr_step,
                         )
 
             # write gp forces
@@ -578,23 +582,33 @@ def run_dft(self):
         f = logging.getLogger(self.output.basename + "log")
         f.info("\nCalling DFT...\n")
 
-        # change from FLARE to DFT calculator
-        self.atoms.calc = deepcopy(self.dft_calc)
+        # avoid duplicating a previous DFT calculation
+        self.dft_calc.reset()
+        self.atoms.calc = self.dft_calc
 
         # Calculate DFT energy, forces, and stress.
         # Note that ASE and QE stresses differ by a minus sign.
         if "forces" in self.dft_calc.implemented_properties:
-            forces = self.atoms.get_forces()
+            if "forces" in self.atoms.calc.results:
+                forces = self.atoms.get_forces()
+            else:
+                forces = None
         else:
             forces = None
 
         if "stress" in self.dft_calc.implemented_properties:
-            stress = self.atoms.get_stress()
+            if "stress" in self.atoms.calc.results:
+                stress = self.atoms.get_stress()
+            else:
+                stress = None
         else:
             stress = None
 
         if "energy" in self.dft_calc.implemented_properties:
-            energy = self.atoms.get_potential_energy()
+            if "energy" in self.atoms.calc.results:
+                energy = self.atoms.get_potential_energy()
+            else:
+                energy = None
         else:
             energy = None
 
@@ -791,14 +805,14 @@ def record_state(self):
                     "ke": self.KE,
                     "pe": self.atoms.get_potential_energy(),
                 },
-                step = self.curr_step,
+                step=self.curr_step,
             )
             if "stds" in self.atoms.calc.results:
                 wandb.log(
                     {
                         "maxunc": np.max(np.abs(self.atoms.calc.results["stds"])),
                     },
-                    step = self.curr_step,
+                    step=self.curr_step,
                 )
 
         if self.md_engine == "Fake" and not self.dft_step:
@@ -827,10 +841,9 @@ def record_state(self):
                         "s_mae": s_mae,
                         "s_mav": s_mav,
                     },
-                    step = self.curr_step,
+                    step=self.curr_step,
                 )
 
-
     def record_dft_data(self, structure, target_atoms):
         structure.info["target_atoms"] = np.array(target_atoms)
         write(self.dft_xyz, structure, append=True)
@@ -849,6 +862,10 @@ def as_dict(self):
         self.gp = None
         self.atoms.calc = None
 
+        # DFT calculator may not be picklable. Temporarily set to None before copying.
+        dft_calc = self.dft_calc
+        self.dft_calc = None
+
         # Deepcopy OTF object.
         dct = deepcopy(dict(vars(self)))
 
@@ -858,6 +875,7 @@ def as_dict(self):
         self.flare_calc = flare_calc
         self.gp = gp
         self.atoms.calc = flare_calc
+        self.dft_calc = dft_calc
 
         # write atoms and flare calculator to separate files
         write(self.atoms_name, self.atoms)
@@ -870,7 +888,7 @@ def as_dict(self):
         try:
             with open(self.dft_name, "wb") as f:
                 pickle.dump(self.dft_calc, f)
-        except AttributeError:
+        except (AttributeError, TypeError):
             with open(self.dft_name + ".json", "w") as f:
                 json.dump(self.dft_calc.todict(), f, cls=NumpyEncoder)
 

flare/md/fake.py

@@ -101,12 +101,21 @@ def step(self):
         ]
 
         self.atoms.calc.reset()
-        gp_energy = self.atoms.get_potential_energy()
         gp_forces = self.atoms.get_forces()
-        gp_stress = self.atoms.get_stress()
-        self.dft_energy = new_atoms.get_potential_energy()
-        self.dft_forces = new_atoms.get_forces()
-        self.dft_stress = new_atoms.get_stress()
+        try:
+            self.dft_energy = new_atoms.get_potential_energy()
+        except:
+            self.dft_energy = None
+
+        try:
+            self.dft_forces = new_atoms.get_forces()
+        except:
+            self.dft_forces = None
+
+        try:
+            self.dft_stress = new_atoms.get_stress()
+        except:
+            self.dft_stress = None
 
         self.curr_step += 1
         self.atoms.info["step"] = self.curr_step
@@ -155,8 +164,7 @@ def calculate(
 
         step = atoms.info.get("step", 0)
 
-        fake_trajectory = read("All_Data.xyz", index=":", format="extxyz")
-        fake_frame = fake_trajectory[step]
+        fake_frame = read("All_Data.xyz", index=step, format="extxyz")
         assert np.allclose(atoms.positions, fake_frame.positions), (
             atoms.positions[0],
             fake_frame.positions[0],

flare/md/lammps.py

@@ -236,12 +236,19 @@ def step(self, std_tolerance, N_steps):
         elif self.params["pair_style"] == "flare":
             self.uncertainty_file = "L_inv_lmp.flare sparse_desc_lmp.flare"
 
-        self.params["model_post"] += [f"reset_timestep {self.nsteps}"]
+        #self.params["model_post"] += [f"reset_timestep {self.nsteps}"]
         self.params["run"] = (
             f"{N_steps} upto every {self.params['dump_period']} "
-            f"\"if '$(c_MaxUnc) > {std_tolerance}' then quit\""
+            f"\"if '$(c_MaxUnc) > {std_tolerance}' then 'write_restart restart_*.dat' quit\""
+            f"\nwrite_restart restart_*.dat"
         )
 
+        print("lammps nsteps =", self.nsteps)
+        if self.nsteps > 0:
+            restart_file = f"restart_{self.nsteps}.dat"
+        else:
+            restart_file = None
+
         lmp_calc, params = get_flare_lammps_calc(
             pair_style="flare",
             potfile=self.potential_file,
@@ -253,6 +260,7 @@ def step(self, std_tolerance, N_steps):
             compute_uncertainty=True,
             thermo_file=self.thermo_file,
             params=self.params,
+            restart=restart_file,
         )
 
         atoms = deepcopy(self.curr_atoms)
@@ -387,6 +395,7 @@ def get_flare_lammps_calc(
     compute_uncertainty=False,
     thermo_file=None,
     params={},
+    restart=None,
 ):
 
     params.setdefault("compute", [])
@@ -437,6 +446,22 @@ def get_flare_lammps_calc(
                 f'thermoprint all print {dump_period} "$(step) $(temp) $(ke) $(pe) $(etotal) $(pxx) $(pyy) $(pzz) $(pyz) $(pxz) $(pxy) $(c_MaxUnc)" append {tmp_dir}/{thermo_file}'
             ]
 
+    if restart is not None:
+        params["model_post"] = [f"""
+clear
+atom_style atomic
+units metal
+boundary p p p
+atom_modify sort 0 0.0
+newton {params["newton"]}
+
+read_restart {restart}
+
+### interactions
+pair_style {params["pair_style"]}
+pair_coeff {params["pair_coeff"][0]}
+"""] + params["model_post"]
+
     # Run lammps with the customized parameters
     if command is None:
         command = os.environ.get("lmp")
@@ -467,6 +492,12 @@ def check_sgp_match(atoms, sgp_calc, logger, specorder, command):
     lmp_stress = atoms.stress
     lmp_stds = atoms.get_array("c_unc")
 
+    # Add back single atom energies
+    if sgp_calc.gp_model.single_atom_energies is not None:
+        for spec in atoms.numbers:
+            coded_spec = sgp_calc.gp_model.species_map[spec]
+            lmp_energy += sgp_calc.gp_model.single_atom_energies[coded_spec]
+
     # subtract the pressure from temperature
     kinetic_stress = get_kinetic_stress(atoms)
     lmp_stress += kinetic_stress