Spaces:
Running
Running
File size: 5,131 Bytes
75ac94f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
from pathlib import Path
import numpy as np
import pandas as pd
from ase.db import connect
from dask.distributed import Client
from dask_jobqueue import SLURMCluster
from prefect import flow, task
from prefect.runtime import task_run
from prefect_dask import DaskTaskRunner
from prefect.cache_policies import INPUTS, TASK_SOURCE
from mlip_arena.models import REGISTRY, MLIPEnum
from mlip_arena.tasks.utils import get_calculator
@task
def load_wbm_structures():
"""
Load the WBM structures from an ASE database file.
Reads structures from 'wbm_structures.db' and yields them as ASE Atoms objects
with additional metadata preserved from the database.
Yields:
ase.Atoms: Individual atomic structures from the WBM database with preserved
metadata in the .info dictionary.
"""
with connect("../wbm_structures.db") as db:
for row in db.select():
yield row.toatoms(add_additional_information=True)
@task(
name="E-V Scan",
task_run_name=lambda: f"{task_run.task_name}: {task_run.parameters['atoms'].get_chemical_formula()} - {task_run.parameters['model'].name}",
cache_policy=TASK_SOURCE + INPUTS,
)
def ev_scan(atoms, model):
"""
Perform an energy-volume scan for a given model and atomic structure.
This function applies uniaxial strain to the structure in all three dimensions,
maintaining the fractional coordinates of atoms, and computes the energy at each
deformation point using the specified model.
Args:
atoms: ASE atoms object containing the structure to analyze.
model: MLIPEnum model to use for the energy calculations.
Returns:
dict: Results dictionary containing:
- method (str): The name of the model used
- id (str): The WBM ID of the structure
- eos (dict): Energy of state data with:
- volumes (list): Volume of the unit cell at each strain point
- energies (list): Computed potential energy at each strain point
Note:
The strain range is fixed at ±20% with 21 evenly spaced points.
Results are also saved as a JSON file in a directory named after the model.
"""
calculator = get_calculator(
model
) # avoid sending entire model over prefect and select freer GPU
wbm_id = atoms.info["key_value_pairs"]["wbm_id"]
c0 = atoms.get_cell()
max_abs_strain = 0.2
npoints = 21
volumes = []
energies = []
for uniaxial_strain in np.linspace(-max_abs_strain, max_abs_strain, npoints):
cloned = atoms.copy()
scale_factor = uniaxial_strain + 1
cloned.set_cell(c0 * scale_factor, scale_atoms=True)
cloned.calc = calculator
volumes.append(cloned.get_volume())
energies.append(cloned.get_potential_energy())
data = {
"method": model.name,
"id": wbm_id,
"eos": {
"volumes": volumes, "energies": energies
}
}
fpath = Path(f"{model.name}") / f"{wbm_id}.json"
fpath.parent.mkdir(exist_ok=True)
df = pd.DataFrame([data])
df.to_json(fpath)
return df
@flow
def submit_tasks():
"""
Create and submit energy-volume scan tasks for subsampled WBM structures and applicable models.
This flow function:
1. Loads all structures from the WBM database
2. Iterates through available models in MLIPEnum
3. Filters models based on their capability to handle the 'wbm_ev' GPU task
4. Submits parallel ev_scan tasks for all valid (structure, model) combinations
5. Collects and returns results from all tasks
Returns:
list: Results from all executed tasks (successful or failed)
"""
futures = []
for atoms in load_wbm_structures():
for model in MLIPEnum:
if "wbm_ev" not in REGISTRY[model.name].get("gpu-tasks", []):
continue
try:
result = ev_scan.submit(atoms, model)
except Exception as e:
print(f"Failed to submit task for {model.name}: {e}")
continue
futures.append(result)
return [f.result(raise_on_failure=False) for f in futures]
nodes_per_alloc = 1
gpus_per_alloc = 1
ntasks = 1
cluster_kwargs = dict(
cores=1,
memory="64 GB",
processes=1,
shebang="#!/bin/bash",
account="m3828",
walltime="00:30: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 wbm_ev",
"-q debug",
f"-N {nodes_per_alloc}",
"-C gpu",
f"-G {gpus_per_alloc}",
"--exclusive",
],
)
cluster = SLURMCluster(**cluster_kwargs)
print(cluster.job_script())
cluster.adapt(minimum_jobs=2, maximum_jobs=2)
client = Client(cluster)
submit_tasks.with_options(
task_runner=DaskTaskRunner(address=client.scheduler.address),
log_prints=True,
)()
|