Speedup import and add regression check for import time (#238)

* Add test framework to monitor module import times with regression tests
* Use time.perf_counter for accurate timing
* Implement lazy imports across multiple modules to improve performance:
  - scipy
  - plotly.figure_factory
  - sklearn
  - pymatgen (Structure, NearNeighbors, PhononDos, PhononBands, Composition)
* Add reference import times for all core modules
* Configure tests to run only on main branch
* Add grace and hard thresholds for import time regression

---------

Co-authored-by: Janosh Riebesell <[email protected]>

Author: DanielYang59

.pre-commit-config.yaml

@@ -2,7 +2,7 @@ ci:
   autoupdate_schedule: quarterly
   skip: [pyright]
 
-default_stages: [commit]
+default_stages: [pre-commit]
 
 default_install_hook_types: [pre-commit, commit-msg]
 
@@ -43,7 +43,7 @@ repos:
     rev: v2.3.0
     hooks:
       - id: codespell
-        stages: [commit, commit-msg]
+        stages: [pre-commit, commit-msg]
         exclude_types: [csv, svg, html, yaml, jupyter]
         args: [--ignore-words-list, "hist,mape,te,nd,fpr", --check-filenames]
 

pymatviz/__init__.py

@@ -14,10 +14,7 @@
 import builtins
 from importlib.metadata import PackageNotFoundError, version
 
-import matplotlib.pyplot as plt
 import plotly.express as px
-import plotly.graph_objects as go
-import plotly.io as pio
 
 from pymatviz import (
     bar,

pymatviz/coordination.py

@@ -1,23 +1,31 @@
 """Visualizations of coordination numbers distributions."""
 
+from __future__ import annotations
+
 import math
 from collections import Counter
-from collections.abc import Callable, Sequence
+from collections.abc import Sequence
 from inspect import isclass
-from typing import Any, Literal
+from typing import TYPE_CHECKING
 
 import numpy as np
 import plotly.graph_objects as go
 from plotly.colors import label_rgb
 from plotly.subplots import make_subplots
 from pymatgen.analysis.local_env import NearNeighbors
-from pymatgen.core import PeriodicSite, Structure
 
 from pymatviz.colors import ELEM_COLORS_JMOL, ELEM_COLORS_VESTA
 from pymatviz.enums import ElemColorScheme, LabelEnum
 from pymatviz.utils import normalize_to_dict
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import Any, Literal
+
+    from pymatgen.core import PeriodicSite, Structure
+
+
 class SplitMode(LabelEnum):
     """How to split the coordination number histogram into subplots."""
 
@@ -57,10 +65,12 @@ def normalize_get_neighbors(
     # Prepare the neighbor-finding strategy
     if isinstance(strategy, int | float):
         return lambda site, structure: structure.get_neighbors(site, strategy)
+
     if isinstance(strategy, NearNeighbors):
         return lambda site, structure: strategy.get_nn_info(
             structure, structure.index(site)
         )
+
     if isclass(strategy) and issubclass(strategy, NearNeighbors):
         nn_instance = strategy()
         return lambda site, structure: nn_instance.get_nn_info(
@@ -417,6 +427,7 @@ def coordination_vs_cutoff_line(
         and {*map(type, strategy)} <= {int, float}
     ):
         cutoff_range = strategy
+
     elif isinstance(strategy, NearNeighbors) or (
         isclass(strategy) and issubclass(strategy, NearNeighbors)
     ):
@@ -428,6 +439,7 @@ def coordination_vs_cutoff_line(
         else:
             raise AttributeError(f"Could not determine cutoff for {nn_instance=}")
         cutoff_range = (0, max_cutoff)
+
     else:
         raise TypeError(
             f"Invalid {strategy=}. Expected float, tuple of floats, NearNeighbors "

pymatviz/process_data.py

@@ -65,6 +65,7 @@ def count_elements(
 
     if is_numeric_dtype(srs):
         pass
+
     elif is_string_dtype(srs) or {*map(type, srs)} <= {str, Composition}:
         # all items are formula strings or Composition objects
         if count_mode == "occurrence":

pymatviz/ptable/ptable_plotly.py

@@ -8,7 +8,6 @@
 import numpy as np
 import pandas as pd
 import plotly.express as px
-import plotly.figure_factory as ff
 
 from pymatviz.enums import ElemCountMode
 from pymatviz.process_data import count_elements
@@ -272,6 +271,8 @@ def ptable_heatmap_plotly(
     zmax = max(non_nan_values) if cscale_range[1] is None else cscale_range[1]
     car_multiplier = 100 if heat_mode == "percent" else 1
 
+    import plotly.figure_factory as ff  # slow import
+
     fig = ff.create_annotated_heatmap(
         car_multiplier * heatmap_values,
         annotation_text=tile_texts,

pymatviz/structure_viz/helpers.py

@@ -6,13 +6,10 @@
 import itertools
 import math
 import warnings
-from collections.abc import Callable, Sequence
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
-import plotly.graph_objects as go
-from pymatgen.analysis.local_env import NearNeighbors
 from pymatgen.core import Composition, Lattice, PeriodicSite, Species, Structure
 
 from pymatviz.colors import ELEM_COLORS_JMOL, ELEM_COLORS_VESTA

pymatviz/structure_viz/mpl.py

@@ -8,9 +8,8 @@
 
 import math
 import warnings
-from collections.abc import Callable, Sequence
 from itertools import product
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING
 
 import matplotlib.pyplot as plt
 import numpy as np

pymatviz/structure_viz/plotly.py

@@ -4,11 +4,9 @@
 
 import math
 import warnings
-from collections.abc import Callable, Sequence
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING
 
 import numpy as np
-import plotly.graph_objects as go
 from plotly.subplots import make_subplots
 from pymatgen.analysis.local_env import CrystalNN, NearNeighbors
 from pymatgen.symmetry.analyzer import SpacegroupAnalyzer

pymatviz/utils.py

@@ -701,6 +701,7 @@ def normalize_to_dict(
     """
     if isinstance(inputs, cls):
         return {"": inputs}
+
     if (
         isinstance(inputs, list | tuple)
         and all(isinstance(obj, cls) for obj in inputs)

pyproject.toml

@@ -33,6 +33,9 @@ dependencies = [
     "pandas[output-formatting]>=2.2",
     "plotly>=5.23",
     "pymatgen>=2024.7.18",
+    # TODO: pmv doesn't actually depend on monty, however latest monty
+    # includes a critical import patch, remove this after pmg bump dep
+    "monty>=2024.10.21",
     "scikit-learn>=1.5",
     "scipy>=1.14",
 ]

tests/performance_tests/test_import_time.py

@@ -0,0 +1,104 @@
+"""
+Test import time of core modules to avoid regression.
+"""
+
+# ruff: noqa: T201 (check for print statement)
+
+from __future__ import annotations
+
+import os
+import subprocess
+import time
+import warnings
+
+import pytest
+
+
+GEN_REF_TIME = False  # switch for generating reference time
+
+# Last update: 2024-10-23
+REF_IMPORT_TIME: dict[str, float] = {
+    "pymatviz": 4085.73,
+    "pymatviz.coordination": 4135.77,
+    "pymatviz.cumulative": 4108.06,
+    "pymatviz.histogram": 4110.41,
+    "pymatviz.phonons": 4109.97,
+    "pymatviz.powerups": 4066.31,
+    "pymatviz.ptable": 4092.35,
+    "pymatviz.rainclouds": 4098.33,
+    "pymatviz.rdf": 4144.26,
+    "pymatviz.relevance": 4126.54,
+    "pymatviz.sankey": 4135.17,
+    "pymatviz.scatter": 4087.62,
+    "pymatviz.structure_viz": 4105.33,
+    "pymatviz.sunburst": 4133.78,
+    "pymatviz.uncertainty": 4179.99,
+    "pymatviz.xrd": 4156.52,
+}
+
+
+@pytest.mark.skipif(
+    not GEN_REF_TIME, reason="Set GEN_REF_TIME to generate reference import time."
+)
+def test_get_ref_import_time() -> None:
+    """A dummy test that would always fail, used to generate copyable reference time."""
+    import_times = {
+        module_name: round(measure_import_time(module_name), 2)
+        for module_name in REF_IMPORT_TIME
+    }
+
+    # Print out the import times in a copyable format
+    print("\nCopyable import time dictionary:")
+    print(import_times)
+
+    pytest.fail("Generate reference import times.")
+
+
+def measure_import_time(module_name: str, repeats: int = 3) -> float:
+    """Measure import time of a module in milliseconds across several runs.
+
+    Args:
+        module_name (str): name of the module to test.
+        count (int): Number of runs to average.
+
+    Returns:
+        float: import time in milliseconds.
+    """
+    total_time = 0.0
+
+    for _ in range(repeats):
+        start_time = time.perf_counter()
+        subprocess.run(["python", "-c", f"import {module_name}"], check=True)  # noqa: S603, S607
+        total_time += time.perf_counter() - start_time
+
+    return total_time / repeats * 1000
+
+
+@pytest.mark.skipif(
+    os.getenv("GITHUB_REF") != "refs/heads/main", reason="Only run on the main branch"
+)
+@pytest.mark.skipif(GEN_REF_TIME, reason="Generating reference import time.")
+def test_import_time(grace_percent: float = 0.20, hard_percent: float = 0.50) -> None:
+    """Test the import time of core modules to avoid regression in performance.
+
+    Args:
+        grace_percentage (float): Maximum allowed percentage increase in import time
+            before a warning is raised.
+        hard_percentage (float): Maximum allowed percentage increase in import time
+            before the test fails.
+    """
+    for module_name, ref_time in REF_IMPORT_TIME.items():
+        current_time = measure_import_time(module_name)
+
+        # Calculate grace and hard thresholds
+        grace_threshold = ref_time * (1 + grace_percent)
+        hard_threshold = ref_time * (1 + hard_percent)
+
+        if current_time > grace_threshold:
+            if current_time > hard_threshold:
+                pytest.fail(f"{module_name} import too slow! {hard_threshold=:.2f} ms")
+            else:
+                warnings.warn(
+                    f"{module_name} import slightly slower: {grace_threshold=:.2f} ms",
+                    stacklevel=2,
+                )

tests/test_coordination.py

@@ -1,6 +1,7 @@
+from __future__ import annotations
+
 import re
-from collections.abc import Sequence
-from typing import Any
+from typing import TYPE_CHECKING
 
 import pytest
 from pymatgen.analysis.local_env import CrystalNN, NearNeighbors, VoronoiNN
@@ -15,6 +16,11 @@
 )
 
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from typing import Any
+
+
 def test_coordination_hist_single_structure(structures: Sequence[Structure]) -> None:
     """Test coordination_hist with a single structure."""
     fig = coordination_hist(structures[0])