From 8bb079bf846e72388fe6d2697fd50d860b8d3032 Mon Sep 17 00:00:00 2001 From: Timur Bazhirov Date: Wed, 20 May 2026 00:14:05 -0700 Subject: [PATCH 1/2] feat: add CHGNet support for Pyodide/JupyterLite - Add chgnet-0.3.8 wheel (stripped Cython, model bundled, 4.3MB) - Add patch_chgnet_deps: stubs nvidia_smi, Cython, palettable - Add CHGNet package group to config.yml - Add relax_structure_with_chgnet.ipynb notebook - Add Playwright test (all 9 cells pass) CHGNet uses pure PyTorch loading (state_dict, no JIT) and pymatgen for neighbor lists (no DGL/torch_geometric). The cygraph.so C extension automatically falls back to Python 'legacy' mode in Pyodide. --- config.yml | 12 + .../relax_structure_with_chgnet.ipynb | 261 ++++++++++++++++++ packages/chgnet-0.3.8-py3-none-any.whl | 3 + .../notebooks_utils/pyodide/packages/torch.py | 49 +++- tests/playwright/test_chgnet_notebook.mjs | 93 +++++++ 5 files changed, 417 insertions(+), 1 deletion(-) create mode 100644 other/experiments/jupyterlite/relax_structure_with_chgnet.ipynb create mode 100644 packages/chgnet-0.3.8-py3-none-any.whl create mode 100644 tests/playwright/test_chgnet_notebook.mjs diff --git a/config.yml b/config.yml index 3ef077674..3e774386e 100644 --- a/config.yml +++ b/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 diff --git a/other/experiments/jupyterlite/relax_structure_with_chgnet.ipynb b/other/experiments/jupyterlite/relax_structure_with_chgnet.ipynb new file mode 100644 index 000000000..19be75f91 --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/packages/chgnet-0.3.8-py3-none-any.whl b/packages/chgnet-0.3.8-py3-none-any.whl new file mode 100644 index 000000000..e5bfd31b1 --- /dev/null +++ b/packages/chgnet-0.3.8-py3-none-any.whl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:623e0ccbdd220323a93bfa562b9bdffe55a80a8d055105e794fcc25e13f6d85c +size 4473561 diff --git a/src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py b/src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py index 15270680b..98a81925d 100644 --- a/src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py +++ b/src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py @@ -959,7 +959,7 @@ def patch_mace_tools(): # ============================================================================== -def apply_all_patches(include_fairchem=False, include_mattersim=False, include_sevennet=False): +def apply_all_patches(include_fairchem=False, include_mattersim=False, include_sevennet=False, include_chgnet=False): """ Apply all torch and model patches for Pyodide in one call. @@ -973,6 +973,9 @@ def apply_all_patches(include_fairchem=False, include_mattersim=False, include_s include_sevennet: If True, also apply SevenNet-specific patches (pandas, tqdm, packaging stubs). Set this when using SevenNet / 7net models. + include_chgnet: If True, also apply CHGNet-specific patches + (nvidia_smi, cython stubs). Set this when + using CHGNet models. """ patch_torch_linalg() patch_torch_compiler() @@ -993,6 +996,9 @@ def apply_all_patches(include_fairchem=False, include_mattersim=False, include_s patch_torch_distributed() patch_sevennet_deps() + if include_chgnet: + patch_chgnet_deps() + print("\n✅ All Pyodide patches applied successfully!") @@ -1417,3 +1423,44 @@ def _noop_script(obj=None, *a, **k): pass print("✓ SevenNet dependency stubs applied") + + +# ============================================================================== +# CHGNet patches +# ============================================================================== + + +def patch_chgnet_deps(): + """ + Stub dependencies required by CHGNet but not needed for inference in Pyodide. + + Stubs: nvidia_smi (GPU memory detection), Cython. + CHGNet's cygraph.so automatically falls back to pure Python 'legacy' mode + when the C extension is unavailable, so no graph library patching is needed. + """ + # --- stub nvidia_smi (GPU memory detection, not needed for CPU inference) --- + if "nvidia_smi" not in sys.modules: + nvidia_mod = types.ModuleType("nvidia_smi") + nvidia_mod.nvmlInit = lambda: None + nvidia_mod.nvmlShutdown = lambda: None + nvidia_mod.nvmlDeviceGetCount = lambda: 0 + nvidia_mod.nvmlDeviceGetHandleByIndex = lambda idx: None + nvidia_mod.nvmlDeviceGetMemoryInfo = lambda handle: None + sys.modules["nvidia_smi"] = nvidia_mod + + # --- stub Cython (build-time dependency, not needed at runtime) --- + for mod_name in ["Cython", "Cython.Build", "cython"]: + if mod_name not in sys.modules: + sys.modules[mod_name] = types.ModuleType(mod_name) + + # --- stub palettable (color palette lib used by pymatgen.util.plotting, + # imported via chgnet.model.dynamics -> pymatgen.analysis.eos) --- + for mod_name in [ + "palettable", + "palettable.colorbrewer", + "palettable.colorbrewer.diverging", + ]: + if mod_name not in sys.modules: + sys.modules[mod_name] = types.ModuleType(mod_name) + + print("✓ CHGNet dependency stubs applied") diff --git a/tests/playwright/test_chgnet_notebook.mjs b/tests/playwright/test_chgnet_notebook.mjs new file mode 100644 index 000000000..eead4c2f4 --- /dev/null +++ b/tests/playwright/test_chgnet_notebook.mjs @@ -0,0 +1,93 @@ +import { chromium } from "playwright"; + +const URL = + "http://localhost:8000/lab/index.html?path=relax_structure_with_chgnet.ipynb"; + +(async () => { + const browser = await chromium.launch({ channel: "chrome", headless: false }); + const page = await browser.newPage(); + + console.log("⏳ Navigating to JupyterLite..."); + await page.goto(URL, { waitUntil: "domcontentloaded", timeout: 60_000 }); + + // Wait for kernel idle + console.log("⏳ Waiting for Pyodide kernel..."); + await page.waitForFunction( + () => { + const ind = document.querySelector(".jp-Notebook-ExecutionIndicator"); + return ind?.dataset?.status === "idle" && document.querySelectorAll(".jp-CodeCell").length > 5; + }, + null, + { timeout: 120_000, polling: 2_000 } + ); + console.log("✅ Kernel idle!"); + + // Click Restart & Run All + await page.locator('button[data-command="runmenu:restart-and-run-all"]').click(); + await page.waitForTimeout(1_000); + try { + await page.locator(".jp-Dialog button", { hasText: /restart/i }).click({ timeout: 3_000 }); + console.log("✅ Restart & Run All confirmed!"); + } catch { + console.log("⚠ No dialog"); + } + + // Monitor with logging + console.log("\n📊 Monitoring..."); + const start = Date.now(); + + while (Date.now() - start < 300_000) { + await page.waitForTimeout(10_000); + + const s = await page.evaluate(() => { + const cells = document.querySelectorAll(".jp-CodeCell"); + let running = 0, done = 0, errors = 0; + cells.forEach((c) => { + const p = c.querySelector(".jp-InputPrompt")?.textContent?.trim() || ""; + if (p === "[*]:") running++; + else if (/\[\d+\]:/.test(p)) done++; + if (c.querySelector(".jp-RenderedText[data-mime-type='application/vnd.jupyter.stderr']")) errors++; + }); + const k = document.querySelector(".jp-Notebook-ExecutionIndicator")?.dataset?.status; + return { running, done, errors, total: cells.length, k }; + }); + + const t = ((Date.now() - start) / 1000) | 0; + console.log(` [${t}s] kernel=${s.k} running=${s.running} done=${s.done}/${s.total} errors=${s.errors}`); + + if (s.k === "idle" && s.running === 0 && s.done > 0) break; + } + + // Final report + console.log("\n" + "=".repeat(60)); + const results = await page.evaluate(() => + Array.from(document.querySelectorAll(".jp-CodeCell")).map((c, i) => { + let t = ""; + for (const o of c.querySelectorAll(".jp-OutputArea-output")) t += o.textContent; + return { cell: i + 1, err: t.includes("Traceback"), out: t }; + }) + ); + + let failures = 0; + for (const r of results) { + if (r.err) { + failures++; + console.log(`Cell ${r.cell} ❌`); + console.log(r.out.substring(0, 2000)); + console.log(); + } else { + const preview = r.out.trim().substring(0, 120).replace(/\n/g, " | "); + console.log(`Cell ${r.cell} ✅${preview ? ": " + preview : ""}`); + } + } + + console.log("=".repeat(60)); + if (failures > 0) { + console.log(`\n⚠ ${failures} cell(s) failed`); + process.exitCode = 1; + } else { + console.log(`\n🎉 ALL ${results.length} CELLS PASSED!`); + } + + await browser.close(); +})(); From d0c5a988b7749f98d10341d205b002742f5ea5c6 Mon Sep 17 00:00:00 2001 From: Timur Bazhirov Date: Wed, 20 May 2026 00:21:44 -0700 Subject: [PATCH 2/2] docs: add CHGNet build instructions to packages README --- packages/README.md | 58 +++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 57 insertions(+), 1 deletion(-) diff --git a/packages/README.md b/packages/README.md index db269077f..03d3256a0 100644 --- a/packages/README.md +++ b/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/`.