From 87f66ae6106f717c3a947bacb4bd9a7490a91392 Mon Sep 17 00:00:00 2001
From: k-harris27 <120191386+k-harris27@users.noreply.github.com>
Date: Wed, 5 Feb 2025 13:49:51 +0000
Subject: [PATCH] Add progress bars for single point properties (#402)

---
 janus_core/calculations/single_point.py | 36 +++++++++++++++++++++----
 janus_core/cli/singlepoint.py           |  1 +
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/janus_core/calculations/single_point.py b/janus_core/calculations/single_point.py
index 7f305a37..e8ce8ea3 100644
--- a/janus_core/calculations/single_point.py
+++ b/janus_core/calculations/single_point.py
@@ -22,7 +22,7 @@
 )
 from janus_core.helpers.mlip_calculators import check_calculator
 from janus_core.helpers.struct_io import output_structs
-from janus_core.helpers.utils import none_to_dict
+from janus_core.helpers.utils import none_to_dict, track_progress
 
 
 class SinglePoint(BaseCalculation):
@@ -69,6 +69,9 @@ class SinglePoint(BaseCalculation):
     write_kwargs
         Keyword arguments to pass to ase.io.write if saving structure with results of
         calculations. Default is {}.
+    enable_progress_bar
+        Whether to show a progress bar when applied to a file containing many
+        structures. Default is False.
 
     Attributes
     ----------
@@ -94,6 +97,7 @@ def __init__(
         properties: MaybeSequence[Properties] = (),
         write_results: bool = False,
         write_kwargs: OutputKwargs | None = None,
+        enable_progress_bar: bool = False,
     ) -> None:
         """
         Read the structure being simulated and attach an MLIP calculator.
@@ -138,12 +142,16 @@ def __init__(
         write_kwargs
             Keyword arguments to pass to ase.io.write if saving structure with results
             of calculations. Default is {}.
+        enable_progress_bar
+            Whether to show a progress bar when applied to a file containing many
+            structures. Default is False.
         """
         read_kwargs, write_kwargs = none_to_dict(read_kwargs, write_kwargs)
 
         self.write_results = write_results
         self.write_kwargs = write_kwargs
         self.log_kwargs = log_kwargs
+        self.enable_progress_bar = enable_progress_bar
 
         # Read full trajectory by default
         read_kwargs.setdefault("index", ":")
@@ -233,7 +241,12 @@ def _get_potential_energy(self) -> MaybeList[float]:
             Potential energy of structure(s).
         """
         if isinstance(self.struct, Sequence):
-            return [struct.get_potential_energy() for struct in self.struct]
+            struct_sequence = self.struct
+            if self.enable_progress_bar:
+                struct_sequence = track_progress(
+                    struct_sequence, "Computing potential energies..."
+                )
+            return [struct.get_potential_energy() for struct in struct_sequence]
 
         return self.struct.get_potential_energy()
 
@@ -247,7 +260,10 @@ def _get_forces(self) -> MaybeList[ndarray]:
             Forces of structure(s).
         """
         if isinstance(self.struct, Sequence):
-            return [struct.get_forces() for struct in self.struct]
+            struct_sequence = self.struct
+            if self.enable_progress_bar:
+                struct_sequence = track_progress(struct_sequence, "Computing forces...")
+            return [struct.get_forces() for struct in struct_sequence]
 
         return self.struct.get_forces()
 
@@ -261,7 +277,12 @@ def _get_stress(self) -> MaybeList[ndarray]:
             Stress of structure(s).
         """
         if isinstance(self.struct, Sequence):
-            return [struct.get_stress() for struct in self.struct]
+            struct_sequence = self.struct
+            if self.enable_progress_bar:
+                struct_sequence = track_progress(
+                    struct_sequence, "Computing stresses..."
+                )
+            return [struct.get_stress() for struct in struct_sequence]
 
         return self.struct.get_stress()
 
@@ -300,7 +321,12 @@ def _get_hessian(self) -> MaybeList[ndarray]:
             Hessian of structure(s).
         """
         if isinstance(self.struct, Sequence):
-            return [self._calc_hessian(struct) for struct in self.struct]
+            struct_sequence = self.struct
+            if self.enable_progress_bar:
+                struct_sequence = track_progress(
+                    struct_sequence, "Computing Hessian..."
+                )
+            return [self._calc_hessian(struct) for struct in struct_sequence]
 
         return self._calc_hessian(self.struct)
 
diff --git a/janus_core/cli/singlepoint.py b/janus_core/cli/singlepoint.py
index 3589bda4..3e46a34a 100644
--- a/janus_core/cli/singlepoint.py
+++ b/janus_core/cli/singlepoint.py
@@ -139,6 +139,7 @@ def singlepoint(
         "attach_logger": True,
         "log_kwargs": log_kwargs,
         "track_carbon": tracker,
+        "enable_progress_bar": True,
     }
 
     # Initialise singlepoint structure and calculator