Env now has force attribute

flare/env.py

@@ -376,3 +376,7 @@ def __str__(self):
         )
 
         return string
+
+    @property
+    def force(self) -> "np.ndarray":
+        return self.structure.forces[self.atom]

flare/gp.py

@@ -54,7 +54,7 @@ class GaussianProcess:
             component ("mc") kernel to use. Defaults to "mc"
         hyps (np.ndarray, optional): Hyperparameters of the GP.
         cutoffs (Dict, optional): Cutoffs of the GP kernel. For simple hyper-
-            parameter setups, formatted like {"twobody":7, "threebody":4.5}, 
+            parameter setups, formatted like {"twobody":7, "threebody":4.5},
             etc.
         hyp_labels (List, optional): List of hyperparameter labels. Defaults
             to None.
@@ -326,7 +326,7 @@ def update_kernel(
     def update_db(
         self,
         struc: Structure,
-        forces: "ndarray",
+        forces: "ndarray" = None,
         custom_range: List[int] = (),
         energy: float = None,
         stress: "ndarray" = None,
@@ -388,7 +388,13 @@ def update_db(
         )
         self.sync_data()
 
-    def add_one_env(self, env: AtomicEnvironment, force, train: bool = False, **kwargs):
+    def add_one_env(
+        self,
+        env: AtomicEnvironment,
+        force: "np.ndarray" = None,
+        train: bool = False,
+        **kwargs,
+    ):
         """Add a single local environment to the training set of the GP.
 
         Args:
@@ -401,7 +407,10 @@ def add_one_env(self, env: AtomicEnvironment, force, train: bool = False, **kwar
                 environment is added.
         """
         self.training_data.append(env)
-        self.training_labels.append(force)
+        if force is None:
+            self.training_labels.append(env.force)
+        else:
+            self.training_labels.append(force)
         self.training_labels_np = np.hstack(self.training_labels)
         self.sync_data()
 
@@ -714,7 +723,7 @@ def predict_local_energy_and_var(self, x_t: AtomicEnvironment):
 
     def predict_efs(self, x_t: AtomicEnvironment):
         """Predict the local energy, forces, and partial stresses of an
-            atomic environment and their predictive variances."""
+        atomic environment and their predictive variances."""
 
         # Kernel vector allows for evaluation of atomic environments.
         if self.parallel and not self.per_atom_par:
@@ -1187,7 +1196,7 @@ def write_model(
 
         else:
             raise ValueError(
-                "Output format not supported: try from " "{}".format(supported_formats)
+                "Output format not supported: try from {}".format(supported_formats)
             )
 
         if len(self.training_data) > split_matrix_size_cutoff:
@@ -1233,14 +1242,13 @@ def from_file(filename: str, format: str = ""):
                         gp_model.alpha = None
                         gp_model.ky_mat_inv = None
                         Warning(
-                            "the covariance matrices are not loaded"
-                            f"it can take extra long time to recompute"
+                            f"the covariance matrices are not loaded, "
+                            f"this can take a long time to recompute"
                         )
 
         else:
             raise ValueError(
-                "Warning: Format unspecieified or file is not "
-                ".json or .pickle format."
+                "Warning: Format unspecieified or file is not .json or .pickle format."
             )
 
         gp_model.check_instantiation()

tests/test_env.py

@@ -67,8 +67,13 @@ def test_env_methods(structure, mask, cutoff, result):
     else:
         mask = None
 
+    structure = deepcopy(structure)
+    structure.forces = np.random.random(size=(len(structure), 3))
+
     env_test = AtomicEnvironment(structure, atom=0, cutoffs=cutoff, cutoffs_mask=mask)
 
+    assert np.array_equal(structure.forces[0], env_test.force)
+
     assert (
         str(env_test)
         == f"Atomic Env. of Type 1 surrounded by {result[0]} atoms of Types [1, 2, 3]"