Env now has force attribute (#229)

* Env now has force attribute

* Modify warning message

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]"