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

[lwalew]

Why

mlipaudit -h takes ~6s. The cost is entirely import time, not argument parsing: main.py imports mlipaudit.benchmarks (needed to build the --benchmarks choices for the parser), and that transitively pulls in the full ML/plotting stack just to print a help string:

Chain	~cost	pulled in via
mlip.models (ESEN network, flax)	~2.2s	Benchmark base / ForceField type hints
mlip.simulation.ase → jax_md.rigid_body	~1.5s	utils.simulation engine
sklearn.metrics	~0.7s	conformer / dihedral / water benchmarks
bare import jax	~0.5s	utils.stability

None of this is needed to list benchmarks or print help — only to actually run a benchmark.

What

Defer the heavy imports to call time:

• benchmark.py, utils/{inference,simulation,stability,unallowed_elements}.py: move mlip.models / mlip.simulation.ase (JAX-MD) / jax / sklearn imports into the functions that use them; type-only references go under TYPE_CHECKING with from __future__ import annotations.
• utils/_ase_engine.py (new): ASESimulationEngineWithCalculator subclasses the heavy ASESimulationEngine, so a class definition can't be lazy. Moved it to its own module imported lazily when a simulation runs; the historical mlipaudit.utils[.simulation].ASESimulationEngineWithCalculator paths still work via module-level __getattr__.
• nudged_elastic_band, conformer_selection, dihedral_scan, water_radial_distribution: defer their mlip / sklearn imports into the run/analyze methods.
• main.py: import launch_app / run_benchmarks inside their command branches.

Tests

The benchmark tests patched the now-lazy symbols at the use-site module (e.g. mlipaudit.utils.simulation.JaxMDSimulationEngine), which no longer exists there. Updated them to patch at the definition site (e.g. mlip.simulation.jax_md.JaxMDSimulationEngine) — the lazy from ... import picks the patch up at call time, and create_autospec in conftest (which binds the real class at import) is unaffected.

Result

• mlipaudit -h: ~5.6s → ~0.9s
• import mlipaudit.benchmarks: ~5.0s → ~1.1s
• All 128 tests pass; ruff / ruff-format / mypy clean.

Notes

• mlipaudit benchmark -h has a pre-existing argparse AssertionError (unrelated to this change; reproduces on main) and is left as-is.

🤖 Generated with Claude Code

[github-actions]

Coverage Report

