From bb45e12e776d0ee51fb965bc5e5d68f96130519e Mon Sep 17 00:00:00 2001 From: Timur Bazhirov Date: Wed, 20 May 2026 18:08:30 -0700 Subject: [PATCH 1/3] feature: nequip poc added --- config.yml | 16 + .../relax_structure_with_nequip.ipynb | 287 +++++++++++++++ .../notebooks_utils/pyodide/packages/torch.py | 330 +++++++++++++++++- tests/playwright/test_nequip_notebook.mjs | 93 +++++ 4 files changed, 725 insertions(+), 1 deletion(-) create mode 100644 other/experiments/jupyterlite/relax_structure_with_nequip.ipynb create mode 100644 tests/playwright/test_nequip_notebook.mjs diff --git a/config.yml b/config.yml index 3e774386e..c2baa2077 100644 --- a/config.yml +++ b/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 diff --git a/other/experiments/jupyterlite/relax_structure_with_nequip.ipynb b/other/experiments/jupyterlite/relax_structure_with_nequip.ipynb new file mode 100644 index 000000000..127d6825e --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py b/src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py index 98a81925d..c4f31881f 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, include_chgnet=False): +def apply_all_patches(include_fairchem=False, include_mattersim=False, include_sevennet=False, include_chgnet=False, include_nequip=False): """ Apply all torch and model patches for Pyodide in one call. @@ -976,6 +976,9 @@ def apply_all_patches(include_fairchem=False, include_mattersim=False, include_s include_chgnet: If True, also apply CHGNet-specific patches (nvidia_smi, cython stubs). Set this when using CHGNet models. + include_nequip: If True, also apply NequIP-specific patches + (lightning, hydra, torchmetrics, e3nn JIT stubs). Set this + when using NequIP models. """ patch_torch_linalg() patch_torch_compiler() @@ -999,6 +1002,9 @@ def apply_all_patches(include_fairchem=False, include_mattersim=False, include_s if include_chgnet: patch_chgnet_deps() + if include_nequip: + patch_nequip_deps() + print("\n✅ All Pyodide patches applied successfully!") @@ -1464,3 +1470,325 @@ def patch_chgnet_deps(): sys.modules[mod_name] = types.ModuleType(mod_name) print("✓ CHGNet dependency stubs applied") + + +# ============================================================================== +# NequIP patches +# ============================================================================== + + +def patch_nequip_deps(): + """ + Stub heavy dependencies required by NequIP but not needed for inference. + + Stubs: hydra, lightning, pytorch_lightning, torchmetrics, lmdb, matscipy. + Patches e3nn and torch.jit for Pyodide compatibility. + Sets NEQUIP_NL=ase to use ASE neighbor lists instead of matscipy. + """ + import os + import torch + + # --- Set environment to use ASE neighbor lists --- + os.environ["NEQUIP_NL"] = "ase" + + # --- lightning_utilities (used by nequip.utils.logger.RankedLogger) --- + _lu = _make_stub_module("lightning_utilities", submodules=[ + "core", "core.rank_zero", + ]) + + def _rank_prefixed_message(msg, rank=None): + return msg + + def _rank_zero_only(fn): + return fn + + sys.modules["lightning_utilities.core.rank_zero"].rank_prefixed_message = _rank_prefixed_message + sys.modules["lightning_utilities.core.rank_zero"].rank_zero_only = _rank_zero_only + + # --- lightning / pytorch_lightning --- + # NequIP uses lightning for training (EMALightningModule, Trainer, etc.) + # Stub all submodules needed by NequIP's import chain + _pl = _make_stub_module("lightning", submodules=[ + "pytorch", "pytorch.utilities", "pytorch.utilities.seed", + "pytorch.utilities.warnings", "pytorch.callbacks", + ]) + + class _IsolateRng: + def __enter__(self): + return self + def __exit__(self, *a): + pass + + def _seed_everything(seed=None, workers=False, **kwargs): + pass + + sys.modules["lightning.pytorch.utilities.seed"].isolate_rng = lambda: _IsolateRng() + sys.modules["lightning.pytorch"].seed_everything = _seed_everything + + # PossibleUserWarning used in nequip.train.lightning + class _PossibleUserWarning(UserWarning): + pass + sys.modules["lightning.pytorch.utilities.warnings"].PossibleUserWarning = _PossibleUserWarning + + # LightningModule used as base class in nequip.train + class _LightningModule(torch.nn.Module): + def __init__(self, *a, **k): + super().__init__() + def log(self, *a, **k): + pass + sys.modules["lightning.pytorch"].LightningModule = _LightningModule + sys.modules["lightning"].pytorch = sys.modules["lightning.pytorch"] + + # Callback for lightning.pytorch.callbacks + sys.modules["lightning.pytorch.callbacks"].Callback = type("Callback", (), {}) + + _make_stub_module("pytorch_lightning", submodules=[ + "utilities", "utilities.seed", + ]) + sys.modules["pytorch_lightning.utilities.seed"].isolate_rng = lambda: _IsolateRng() + + # --- torchmetrics (Metric is used as base class by DataStatisticsManager) --- + _tm = _make_stub_module("torchmetrics") + + class _Metric(torch.nn.Module): + def __init__(self, *a, **k): + super().__init__() + + def add_state(self, name, default=None, dist_reduce_fx=None, **k): + if default is not None: + setattr(self, name, default) + _tm.Metric = _Metric + _tm.MeanMetric = _Metric + + # --- packaging (needed by nequip.__init__) --- + if "packaging" not in sys.modules: + try: + import packaging.version # noqa: F401 + except (ImportError, ModuleNotFoundError): + pkg = _make_stub_module("packaging", submodules=["version"]) + + class _Version: + def __init__(self, v): + self._v = str(v) + parts = self._v.split(".") + self.major = int(parts[0]) if parts else 0 + self.minor = int(parts[1]) if len(parts) > 1 else 0 + + def __lt__(self, other): + return (self.major, self.minor) < (other.major, other.minor) + + def __ge__(self, other): + return not self.__lt__(other) + + def __repr__(self): + return f"Version('{self._v}')" + + sys.modules["packaging.version"].Version = _Version + sys.modules["packaging.version"].parse = lambda v: _Version(v) + + # --- lmdb --- + if "lmdb" not in sys.modules: + sys.modules["lmdb"] = types.ModuleType("lmdb") + + # --- hydra (NequIP uses hydra.utils.instantiate for ZBL pair potential) --- + if "hydra" not in sys.modules: + _make_stub_module("hydra", submodules=[ + "core", "core.global_hydra", "utils", + "_internal", "_internal.instantiate", "_internal.instantiate._instantiate2", + ]) + + def _hydra_instantiate(config, *args, _recursive_=True, **kwargs): + import importlib + if isinstance(config, dict) and "_target_" in config: + target = config["_target_"] + mod_path, cls_name = target.rsplit(".", 1) + mod = importlib.import_module(mod_path) + cls = getattr(mod, cls_name) + pos_args = list(args) + list(config.get("_args_", [])) + cfg = {k: v for k, v in config.items() if not k.startswith("_")} + if _recursive_: + for k, v in cfg.items(): + if isinstance(v, dict) and "_target_" in v: + cfg[k] = _hydra_instantiate(v) + elif isinstance(v, list): + cfg[k] = [_hydra_instantiate(i) if isinstance(i, dict) and "_target_" in i else i for i in v] + pos_args = [_hydra_instantiate(a) if isinstance(a, dict) and "_target_" in a else a for a in pos_args] + return cls(*pos_args, **{**cfg, **kwargs}) + return config + + sys.modules["hydra.utils"].instantiate = _hydra_instantiate + sys.modules["hydra"].utils = sys.modules["hydra.utils"] + sys.modules["hydra._internal.instantiate._instantiate2"].InstantiationException = type( + "InstantiationException", (Exception,), {} + ) + + # get_method / get_class used by nequip.train.lightning + def _hydra_get_target(target_str): + import importlib + mod_path, name = target_str.rsplit(".", 1) + return getattr(importlib.import_module(mod_path), name) + + sys.modules["hydra.utils"].get_method = _hydra_get_target + sys.modules["hydra.utils"].get_class = _hydra_get_target + + # --- omegaconf --- + if "omegaconf" not in sys.modules: + omegaconf_mod = _make_stub_module("omegaconf") + + class _DictConfig(dict): + pass + + class _ListConfig(list): + pass + + omegaconf_mod.DictConfig = _DictConfig + omegaconf_mod.ListConfig = _ListConfig + omegaconf_mod.OmegaConf = type( + "OmegaConf", + (), + { + "to_container": staticmethod(lambda cfg, **k: dict(cfg) if isinstance(cfg, dict) else cfg), + "create": staticmethod(lambda d: _DictConfig(d) if isinstance(d, dict) else d), + "register_new_resolver": staticmethod(lambda name, func, **k: None), + }, + ) + + # --- matscipy (use ASE neighbor lists) --- + if "matscipy" not in sys.modules: + _matscipy = types.ModuleType("matscipy") + _matscipy.__path__ = [] + _matscipy.__package__ = "matscipy" + _matscipy_neighbours = types.ModuleType("matscipy.neighbours") + _matscipy_neighbours.neighbour_list = _matscipy_neighbour_list_compat + _matscipy.neighbours = _matscipy_neighbours + sys.modules["matscipy"] = _matscipy + sys.modules["matscipy.neighbours"] = _matscipy_neighbours + + # --- patch e3nn to skip JIT compilation --- + try: + import e3nn + e3nn._SO3_INITIALIZED = True + if hasattr(e3nn, '_OPT_DEFAULTS'): + e3nn._OPT_DEFAULTS["jit_mode"] = "eager" + except Exception: + pass + + # --- patch torch.jit for e3nn --- + if not hasattr(torch.jit, '_original_script'): + _orig_script = torch.jit.script + + def _noop_script(obj=None, *a, **k): + if obj is not None: + return obj + return lambda fn: fn + + torch.jit.script = _noop_script + + # --- patch e3nn compile_mode decorator --- + try: + import e3nn.util.jit + e3nn.util.jit.compile_mode = lambda mode: lambda cls: cls + except Exception: + pass + + # --- tqdm --- + if "tqdm" not in sys.modules: + tqdm_mod = _make_stub_module("tqdm", submodules=["auto", "std"]) + + def _tqdm_passthrough(iterable=None, *a, **k): + return iterable if iterable is not None else iter([]) + + tqdm_mod.tqdm = _tqdm_passthrough + sys.modules["tqdm.auto"].tqdm = _tqdm_passthrough + sys.modules["tqdm.std"].tqdm = _tqdm_passthrough + + print("✓ NequIP dependency stubs applied") + + +def load_nequip_model(checkpoint_path): + """ + Load a NequIP model from a config+state_dict checkpoint file. + + This bypasses torch.package (which doesn't work in Pyodide) by + rebuilding the model architecture from saved config parameters + and loading the state_dict directly. + + Args: + checkpoint_path: Path to the .pth file containing config + state_dict. + Expected keys: type_names, r_max, irreps_edge_sh, + type_embed_num_features, feature_irreps_hidden, + radial_mlp_depth, radial_mlp_width, avg_num_neighbors, + per_type_energy_scales, per_type_energy_shifts, + has_zbl, state_dict. + + Returns: + A NequIP GraphModel ready for inference. + """ + import torch + + data = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + + # Override set_global_state with Pyodide-safe version + # The real one calls torch.jit.set_fusion_strategy, torch.multiprocessing, + # and e3nn.set_optimization_defaults(jit_script_fx=True), which don't work in WASM. + import nequip.utils.global_state as _gs + + def _pyodide_set_global_state(allow_tf32=False, warn_on_override=False): + if not _gs._GLOBAL_STATE_INITIALIZED: + torch.set_default_dtype(torch.float64) + try: + import e3nn + e3nn.set_optimization_defaults( + specialized_code=True, + optimize_einsums=True, + jit_script_fx=False, # Disabled for Pyodide + ) + except Exception: + pass + _gs._GLOBAL_STATE_INITIALIZED = True + _gs._latest_global_config["allow_tf32"] = allow_tf32 + + _gs.set_global_state = _pyodide_set_global_state + + # Initialize NequIP global state + _gs.set_global_state(allow_tf32=False) + + # Build the model from config + from nequip.model import FullNequIPGNNModel + + model = FullNequIPGNNModel( + seed=0, + model_dtype="float32", + r_max=data["r_max"], + type_names=data["type_names"], + irreps_edge_sh=data["irreps_edge_sh"], + type_embed_num_features=data["type_embed_num_features"], + feature_irreps_hidden=data["feature_irreps_hidden"], + radial_mlp_depth=data["radial_mlp_depth"], + radial_mlp_width=data["radial_mlp_width"], + avg_num_neighbors=data["avg_num_neighbors"], + per_type_energy_scales=data["per_type_energy_scales"], + per_type_energy_shifts=data["per_type_energy_shifts"], + polynomial_cutoff_p=data.get("polynomial_cutoff_p", 6), + ) + + # Add ZBL pair potential if needed (must come BEFORE total_energy_sum) + if data.get("has_zbl", False): + from nequip.nn.pair_potential import ZBL + seq_net = model.model.func + zbl = ZBL( + type_names=data["type_names"], + chemical_species=data["type_names"], + units="metal", + irreps_in=seq_net.irreps_out, + ) + seq_net.insert( + name="pair_potential", module=zbl, before="total_energy_sum" + ) + + # Load the state_dict + model.load_state_dict(data["state_dict"], strict=True) + model.eval() + + print(f"✓ NequIP model loaded ({sum(p.numel() for p in model.parameters()):,} parameters)") + return model diff --git a/tests/playwright/test_nequip_notebook.mjs b/tests/playwright/test_nequip_notebook.mjs new file mode 100644 index 000000000..d1d38eaf3 --- /dev/null +++ b/tests/playwright/test_nequip_notebook.mjs @@ -0,0 +1,93 @@ +import { chromium } from "playwright"; + +const URL = + "http://localhost:8000/lab/index.html?path=relax_structure_with_nequip.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 b4f262d0bfa5b8807c37f4205c51b4a7af2ec67f Mon Sep 17 00:00:00 2001 From: Timur Bazhirov Date: Wed, 20 May 2026 18:16:36 -0700 Subject: [PATCH 2/3] chore: git attributes file --- .gitattributes | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitattributes b/.gitattributes index 7f638c094..a32ccd90f 100644 --- a/.gitattributes +++ b/.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 From 6ce013157f5b5055b36181df4874a3b101581471 Mon Sep 17 00:00:00 2001 From: Timur Bazhirov Date: Wed, 20 May 2026 18:18:10 -0700 Subject: [PATCH 3/3] chore: whl, pth --- packages/README.md | 105 ++++++++++++++++++++- packages/models/nequip-oam-s-config-sd.pth | 3 + packages/nequip-0.15.0-py3-none-any.whl | 3 + 3 files changed, 108 insertions(+), 3 deletions(-) create mode 100644 packages/models/nequip-oam-s-config-sd.pth create mode 100644 packages/nequip-0.15.0-py3-none-any.whl diff --git a/packages/README.md b/packages/README.md index 03d3256a0..1ce5a7632 100644 --- a/packages/README.md +++ b/packages/README.md @@ -193,12 +193,111 @@ EOF --- +#### `nequip-0.15.0-py3-none-any.whl` (254 KB) + +**Source**: [mir-group/nequip](https://github.com/mir-group/nequip) v0.15.0 (PyPI) + +**Why custom**: The upstream NequIP wheel declares heavy training dependencies (`hydra-core`, `lightning`, `torchmetrics`, `lmdb`) that are not needed for inference and would fail to install in Pyodide. The custom wheel strips these from `Requires-Dist` metadata. + +**How to reproduce**: + +```bash +# 1. Create a venv with build tools +python3 -m venv /tmp/nequip_build && source /tmp/nequip_build/bin/activate +pip install build wheel setuptools + +# 2. Download and extract source +pip download nequip==0.15.0 --no-deps --no-binary :all: +tar xf nequip-0.15.0.tar.gz && cd nequip-0.15.0 + +# 3. Strip heavy dependencies from pyproject.toml +# Remove these lines from [project].dependencies: +# "hydra-core", "lightning", "torchmetrics>=1.6.0", +# "lmdb", "tqdm", "requests" +# Keep: torch, numpy, matscipy, ase, e3nn, pyyaml + +# 4. Build the wheel +python -m build --wheel +# Output: dist/nequip-0.15.0-py3-none-any.whl (~254 KB) +``` + +**Runtime patches**: Requires `apply_all_patches(include_nequip=True)` in [torch.py](../src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py) which stubs `hydra` (with working `instantiate`), `lightning`, `pytorch_lightning`, `torchmetrics`, `lmdb`, `matscipy` (uses ASE neighbor lists via `NEQUIP_NL=ase`), and patches `e3nn` for eager mode (no JIT). + +--- + ## Models The `models/` subdirectory contains pretrained model checkpoint files: -| File | Model | Size | -|------|-------|------| -| `mattersim-v1.0.0-1M.pth` | MatterSim M3GNet (1M params) | ~4 MB | +| File | Model | Size | Notes | +|------|-------|------|-------| +| `mattersim-v1.0.0-1M.pth` | MatterSim M3GNet (1M params) | ~4 MB | Direct checkpoint | +| `nequip-oam-s-config-sd.pth` | NequIP-OAM-S (618K params) | ~2.4 MB | Config + state_dict extracted from `.nequip.zip` package | > **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/`. + +### NequIP-OAM-S model extraction + +The `nequip-oam-s-config-sd.pth` file contains the model architecture config and state_dict extracted from the NequIP package format (`.nequip.zip`). This is necessary because NequIP's standard loading paths use `torch.package.PackageImporter` or `torch.jit.load`, neither of which work in Pyodide. + +**How to reproduce**: + +```bash +# 1. Set up venv with NequIP +python3 -m venv /tmp/nequip_venv && source /tmp/nequip_venv/bin/activate +pip install nequip==0.15.0 + +# 2. Download the NequIP-OAM-S model from Zenodo +wget https://zenodo.org/records/18775904/files/NequIP-OAM-S-0.1.nequip.zip + +# 3. Extract config + state_dict +python3 << 'PYEOF' +import torch, io, yaml + +pkg_path = "NequIP-OAM-S-0.1.nequip.zip" +imp = torch.package.PackageImporter(pkg_path) + +# Patch for CPU-only loading +orig = torch.storage._load_from_bytes +def _load_cpu(b): + return torch.load(io.BytesIO(b), map_location="cpu", weights_only=False) +torch.storage._load_from_bytes = _load_cpu +pkg_model = imp.load_pickle(package="model", resource="eager_model.pkl", map_location="cpu") +torch.storage._load_from_bytes = orig + +# Extract state_dict and metadata +inner_sd = pkg_model.sole_model.state_dict() +metadata = yaml.safe_load(imp.load_text(package="model", resource="package_metadata.txt")) +type_names = [metadata["atom_types"][i] for i in range(len(metadata["atom_types"]))] + +# Extract per-type energy scales and shifts +scales_t = pkg_model.sole_model.model.func.per_type_energy_scale_shift.scales.data +shifts_t = pkg_model.sole_model.model.func.per_type_energy_scale_shift.shifts.data +scales_dict = {tn: scales_t[i, 0].item() for i, tn in enumerate(type_names)} +shifts_dict = {tn: shifts_t[i, 0].item() for i, tn in enumerate(type_names)} + +# Extract avg_num_neighbors from scatter_norm_factor +snf = pkg_model.sole_model.model.func.layer0_convnet.conv.scatter_norm_factor +avg_nn = 1.0 / (snf ** 2) + +# Save config + state_dict +torch.save({ + "type_names": type_names, + "r_max": 4.5, + "irreps_edge_sh": "1x0e+1x1e", + "type_embed_num_features": 32, + "feature_irreps_hidden": ["128x0e+64x1e", "128x0e"], + "radial_mlp_depth": [1, 1], + "radial_mlp_width": [128, 128], + "avg_num_neighbors": avg_nn, + "per_type_energy_scales": scales_dict, + "per_type_energy_shifts": shifts_dict, + "has_zbl": True, + "state_dict": inner_sd, + "metadata": metadata, +}, "nequip-oam-s-config-sd.pth") +PYEOF +# Output: nequip-oam-s-config-sd.pth (~2.4 MB) +``` + +The model is then loaded at runtime using `load_nequip_model()` from [torch.py](../src/py/mat3ra/notebooks_utils/pyodide/packages/torch.py), which rebuilds the architecture using `FullNequIPGNNModel` + manually adds the ZBL pair potential, then loads the state_dict. diff --git a/packages/models/nequip-oam-s-config-sd.pth b/packages/models/nequip-oam-s-config-sd.pth new file mode 100644 index 000000000..ebaf076c5 --- /dev/null +++ b/packages/models/nequip-oam-s-config-sd.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e547801c211ebdef3750e9dfc1baa8b1919c7d4e66624e9f316e50875e445eea +size 2495009 diff --git a/packages/nequip-0.15.0-py3-none-any.whl b/packages/nequip-0.15.0-py3-none-any.whl new file mode 100644 index 000000000..2462ad72a --- /dev/null +++ b/packages/nequip-0.15.0-py3-none-any.whl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0b3ffe7b2c95dd4954515c150e374890a49dd091b1c8db255772601d037df6e3 +size 260116