diff --git a/flare/env.py b/flare/env.py index c99296e59..ee1331353 100644 --- a/flare/env.py +++ b/flare/env.py @@ -376,3 +376,7 @@ def __str__(self): ) return string + + @property + def force(self) -> "np.ndarray": + return self.structure.forces[self.atom] diff --git a/flare/gp.py b/flare/gp.py index 7ba920334..e99fd74d8 100644 --- a/flare/gp.py +++ b/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() diff --git a/tests/test_env.py b/tests/test_env.py index ef98ad4fb..9bf728f4c 100644 --- a/tests/test_env.py +++ b/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]"