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 3 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
54 changes: 40 additions & 14 deletions mlip_arena/flows/diatomics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,16 +15,27 @@
from tqdm.auto import tqdm

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


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


@task(
task_run_name=lambda: (
f"{task_run.task_name}: {task_run.parameters['symbol']} - {task_run.parameters.get('calculator', 'Unknown')}"
),
task_run_name=_generate_task_run_name,
)
def homonuclear_diatomic(
symbol: str, calculator: str | MLIPEnum | BaseCalculator, calculator_kwargs: dict | None, out_dir: Path
symbol: str,
calculator: str | MLIPEnum | BaseCalculator | None = None,
calculator_kwargs: dict | None = None,
dispersion: bool = False,
dispersion_kwargs: dict | None = None,
out_dir: Path | None = None,
):
Comment on lines +38 to 39

@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 output-directory parameters were introduced without matching None handling. This creates direct runtime crashes when these tasks are invoked without explicit directory arguments.

  • mlip_arena/flows/diatomics.py#L38-L39: either make out_dir required again or initialize a default Path before first use.
  • mlip_arena/flows/stability.py#L38-L39: initialize run_dir before building traj_file in nvt_heat_one.
  • mlip_arena/flows/stability.py#L71-L72: initialize run_dir before building traj_file in npt_compress_one.
📍 Affects 2 files
  • mlip_arena/flows/diatomics.py#L38-L39 (this comment)
  • mlip_arena/flows/stability.py#L38-L39
  • mlip_arena/flows/stability.py#L71-L72
