Spaces:
Running
Running
File size: 4,638 Bytes
9d1a2a5 |
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 |
from datetime import timedelta
from typing import Union
# import covalent as ct
import numpy as np
import pandas as pd
import torch
from ase import Atoms
from ase.calculators.calculator import Calculator
from ase.data import chemical_symbols
from dask.distributed import Client
from dask_jobqueue import SLURMCluster
from prefect import flow, task
from prefect.tasks import task_input_hash
from prefect_dask import DaskTaskRunner
from mlip_arena.models import MLIPCalculator
from mlip_arena.models.utils import EXTMLIPEnum, MLIPMap, external_ase_calculator
cluster_kwargs = {
"cores": 4,
"memory": "64 GB",
"shebang": "#!/bin/bash",
"account": "m3828",
"walltime": "00:10:00",
"job_mem": "0",
"job_script_prologue": ["source ~/.bashrc"],
"job_directives_skip": ["-n", "--cpus-per-task"],
"job_extra_directives": ["-q debug", "-C gpu"],
}
cluster = SLURMCluster(**cluster_kwargs)
cluster.scale(jobs=10)
client = Client(cluster)
@task(cache_key_fn=task_input_hash, cache_expiration=timedelta(hours=1))
def calculate_single_diatomic(
calculator_name: str | EXTMLIPEnum,
calculator_kwargs: dict | None,
atom1: str,
atom2: str,
rmin: float = 0.1,
rmax: float = 6.5,
npts: int = int(1e3),
):
calculator_kwargs = calculator_kwargs or {}
if isinstance(calculator_name, EXTMLIPEnum) and calculator_name in EXTMLIPEnum:
calc = external_ase_calculator(calculator_name, **calculator_kwargs)
elif calculator_name in MLIPMap:
calc = MLIPMap[calculator_name](**calculator_kwargs)
a = 2 * rmax
rs = np.linspace(rmin, rmax, npts)
e = np.zeros_like(rs)
f = np.zeros_like(rs)
da = atom1 + atom2
for i, r in enumerate(rs):
positions = [
[0, 0, 0],
[r, 0, 0],
]
# Create the unit cell with two atoms
atoms = Atoms(da, positions=positions, cell=[a, a, a])
atoms.calc = calc
e[i] = atoms.get_potential_energy()
f[i] = np.inner(np.array([1, 0, 0]), atoms.get_forces()[1])
return {"r": rs, "E": e, "F": f, "da": da}
@flow
def calculate_multiple_diatomics(calculator_name, calculator_kwargs):
futures = []
for symbol in chemical_symbols:
if symbol == "X":
continue
futures.append(
calculate_single_diatomic.submit(
calculator_name, calculator_kwargs, symbol, symbol
)
)
return [i for future in futures for i in future.result()]
@flow(task_runner=DaskTaskRunner(address=client.scheduler.address), log_prints=True)
def calculate_homonuclear_diatomics(calculator_name, calculator_kwargs):
curves = calculate_multiple_diatomics(calculator_name, calculator_kwargs)
pd.DataFrame(curves).to_csv(f"homonuclear-diatomics-{calculator_name}.csv")
# with plt.style.context("default"):
# SMALL_SIZE = 6
# MEDIUM_SIZE = 8
# LARGE_SIZE = 10
# LINE_WIDTH = 1
# plt.rcParams.update(
# {
# "pgf.texsystem": "pdflatex",
# "font.family": "sans-serif",
# "text.usetex": True,
# "pgf.rcfonts": True,
# "figure.constrained_layout.use": True,
# "axes.labelsize": MEDIUM_SIZE,
# "axes.titlesize": MEDIUM_SIZE,
# "legend.frameon": False,
# "legend.fontsize": MEDIUM_SIZE,
# "legend.loc": "best",
# "lines.linewidth": LINE_WIDTH,
# "xtick.labelsize": SMALL_SIZE,
# "ytick.labelsize": SMALL_SIZE,
# }
# )
# fig, ax = plt.subplots(layout="constrained", figsize=(3, 2), dpi=300)
# color = "tab:red"
# ax.plot(rs, e, color=color, zorder=1)
# ax.axhline(ls="--", color=color, alpha=0.5, lw=0.5 * LINE_WIDTH)
# ylo, yhi = ax.get_ylim()
# ax.set(xlabel=r"r [$\AA]$", ylim=(max(-7, ylo), min(5, yhi)))
# ax.set_ylabel(ylabel="E [eV]", color=color)
# ax.tick_params(axis="y", labelcolor=color)
# ax.text(0.8, 0.85, da, fontsize=LARGE_SIZE, transform=ax.transAxes)
# color = "tab:blue"
# at = ax.twinx()
# at.plot(rs, f, color=color, zorder=0, lw=0.5 * LINE_WIDTH)
# at.axhline(ls="--", color=color, alpha=0.5, lw=0.5 * LINE_WIDTH)
# ylo, yhi = at.get_ylim()
# at.set(
# xlabel=r"r [$\AA]$",
# ylim=(max(-20, ylo), min(20, yhi)),
# )
# at.set_ylabel(ylabel="F [eV/$\AA$]", color=color)
# at.tick_params(axis="y", labelcolor=color)
# plt.show()
if __name__ == "__main__":
calculate_homonuclear_diatomics(
EXTMLIPEnum.MACE, dict(model="medium", device="cuda")
)
|