Skip to content

Commit

Permalink
add notebooks
Browse files Browse the repository at this point in the history
  • Loading branch information
@chiang-yuan
chiang-yuan committed Jul 7, 2024
1 parent d72faca commit 5225959
Show file tree
Hide file tree
Showing 3 changed files with 1,033 additions and 0 deletions.
243 changes: 243 additions & 0 deletions mlip_arena/tasks/diatomics/alignn/run.ipynb
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,243 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "3200850a-b8fb-4f50-9815-16ae8da0f942",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from ase import Atoms, Atom\n",
"from ase.io import read, write\n",
"from ase.data import chemical_symbols, covalent_radii, vdw_alvarez\n",
"from ase.parallel import paropen as open\n",
"\n",
"from pathlib import Path\n",
"import os\n",
"import numpy as np\n",
"from pymatgen.core import Element\n",
"from tqdm.auto import tqdm\n",
"import pandas as pd\n",
"\n",
"\n",
"from alignn.ff.ff import AlignnAtomwiseCalculator,default_path\n",
"\n",
"\n",
"model_path = default_path()\n",
"calc = AlignnAtomwiseCalculator(path=model_path, device='cuda')\n",
"\n",
"model_name = 'ALIGNN'\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "90887faa-1601-4c4c-9c44-d16731471d7f",
"metadata": {
"scrolled": true,
"tags": []
},
"outputs": [],
"source": [
"\n",
"\n",
"for symbol in tqdm(chemical_symbols):\n",
" \n",
" s = set([symbol])\n",
" \n",
" if 'X' in s:\n",
" continue\n",
" \n",
" try:\n",
" atom = Atom(symbol)\n",
" rmin = covalent_radii[atom.number] * 0.95\n",
" rvdw = vdw_alvarez.vdw_radii[atom.number] if atom.number < len(vdw_alvarez.vdw_radii) else np.nan \n",
" rmax = 3.1 * rvdw if not np.isnan(rvdw) else 6\n",
" rstep = 0.01 #if rmin < 1 else 0.4\n",
"\n",
" a = 2 * rmax\n",
"\n",
" npts = int((rmax - rmin)/rstep)\n",
"\n",
" rs = np.linspace(rmin, rmax, npts)\n",
" e = np.zeros_like(rs)\n",
"\n",
" da = symbol + symbol\n",
"\n",
" out_dir = Path(str(da))\n",
"\n",
" os.makedirs(out_dir, exist_ok=True)\n",
"\n",
" skip = 0\n",
" \n",
" element = Element(symbol)\n",
" \n",
" try:\n",
" m = element.valence[1]\n",
" if element.valence == (0, 2):\n",
" m = 0\n",
" except:\n",
" m = 0\n",
" \n",
" \n",
" r = rs[0]\n",
" \n",
" positions = [\n",
" [a/2-r/2, a/2, a/2],\n",
" [a/2+r/2, a/2, a/2],\n",
" ]\n",
" \n",
" traj_fpath = out_dir / f\"{model_name}.extxyz\"\n",
"\n",
" if traj_fpath.exists():\n",
" traj = read(traj_fpath, index=\":\")\n",
" skip = len(traj)\n",
" atoms = traj[-1]\n",
" else:\n",
" # Create the unit cell with two atoms\n",
" atoms = Atoms(\n",
" da, \n",
" positions=positions,\n",
" # magmoms=magmoms,\n",
" cell=[a, a+0.001, a+0.002], \n",
" pbc=True\n",
" )\n",
" \n",
" print(atoms)\n",
"\n",
" calc = calc\n",
"\n",
" atoms.calc = calc\n",
"\n",
" for i, r in enumerate(tqdm(rs)):\n",
"\n",
" if i < skip:\n",
" continue\n",
"\n",
" positions = [\n",
" [a/2-r/2, a/2, a/2],\n",
" [a/2+r/2, a/2, a/2],\n",
" ]\n",
" \n",
" # atoms.set_initial_magnetic_moments(magmoms)\n",
" \n",
" atoms.set_positions(positions)\n",
"\n",
" e[i] = atoms.get_potential_energy()\n",
" \n",
" atoms.calc.results.update({\n",
" \"forces\": atoms.get_forces()\n",
" })\n",
"\n",
" write(traj_fpath, atoms, append=\"a\")\n",
" except Exception as e:\n",
" print(e)\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a0ac2c09-370b-4fdd-bf74-ea5c4ade0215",
"metadata": {},
"outputs": [],
"source": [
"\n",
"\n",
"df = pd.DataFrame(columns=['name', 'method', 'R', 'E', 'F', 'S^2'])\n",
"\n",
"for symbol in tqdm(chemical_symbols):\n",
" \n",
" da = symbol + symbol\n",
" \n",
" out_dir = Path(da)\n",
" \n",
" traj_fpath = out_dir / f\"{model_name}.extxyz\"\n",
"\n",
"\n",
" if traj_fpath.exists():\n",
" traj = read(traj_fpath, index=\":\")\n",
" else:\n",
" continue\n",
" \n",
" Rs, Es, Fs, S2s = [], [], [], []\n",
" for atoms in traj:\n",
" \n",
" vec = atoms.positions[1] - atoms.positions[0]\n",
" r = np.linalg.norm(vec)\n",
" e = atoms.get_potential_energy()\n",
" f = np.inner(vec/r, atoms.get_forces()[1])\n",
" # s2 = np.mean(np.power(atoms.get_magnetic_moments(), 2))\n",
" \n",
" Rs.append(r)\n",
" Es.append(e)\n",
" Fs.append(f)\n",
" # S2s.append(s2)\n",
" \n",
" data = {\n",
" 'name': da,\n",
" 'method': 'ALIGNN',\n",
" 'R': Rs,\n",
" 'E': Es,\n",
" 'F': Fs,\n",
" 'S^2': S2s\n",
" }\n",
"\n",
" df = pd.concat([df, pd.DataFrame([data])], ignore_index=True)\n",
"\n",
"json_fpath = 'homonuclear-diatomics.json'\n",
"\n",
"df.to_json(json_fpath, orient='records') "
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e0dd4367-3dca-440f-a7a9-7fdd84183f2c",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"df"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4e6ae884-89f3-43f2-8fd9-19bf00c91566",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "mlip-arena",
"language": "python",
"name": "mlip-arena"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.8"
},
"widgets": {
"application/vnd.jupyter.widget-state+json": {
"state": {},
"version_major": 2,
"version_minor": 0
}
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Loading

0 comments on commit 5225959

Please sign in to comment.