Spaces:
Running
Running
File size: 3,838 Bytes
b7d94da |
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 |
import itertools
from pathlib import Path
import numpy as np
from ase import Atom, Atoms
from ase.calculators.calculator import BaseCalculator
from ase.data import chemical_symbols, covalent_radii, vdw_alvarez
from ase.io import read, write
from prefect import flow, task
from tqdm.auto import tqdm
from mlip_arena.models import REGISTRY, MLIPEnum
from mlip_arena.tasks.utils import get_calculator
@task
def homonuclear_diatomics(symbol: str, calculator: BaseCalculator, out_dir: Path):
"""
Calculate potential energy curves for homonuclear diatomic molecules.
This function computes the potential energy of a diatomic molecule (two atoms of
the same element) across a range of interatomic distances. The distance range is
automatically determined from the covalent and van der Waals radii of the element.
Args:
symbol: Chemical symbol of the atom (e.g., 'H', 'O', 'Fe')
calculator: ASE calculator object used to compute the potential energies. Could be VASP, MLIP, etc.
Returns:
None: Results are saved as trajectory files in a directory structure:
/{model_family}/{element_pair}/{model_name}.extxyz
Note:
- Minimum distance is set to 0.9× the covalent radius
- Maximum distance is set to 3.1× the van der Waals radius (or 6 Å if unknown)
- Distance step size is fixed at 0.01 Å
- If an existing trajectory file is found, the calculation will resume from where it left off
- The atoms are placed in a periodic box large enough to avoid self-interaction
"""
atom = Atom(symbol)
rmin = 0.9 * covalent_radii[atom.number]
rvdw = (
vdw_alvarez.vdw_radii[atom.number]
if atom.number < len(vdw_alvarez.vdw_radii)
else np.nan
)
rmax = 3.1 * rvdw if not np.isnan(rvdw) else 6
rstep = 0.01
npts = int((rmax - rmin) / rstep)
rs = np.linspace(rmin, rmax, npts)
es = np.zeros_like(rs)
da = symbol + symbol
out_dir.mkdir(parents=True, exist_ok=True)
skip = 0
a = 5 * rmax
r = rs[0]
positions = [
[a / 2 - r / 2, a / 2, a / 2],
[a / 2 + r / 2, a / 2, a / 2],
]
traj_fpath = out_dir / f"{da!s}.extxyz"
if traj_fpath.exists():
traj = read(traj_fpath, index=":")
skip = len(traj)
atoms = traj[-1]
else:
# Create the unit cell with two atoms
atoms = Atoms(
da,
positions=positions,
# magmoms=magmoms,
cell=[a, a + 0.001, a + 0.002],
pbc=False,
)
atoms.calc = calculator
for i, r in enumerate(tqdm(rs)):
if i < skip:
continue
positions = [
[a / 2 - r / 2, a / 2, a / 2],
[a / 2 + r / 2, a / 2, a / 2],
]
# atoms.set_initial_magnetic_moments(magmoms)
atoms.set_positions(positions)
es[i] = atoms.get_potential_energy()
write(traj_fpath, atoms, append="a")
@flow
def submit_homonuclear_diatomics():
futures = []
for symbol, model in itertools.product(
chemical_symbols[1:],
MLIPEnum,
):
if "homonuclear-diatomics" not in REGISTRY[model.name].get("gpu-tasks", []):
continue
out_dir = Path(__file__).parent / model.name
calculator = get_calculator(model)
# if not (out_dir / "homonuclear-diatomics.json").exists():
future = homonuclear_diatomics.submit(
symbol,
calculator,
out_dir=out_dir,
)
futures.append(future)
return [f.result(raise_on_failure=False) for f in futures]
if __name__ == "__main__":
submit_homonuclear_diatomics.with_options(
# task_runner=DaskTaskRunner(address=client.scheduler.address),
log_prints=True,
)()
|