Add progress bar to phonon calculations (#319)

* Add rich progress bar to Phonons.calc_force_constants

* Add rich as dependency

---------

Co-authored-by: Turanyi <[email protected]>

Author: RastislavTuranyi

janus_core/calculations/phonons.py

@@ -20,7 +20,7 @@
     PathLike,
     PhononCalcs,
 )
-from janus_core.helpers.utils import none_to_dict, write_table
+from janus_core.helpers.utils import none_to_dict, track_progress, write_table
 
 
 class Phonons(BaseCalculation):
@@ -90,6 +90,8 @@ class Phonons(BaseCalculation):
     file_prefix : Optional[PathLike]
         Prefix for output filenames. Default is inferred from chemical formula of the
         structure.
+    enable_progress_bar : bool
+        Whether to show a progress bar during phonon calculations. Default is False.
 
     Attributes
     ----------
@@ -152,6 +154,7 @@ def __init__(
         write_results: bool = True,
         write_full: bool = True,
         file_prefix: Optional[PathLike] = None,
+        enable_progress_bar: bool = False,
     ) -> None:
         """
         Initialise Phonons class.
@@ -219,6 +222,8 @@ def __init__(
         file_prefix : Optional[PathLike]
             Prefix for output filenames. Default is inferred from structure name, or
             chemical formula of the structure.
+        enable_progress_bar : bool
+            Whether to show a progress bar during phonon calculations. Default is False.
         """
         (read_kwargs, minimize_kwargs) = none_to_dict((read_kwargs, minimize_kwargs))
 
@@ -235,6 +240,7 @@ def __init__(
         self.plot_to_file = plot_to_file
         self.write_results = write_results
         self.write_full = write_full
+        self.enable_progress_bar = enable_progress_bar
 
         # Ensure supercell is a valid list
         self.supercell = [supercell] * 3 if isinstance(supercell, int) else supercell
@@ -363,6 +369,11 @@ def calc_force_constants(
         phonon.generate_displacements(distance=self.displacement)
         disp_supercells = phonon.supercells_with_displacements
 
+        if self.enable_progress_bar:
+            disp_supercells = track_progress(
+                disp_supercells, "Computing displacements..."
+            )
+
         phonon.forces = [
             self._calc_forces(supercell)
             for supercell in disp_supercells

janus_core/cli/phonons.py

@@ -238,6 +238,7 @@ def phonons(
         "write_results": True,
         "write_full": write_full,
         "file_prefix": file_prefix,
+        "enable_progress_bar": True,
     }
 
     # Initialise phonons

janus_core/helpers/utils.py

@@ -6,11 +6,19 @@
 from io import StringIO
 import logging
 from pathlib import Path
-from typing import Any, Literal, Optional, TextIO, get_args
+from typing import Any, Literal, Optional, TextIO, Union, get_args
 
 from ase import Atoms
 from ase.io import read, write
 from ase.io.formats import filetype
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    TextColumn,
+    TimeRemainingColumn,
+)
+from rich.style import Style
 from spglib import get_spacegroup
 
 from janus_core.helpers.janus_types import (
@@ -674,3 +682,44 @@ def _dump_csv(
 
     for cols in zip(*columns.values()):
         print(",".join(map(format, cols, formats)), file=file)
+
+
+def track_progress(sequence: Union[Sequence, Iterable], description: str) -> Iterable:
+    """
+    Track the progress of iterating over a sequence.
+
+    This is done by displaying a progress bar in the console using the rich library.
+    The function is an iterator over the sequence, updating the progress bar each
+    iteration.
+
+    Parameters
+    ----------
+    sequence : Iterable
+        The sequence to iterate over. Must support "len".
+    description : str
+        The text to display to the left of the progress bar.
+
+    Yields
+    ------
+    Iterable
+        An iterable of the values in the sequence.
+    """
+    text_column = TextColumn("{task.description}")
+    bar_column = BarColumn(
+        bar_width=None,
+        complete_style=Style(color="#FBBB10"),
+        finished_style=Style(color="#E38408"),
+    )
+    completion_column = MofNCompleteColumn()
+    time_column = TimeRemainingColumn()
+    progress = Progress(
+        text_column,
+        bar_column,
+        completion_column,
+        time_column,
+        expand=True,
+        auto_refresh=False,
+    )
+
+    with progress:
+        yield from progress.track(sequence, description=description)

pyproject.toml

@@ -33,6 +33,7 @@ numpy = "^1.26.4"
 phonopy = "^2.23.1"
 python = "^3.9"
 pyyaml = "^6.0.1"
+rich = "^13.9.1"
 seekpath = "^1.9.7"
 spglib = "^2.3.0"
 torch = ">= 2.1, <= 2.2" # Range required for dgl