Refactor CI test environments

.github/workflows/ci.yaml

@@ -13,33 +13,31 @@ jobs:
   lint:
     runs-on: ubuntu-latest
     steps:
-    - name: Checkout PR with full history
+    - name: Checkout code
       uses: actions/checkout@v4
-      with:
-        fetch-depth: 0
 
     - name: Install uv
       uses: astral-sh/setup-uv@v6
       with:
         enable-cache: true
 
     - name: Set up Python
-      uses: actions/setup-python@v5
-      with:
-        python-version: "3.12"
+      run: uv python install 3.12
 
-    - name: Run pre-commit
-      run: uvx pre-commit run --all-files --show-diff-on-failure
-      env:
-        PRE_COMMIT_USE_UV: 1
+    - name: Run Ruff (Lint & Format)
+      # This assumes ruff is in your pyproject.toml or you use uvx
+      run: |
+        uvx ruff check .
+        uvx ruff format --check .
 
   test:
     needs: lint
     runs-on: ubuntu-latest
 
     strategy:
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.11", "3.12", "3.13"]
+        test-group: ["default", "nequip", "sevennet", "mace", "fairchem"]
 
     steps:
     - name: Checkout PR with full history
@@ -65,7 +63,7 @@ jobs:
        swap-size-gb: 12
 
     - name: Install dependencies
-      run: bash scripts/install-linux.sh
+      run: bash scripts/install-linux.sh ${{ matrix.test-group }}
 
     - name: List dependencies
       run: pip list
@@ -80,7 +78,12 @@ jobs:
       env:
         PREFECT_API_KEY: ${{ secrets.PREFECT_API_KEY }}
         PREFECT_API_URL: ${{ secrets.PREFECT_API_URL }}
-      run: pytest -vra -n 5 --dist=loadscope --cov=mlip_arena --cov-report=xml tests
+      run: |
+        if [ "${{ matrix.test-group }}" == "default" ]; then
+          pytest -vra --cov=mlip_arena --cov-report=xml tests -m "not (nequip or sevennet or mace or fairchem)"
+        else
+          pytest -vra --cov=mlip_arena --cov-report=xml tests -m ${{ matrix.test-group }}
+        fi
 
     - name: Upload coverage reports to Codecov
       uses: codecov/codecov-action@v4

benchmarks/asymptotes.py

@@ -24,6 +24,8 @@
 from mlip_arena.flows.eos_bulk import run_db as EOSFlow
 from mlip_arena.flows.ev import run_db as EVFlow
 
+from mlip_arena.models import REGISTRY, MLIPEnum
+
 
 @flow
 def asymptotic_behaviors(calculator: str | BaseCalculator, calculator_kwargs: dict | None = None):
