Add simple progress bar to simulations

diffsims/generators/simulation_generator.py

@@ -19,7 +19,9 @@
 """Kinematic Diffraction Simulation Generator."""
 
 from typing import Union, Sequence
+
 import numpy as np
+from tqdm import tqdm
 
 from orix.quaternion import Rotation
 from orix.crystal_map import Phase
@@ -139,6 +141,7 @@ def calculate_diffraction2d(
         max_excitation_error: float = 1e-2,
         shape_factor_width: float = None,
         debye_waller_factors: dict = None,
+        show_progressbar: bool = False,
     ):
         """Calculates the diffraction pattern for one or more phases given a list
         of rotations for each phase.
@@ -166,6 +169,8 @@ def calculate_diffraction2d(
             control. If not set will be set equal to max_excitation_error.
         debye_waller_factors
             Maps element names to their temperature-dependent Debye-Waller factors.
+        show_progressbar
+            If True, display a progressbar. Defaults to False
 
         Returns
         -------
@@ -196,7 +201,13 @@ def calculate_diffraction2d(
                 include_zero_vector=with_direct_beam,
             )
             phase_vectors = []
-            for rot in rotate:
+
+            # Progress bar setup
+            rotate_iter = rotate
+            if show_progressbar:
+                rotate_iter = tqdm(rotate_iter, desc=p.name, total=rotate.size)
+
+            for rot in rotate_iter:
                 # Calculate the reciprocal lattice vectors that intersect the Ewald sphere.
                 (
                     intersected_vectors,

diffsims/tests/generators/test_simulation_generator.py

@@ -341,3 +341,63 @@ def test_same_simulation_results():
     )
     old_data = np.load(FILE1)
     np.testing.assert_allclose(new_data, old_data, atol=1e-8)
+
+
+def test_calculate_diffraction2d_progressbar_single_phase(capsys):
+    gen = SimulationGenerator()
+    phase = make_phase()
+    phase.name = "test phase"
+    rots = Rotation.random(10)
+
+    # no pbar
+    sims = gen.calculate_diffraction2d(phase, rots, show_progressbar=False)
+
+    # Note: tqdm outputs to stderr by default
+    captured = capsys.readouterr()
+    assert captured.err == ""
+
+    # with pbar
+    sims = gen.calculate_diffraction2d(phase, rots, show_progressbar=True)
+
+    captured = capsys.readouterr()
+    expected = "test phase: 100%|██████████| 10/10"  # also some more, but that is compute-time dependent
+    # ignore possible flushing
+    captured = captured.err.split("\r")[-1]
+    assert captured[: len(expected)] == expected
+
+
+def test_calculate_diffraction2d_progressbar_multi_phase(capsys):
+    gen = SimulationGenerator()
+    phase1 = make_phase()
+    phase1.name = "A"
+    phase2 = make_phase()
+    phase2.name = "B"
+    rots = Rotation.random(10)
+
+    # no pbar
+    sims = gen.calculate_diffraction2d(
+        [phase1, phase2], [rots, rots], show_progressbar=False
+    )
+
+    # Note: tqdm outputs to stderr by default
+    captured = capsys.readouterr()
+    assert captured.err == ""
+
+    # with pbar
+    sims = gen.calculate_diffraction2d(
+        [phase1, phase2], [rots, rots], show_progressbar=True
+    )
+
+    captured = capsys.readouterr()
+    expected1 = "A: 100%|██████████| 10/10 "
+    expected2 = "B: 100%|██████████| 10/10 "
+    # Find the correct output in the stream, i.e. final line containing the name of the phase
+    captured1 = ""
+    captured2 = ""
+    for line in captured.err.split("\r"):
+        if "A" in line:
+            captured1 = line
+        if "B" in line:
+            captured2 = line
+    assert captured1[: len(expected1)] == expected1
+    assert captured2[: len(expected2)] == expected2