Display training/reconstruction steps

pycroscopy · Apr 19, 2023 · bce2742 · bce2742
1 parent 7f718cc
commit bce2742
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/atomai/models/dklgp/gpr.py b/atomai/models/dklgp/gpr.py
@@ -104,8 +104,11 @@ def reconstruct(self, sparse_image: np.ndarray,
         X_train, y_train, X_full = prepare_gp_input(sparse_image)
         if not lengthscale_constraints:
             lengthscale_constraints = get_lengthscale_constraints(X_full)
+        print("Model training .../n")
         self.fit(X_train, y_train, training_cycles,
                  lengthscale_constraints=lengthscale_constraints,
                  grid_points_ratio=grid_points_ratio, **kwargs)
+        print('\n\rPerforming reconstruction... ', end="")
         reconstruction = self.predict(X_full, **kwargs)
+        print("Done")
         return reconstruction.view(sparse_image.shape).cpu().numpy()