File	Stmts	Miss	Cover	Missing
src/mlipaudit
__init__.py	2	0	100%
benchmark.py	78	7	91%	154, 173, 239, 245–246, 256, 270
exceptions.py	2	0	100%
io.py	100	25	75%	117, 121, 146–150, 152, 154, 171–172, 174–175, 192–193, 195, 213–221
io_helpers.py	52	0	100%
run_mode.py	5	0	100%
scoring.py	29	7	75%	101–102, 104, 107, 111–112, 114
src/mlipaudit/benchmarks
__init__.py	29	0	100%
src/mlipaudit/benchmarks/bond_length_distribution
bond_length_distribution.py	95	3	96%	208, 211, 276
src/mlipaudit/benchmarks/conformer_selection
conformer_selection.py	132	5	96%	253, 256, 304, 309, 349
src/mlipaudit/benchmarks/dihedral_scan
dihedral_scan.py	116	5	95%	221, 224, 258, 261, 303
src/mlipaudit/benchmarks/folding_stability
folding_stability.py	106	4	96%	212, 270, 275, 339
helpers.py	35	0	100%
src/mlipaudit/benchmarks/noncovalent_interactions
noncovalent_interactions.py	168	10	94%	245, 248, 272, 283, 407, 424, 441, 525–527
src/mlipaudit/benchmarks/nudged_elastic_band
engine.py	126	94	25%	62, 81–82, 84–89, 91, 93, 95–96, 98–99, 101–103, 112–113, 115–116, 120–121, 123, 130, 132, 134–136, 138–139, 141, 143–145, 147–151, 153, 159–163, 165–167, 169–171, 173–175, 177–180, 182, 185, 189–190, 192, 194–195, 197, 201–202, 205–210, 212–213, 229, 234, 236–238, 240, 242–243, 246–247, 253, 256–257, 261–262, 267
nudged_elastic_band.py	124	6	95%	302–303, 310, 322, 328, 335
src/mlipaudit/benchmarks/reactivity
reactivity.py	117	3	97%	239–240, 272
src/mlipaudit/benchmarks/reference_geometry_stability
reference_geometry_stability.py	122	3	97%	233, 305, 328
src/mlipaudit/benchmarks/ring_planarity
ring_planarity.py	99	6	93%	229, 232, 264, 267–268, 287
src/mlipaudit/benchmarks/sampling
helpers.py	64	4	93%	46, 67, 97, 129
sampling.py	225	10	95%	303, 344, 385, 388–390, 440, 707, 825–826
src/mlipaudit/benchmarks/scaling
scaling.py	103	34	66%	128–130, 138–139, 141–144, 154–156, 193–198, 201–203, 209–210, 212–213, 215–216, 219–220, 222, 243, 252, 271, 284
src/mlipaudit/benchmarks/solvent_radial_distribution
solvent_radial_distribution.py	101	8	92%	189–190, 192, 231, 236–237, 302, 325
src/mlipaudit/benchmarks/stability
stability.py	150	18	88%	166, 169, 201–202, 205–206, 208, 210, 213, 215, 431–432, 442, 478, 487, 522, 537, 548
src/mlipaudit/benchmarks/tautomers
tautomers.py	92	5	94%	182, 208, 211–212, 231
src/mlipaudit/benchmarks/water_radial_distribution
water_radial_distribution.py	94	2	97%	190, 272
src/mlipaudit/ui
__init__.py	16	0	100%
bond_length_distribution.py	58	11	81%	89–90, 95–96, 141–142, 147, 150, 176, 188, 192
conformer_selection.py	100	7	93%	133–134, 139–140, 195, 320, 324
dihedral_scan.py	115	12	89%	141–142, 147–148, 259–260, 302, 313, 353, 355, 365, 369
folding_stability.py	96	7	92%	44, 226–227, 232–233, 406, 410
leaderboard.py	80	10	87%	67, 74, 78, 85, 287–289, 291, 293, 303
noncovalent_interactions.py	152	9	94%	205–206, 211–212, 392, 397, 432, 454, 458
nudged_elastic_band.py	55	39	29%	42–46, 50–51, 53, 66, 68, 75, 82, 89–90, 93, 97–99, 101, 103–105, 107–109, 111–113, 115, 117, 119, 121–122, 124, 135, 137, 149, 159, 163
page_wrapper.py	14	2	85%	36, 46
reactivity.py	59	6	89%	96–97, 102–103, 173, 177
reference_geometry_stability.py	66	9	86%	62, 131–132, 137–138, 199, 226, 236, 240
ring_planarity.py	56	9	83%	71–72, 77–78, 123–124, 160, 170, 174
sampling.py	87	9	89%	85, 129–130, 135–136, 166, 234, 265, 269
scaling.py	52	6	88%	135–136, 147–148, 162, 166
solvent_radial_distribution.py	83	20	75%	117–118, 123–124, 150, 156, 162–166, 168–169, 175, 178, 190, 197, 199, 211, 215
stability.py	60	10	83%	106–107, 112–113, 134–135, 138, 148, 177, 181
tautomers.py	62	6	90%	82–83, 88–89, 194, 198
utils.py	106	19	82%	69, 73–74, 130, 192–195, 201, 251, 269–270, 273–274, 277, 287–289, 291
water_radial_distribution.py	98	6	93%	132–133, 138–139, 279, 283
src/mlipaudit/utils
__init__.py	9	4	55%	28–29, 33–34
_ase_engine.py	27	0	100%
inference.py	29	7	75%	66–68, 77–78, 81, 84
simulation.py	42	6	85%	45, 49, 122, 160–162
stability.py	31	5	83%	60, 67, 91, 102, 106
trajectory_helpers.py	34	0	100%
unallowed_elements.py	11	0	100%
TOTAL	3964	478	87%

Tests	Skipped	Failures	Errors	Time
128	0 💤	0 ❌	0 🔥	13.632s ⏱️