Merge pull request #233 from viljarjf/simulation-progressbar

Add simple progress bar to simulations

Author: CSSFrancis

.github/workflows/build.yml

@@ -40,13 +40,13 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, windows-latest, macos-latest]
-        python-version: ['3.10', '3.11']
+        python-version: ['3.10', '3.11', '3.12', '3.13']
         include:
           # Oldest supported version of main dependencies
           - os: ubuntu-latest
-            python-version: 3.8
+            python-version: 3.9
             OLDEST_SUPPORTED_VERSION: true
-            DEPENDENCIES: diffpy.structure==3.0.2 matplotlib==3.5 numpy==1.17.3 orix==0.12.1 scipy==1.8 tqdm==4.9
+            DEPENDENCIES: diffpy.structure==3.0.2 matplotlib==3.5 numpy==1.17.3 orix==0.12.1 scipy==1.8 tqdm==4.61.2
             LABEL: -oldest
     steps:
       - uses: actions/checkout@v4

.zenodo.json

@@ -36,6 +36,9 @@
       "name":"Niels Cautaerts",
       "orcid": "0000-0002-6402-9879"
     },
+    {
+      "name": "Viljar Johan Femoen"
+    },
     {
       "name":"Isabel Wood"
     },

CHANGELOG.rst

@@ -12,6 +12,7 @@ Unreleased
 
 Added
 -----
+- Progressbar for SimulationGenerator.calculate_diffraction2d in (#233)
 
 Changed
 -------

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/release_info.py

@@ -14,6 +14,7 @@
     "Ben Martineau",
     "Niels Cautaerts",
     "Joonatan Laulainen",
+    "Viljar Johan Femoen",
     "Isabel Wood",
     "Sean Collins",
     "Stef Smeets",

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