perf: lazy-load heavy deps to speed up CLI startup

Importing mlipaudit.benchmarks (which the CLI does to build the argument
parser) eagerly pulled in mlip.models, JAX/JAX-MD and scikit-learn, making
even `mlipaudit -h` take ~6s. These libraries are only needed when a
benchmark is actually run.

Defer those imports to call time:
- benchmark.py / utils.{inference,simulation,stability,unallowed_elements}:
  move mlip.models / mlip.simulation.ase (JAX-MD) / jax / sklearn imports
  into the functions that use them; type-only refs go under TYPE_CHECKING
  with `from __future__ import annotations`.
- Move ASESimulationEngineWithCalculator into utils/_ase_engine.py so the
  heavy ASE-engine base class is only imported when a simulation runs; keep
  the historical import paths working via module-level __getattr__.
- nudged_elastic_band / conformer_selection / dihedral_scan /
  water_radial_distribution: defer their mlip / sklearn imports.
- main.py: import launch_app / run_benchmarks inside their command branches.

Tests patched the now-lazy symbols at the use-site module; update them to
patch at the definition site (e.g. mlip.simulation.jax_md.JaxMDSimulationEngine),
which the lazy `from ... import` picks up at call time.

`mlipaudit -h` drops from ~5.6s to ~0.9s; importing mlipaudit.benchmarks
drops from ~5.0s to ~1.1s. All 128 tests pass.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: lwalew

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## Unreleased
+
+- Speed up CLI startup (e.g. `mlipaudit -h`) from ~6s to ~1s by importing the
+  heavy `mlip`/JAX/JAX-MD/scikit-learn dependencies lazily, only when a benchmark
+  is actually run, rather than at import time of the benchmark modules.
+
 ## Release 0.1.3
 
 - Populate `atoms.info["charge"]` and `atoms.info["spin"]` on every benchmark's

src/mlipaudit/benchmark.py

@@ -12,21 +12,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import os
 import zipfile
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Any, Literal, TypeAlias
+from typing import TYPE_CHECKING, Any, Literal, TypeAlias
 
 from ase import Atom
 from ase.calculators.calculator import Calculator as ASECalculator
 from huggingface_hub import hf_hub_download
-from mlip.models import ForceField
 from pydantic import BaseModel, Field
 
 from mlipaudit.exceptions import ChemicalElementsMissingError
 from mlipaudit.run_mode import RunMode
 
+if TYPE_CHECKING:
+    # `mlip.models` pulls in the full model/JAX stack (~2s); importing it lazily
+    # keeps benchmark classes cheap to import (e.g. for `mlipaudit -h`).
+    from mlip.models import ForceField
+
 RunModeAsString: TypeAlias = Literal["dev", "fast", "standard"]
 
 #: Default total charge applied to `atoms.info["charge"]` when a benchmark's
