improve deps and imports

mlip_arena/models/__init__.py

@@ -42,7 +42,7 @@ for model, metadata in REGISTRY.items():
             f"{__package__}.{metadata['module']}.{metadata['family']}"
         )
         MLIPMap[model] = getattr(module, metadata["class"])
-    except (ModuleNotFoundError, AttributeError, ValueError, ImportError) as e:
+    except (ModuleNotFoundError, AttributeError, ValueError, ImportError, Exception) as e:
         logger.warning(e)
         continue
 

mlip_arena/tasks/__init__.py

@@ -6,6 +6,13 @@ from huggingface_hub import HfApi, HfFileSystem, hf_hub_download
 # from mlip_arena.models import MLIP
 # from mlip_arena.models import REGISTRY as MODEL_REGISTRY
 
+try:
+    from prefect.logging import get_run_logger
+
+    logger = get_run_logger()
+except (ImportError, RuntimeError):
+    from loguru import logger
+
 try:
     from .elasticity import run as ELASTICITY
     from .eos import run as EOS
@@ -16,8 +23,9 @@ try:
     from .phonon import run as PHONON
 
     __all__ = ["OPT", "EOS", "MD", "NEB", "NEB_FROM_ENDPOINTS", "ELASTICITY", "PHONON"]
-except ImportError:
-    pass
+except (ImportError, TypeError) as e:
+    logger.warning(e)
+
 
 with open(Path(__file__).parent / "registry.yaml", encoding="utf-8") as f:
     REGISTRY = yaml.safe_load(f)

mlip_arena/tasks/utils.py

@@ -8,7 +8,6 @@ import torch
 from ase import units
 from ase.calculators.calculator import BaseCalculator
 from ase.calculators.mixing import SumCalculator
-from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
 
 from mlip_arena.models import MLIPEnum
 
@@ -102,6 +101,13 @@ def get_calculator(
     dispersion_kwargs.update({"device": device})
 
     if dispersion:
+        try:
+            from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
+        except ImportError as e:
+            raise ImportError(
+                "torch_dftd is required for dispersion but is not installed."
+            ) from e
+
         disp_calc = TorchDFTD3Calculator(
             **dispersion_kwargs,
         )
@@ -112,5 +118,5 @@ def get_calculator(
         if dispersion_kwargs:
             logger.info(pformat(dispersion_kwargs))
 
-    assert isinstance(calc, BaseCalculator)    
+    assert isinstance(calc, BaseCalculator)
     return calc

pyproject.toml

@@ -10,7 +10,7 @@ authors=[
 ]
 description="Fair and transparent benchmark of machine learning interatomic potentials (MLIPs), beyond error-based regression metrics"
 readme=".github/README.md"
-requires-python=">=3.10"
+requires-python=">=3.9"
 keywords=[
     "pytorch",
     "machine-learning-interatomic-potentials",
@@ -27,6 +27,7 @@ classifiers=[
     "Programming Language :: Python :: 3 :: Only",
 ]
 dependencies=[
+    "loguru",
     "ase",
     "pymatgen",
     "torch",