🤖 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/diatomics.py` around lines 38 - 39, Optional
output-directory parameters were introduced without proper `None` handling,
causing runtime crashes when tasks are invoked without explicit directory
arguments. Fix this across three locations: In mlip_arena/flows/diatomics.py
(lines 38-39), either make the `out_dir` parameter required again or initialize
a default `Path` object before any first use of `out_dir`. In
mlip_arena/flows/stability.py (lines 38-39) within the `nvt_heat_one` function,
initialize `run_dir` with a default `Path` value before it is used to construct
`traj_file`. In mlip_arena/flows/stability.py (lines 71-72) within the
`npt_compress_one` function, initialize `run_dir` with a default `Path` value
before it is used to construct `traj_file`. Ensure all three locations properly
handle the `None` case by providing sensible defaults before any path
construction operations.

"""Calculate the potential energy curve for single homonuclear diatomic
molecule.
Expand All @@ -36,6 +47,10 @@ def homonuclear_diatomic(
Args:
symbol: Chemical symbol of the atom (e.g., 'H', 'O', 'Fe')
calculator: An ASE calculator instance or a string/MLIPEnum to create one
calculator_kwargs (dict, optional): Keyword arguments to pass to the calculator. Defaults to None.
dispersion (bool, optional): Whether to use dispersion correction. Defaults to False.
dispersion_kwargs (dict, optional): Keyword arguments for dispersion correction.
out_dir: Path to directory to write output.

Returns:
None: Results are saved as trajectory files.
Expand Down Expand Up @@ -88,7 +103,12 @@ def homonuclear_diatomic(
pbc=False,
)

atoms.calc = get_calculator(calculator=calculator, calculator_kwargs=calculator_kwargs)
atoms.calc = get_calculator(
calculator=calculator,
calculator_kwargs=calculator_kwargs,
dispersion=dispersion,
dispersion_kwargs=dispersion_kwargs,
)

for i, r in enumerate(tqdm(rs)):
if i < skip:
Expand All @@ -104,6 +124,8 @@ def homonuclear_diatomic(
es[i] = atoms.get_potential_energy()
write(traj_fpath, atoms, append="a")

atoms.calc = None


def analyze(out_dir: Path):
"""Analyze potential energy curves (PEC) for diatomic molecules in a
Expand Down Expand Up @@ -257,24 +279,26 @@ def analyze(out_dir: Path):

@flow
def homonuclear_diatomics(
calculator: BaseCalculator | str, calculator_kwargs: dict | None, run_dir: Path | None = None
calculator: BaseCalculator | str,
calculator_kwargs: dict | None = None,
dispersion: bool = False,
dispersion_kwargs: dict | None = None,
run_dir: Path | None = None,
):
"""Run homonuclear diatomic calculations for all elements using a specific
model.

Args:
calculator (BaseCalculator | str): The model or ASE calculator to use.
calculator_kwargs (dict, optional): Keyword arguments to pass to the calculator. Defaults to None.
dispersion (bool, optional): Whether to use dispersion correction. Defaults to False.
dispersion_kwargs (dict, optional): Keyword arguments for dispersion correction.
run_dir (Path, optional): Directory to save outputs. Defaults to None.

Returns:
list: List of results from the distributed homonuclear_diatomic tasks.
"""
if isinstance(calculator, BaseCalculator):
model_name = calculator.__class__.__name__
elif isinstance(calculator, str) and hasattr(MLIPEnum, calculator):
model_name = calculator
else:
raise ValueError(f"Unsupported model: {calculator}")
model_name = resolve_calculator_name(calculator)

family = REGISTRY[model_name]["family"] if hasattr(MLIPEnum, model_name) else "custom"

Expand All @@ -284,8 +308,10 @@ def homonuclear_diatomics(
for symbol in chemical_symbols[1:]:
future = homonuclear_diatomic.submit(
symbol,
calculator,
calculator=calculator,
calculator_kwargs=calculator_kwargs,
dispersion=dispersion,
dispersion_kwargs=dispersion_kwargs,
out_dir=out_dir,
)
futures.append(future)
Expand Down
63 changes: 45 additions & 18 deletions mlip_arena/flows/eos_bulk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
from mlip_arena.models import MLIPEnum
from mlip_arena.tasks.eos import run as EOS
from mlip_arena.tasks.optimize import run as OPT
from mlip_arena.tasks.utils import get_calculator
from mlip_arena.tasks.utils import resolve_calculator_name

if TYPE_CHECKING:
from ase import Atoms
Expand Down Expand Up @@ -61,30 +61,48 @@ def calculate_metrics(res_eos: dict, b0: float, atoms: Atoms, model_name: str, s
}


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(
name="EOS bulk",
task_run_name=lambda: (
f"{task_run.task_name}: {task_run.parameters['atoms'].get_chemical_formula()} - {task_run.parameters['model_name']}"
),
task_run_name=_generate_task_run_name,
cache_policy=TASK_SOURCE + INPUTS,
)
def run(atoms: Atoms, model_name: str, calculator: str | BaseCalculator, calculator_kwargs: dict | None = None):
def run(
atoms: Atoms,
calculator: str | MLIPEnum | BaseCalculator | None = None,
calculator_kwargs: dict | None = None,
dispersion: bool = False,
dispersion_kwargs: dict | None = None,
):
"""Run EOS bulk task for a single structure and model.

Args:
atoms (Atoms): ASE Atoms structure.
model_name (str): Human-readable name of the model.
calculator (str | BaseCalculator): The model or ASE calculator.
calculator_kwargs (dict | None, optional): Arguments for the ASE calculator.
calculator (str | MLIPEnum | BaseCalculator, optional): The model or ASE calculator.
calculator_kwargs (dict, optional): Arguments for the ASE calculator. Defaults to None.
dispersion (bool, optional): Whether to use dispersion correction. Defaults to False.
dispersion_kwargs (dict, optional): Keyword arguments for dispersion correction.

Returns:
pd.DataFrame: A DataFrame containing the raw EOS results.
"""
calculator = calculator if isinstance(calculator, BaseCalculator) else get_calculator(calculator, calculator_kwargs)
model_name = resolve_calculator_name(calculator)

result = OPT(
atoms,
calculator,
calculator=calculator,
calculator_kwargs=calculator_kwargs,
dispersion=dispersion,
dispersion_kwargs=dispersion_kwargs,
optimizer="FIRE",
criterion=dict(
fmax=0.1,
Expand All @@ -93,6 +111,9 @@ def run(atoms: Atoms, model_name: str, calculator: str | BaseCalculator, calcula
result = EOS(
atoms=result["atoms"],
calculator=calculator,
calculator_kwargs=calculator_kwargs,
dispersion=dispersion,
dispersion_kwargs=dispersion_kwargs,
optimizer="FIRE",
npoints=21,
max_abs_strain=0.2,
Expand All @@ -110,27 +131,27 @@ def run(atoms: Atoms, model_name: str, calculator: str | BaseCalculator, calcula
def run_db(
calculator: str | BaseCalculator,
calculator_kwargs: dict | None = None,
dispersion: bool = False,
dispersion_kwargs: dict | None = None,
run_dir: Path | None = None,
dataset: str = "atomind/mlip-arena",
dataset_file: str = "wbm_subset.db",
):
"""Run bulk EOS calculations over a database of structures.

Args:
model (str | BaseCalculator): Model name or ASE calculator.
calculator (str | BaseCalculator): Model name or ASE calculator.
calculator_kwargs (dict, optional): Arguments for the ASE calculator. Defaults to None.
dispersion (bool, optional): Whether to use dispersion correction. Defaults to False.
dispersion_kwargs (dict, optional): Keyword arguments for dispersion correction.
run_dir (Path, optional): Directory to save outputs (parquet files). Defaults to None.
dataset (str, optional): HuggingFace dataset ID. Defaults to "atomind/mlip-arena".
dataset_file (str, optional): Database filename in the dataset. Defaults to "wbm_subset.db".

Returns:
pd.DataFrame: A DataFrame containing analyzed results for all structures in the database.
"""
if isinstance(calculator, BaseCalculator):
model_name = calculator.__class__.__name__
elif isinstance(calculator, str) and hasattr(MLIPEnum, calculator):
model_name = calculator
else:
raise ValueError(f"Unsupported model: {calculator}")
model_name = resolve_calculator_name(calculator)

out_dir = run_dir if run_dir is not None else Path.cwd() / "eos_bulk" / model_name
out_dir.mkdir(parents=True, exist_ok=True)
Expand All @@ -142,7 +163,13 @@ def run_db(
with connect(db_path) as db:
for row in db.select():
atoms = row.toatoms(add_additional_information=True)
future = run.submit(atoms, model_name, calculator, calculator_kwargs)
future = run.submit(
atoms,
calculator=calculator,
calculator_kwargs=calculator_kwargs,
dispersion=dispersion,
dispersion_kwargs=dispersion_kwargs,
)
futures.append(future)

results = [f.result(raise_on_failure=False) for f in futures]
Expand Down
77 changes: 51 additions & 26 deletions mlip_arena/flows/stability.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@

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

if TYPE_CHECKING:
pass
Expand All @@ -28,28 +28,24 @@
HF_TOKEN = os.environ.get("HF_TOKEN", None)


def resolve_model_name(model: MLIPEnum | BaseCalculator | str) -> str:
"""Resolve a string model name from MLIPEnum, BaseCalculator, or string."""
if isinstance(model, str):
return model
if isinstance(model, MLIPEnum):
return model.name
if isinstance(model, type):
return model.__name__
if hasattr(model, "__class__"):
return model.__class__.__name__
return str(model)


@task(cache_policy=TASK_SOURCE + INPUTS)
def nvt_heat_one(atoms: Atoms, model: MLIPEnum | BaseCalculator | str, run_dir: Path):
def nvt_heat_one(
atoms: Atoms,
calculator: str | MLIPEnum | BaseCalculator | None = None,
calculator_kwargs: dict | None = None,
dispersion: bool = False,
dispersion_kwargs: dict | None = None,
run_dir: Path | None = None,
):
"""Run a 10 ps NVT MD simulation with linear heating schedule."""
calculator = get_calculator(model)
model_name = resolve_model_name(model)
model_name = resolve_calculator_name(calculator)

return MD.with_options(refresh_cache=True)(
atoms=atoms,
calculator=calculator,
calculator_kwargs=calculator_kwargs,
dispersion=dispersion,
dispersion_kwargs=dispersion_kwargs,
ensemble="nvt",
dynamics="nose-hoover",
time_step=None,
Expand All @@ -66,14 +62,23 @@ def nvt_heat_one(atoms: Atoms, model: MLIPEnum | BaseCalculator | str, run_dir:


@task(cache_policy=TASK_SOURCE + INPUTS)
def npt_compress_one(atoms: Atoms, model: MLIPEnum | BaseCalculator | str, run_dir: Path):
def npt_compress_one(
atoms: Atoms,
calculator: str | MLIPEnum | BaseCalculator | None = None,
calculator_kwargs: dict | None = None,
dispersion: bool = False,
dispersion_kwargs: dict | None = None,
run_dir: Path | None = None,
):
"""Run a 10 ps NPT MD simulation with linear pressure ramp."""
calculator = get_calculator(model)
model_name = resolve_model_name(model)
model_name = resolve_calculator_name(calculator)

return MD.with_options(timeout_seconds=600, retries=2, refresh_cache=True)(
atoms=atoms,
calculator=calculator,
calculator_kwargs=calculator_kwargs,
dispersion=dispersion,
dispersion_kwargs=dispersion_kwargs,
ensemble="npt",
dynamics="nose-hoover",
time_step=None,
Expand All @@ -89,12 +94,15 @@ def npt_compress_one(atoms: Atoms, model: MLIPEnum | BaseCalculator | str, run_d

@flow
def heating(
model: MLIPEnum | BaseCalculator | str,
calculator: str | MLIPEnum | BaseCalculator | None = None,
calculator_kwargs: dict | None = None,
dispersion: bool = False,
dispersion_kwargs: dict | None = None,
run_dir: Path | None = None,
hf_token: str | None = HF_TOKEN,
):
"""Prefect flow to run NVT heating tasks for many database structures."""
model_name = resolve_model_name(model)
model_name = resolve_calculator_name(calculator)

family = REGISTRY[model_name]["family"] if hasattr(MLIPEnum, model_name) else "custom"
out_dir = run_dir if run_dir is not None else Path.cwd() / "stability" / family
Expand All @@ -107,7 +115,14 @@ def heating(
break
future = nvt_heat_one.with_options(
timeout_seconds=600, retries=2, persist_result=False, refresh_cache=True
).submit(atoms.copy(), model, out_dir)
).submit(
atoms=atoms.copy(),
calculator=calculator,
calculator_kwargs=calculator_kwargs,
dispersion=dispersion,
dispersion_kwargs=dispersion_kwargs,
run_dir=out_dir,
)
futures.append(future)

wait(futures)
Expand All @@ -117,12 +132,15 @@ def heating(

@flow
def compression(
model: MLIPEnum | BaseCalculator | str,
calculator: str | MLIPEnum | BaseCalculator | None = None,
calculator_kwargs: dict | None = None,
dispersion: bool = False,
dispersion_kwargs: dict | None = None,
run_dir: Path | None = None,
hf_token: str | None = HF_TOKEN,
):
"""Prefect flow to run NPT compression tasks for many database structures."""
model_name = resolve_model_name(model)
model_name = resolve_calculator_name(calculator)

family = REGISTRY[model_name]["family"] if hasattr(MLIPEnum, model_name) else "custom"
out_dir = run_dir if run_dir is not None else Path.cwd() / "stability" / family
Expand All @@ -135,7 +153,14 @@ def compression(
break
future = npt_compress_one.with_options(
timeout_seconds=600, retries=2, persist_result=False, refresh_cache=True
).submit(atoms.copy(), model, out_dir)
).submit(
atoms=atoms.copy(),
calculator=calculator,
calculator_kwargs=calculator_kwargs,
dispersion=dispersion,
dispersion_kwargs=dispersion_kwargs,
run_dir=out_dir,
)
futures.append(future)

wait(futures)
Expand Down
Loading
Loading