@@ -32,32 +34,33 @@ def asymptotic_behaviors(calculator: str | BaseCalculator, calculator_kwargs: di
 
     if isinstance(calculator, BaseCalculator):
         model_name = calculator.__class__.__name__
-    else:
+    elif isinstance(calculator, str) and hasattr(MLIPEnum, calculator):
         model_name = calculator
+    else:
+        raise ValueError(f"Unsupported model: {calculator}")
 
-    from mlip_arena.tasks.utils import get_calculator
-
-    calc_instance = get_calculator(calculator, calculator_kwargs)
+    family = REGISTRY[model_name]["family"] if hasattr(MLIPEnum, model_name) else "custom"
 
     # 1. Diatomics
     print(f"Starting homonuclear diatomics benchmark for {model_name}...")
-    out_dir_diatomics = Path(__file__).parent / "diatomics"
+
+    out_dir_diatomics = Path(__file__).parent / "diatomics" / family / model_name
     homonuclear_diatomics.with_options(name=f"diatomics-{model_name}", task_runner=parent_task_runner)(
-        model=calc_instance, run_dir=out_dir_diatomics
+        calculator=calculator, calculator_kwargs=calculator_kwargs, run_dir=out_dir_diatomics
     )
 
     # 2. EOS Bulk
     print(f"Starting EOS bulk benchmark for {model_name}...")
     out_dir_eos = Path(__file__).parent / "eos_bulk"
     EOSFlow.with_options(name=f"eos_bulk-{model_name}", task_runner=parent_task_runner)(
-        model=calc_instance, run_dir=out_dir_eos, dataset_file="wbm_subset.db"
+        calculator=calculator, calculator_kwargs=calculator_kwargs, run_dir=out_dir_eos, dataset_file="wbm_subset.db"
     )
 
     # 3. E-V
     print(f"Starting E-V scan benchmark for {model_name}...")
     out_dir_ev = Path(__file__).parent / "ev"
     EVFlow.with_options(name=f"ev-{model_name}", task_runner=parent_task_runner)(
-        model=calc_instance, run_dir=out_dir_ev, dataset_file="wbm_subset.db"
+        calculator=calculator, calculator_kwargs=calculator_kwargs, run_dir=out_dir_ev, dataset_file="wbm_subset.db"
     )
 
 

benchmarks/stability.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from ase.calculators.calculator import BaseCalculator
+from loguru import logger
+from prefect import flow
+from prefect.context import FlowRunContext
+
+from mlip_arena.flows.stability import compression, heating
+from mlip_arena.models import REGISTRY, MLIPEnum
+from mlip_arena.tasks.stability.analysis import gather_results
+from mlip_arena.tasks.stability.data import get_atoms_from_db
+
+
+@flow
+def random_mixtures(calculator: str | BaseCalculator, calculator_kwargs: dict | None = None):
+    ctx = FlowRunContext.get()
+    parent_task_runner = ctx.task_runner if ctx else None
+
+    if isinstance(calculator, BaseCalculator):
+        model_name = calculator.__class__.__name__
+        model = calculator
+        family = "custom"
+    elif isinstance(calculator, str) and hasattr(MLIPEnum, calculator):
+        model_name = calculator
+        model = MLIPEnum[calculator]
+        family = REGISTRY[model_name]["family"]
+    else:
+        raise ValueError(f"Unsupported model: {calculator}")
+
+    run_dir = Path(__file__).parent / "stability" / family
+
+    # 1. Run heating
+    print(f"Starting stability heating simulation for {model_name}...")
+    heating.with_options(
+        name=f"stability-heating-{model_name}",
+        task_runner=parent_task_runner,
+        log_prints=True,
+    )(model, run_dir)
+
+    # 2. Run compression
+    print(f"Starting stability compression simulation for {model_name}...")
+    compression.with_options(
+        name=f"stability-compression-{model_name}",
+        task_runner=parent_task_runner,
+        log_prints=True,
+    )(model, run_dir)
+
+    # 3. Analysis / gather results
+    print(f"Analyzing stability results for {model_name}...")
+    compositions = []
+    for atoms in get_atoms_from_db("random-mixture.db"):
+        if len(atoms) == 0:
+            continue
+        compositions.append(atoms.get_chemical_formula())
+
+    try:
+        df = gather_results(run_dir, prefix=model_name, run_type="nvt")
+        df = df[df["formula"].isin(compositions[:120])].copy()  # tentatively we only take the first 120 structures
+        assert len(df) > 0
+        df.to_parquet(run_dir / f"{model_name}-heating.parquet", index=False)
+    except Exception as e:
+        logger.warning(f"Error processing model {model_name} (heating): {e}")
+
+    try:
+        df = gather_results(run_dir, prefix=model_name, run_type="npt")
+        df = df[df["formula"].isin(compositions[:80])].copy()  # tentatively we only take the first 80 structures
+        assert len(df) > 0
+        df.to_parquet(run_dir / f"{model_name}-compression.parquet", index=False)
+    except Exception as e:
+        logger.warning(f"Error processing model {model_name} (compression): {e}")
+
+
+if __name__ == "__main__":
+    from dask.distributed import Client
+    from dask_jobqueue import SLURMCluster
+    from prefect_dask import DaskTaskRunner
+
+    nodes_per_alloc = 1
+    gpus_per_alloc = 4
+
+    cluster_kwargs = dict(
+        cores=1,
+        memory="64 GB",
+        processes=1,
+        shebang="#!/bin/bash",
+        account="matgen",
+        walltime="04:00:00",
+        job_mem="0",
+        job_script_prologue=[
+            "source ~/.bashrc",
+            "module load python",
+            "source activate /pscratch/sd/c/cyrusyc/.conda/mlip-arena",
+        ],
+        job_directives_skip=["-n", "--cpus-per-task", "-J"],
+        job_extra_directives=[
+            "-J arena-stability",
+            "-q preempt",
+            "--time-min=00:30:00",
+            "--comment=12:00:00",
+            f"-N {nodes_per_alloc}",
+            "-C gpu",
+            f"-G {gpus_per_alloc}",
+        ],
+    )
+
+    cluster = SLURMCluster(**cluster_kwargs)
+    print(cluster.job_script())
+    cluster.adapt(minimum_jobs=10, maximum_jobs=50)
+    client = Client(cluster)
+
+    for model in MLIPEnum:
+        random_mixtures.with_options(
+            task_runner=DaskTaskRunner(address=client.scheduler.address),
+            log_prints=True,
+        )(model.name)

benchmarks/submit.py

@@ -2,14 +2,14 @@
 from dask.distributed import Client
 from dask_jobqueue import SLURMCluster
 from prefect_dask import DaskTaskRunner
-from shifts import distribution_shifts
+from stability import random_mixtures
 
 # ==============================================================================
 # 1. JOB CONFIGURATION
 # ==============================================================================
 
 # Example A: Registered string model (e.g., "MACE-MP(M)", "CHGNet")
-calculator = "NequIP-OAM-L"
+calculator = "PET-OAM"
 
 # Example B: Custom ASE Calculator class and arguments (e.g., custom MLIP)
 # from ase.calculators.lj import LennardJones
@@ -77,7 +77,17 @@
     # calculator_kwargs=calculator_kwargs # Uncomment for custom ASE Calculator class
 )
 
-distribution_shifts.with_options(
+# distribution_shifts.with_options(
+#     task_runner=DaskTaskRunner(address=client.scheduler.address),
+#     log_prints=True,
+#     persist_result=False,
+# )(
+#     calculator=calculator,
+#     # calculator_kwargs=calculator_kwargs # Uncomment for custom ASE Calculator class
+# )
+
+
+random_mixtures.with_options(
     task_runner=DaskTaskRunner(address=client.scheduler.address),
     log_prints=True,
     persist_result=False,
@@ -86,4 +96,4 @@
     # calculator_kwargs=calculator_kwargs # Uncomment for custom ASE Calculator class
 )
 
-# TODO: Add Stability and Reactivity
+# TODO: Add Reactivity