experiment: support for uma

.gitattributes

@@ -17,4 +17,5 @@ images/*.png filter=lfs diff=lfs merge=lfs -text
 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
 packages filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -110,3 +110,4 @@ examples/assets/Molecules_Dataset_Collection
 pw_scf.*
 examples/**/*.html
 .DS_Store
+node_modules

config.yml

@@ -101,3 +101,24 @@ notebooks:
       - ssl
       - h5py
       - lmdb
+  - name: uma
+    packages_pyodide:
+      # antlr4 local wheel (required by omegaconf, PyPI version has ATN mismatch)
+      - emfs:/drive/packages/antlr4_python3_runtime-4.9.3-py3-none-any.whl
+      # Packages with dependencies
+      - opt_einsum
+      - orjson
+      - pyyaml
+      - sqlite3
+      - omegaconf
+      - hydra-core
+      # 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 \u2014 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 = \"omat\"  # Task name for the UMA model (e.g. \"omat\", \"oc20\", \"omol\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1e9d9e76d619cd91",
+   "metadata": {
+    "trusted": true
+   },
+   "outputs": [],
+   "source": [
+    "UMA_MODEL_PATH = \"/drive/packages/models/uma-s-1p1-int8.pt\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4e89f2d820acb1ed",
+   "metadata": {},
+   "source": [
+    "## 2. Install Packages"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ece62358982306f2",
+   "metadata": {
+    "trusted": true
+   },
+   "outputs": [],
+   "source": [
+    "from mat3ra.notebooks_utils.packages import install_packages\n",
+    "\n",
+    "await install_packages(\"made|api_examples|torch|uma\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c7f8e5e6a1b34d90",
+   "metadata": {
+    "trusted": true
+   },
+   "outputs": [],
+   "source": [
+    "from mat3ra.notebooks_utils.pyodide.packages.torch import apply_all_patches\n",
+    "\n",
+    "apply_all_patches(include_fairchem=True)"
+   ]
+  },
+  {
+   "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. Load UMA Model and Create Calculator"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a1b2c3d4e5f60001",
+   "metadata": {
+    "trusted": true
+   },
+   "outputs": [],
+   "source": [
+    "from fairchem.core import FAIRChemCalculator\n",
+    "from fairchem.core.units.mlip_unit import load_predict_unit\n",
+    "\n",
+    "predictor = load_predict_unit(UMA_MODEL_PATH, device=\"cpu\")\n",
+    "calculator = FAIRChemCalculator(predictor, task_name=UMA_TASK)\n",
+    "\n",
+    "print(f\"UMA model loaded from {UMA_MODEL_PATH}\")\n",
+    "print(f\"Task: {UMA_TASK}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d4e5f6a7b8c90002",
+   "metadata": {},
+   "source": [
+    "### 4.2. Relax with UMA"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3d8746a77f71bab5",
+   "metadata": {
+    "jupyter": {
+     "is_executing": true
+    },
+    "trusted": true
+   },
+   "outputs": [],
+   "source": [
+    "import plotly.graph_objs as go\nfrom IPython.display import display\nfrom plotly.subplots import make_subplots\n\nfrom mat3ra.made.tools.convert import to_ase\nfrom ase.optimize import BFGS\n\nase_structure = to_ase(structure)\nase_structure.set_calculator(calculator)\ndyn = BFGS(ase_structure)\n\nsteps = []\nenergies = []\n\nfig = make_subplots(rows=1, cols=1, specs=[[{\"type\": \"scatter\"}]])\nscatter = go.Scatter(x=[], y=[], mode=\"lines+markers\", name=\"Energy\")\nfig.add_trace(scatter)\nfig.update_layout(title_text=\"Real-time Optimization Progress\", xaxis_title=\"Step\", yaxis_title=\"Energy (eV)\")\n\ntry:\n    f = go.FigureWidget(fig)\nexcept ImportError:\n    f = go.Figure(fig)\ndisplay(f)\n\n\ndef 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    if hasattr(f, \"batch_update\"):\n        with f.batch_update():\n            f.data[0].x = steps\n            f.data[0].y = energies\n    else:\n        f.data[0].x = steps\n        f.data[0].y = energies\n\n\ndyn.attach(plotly_callback, interval=1)\ndyn.run(fmax=RELAXATION_PARAMETERS[\"FMAX\"])\n\nase_original_structure = to_ase(structure)\nase_original_structure.set_calculator(calculator)\nase_final_structure = ase_structure\n\noriginal_energy = ase_original_structure.get_total_energy()\nrelaxed_energy = ase_structure.get_total_energy()\n\nprint(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} \u00c5\")\n",
+    "print(\n",
+    "    f\"Interlayer distance after relaxation:  {get_average_interlayer_distance(material_relaxed, SUBSTRATE_TAG, FILM_TAG):.4f} \u00c5\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2f60fdb73e44c09c",
+   "metadata": {},
+   "source": [
+    "## References\n",
+    "\n",
+    "[1] FAIRChem: https://github.com/FAIR-Chem/fairchem  \n",
+    "[2] UMA \u2014 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
+}