promote run deps, bump prefect and mace versions (#40)

* promote run deps, bump prefect and mace versions

* reorganize pyproject toml

* test caching

* disable granular cache in eos

mlip_arena/tasks/eos.py

@@ -54,8 +54,6 @@ def run(
     max_abs_strain: float = 0.1,
     npoints: int = 11,
     concurrent: bool = True,
-    persist_opt: bool = True,
-    cache_opt: bool = True,
 ) -> dict[str, Any] | State:
     """
     Compute the equation of state (EOS) for the given atoms and calculator.
@@ -80,12 +78,7 @@ def run(
         A dictionary containing the EOS data, bulk modulus, equilibrium volume, and equilibrium energy if successful. Otherwise, a prefect state object.
     """
 
-    OPT_ = OPT.with_options(
-        refresh_cache=not cache_opt,
-        persist_result=persist_opt,
-    )
-
-    state = OPT_(
+    state = OPT(
         atoms=atoms,
         calculator_name=calculator_name,
         calculator_kwargs=calculator_kwargs,
@@ -119,7 +112,7 @@ def run(
             atoms = relaxed.copy()
             atoms.set_cell(c0 * f, scale_atoms=True)
 
-            future = OPT_.submit(
+            future = OPT.submit(
                 atoms=atoms,
                 calculator_name=calculator_name,
                 calculator_kwargs=calculator_kwargs,
@@ -145,7 +138,7 @@ def run(
             atoms = relaxed.copy()
             atoms.set_cell(c0 * f, scale_atoms=True)
 
-            state = OPT_(
+            state = OPT(
                 atoms=atoms,
                 calculator_name=calculator_name,
                 calculator_kwargs=calculator_kwargs,

pyproject.toml

@@ -4,7 +4,7 @@ build-backend="flit_core.buildapi"
 
 [project]
 name="mlip-arena"
-version="0.0.1a1"
+version="0.0.1a2"
 authors=[
     {name="Yuan Chiang", email="[email protected]"},
 ]
@@ -21,37 +21,33 @@ keywords=[
 classifiers=[
     "Development Status :: 1 - Planning",
     "Programming Language :: Python",
+    "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3 :: Only",
 ]
 dependencies=[
     "ase",
+    "pymatgen",
     "torch",
-    "torch_dftd>=0.4.0",
     "huggingface_hub",
     "datasets",
-    "torch-geometric",
     "safetensors",
-    "pymatgen",
-]
-
-[project.optional-dependencies]
-run = [
-    "prefect>=3.0.4,<3.1.0",
+    "prefect>3.1.0",
     "prefect-dask",
     "dask",
     "dask_jobqueue",
     "tables",
 ]
+
+[project.optional-dependencies]
 app = [
     "streamlit==1.38.0",
     "plotly",
-    "bokeh==2.4.3",
-    "statsmodels"
 ]
 test = [
     "torch==2.2.0",
+    "torch_dftd==0.4.0",
     "nvidia-ml-py==12.560.30",
     "e3nn==0.5.0",
     "matgl==1.1.2",
@@ -65,11 +61,11 @@ test = [
     "mattersim==1.0.0rc9",
     "pytest",
     "pytest-xdist",
-    "prefect>=3.0.4,<3.1.0",
+    "prefect==3.1.11",
     "pymatgen@git+https://github.com/chiang-yuan/pymatgen" # hotfix for deprecated pymatgen property https://github.com/materialsproject/pymatgen/pull/4242
 ]
 mace = [
-    "mace-torch==0.3.4",
+    "mace-torch==0.3.9",
 ]
 deepmd = [
     "deepmd-kit@git+https://github.com/deepmodeling/[email protected]"

tests/test_eos.py

@@ -9,15 +9,13 @@ from mlip_arena.models import MLIPEnum
 from mlip_arena.tasks.eos import run as EOS
 
 
-@flow
+@flow(persist_result=True)
 def single_eos_flow(calculator_name, concurrent=True, cache=False):
     atoms = bulk("Cu", "fcc", a=3.6)
 
-    EOS_ = EOS.with_options(
+    return EOS.with_options(
         refresh_cache=not cache,
-    )
-    
-    return EOS_(
+    )(
         atoms=atoms,
         calculator_name=calculator_name,
         calculator_kwargs={},
@@ -32,7 +30,6 @@ def single_eos_flow(calculator_name, concurrent=True, cache=False):
         max_abs_strain=0.1,
         npoints=6,
         concurrent=concurrent,
-        cache_opt=cache
     )