Skip to content

Calculator instantiation inside Prefect task #91

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chiang-yuan merged 9 commits into from
Jun 16, 2026
Merged

Calculator instantiation inside Prefect task #91

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
244a798
refactor(tasks): unify calculator argument signatures in core tasks
chiang-yuan Jun 15, 2026
bb7bfb4
refactor(flows): update flows and tasks to use unified calculator par…
chiang-yuan Jun 15, 2026
4fba6da
test(prefect): add regression test for NequIP calculator serialization
chiang-yuan Jun 15, 2026
75dd513
refactor(flows/tasks): update EV scan and Conservation tasks/flows wi…
chiang-yuan Jun 15, 2026
c7a072f
test(refactor): update test suite call sites for unified calculator s…
chiang-yuan Jun 15, 2026
168ca91
refactor(benchmarks): update benchmarks and submission flows to use u…
chiang-yuan Jun 15, 2026
da8728b
fix(tasks): restrict calculator clearing to end of MD runs and add te…
chiang-yuan Jun 15, 2026
493753d
ci: include test group in Hugging Face trial-branch name
chiang-yuan Jun 15, 2026
a957d9a
test: add unit tests for eos_bulk, ev, and diatomics flows
chiang-yuan Jun 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 6 additions & 9 deletions benchmarks/c2db/run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@

from mlip_arena.models import MLIPEnum
from mlip_arena.tasks import ELASTICITY, OPT, PHONON
from mlip_arena.tasks.utils import get_calculator

