experiment/chgnet

config.yml

@@ -161,3 +161,15 @@ notebooks:
       - ssl
       - h5py
       - lmdb
+  - name: chgnet
+    packages_pyodide:
+      # Packages with dependencies
+      - pyyaml
+      - setuptools
+      # Packages without dependencies
+      - nodeps:ase
+      - nodeps:monty
+      # CHGNet wheel (stripped Cython extension + old model, uses legacy Python mode)
+      - emfs:/drive/packages/chgnet-0.3.8-py3-none-any.whl
+      # Stubbed packages (patched by torch_pyodide with include_chgnet=True)
+      - ssl

other/experiments/jupyterlite/relax_structure_with_chgnet.ipynb

@@ -0,0 +1,261 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Relax Structure with CHGNet \u2014 Crystal Hamiltonian Graph Network\n",
+    "\n",
+    "This notebook demonstrates structural relaxation using **CHGNet**,\n",
+    "a graph neural network interatomic potential that models the universal potential\n",
+    "energy surface of the Materials Project (MP) dataset.\n",
+    "\n",
+    "CHGNet predicts energies, forces, stresses, and magnetic moments,\n",
+    "making it well-suited for ionic and magnetic systems.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Set Input Parameters\n",
+    "### 1.1. Structure and Relaxation\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "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",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from mat3ra.notebooks_utils.packages import install_packages\n",
+    "\n",
+    "await install_packages(\"made|api_examples|torch|chgnet\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from mat3ra.notebooks_utils.pyodide.packages.torch import apply_all_patches\n",
+    "\n",
+    "apply_all_patches(include_chgnet=True)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Load Materials\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "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",
+   "metadata": {},
+   "execution_count": null,
+   "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 CHGNet Model and Create Calculator\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from chgnet.model.model import CHGNet\n",
+    "from chgnet.model.dynamics import CHGNetCalculator\n",
+    "\n",
+    "# Load the pretrained CHGNet v0.3.0 model\n",
+    "chgnet = CHGNet.load(model_name=\"0.3.0\", use_device=\"cpu\")\n",
+    "calculator = CHGNetCalculator(model=chgnet, use_device=\"cpu\")\n",
+    "\n",
+    "print(f\"CHGNet v0.3.0 loaded ({sum(p.numel() for p in chgnet.parameters()):,} parameters)\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 4.2. Relax with CHGNet\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "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 = (f**2).sum(axis=1).max()**0.5\n",
+    "    steps.append(len(steps))\n",
+    "    energies.append(e)\n",
+    "    forces_max.append(fmax)\n",
+    "    print(f\"Step {len(steps)-1:3d}: E={e:.4f} eV  Fmax={fmax:.4f} eV/\u00c5\")\n",
+    "\n",
+    "dyn.attach(log_step, interval=1)\n",
+    "log_step()  # log initial state\n",
+    "dyn.run(fmax=RELAXATION_PARAMETERS[\"FMAX\"], steps=200)\n",
+    "\n",
+    "ase_final_structure = ase_structure.copy()\n",
+    "\n",
+    "# Plot 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 (eV)\"), row=1, col=1)\n",
+    "fig.add_trace(go.Scatter(x=steps, y=forces_max, mode=\"lines+markers\", name=\"Fmax (eV/\u00c5)\"), row=1, col=2)\n",
+    "fig.update_layout(height=350, showlegend=False)\n",
+    "fig.show()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Analyze Results\n",
+    "### 5.1. View Structure Before and After Relaxation\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "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 + \" (CHGNet 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",
+   "metadata": {},
+   "execution_count": null,
+   "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] CHGNet: https://github.com/CederGroupHub/chgnet\n",
+    "\n",
+    "[2] Bowen Deng et al., \"CHGNet as a pretrained universal neural network potential for charge-informed atomistic modelling,\" Nature Machine Intelligence (2023)\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

packages/README.md

@@ -137,6 +137,62 @@ EOF
 
 ---
 
+#### `chgnet-0.3.8-py3-none-any.whl` (4.3 MB)
+
+**Source**: [CederGroupHub/chgnet](https://github.com/CederGroupHub/chgnet) v0.3.8 (PyPI)
+
+**Why custom**: The upstream wheel includes a Cython-compiled graph converter extension (`cygraph.so`) that cannot run in Pyodide's WASM runtime. CHGNet has a built-in `algorithm="legacy"` fallback that uses pure Python when the C extension is unavailable, so we simply strip the `.so` file.
+
+**How to reproduce**:
+
+```bash
+# 1. Download the upstream wheel
+pip download chgnet==0.3.8 --no-deps
+# chgnet-0.3.8-*.whl (~4.5 MB)
+
+# 2. Extract, strip the Cython extension, and repack
+python3 << 'EOF'
+import zipfile, os, shutil
+
+src_whl = next(f for f in os.listdir('.') if f.startswith('chgnet-0.3.8'))
+work_dir = "chgnet_work"
+
+# Extract
+with zipfile.ZipFile(src_whl, 'r') as z:
+    z.extractall(work_dir)
+
+# Remove the Cython .so extension (runtime falls back to pure Python)
+pkg = os.path.join(work_dir, "chgnet", "graph")
+for f in os.listdir(pkg):
+    if f.startswith("cygraph") and f.endswith(".so"):
+        os.remove(os.path.join(pkg, f))
+        print(f"Removed: {f}")
+
+# Repack
+out_whl = "chgnet-0.3.8-py3-none-any.whl"
+with zipfile.ZipFile(out_whl, 'w', zipfile.ZIP_DEFLATED) as zf:
+    for root, dirs, files in os.walk(work_dir):
+        for f in files:
+            full = os.path.join(root, f)
+            arcname = os.path.relpath(full, work_dir)
+            zf.write(full, arcname)
+
+shutil.rmtree(work_dir)
+EOF
+# Output: chgnet-0.3.8-py3-none-any.whl (~4.3 MB)
+```
+
+**What was removed**:
+- `chgnet/graph/cygraph.*.so` — Cython-compiled graph converter (runtime falls back to `algorithm="legacy"`)
+
+**What was kept**:
+- All Python source code
+- Pretrained model v0.3.0 checkpoint (~4 MB, bundled in `chgnet/pretrained/`)
+
+**Runtime patches**: Requires `apply_all_patches(include_chgnet=True)` in [torch.py](../src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py) which stubs `nvidia_smi` (GPU memory detection), `Cython` (build-time dep), and `palettable` (color palette lib imported via `pymatgen.util.plotting`).
+
+---
+
 ## Models
 
 The `models/` subdirectory contains pretrained model checkpoint files:
@@ -145,4 +201,4 @@ The `models/` subdirectory contains pretrained model checkpoint files:
 |------|-------|------|
 | `mattersim-v1.0.0-1M.pth` | MatterSim M3GNet (1M params) | ~4 MB |
 
-> **Note**: The SevenNet 7net-0 model is bundled inside the `sevenn` wheel itself under `pretrained_potentials/`.
+> **Note**: The SevenNet 7net-0 model is bundled inside the `sevenn` wheel under `pretrained_potentials/`. The CHGNet v0.3.0 model is bundled inside the `chgnet` wheel under `chgnet/pretrained/`.