Spaces:
Running
Running
File size: 5,298 Bytes
fdf446a |
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 165 166 167 168 |
import pickle
from functools import partial
from pathlib import Path
from ase import Atoms
from prefect import Task, flow, task
from prefect.client.schemas.objects import TaskRun
from prefect.results import ResultRecord
from prefect.states import State
from mlip_arena.models import REGISTRY, MLIPEnum
from mlip_arena.tasks.eos import run as EOS
from mlip_arena.tasks.neb import run_from_endpoints as NEB
from mlip_arena.tasks.vacancy_migration.input import get_fcc_pristine, get_hcp_pristine
MP_API_KEY = None
test_models = ["MACE-MP(M)", "MatterSim", "ORBv2", "CHGNet", "M3GNet", "SevenNet"]
def save_to_pickle(
tsk: Task, run: TaskRun, state: State, crystal: str
):
result = run.state.result(raise_on_failure=False)
pristine = result["pristine"]
calculator_name = result["calculator_name"]
calculator_name = calculator_name.name if isinstance(calculator_name, MLIPEnum) else calculator_name
family_path = Path(REGISTRY[calculator_name]["family"])
family_path.mkdir(parents=True, exist_ok=True)
with open(family_path / f"{calculator_name}-{crystal}-{pristine.get_chemical_formula()}.pkl", "wb") as f:
pickle.dump(result, f)
# with open(family_path / f"{crystal}-{pristine.get_chemical_formula()}.json", 'w') as f:
# json.dump(result, f)
@task
def calculate_vacancy_migration(
pristine: Atoms,
istart: int,
iend: int,
calculator_name: MLIPEnum | str,
optimizer: str,
criterion: dict = {}
):
eos = EOS.with_options(refresh_cache=True, persist_result=True)(
atoms=pristine,
calculator_name=calculator_name,
optimizer=optimizer,
criterion=criterion,
concurrent=False,
)
if isinstance(eos, ResultRecord):
eos = eos.result
if isinstance(eos, dict):
pristine = eos["atoms"]
else:
return eos
atoms = pristine.copy()
del atoms[istart]
start = atoms.copy()
atoms = pristine.copy()
del atoms[iend]
end = atoms.copy()
neb = NEB.with_options(refresh_cache=True, persist_result=True)(
start, end, n_images=7,
calculator_name=calculator_name,
optimizer=optimizer,
criterion=criterion,
relax_end_points=True
)
e_defect = 0.5 * (neb["images"][0].get_potential_energy() + neb["images"][-1].get_potential_energy())
e_pristine = pristine.get_potential_energy()
e_vacform = e_defect - (len(neb["images"][0]) / len(pristine)) * e_pristine
e_vacmig = neb["barrier"][0]
asymmetry = abs(neb["barrier"][1] / e_vacmig)
# TODO: temporary solution to pickling problem of mattersim
pristine.calc = None
for image in neb["images"]:
image.calc = None
eos["atoms"].calc = None
return {
"pristine": pristine,
"calculator_name": calculator_name,
"e_vacform": e_vacform,
"e_vacmig": e_vacmig,
"asymmetry": asymmetry,
"neb": neb,
"eos": eos
}
@flow(persist_result=True, result_serializer="pickle")
def run_fcc():
futures = []
for atoms in get_fcc_pristine(MP_API_KEY):
for model in MLIPEnum:
if model.name not in test_models:
continue
try:
result = calculate_vacancy_migration.with_options(
refresh_cache=True, persist_result=True,
on_completion=[partial(save_to_pickle, crystal="fcc")]
)(
pristine=atoms,
istart=0,
iend=1,
calculator_name=model,
optimizer="BFGS",
criterion=dict(fmax=0.05, steps=500),
)
except Exception:
continue
futures.append(result)
return futures
# wait(futures)
# return [f.result(raise_on_failure=False) for f in futures if f.state.is_completed()]
@flow(persist_result=True, result_serializer="pickle")
def run_hcp():
futures = []
for i, atoms in enumerate(get_hcp_pristine(MP_API_KEY)):
if i <= 30:
continue
for model in MLIPEnum:
if model.name not in test_models:
continue
try:
result = calculate_vacancy_migration.with_options(
refresh_cache=True, persist_result=True,
on_completion=[partial(save_to_pickle, crystal="hcp")]
)(
pristine=atoms,
istart=0,
iend=1,
calculator_name=model,
optimizer="BFGS",
criterion=dict(fmax=0.05, steps=500),
)
# calculator_name = model.name if isinstance(model, MLIPEnum) else model
# family_path = Path(REGISTRY[calculator_name]['family'])
# family_path.mkdir(parents=True, exist_ok=True)
# with open(family_path / f"{'hcp'}-{atoms.get_chemical_formula()}.pkl", 'wb') as f:
# pickle.dump(result, f)
except Exception:
continue
futures.append(result)
return futures
# wait(futures)
# return [f.result(raise_on_failure=False) for f in futures if f.state.is_completed()]
|