experiment/nequip

.gitattributes

@@ -18,4 +18,5 @@ examples/assets/bash_workflow_template.json filter=lfs diff=lfs merge=lfs -text
 *.whl filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
 packages filter=lfs diff=lfs merge=lfs -text

config.yml

@@ -173,3 +173,19 @@ notebooks:
       - emfs:/drive/packages/chgnet-0.3.8-py3-none-any.whl
       # Stubbed packages (patched by torch_pyodide with include_chgnet=True)
       - ssl
+  - name: nequip
+    packages_pyodide:
+      # Packages with dependencies
+      - opt_einsum
+      - pyyaml
+      - setuptools
+      # Packages without dependencies
+      - nodeps:opt_einsum_fx
+      - nodeps:e3nn>=0.5
+      - nodeps:ase
+      # NequIP wheel (stripped of heavy deps: hydra, lightning, torchmetrics)
+      - emfs:/drive/packages/nequip-0.15.0-py3-none-any.whl
+      # Stubbed packages (patched by torch_pyodide with include_nequip=True)
+      - ssl
+      - h5py
+      - lmdb

other/experiments/jupyterlite/relax_structure_with_nequip.ipynb

@@ -0,0 +1,287 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Relax Structure with NequIP \u2014 E(3)-Equivariant Neural Network Potential\n",
+    "\n",
+    "This notebook demonstrates structural relaxation using **NequIP**,\n",
+    "an E(3)-equivariant message passing neural network interatomic potential.\n",
+    "\n",
+    "NequIP achieves high accuracy with small model sizes by using equivariant\n",
+    "features, making it efficient for materials simulations.\n",
+    "\n",
+    "We use the **NequIP-OAM-S** foundation model, trained on OMat24 + sAlex + MPTrj datasets.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Set Input Parameters\n",
+    "### 1.1. Structure and Relaxation\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "FOLDER = \"uploads\"\n",
+    "STRUCTURE_NAME = \"Interface\"  # Name of the structure to load from local file\n",
+    "\n",
+    "RELAXATION_PARAMETERS = {\n",
+    "    \"FMAX\": 0.05,\n",
+    "}\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Install Packages\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from mat3ra.notebooks_utils.packages import install_packages\n",
+    "\n",
+    "await install_packages(\"made|api_examples|torch|nequip\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from mat3ra.notebooks_utils.pyodide.packages.torch import apply_all_patches\n",
+    "\n",
+    "apply_all_patches(include_nequip=True)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Load Materials\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from mat3ra.made.material import Material\n",
+    "from mat3ra.notebooks_utils.material import load_material_from_folder\n",
+    "from mat3ra.standata.materials import Materials\n",
+    "\n",
+    "structure = load_material_from_folder(FOLDER, STRUCTURE_NAME) or Material.create(\n",
+    "    Materials.get_by_name_first_match(STRUCTURE_NAME))\n",
+    "\n",
+    "print(f\"Structure: {structure.name}\")\n",
+    "print(f\"Formula: {structure.formula}\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 3.1. Visualize Input Structure\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from mat3ra.notebooks_utils.ipython.entity.material.visualize import ViewersEnum, visualize_materials as visualize\n",
+    "\n",
+    "visualize(structure, repetitions=[1, 1, 1], rotation=\"-90x\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Apply Relaxation\n",
+    "### 4.1. Load NequIP Model and Create Calculator\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from mat3ra.notebooks_utils.pyodide.packages.torch import load_nequip_model\n",
+    "from nequip.ase import NequIPCalculator\n",
+    "from nequip.data.transforms import ChemicalSpeciesToAtomTypeMapper, NeighborListTransform\n",
+    "\n",
+    "# Load the NequIP-OAM-S model from config + state_dict\n",
+    "nequip_model = load_nequip_model(\"/drive/packages/models/nequip-oam-s-config-sd.pth\")\n",
+    "\n",
+    "# Build the ASE calculator with proper transforms\n",
+    "r_max = float(nequip_model.metadata[\"r_max\"])\n",
+    "type_names = nequip_model.metadata[\"type_names\"].split(\" \")\n",
+    "\n",
+    "calculator = NequIPCalculator(\n",
+    "    model=nequip_model,\n",
+    "    device=\"cpu\",\n",
+    "    transforms=[\n",
+    "        ChemicalSpeciesToAtomTypeMapper(type_names),\n",
+    "        NeighborListTransform(r_max=r_max),\n",
+    "    ],\n",
+    ")\n",
+    "\n",
+    "print(f\"NequIP-OAM-S calculator ready (r_max={r_max} \u00c5)\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 4.2. Relax with NequIP\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import plotly.graph_objs as go\n",
+    "from IPython.display import display\n",
+    "from plotly.subplots import make_subplots\n",
+    "\n",
+    "from mat3ra.made.tools.convert import to_ase\n",
+    "from ase.optimize import BFGS\n",
+    "\n",
+    "ase_structure = to_ase(structure)\n",
+    "ase_structure.calc = calculator\n",
+    "dyn = BFGS(ase_structure)\n",
+    "\n",
+    "steps = []\n",
+    "energies = []\n",
+    "forces_max = []\n",
+    "\n",
+    "# Store original structure\n",
+    "ase_original_structure = ase_structure.copy()\n",
+    "\n",
+    "def log_step():\n",
+    "    e = ase_structure.get_potential_energy()\n",
+    "    f = ase_structure.get_forces()\n",
+    "    fmax = max(((f**2).sum(axis=1) ** 0.5))\n",
+    "    step = dyn.nsteps\n",
+    "    steps.append(step)\n",
+    "    energies.append(e)\n",
+    "    forces_max.append(fmax)\n",
+    "    print(f\"  Step {step:3d}: E = {e:12.6f} eV | Fmax = {fmax:.6f} eV/\u00c5\")\n",
+    "\n",
+    "dyn.attach(log_step)\n",
+    "\n",
+    "print(f\"Starting relaxation (FMAX = {RELAXATION_PARAMETERS['FMAX']} eV/\u00c5)...\")\n",
+    "dyn.run(fmax=RELAXATION_PARAMETERS[\"FMAX\"], steps=200)\n",
+    "ase_final_structure = ase_structure.copy()\n",
+    "\n",
+    "# Plot energy and forces convergence\n",
+    "fig = make_subplots(rows=1, cols=2, subplot_titles=(\"Energy\", \"Max Force\"))\n",
+    "fig.add_trace(go.Scatter(x=steps, y=energies, mode=\"lines+markers\", name=\"Energy\"), row=1, col=1)\n",
+    "fig.add_trace(go.Scatter(x=steps, y=forces_max, mode=\"lines+markers\", name=\"Fmax\"), row=1, col=2)\n",
+    "fig.add_hline(y=RELAXATION_PARAMETERS[\"FMAX\"], line_dash=\"dash\", line_color=\"red\", row=1, col=2)\n",
+    "fig.update_xaxes(title_text=\"Step\", row=1, col=1)\n",
+    "fig.update_xaxes(title_text=\"Step\", row=1, col=2)\n",
+    "fig.update_yaxes(title_text=\"Energy (eV)\", row=1, col=1)\n",
+    "fig.update_yaxes(title_text=\"Fmax (eV/\u00c5)\", row=1, col=2)\n",
+    "fig.update_layout(height=400, showlegend=False)\n",
+    "display(fig)\n",
+    "\n",
+    "print(f\"\\nRelaxation converged in {len(steps)} steps\")\n",
+    "print(f\"Final energy: {energies[-1]:.6f} eV\")\n",
+    "print(f\"Final Fmax:   {forces_max[-1]:.6f} eV/\u00c5\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Analyze Results\n",
+    "### 5.1. View Structure Before and After Relaxation\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from mat3ra.made.tools.convert import from_ase\n",
+    "\n",
+    "material_original = Material.create(from_ase(ase_original_structure))\n",
+    "material_relaxed = Material.create(from_ase(ase_final_structure))\n",
+    "material_original.name = structure.name\n",
+    "material_relaxed.name = structure.name + \" (NequIP Relaxed)\"\n",
+    "\n",
+    "visualize([material_original, material_relaxed], repetitions=[1, 1, 1], rotation=\"-90x\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 5.2. Output interlayer distance before and after relaxation\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from mat3ra.made.tools.analyze.other import get_average_interlayer_distance\n",
+    "\n",
+    "SUBSTRATE_TAG = 0\n",
+    "FILM_TAG = 1\n",
+    "\n",
+    "print(\n",
+    "    f\"Interlayer distance before relaxation: {get_average_interlayer_distance(material_original, SUBSTRATE_TAG, FILM_TAG):.4f} \u00c5\")\n",
+    "print(\n",
+    "    f\"Interlayer distance after relaxation:  {get_average_interlayer_distance(material_relaxed, SUBSTRATE_TAG, FILM_TAG):.4f} \u00c5\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## References\n",
+    "\n",
+    "[1] NequIP: https://github.com/mir-group/nequip\n",
+    "\n",
+    "[2] Simon Batzner et al., \"E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials,\" Nature Communications (2022)\n",
+    "\n",
+    "[3] NequIP-OAM Foundation Models: https://zenodo.org/records/18775904\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}