@@ -139,6 +145,8 @@ def __init__(
 
         self.force_field = force_field
 
+        from mlip.models import ForceField  # noqa: PLC0415
+
         if not (
             isinstance(self.force_field, ForceField)
             or isinstance(self.force_field, ASECalculator)

src/mlipaudit/benchmarks/conformer_selection/conformer_selection.py

@@ -12,19 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import functools
 import logging
 import os
 import statistics
-from typing import Literal, TypeAlias
+from typing import TYPE_CHECKING, Literal, TypeAlias
 
 import numpy as np
 from ase import Atoms, units
 from ase.calculators.calculator import Calculator as ASECalculator
-from mlip.models import ForceField
 from pydantic import BaseModel, Field, NonNegativeFloat, TypeAdapter
 from scipy.stats import spearmanr
-from sklearn.metrics import mean_absolute_error, root_mean_squared_error
 
 from mlipaudit.benchmark import (
     DEFAULT_CHARGE,
@@ -38,6 +38,10 @@
 from mlipaudit.scoring import compute_benchmark_score
 from mlipaudit.utils import run_inference
 
+if TYPE_CHECKING:
+    # `mlip.models` pulls in the heavy model/JAX stack; only needed for typing here.
+    from mlip.models import ForceField
+
 DatasetName: TypeAlias = Literal["wiggle150", "folmsbee"]
 
 logger = logging.getLogger("mlipaudit")
@@ -279,6 +283,11 @@ def analyze(self) -> ConformerSelectionResult:
         Raises:
             RuntimeError: If called before `run_model()`.
         """
+        from sklearn.metrics import (  # noqa: PLC0415
+            mean_absolute_error,
+            root_mean_squared_error,
+        )
+
         if self.model_output is None:
             raise RuntimeError("Must call run_model() first.")
 

src/mlipaudit/benchmarks/dihedral_scan/dihedral_scan.py

@@ -21,7 +21,6 @@
 from ase import Atoms, units
 from pydantic import BaseModel, Field, NonNegativeFloat, TypeAdapter
 from scipy.stats import pearsonr
-from sklearn.metrics import mean_absolute_error, root_mean_squared_error
 
 from mlipaudit.benchmark import (
     DEFAULT_CHARGE,
@@ -331,6 +330,11 @@ def _compute_metrics(
         ref_energy_profile: np.ndarray,
         predicted_energy_profile: np.ndarray,
     ) -> tuple[float, float, float, float, float]:
+        from sklearn.metrics import (  # noqa: PLC0415
+            mean_absolute_error,
+            root_mean_squared_error,
+        )
+
         mae = mean_absolute_error(ref_energy_profile, predicted_energy_profile)
         rmse = root_mean_squared_error(ref_energy_profile, predicted_energy_profile)
 

src/mlipaudit/benchmarks/nudged_elastic_band/nudged_elastic_band.py

@@ -12,14 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import functools
 import logging
+from typing import TYPE_CHECKING
 
 import numpy as np
 from ase import Atoms
 from mlip.simulation import SimulationState
-from mlip.simulation.ase import ASESimulationEngine
-from mlip.simulation.configs import ASESimulationConfig
 from pydantic import BaseModel, ConfigDict, TypeAdapter
 
 from mlipaudit.benchmark import (
@@ -29,12 +30,16 @@
     BenchmarkResult,
     ModelOutput,
 )
-from mlipaudit.benchmarks.nudged_elastic_band.engine import (
-    NEBSimulationConfig,
-    NEBSimulationEngine,
-)
 from mlipaudit.run_mode import RunMode
 
+if TYPE_CHECKING:
+    # Used only in type annotations. The matching runtime imports happen lazily inside
+    # the run methods so that importing this benchmark module (e.g. for the CLI
+    # benchmark listing) does not pull in the heavy mlip/JAX-MD stack.
+    from mlip.simulation.configs import ASESimulationConfig
+
+    from mlipaudit.benchmarks.nudged_elastic_band.engine import NEBSimulationConfig
+
 logger = logging.getLogger("mlipaudit")
 
 NEB_DATASET_FILENAME = "grambow_dataset_neb.json"
@@ -230,6 +235,13 @@ class NudgedElasticBandBenchmark(Benchmark):
 
     def run_model(self) -> None:
         """Run the NEB calculation."""
+        # Imported lazily (pulls in the heavy mlip/JAX-MD stack via the NEB engine).
+        from mlip.simulation.configs import ASESimulationConfig  # noqa: PLC0415
+
+        from mlipaudit.benchmarks.nudged_elastic_band.engine import (  # noqa: PLC0415
+            NEBSimulationConfig,
+        )
+
         self.model_output = NEBModelOutput(
             simulation_states=[],
         )
@@ -364,6 +376,8 @@ def _run_minimization(
             atoms_initial_em: The initial atoms after energy minimization.
             atoms_final_em: The final atoms after energy minimization.
         """
+        from mlip.simulation.ase import ASESimulationEngine  # noqa: PLC0415
+
         em_engine_initial = ASESimulationEngine(initial_atoms, ff, em_config)
         em_engine_initial.run()
 
@@ -398,6 +412,10 @@ def _run_neb(
         Returns:
             neb_engine_climb: The nudged elastic band engine with climbing image method.
         """
+        from mlipaudit.benchmarks.nudged_elastic_band.engine import (  # noqa: PLC0415
+            NEBSimulationEngine,
+        )
+
         neb_engine = NEBSimulationEngine(
             initial_atoms, final_atoms, ff, neb_config, transition_state=ts_atoms
         )

src/mlipaudit/benchmarks/water_radial_distribution/water_radial_distribution.py

@@ -22,7 +22,6 @@
 from ase.io import read as ase_read
 from mlip.simulation import SimulationState
 from pydantic import ConfigDict, NonNegativeFloat
-from sklearn.metrics import mean_absolute_error, root_mean_squared_error
 
 from mlipaudit.benchmark import (
     DEFAULT_CHARGE,
@@ -177,6 +176,11 @@ def analyze(self) -> WaterRadialDistributionResult:
         Raises:
             RuntimeError: If called before `run_model()`.
         """
+        from sklearn.metrics import (  # noqa: PLC0415
+            mean_absolute_error,
+            root_mean_squared_error,
+        )
+
         if self.model_output is None:
             raise RuntimeError("Must call run_model() first.")
 

src/mlipaudit/main.py

@@ -17,13 +17,11 @@
 from argparse import ArgumentParser, Namespace, RawDescriptionHelpFormatter
 
 import mlipaudit
-from mlipaudit.app import launch_app
 from mlipaudit.benchmark import Benchmark
 from mlipaudit.benchmarks import (
     BENCHMARK_NAMES,
     BENCHMARKS,
 )
-from mlipaudit.benchmarks_cli import run_benchmarks
 from mlipaudit.run_mode import RunMode
 
 logger = logging.getLogger("mlipaudit")
@@ -202,6 +200,8 @@ def main():
         mlip_logger.setLevel(logging.WARNING)
 
     if args.command == "benchmark":
+        from mlipaudit.benchmarks_cli import run_benchmarks  # noqa: PLC0415
+
         benchmarks_to_run = _get_benchmarks_to_run(args)
         run_benchmarks(
             model_paths=args.models,
@@ -213,6 +213,8 @@ def main():
             log_timings=args.log_timings,
         )
     elif args.command == "gui":
+        from mlipaudit.app import launch_app  # noqa: PLC0415
+
         launch_app(args.results_dir, args.is_public)
     else:
         parser.print_help()

src/mlipaudit/utils/__init__.py

@@ -13,12 +13,22 @@
 # limitations under the License.
 
 from mlipaudit.utils.inference import run_inference
-from mlipaudit.utils.simulation import (
-    ASESimulationEngineWithCalculator,
-    run_simulation,
-)
+from mlipaudit.utils.simulation import run_simulation
 from mlipaudit.utils.trajectory_helpers import (
     create_ase_trajectory_from_simulation_state,
     create_mdtraj_trajectory_from_simulation_state,
 )
 from mlipaudit.utils.unallowed_elements import skip_unallowed_elements
+
+
+def __getattr__(name: str):
+    # `ASESimulationEngineWithCalculator` is built lazily (it derives from a heavy
+    # mlip base class). Re-export it via PEP 562 so importing `mlipaudit.utils` does
+    # not pull in the JAX-MD stack unless the class is actually used.
+    if name == "ASESimulationEngineWithCalculator":
+        from mlipaudit.utils._ase_engine import (  # noqa: PLC0415
+            ASESimulationEngineWithCalculator,
+        )
+
+        return ASESimulationEngineWithCalculator
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/mlipaudit/utils/_ase_engine.py

@@ -0,0 +1,79 @@
+# Copyright 2025 InstaDeep Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Custom ASE simulation engine.
+
+This module is kept separate from :mod:`mlipaudit.utils.simulation` because
+importing it pulls in the heavy ``mlip.simulation.ase`` (and therefore JAX-MD)
+stack. It is only imported lazily, when a simulation is actually run, so that
+merely importing the benchmark modules (e.g. for the CLI) stays fast.
+"""
+
+import logging
+from typing import Callable
+
+import ase
+from ase.calculators.calculator import Calculator as ASECalculator
+from mlip.simulation import SimulationState
+from mlip.simulation.ase import ASESimulationEngine
+from mlip.simulation.configs import ASESimulationConfig
+from mlip.simulation.enums import SimulationType
+from mlip.simulation.temperature_scheduling import get_temperature_schedule
+
+logger = logging.getLogger("mlipaudit")
+
+
+class ASESimulationEngineWithCalculator(ASESimulationEngine):
+    """Class derived from mlip's ASE simulation engine but allowing for a passed
+    ASE calculator object.
+    """
+
+    def __init__(
+        self,
+        atoms: ase.Atoms,
+        ase_calculator: ASECalculator,
+        config: ASESimulationConfig,
+    ) -> None:
+        """Overridden constructor that takes in an ASE calculator instead of an
+        mlip force field class.
+
+        Args:
+            atoms: The ASE atoms.
+            ase_calculator: The ASE calculator to use in the simulation.
+            config: The simulation config.
+        """
+        self.state = SimulationState()
+        self.loggers: list[Callable[[SimulationState], None]] = []
+
+        logger.debug("Initialization of simulation begins...")
+        self._config = config
+        self.atoms = atoms
+        self.atoms.center()
+        positions = atoms.get_positions()
+        self._num_atoms = positions.shape[0]
+        self.state.atomic_numbers = atoms.numbers
+
+        self._init_box()
+
+        self.is_md_simulation = self._config.simulation_type == SimulationType.MD
+        self.is_npt_simulation = (
+            self.is_md_simulation and self._config.md_integrator.ensemble == "npt"
+        )
+
+        self.model_calculator = ase_calculator
+
+        self._temperature_schedule = get_temperature_schedule(
+            self._config.temperature_schedule_config, self._config.num_steps
+        )
+
+        logger.debug("Initialization of simulation completed.")

src/mlipaudit/utils/inference.py

@@ -11,13 +11,19 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
 
 import ase
 from ase.calculators.calculator import Calculator as ASECalculator
-from mlip.inference import run_batched_inference
-from mlip.models import ForceField
-from mlip.typing import Prediction
+
+if TYPE_CHECKING:
+    # These pull in the heavy mlip/JAX stack; imported lazily inside `run_inference`
+    # so that merely importing this module stays cheap.
+    from mlip.models import ForceField
+    from mlip.typing import Prediction
 
 logger = logging.getLogger("mlipaudit")
 
@@ -45,6 +51,12 @@ def run_inference(
     Raises:
         ValueError: If force field type is not compatible.
     """
+    # Imported lazily to keep importing this module cheap (these pull in the heavy
+    # mlip/JAX stack, only needed when inference is actually run).
+    from mlip.inference import run_batched_inference  # noqa: PLC0415
+    from mlip.models import ForceField  # noqa: PLC0415
+    from mlip.typing import Prediction  # noqa: PLC0415
+
     if isinstance(force_field, ForceField):
         try:
             predictions = run_batched_inference(