select_models = [
"ALIGNN",
Expand Down Expand Up @@ -94,11 +93,9 @@ def run_one(model, row):
atoms = row.toatoms()
# print(data := row.key_value_pairs)

calc = get_calculator(model)

result_opt = OPT(
atoms,
calc,
atoms=atoms,
calculator=model,
optimizer="FIRE",
criterion=dict(fmax=0.05, steps=500),
symmetry=True,
Expand All @@ -107,8 +104,8 @@ def run_one(model, row):
atoms = result_opt["atoms"]

result_elastic = ELASTICITY(
atoms,
calc,
atoms=atoms,
calculator=model,
optimizer="FIRE",
criterion=dict(fmax=0.05, steps=500),
pre_relax=False,
Expand All @@ -123,8 +120,8 @@ def run_one(model, row):
np.savez(outdir / "elastic.npz", tensor=elastic_tensor, eigenvalues=eigenvalues)

result_phonon = PHONON(
atoms,
calc,
atoms=atoms,
calculator=model,
supercell_matrix=(2, 2, 1),
outdir=outdir,
)
Expand Down
21 changes: 10 additions & 11 deletions benchmarks/eos_bulk/run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
from prefect_dask import DaskTaskRunner

from mlip_arena.models import REGISTRY, MLIPEnum
from mlip_arena.tasks.utils import get_calculator


@task
Expand Down Expand Up @@ -64,26 +63,26 @@ def eos_bulk(atoms: Atoms, model: MLIPEnum):
from mlip_arena.tasks.eos import run as EOS
from mlip_arena.tasks.optimize import run as OPT

calculator = get_calculator(model) # avoid sending entire model over prefect and select freer GPU

result = OPT.with_options(
refresh_cache=True,
)(
atoms,
calculator,
atoms=atoms,
calculator=model,
optimizer="FIRE",
criterion=dict(
fmax=0.1,
),
)
result = EOS.with_options(
refresh_cache=True,
# on_completion=[functools.partial(
# save_result,
# model_name=model.name,
# id=atoms.info["key_value_pairs"]["wbm_id"],
# )],
)(atoms=result["atoms"], calculator=calculator, optimizer="FIRE", npoints=21, max_abs_strain=0.2, concurrent=False)
)(
atoms=result["atoms"],
calculator=model,
optimizer="FIRE",
npoints=21,
max_abs_strain=0.2,
concurrent=False,
)

result["method"] = model.name
result["id"] = atoms.info["key_value_pairs"]["wbm_id"]
Expand Down
4 changes: 2 additions & 2 deletions benchmarks/mof/classification/classification.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@

from mlip_arena.models import MLIPEnum
from mlip_arena.tasks.mof.flow import widom_insertion
from mlip_arena.tasks.utils import get_calculator


def load_row_from_df(fpath: str):
Expand Down Expand Up @@ -102,7 +101,8 @@ def run_one(model, row, gas):
)(
structure=row["structure"],
gas=gas,
calculator=get_calculator(model, dispersion=True),
calculator=model.name,
dispersion=True,
criterion=dict(fmax=0.05, steps=50),
init_structure_optimize_loops=10,
)
Expand Down
15 changes: 5 additions & 10 deletions benchmarks/submit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,11 +23,6 @@ def asymptotic_behaviors(calculator: str | BaseCalculator, calculator_kwargs: di

model_name = get_model_name(calculator)

from mlip_arena.tasks.utils import get_calculator

# Use string directly for registered models to avoid client-side unpickleable PyTorch objects
model = calculator if isinstance(calculator, str) else get_calculator(calculator, calculator_kwargs)

# 1. Diatomics
print(f"Starting homonuclear diatomics benchmark for {model_name}...")
from mlip_arena.models import REGISTRY, MLIPEnum
Expand All @@ -41,21 +36,21 @@ def asymptotic_behaviors(calculator: str | BaseCalculator, calculator_kwargs: di
family = family.lower()
out_dir_diatomics = Path(__file__).parent / "diatomics" / family / model_name
homonuclear_diatomics.with_options(name=f"diatomics-{model_name}", task_runner=parent_task_runner)(
calculator=model, calculator_kwargs=calculator_kwargs, 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)(
calculator=model, calculator_kwargs=calculator_kwargs, 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)(
calculator=model, calculator_kwargs=calculator_kwargs, 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"
)


Expand Down Expand Up @@ -102,12 +97,12 @@ def stability(calculator: str | BaseCalculator, calculator_kwargs: dict | None =
heating.with_options(
name=f"stability-heating-{model_name}",
task_runner=parent_task_runner,
)(calculator, run_dir_stability)
)(calculator=calculator, calculator_kwargs=calculator_kwargs, run_dir=run_dir_stability)

compression.with_options(
name=f"stability-compression-{model_name}",
task_runner=parent_task_runner,
)(calculator, run_dir_stability)
)(calculator=calculator, calculator_kwargs=calculator_kwargs, run_dir=run_dir_stability)

# 3. Analysis / gather results
print(f"Analyzing stability results for {model_name}...")
Expand Down
8 changes: 4 additions & 4 deletions benchmarks/vacancy_migration/run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ def calculate_vacancy_migration(
):
eos = EOS.with_options(refresh_cache=True, persist_result=True)(
atoms=pristine,
calculator_name=calculator_name,
calculator=calculator_name,
optimizer=optimizer,
criterion=criterion,
concurrent=False,
Expand All @@ -64,10 +64,10 @@ def calculate_vacancy_migration(
end = atoms.copy()

neb = NEB.with_options(refresh_cache=True, persist_result=True)(
start,
end,
start=start,
end=end,
n_images=7,
calculator_name=calculator_name,
calculator=calculator_name,
optimizer=optimizer,
criterion=criterion,
relax_end_points=True,
Expand Down
93 changes: 62 additions & 31 deletions mlip_arena/flows/conservation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@

from mlip_arena.models import MLIPEnum
from mlip_arena.tasks.md import run as MD
from mlip_arena.tasks.utils import BaseCalculator, get_calculator, logger
from mlip_arena.tasks.utils import BaseCalculator, logger, resolve_calculator_name

try:
from quests.descriptor import get_descriptors
Expand Down Expand Up @@ -139,21 +139,35 @@ def get_trajectory_entropy(
return delta_entropy(descriptors, ref_descriptors, h=h), structures


def _generate_task_run_name():
task_name = task_run.task_name
parameters = task_run.parameters
atoms = parameters["atoms"]
calculator_name = resolve_calculator_name(parameters.get("calculator"))
return f"{task_name}: {atoms.get_chemical_formula()} - {calculator_name}"


@task(
task_run_name=lambda: (
f"{task_run.task_name}: {task_run.parameters['atoms'].get_chemical_formula()} - {task_run.parameters['calculator']}"
),
name="NVE MD",
task_run_name=_generate_task_run_name,
cache_policy=TASK_SOURCE + INPUTS,
)
def run_nve_md(
atoms: Atoms, calculator: MLIPEnum | BaseCalculator | str, calculator_kwargs: dict | None, traj_file: Path
atoms: Atoms,
calculator: str | MLIPEnum | BaseCalculator | None = None,
calculator_kwargs: dict | None = None,
dispersion: bool = False,
dispersion_kwargs: dict | None = None,
traj_file: Path | None = None,
):
"""Run NVE molecular dynamics simulation.

Args:
atoms (Atoms): ASE Atoms structure.
calculator (MLIPEnum | BaseCalculator | str): Model or calculator to use.
calculator_kwargs (dict, optional): Kwargs for calculator initialization.
calculator (str | MLIPEnum | BaseCalculator, optional): Model or calculator to use.
calculator_kwargs (dict, optional): Kwargs for calculator initialization. Defaults to None.
dispersion (bool, optional): Whether to use dispersion correction. Defaults to False.
dispersion_kwargs (dict, optional): Keyword arguments for dispersion correction.
traj_file (Path): Path to save the trajectory file.

Returns:
Expand All @@ -163,9 +177,10 @@ def run_nve_md(
refresh_cache=True,
)(
atoms,
calculator=get_calculator(
calculator, calculator_kwargs
), # wrap calculater inside task for separate calculator instances
calculator=calculator,
calculator_kwargs=calculator_kwargs,
dispersion=dispersion,
dispersion_kwargs=dispersion_kwargs,
ensemble="nve",
dynamics="velocityverlet",
time_step=1.0, # fs
Expand All @@ -179,12 +194,20 @@ def run_nve_md(


def run_simulations(
calculator: MLIPEnum | BaseCalculator | str, calculator_kwargs: dict | None, structures: list[Atoms], out_dir: Path
calculator: str | MLIPEnum | BaseCalculator | None = None,
calculator_kwargs: dict | None = None,
dispersion: bool = False,
dispersion_kwargs: dict | None = None,
structures: list[Atoms] = None,
out_dir: Path = None,
):
Comment on lines +197 to 203

@coderabbitai coderabbitai Bot Jun 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Optional parameters are treated as required and can crash at runtime.

structures, out_dir, input_path, and reference_path are defaulted to None, but the code dereferences them unconditionally (out_dir.mkdir(...), read(input_path, ...), parameters["reference_path"].stem). This creates AttributeError/TypeError paths and makes the flow contract misleading.

Suggested fix 
 def _generate_flow_run_name():
     name = flow_run.flow_name
     parameters = flow_run.parameters

     calculator = parameters.get("calculator")
     model_name = resolve_calculator_name(calculator)
-    reference_path = parameters["reference_path"].stem
+    reference_path = parameters.get("reference_path")
+    reference_stem = Path(reference_path).stem if reference_path is not None else "unknown-reference"

-    return f"{name}: {model_name} - {reference_path}"
+    return f"{name}: {model_name} - {reference_stem}"
 def run_simulations(
     calculator: str | MLIPEnum | BaseCalculator | None = None,
     calculator_kwargs: dict | None = None,
     dispersion: bool = False,
     dispersion_kwargs: dict | None = None,
     structures: list[Atoms] = None,
     out_dir: Path = None,
 ):
+    if structures is None or out_dir is None:
+        raise ValueError("`structures` and `out_dir` are required.")
 def differential_entropy_along_nve_trajectory(
@@
 ):
+    if input_path is None or reference_path is None:
+        raise ValueError("`input_path` and `reference_path` are required.")
+    if start_idx is None or end_idx is None or step is None:
+        raise ValueError("`start_idx`, `end_idx`, and `step` are required.")

Also applies to: 250-257, 266-274, 308-308

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@mlip_arena/flows/conservation.py` around lines 197 - 203, The parameters
structures, out_dir, input_path, and reference_path are defined with default
values of None but are dereferenced unconditionally throughout the code (e.g.,
out_dir.mkdir(), read(input_path, ...), parameters["reference_path"].stem),
causing AttributeError or TypeError at runtime. Either remove the None defaults
to make these parameters required, or add explicit None checks before each use
of these variables. Apply this fix across all affected locations in
mlip_arena/flows/conservation.py (the function definition at lines 197-203, and
the dereferencing sites at lines 250-257, 266-274, and 308-308) to ensure the
function contract matches its actual requirements.

"""Runs simulations on a list of structures.

Parameters:
calculator (MLIPEnum | BaseCalculator | str): Model to use.
calculator (str | MLIPEnum | BaseCalculator, optional): Model to use.
calculator_kwargs (dict, optional): Keyword arguments for calculator initialization.
dispersion (bool, optional): Whether to use dispersion correction.
dispersion_kwargs (dict, optional): Keyword arguments for dispersion correction.
structures (list[ase.Atoms]): List of structures to simulate.
out_dir (str): Directory to save the simulation trajectories to.

Expand Down Expand Up @@ -214,6 +237,8 @@ def run_simulations(
supercell_atoms,
calculator=calculator,
calculator_kwargs=calculator_kwargs,
dispersion=dispersion,
dispersion_kwargs=dispersion_kwargs,
traj_file=out_dir / f"{i}.traj",
)
futures.append(future)
Expand All @@ -226,24 +251,27 @@ def _generate_flow_run_name():
name = flow_run.flow_name
parameters = flow_run.parameters

model = parameters["model"]
calculator = parameters.get("calculator")
model_name = resolve_calculator_name(calculator)
reference_path = parameters["reference_path"].stem

return f"{name}: {model} - {reference_path}"
return f"{name}: {model_name} - {reference_path}"


@flow(
name="Entropy along trajectory",
flow_run_name=_generate_flow_run_name,
)
def differential_entropy_along_nve_trajectory(
calculator: MLIPEnum | BaseCalculator | str,
calculator_kwargs: dict | None,
input_path: Path,
reference_path: Path,
start_idx: int,
end_idx: int,
step: int,
calculator: str | MLIPEnum | BaseCalculator | None = None,
calculator_kwargs: dict | None = None,
dispersion: bool = False,
dispersion_kwargs: dict | None = None,
input_path: Path | None = None,
reference_path: Path | None = None,
start_idx: int | None = None,
end_idx: int | None = None,
step: int | None = None,
k: int = 32,
cutoff: float = 5.0,
h: float = 0.015,
Expand All @@ -253,7 +281,10 @@ def differential_entropy_along_nve_trajectory(
trajectory with respect to a reference dataset.

Arguments:
model (MLIPEnum | BaseCalculator | str): Model to use.
calculator (str | MLIPEnum | BaseCalculator, optional): Model to use.
calculator_kwargs (dict, optional): Keyword arguments to pass to the calculator.
dispersion (bool, optional): Whether to use dispersion correction.
dispersion_kwargs (dict, optional): Keyword arguments for dispersion correction.
input_path (Path): Path to the directory containing the trajectory files.
reference_path (Path): Path to the file containing the descriptors of the full dataset of structures without the subset.
start_idx (int): Starting index of the subset of structures to select from each trajectory.
Expand All @@ -269,20 +300,20 @@ def differential_entropy_along_nve_trajectory(
or (np.ndarray): an (H,M) matrix if 'h' is a vector of length H
sampled_structures (list[ase.Atoms]): List of structures selected from the trajectory.
"""
if isinstance(calculator, MLIPEnum):
model_name = calculator.name
elif isinstance(calculator, BaseCalculator):
model_name = calculator.__class__.__name__
elif isinstance(calculator, str) and hasattr(MLIPEnum, calculator):
model_name = calculator
else:
raise ValueError(f"Unsupported calculator: {calculator}")
model_name = resolve_calculator_name(calculator)

# Run simulations
out_dir = work_dir / model_name if work_dir is not None else Path.cwd() / model_name

structures = read(input_path, index=":")
run_simulations(calculator, calculator_kwargs, structures, out_dir)
run_simulations(
calculator=calculator,
calculator_kwargs=calculator_kwargs,
dispersion=dispersion,
dispersion_kwargs=dispersion_kwargs,
structures=structures,
out_dir=out_dir,
)

# Get entropy for structures along trajectories
dH, sampled_structures = get_trajectory_entropy(
Expand Down
Loading
Loading