experiment: support for uma

config.yml

@@ -101,3 +101,20 @@ notebooks:
       - ssl
       - h5py
       - lmdb
+  - name: uma
+    packages_pyodide:
+      # Packages with dependencies
+      - opt_einsum
+      - orjson
+      - pyyaml
+      - sqlite3
+      # Packages without dependencies (using nodeps: prefix)
+      - nodeps:opt_einsum_fx
+      - nodeps:e3nn>=0.5
+      - nodeps:ase
+      - nodeps:monty
+      - nodeps:fairchem-core
+      # Stubbed packages (will be patched by torch_pyodide with include_fairchem=True)
+      - ssl
+      - h5py
+      - lmdb

other/experiments/jupyterlite/relax_structure_with_uma.ipynb

@@ -0,0 +1,300 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "9515ed910c085db8",
+   "metadata": {},
+   "source": [
+    "# Relax Structure with FAIRChem UMA — Universal Machine-learning Force Field\n",
+    "\n",
+    "Use FAIRChem's [UMA](https://github.com/FAIR-Chem/fairchem) (Universal Model for Atoms) to relax crystal structures using machine-learned interatomic potentials.\n",
+    "\n",
+    "<h2 style=\"color:green\">Usage</h2>\n",
+    "\n",
+    "1. Drop the materials files into the \"uploads\" folder in the JupyterLab file browser\n",
+    "1. Set Input Parameters below or use the default values\n",
+    "1. Click \"Run\" > \"Run All\" to run all cells\n",
+    "1. Wait for the run to complete. Scroll down to view cell results.\n",
+    "1. Review the relaxation plot and modify parameters as needed\n",
+    "\n",
+    "## Methodology\n",
+    "\n",
+    "1. Load materials from JSON files and create structure via `mat3ra-made`\n",
+    "2. Install FAIRChem UMA and apply Pyodide patches\n",
+    "3. Convert to ASE atoms with `to_ase()`\n",
+    "4. Relax the structure with FAIRChem UMA and visualize convergence\n",
+    "5. Compute relaxation energy"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4737d145950b1cc8",
+   "metadata": {},
+   "source": [
+    "## 1. Set Input Parameters\n",
+    "### 1.1. Structure and Relaxation Parameters\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "57dc565952aa2e4e",
+   "metadata": {
+    "trusted": true
+   },
+   "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",
+    "UMA_TASK = \"s2ef\"  # \"s2ef\" (structure to energy/forces) or \"is2re\" (initial structure to relaxed energy)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4e89f2d820acb1ed",
+   "metadata": {},
+   "source": [
+    "## 2. Install Packages"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": [
+    "from mat3ra.notebooks_utils.packages import install_packages\n",
+    "\n",
+    "await install_packages(\"made|api_examples|torch|uma\")\n",
+    "\n",
+    "    from mat3ra.notebooks_utils.pyodide.packages.torch import apply_all_patches\n",
+    "\n",
+    "    apply_all_patches(include_fairchem=True)"
+   ],
+   "id": "ece62358982306f2"
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f89d3c98ddce2ab5",
+   "metadata": {},
+   "source": [
+    "## 3. Load Materials"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8fd400dace70549e",
+   "metadata": {
+    "trusted": true
+   },
+   "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))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "42f12abf6b65aa2c",
+   "metadata": {},
+   "source": [
+    "### 3.1. Visualize Input Structure"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fcbb4e6c1de21233",
+   "metadata": {
+    "trusted": true
+   },
+   "outputs": [],
+   "source": [
+    "from mat3ra.notebooks_utils.ipython.entity.material.visualize import ViewersEnum, visualize_materials as visualize\n",
+    "\n",
+    "visualize([{\"material\": structure, \"title\": structure.name}], viewer=ViewersEnum.wave)\n",
+    "visualize(structure, repetitions=[1, 1, 1], rotation=\"-90x\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e64688fc18c49bb6",
+   "metadata": {},
+   "source": [
+    "## 4. Apply Relaxation\n",
+    "### 4.1. Relax with FAIRChem UMA"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3d8746a77f71bab5",
+   "metadata": {
+    "jupyter": {
+     "is_executing": true
+    },
+    "trusted": true
+   },
+   "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",
+    "from fairchem.core import FAIRChemCalculator, pretrained_mlip\n",
+    "\n",
+    "# Load UMA model\n",
+    "predictor = pretrained_mlip.get_predict_unit(\"uma-s-1\", device=\"cpu\")\n",
+    "calculator = FAIRChemCalculator(predictor, task_name=UMA_TASK)\n",
+    "\n",
+    "ase_structure = to_ase(structure)\n",
+    "ase_structure.set_calculator(calculator)\n",
+    "dyn = BFGS(ase_structure)\n",
+    "\n",
+    "steps = []\n",
+    "energies = []\n",
+    "\n",
+    "fig = make_subplots(rows=1, cols=1, specs=[[{\"type\": \"scatter\"}]])\n",
+    "scatter = go.Scatter(x=[], y=[], mode=\"lines+markers\", name=\"Energy\")\n",
+    "fig.add_trace(scatter)\n",
+    "fig.update_layout(title_text=\"Real-time Optimization Progress\", xaxis_title=\"Step\", yaxis_title=\"Energy (eV)\")\n",
+    "\n",
+    "f = go.FigureWidget(fig)\n",
+    "display(f)\n",
+    "\n",
+    "\n",
+    "def plotly_callback():\n",
+    "    step = dyn.nsteps\n",
+    "    energy = ase_structure.get_total_energy()\n",
+    "    steps.append(step)\n",
+    "    energies.append(energy)\n",
+    "    print(f\"Step: {step}, Energy: {energy:.4f} eV\")\n",
+    "    with f.batch_update():\n",
+    "        f.data[0].x = steps\n",
+    "        f.data[0].y = energies\n",
+    "\n",
+    "\n",
+    "dyn.attach(plotly_callback, interval=1)\n",
+    "dyn.run(fmax=RELAXATION_PARAMETERS[\"FMAX\"])\n",
+    "\n",
+    "ase_original_structure = to_ase(structure)\n",
+    "ase_original_structure.set_calculator(calculator)\n",
+    "ase_final_structure = ase_structure\n",
+    "\n",
+    "original_energy = ase_original_structure.get_total_energy()\n",
+    "relaxed_energy = ase_structure.get_total_energy()\n",
+    "\n",
+    "print(f\"The final energy is {float(relaxed_energy):.3f} eV.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "abfa372909a96bf8",
+   "metadata": {},
+   "source": [
+    "## 5. Analyze Results\n",
+    "### 5.1. View Structure Before and After Relaxation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9565d0931b198f63",
+   "metadata": {
+    "trusted": true
+   },
+   "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 + \" (UMA Relaxed)\"\n",
+    "\n",
+    "visualize(\n",
+    "    [\n",
+    "        {\"material\": material_original, \"title\": material_original.name},\n",
+    "        {\"material\": material_relaxed, \"title\": material_relaxed.name},\n",
+    "    ],\n",
+    "    viewer=ViewersEnum.wave,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e4b49774283e5517",
+   "metadata": {},
+   "source": [
+    "### 5.2. Output interlayer distance before and after relaxation\n",
+    "This requires labels for substrate and film present in the interface structure, which is already done if interface was created with `mat3ra-made`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6dd00402bc2e9d59",
+   "metadata": {
+    "trusted": true
+   },
+   "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} Å\")\n",
+    "print(\n",
+    "    f\"Interlayer distance after relaxation:  {get_average_interlayer_distance(material_relaxed, SUBSTRATE_TAG, FILM_TAG):.4f} Å\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2f60fdb73e44c09c",
+   "metadata": {},
+   "source": [
+    "## References\n",
+    "\n",
+    "[1] FAIRChem: https://github.com/FAIR-Chem/fairchem  \n",
+    "[2] UMA — Universal Machine-learning Force Field for Atomistic Systems: https://arxiv.org/abs/2410.22570  \n",
+    "[3] mat3ra-made interface builder: https://github.com/Exabyte-io/made  "
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python (Pyodide)",
+   "language": "python",
+   "name": "python"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "python",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}