From 2136cd680e3d537a2bdbb44aa1817218cfa54924 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Mon, 31 Aug 2026 14:40:08 -0700 Subject: [PATCH 1/3] SOF-8039: AFIR reaction path with MACE through NVIDIA ALCHEMI Adds the notebook that dispatches the platform workflow and the script that workflow runs. reaction_path_afir_alchemi.ipynb selects a material, looks the workflow up on the platform by id, submits the job, and reads back the published results. The workflow is not built here, so it stays editable in the Workflow Designer. scripts/mace_afir_alchemi.py is the unit's script: it takes its molecule from the workflow's IO unit, runs reactant relaxation, an AFIR force ramp, a dimer refinement of the transition state and an imaginary-mode check, then writes the plots, tables and structures the unit declares as file_content results. Energies and forces come from MACE evaluated through the NVIDIA ALCHEMI toolkit via an ASE calculator adapter, since nvalchemi ships none of its own. Co-Authored-By: Claude Opus 5 (1M context) --- .../reaction_path_afir_alchemi.ipynb | 512 ++++++++++++++++++ .../scripts/afir_alchemi_requirements.txt | 28 + .../workflows/scripts/mace_afir_alchemi.py | 387 +++++++++++++ 3 files changed, 927 insertions(+) create mode 100644 other/materials_designer/workflows/reaction_path_afir_alchemi.ipynb create mode 100644 other/materials_designer/workflows/scripts/afir_alchemi_requirements.txt create mode 100644 other/materials_designer/workflows/scripts/mace_afir_alchemi.py diff --git a/other/materials_designer/workflows/reaction_path_afir_alchemi.ipynb b/other/materials_designer/workflows/reaction_path_afir_alchemi.ipynb new file mode 100644 index 00000000..ce4509e0 --- /dev/null +++ b/other/materials_designer/workflows/reaction_path_afir_alchemi.ipynb @@ -0,0 +1,512 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Reaction path with AFIR + MACE through NVIDIA ALCHEMI, on a platform GPU\n", + "\n", + "Runs the Claisen-rearrangement search of an input molecule as a platform job: reactant relaxation, an\n", + "artificial-force (AFIR) push along the forming bond, a dimer refinement of the transition state, and a check\n", + "for the imaginary mode. Energies and forces come from MACE evaluated through the\n", + "[NVIDIA ALCHEMI toolkit](https://github.com/NVIDIA/nvalchemi-toolkit).\n", + "\n", + "The workflow itself lives on the platform and is looked up by name, so it stays editable in the Workflow\n", + "Designer. This notebook selects the material, submits the job, shows the published results, and saves the\n", + "transition state back as a material.\n", + "\n", + "The same science runs on a laptop in `local/reaction_path_afir_alchemi.ipynb`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Set up the environment and parameters" + ] + }, + { + "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\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 1.2. Set parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from datetime import datetime\n", + "\n", + "from mat3ra.ide.compute import QueueName\n", + "\n", + "# 2. Auth and organization parameters\n", + "ORGANIZATION_NAME = None # set to use an organization as the owner, otherwise the personal account\n", + "\n", + "# 3. Material parameters\n", + "FOLDER = \"../uploads\"\n", + "MATERIAL_NAME = \"allyl vinyl ether\" # loaded from FOLDER, otherwise from Materials Standata\n", + "\n", + "# 4. Workflow parameters — the workflow is looked up on the platform, not built here\n", + "WORKFLOW_NAME = \"AFIR ALCHEMI - Material IO (WIP)\"\n", + "\n", + "# Files the workflow's execution unit declares as file_content results\n", + "RESULT_FILES = {\n", + " \"results.csv\": \"csv\",\n", + " \"afir_energy_profile.png\": \"image\",\n", + " \"afir_bond_distances.png\": \"image\",\n", + " \"afir_path.csv\": \"csv\",\n", + " \"structures.json\": \"text\",\n", + " \"transition_state.json\": \"text\",\n", + " \"transition_state.xyz\": \"text\",\n", + "}\n", + "\n", + "# 5. Compute parameters\n", + "CLUSTER_NAME = \"cluster-003\"\n", + "QUEUE_NAME = QueueName.GSF # H100 nodes on cluster-003 (GOF on-demand, GSF spot)\n", + "PPN = 40\n", + "\n", + "# 6. Job parameters\n", + "timestamp = datetime.now().strftime(\"%Y-%m-%d %H:%M\")\n", + "POLL_INTERVAL = 30 # seconds" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Authenticate and initialize API client" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.auth import authenticate\n", + "\n", + "\n", + "await authenticate()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.2. Initialize API client" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.api_client import APIClient\n", + "\n", + "client = APIClient.authenticate()\n", + "client" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.3. Select account" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "selected_account = client.my_account\n", + "\n", + "if ORGANIZATION_NAME:\n", + " selected_account = client.get_account(name=ORGANIZATION_NAME)\n", + "\n", + "ACCOUNT_ID = selected_account.id\n", + "print(f\"✅ Selected account ID: {ACCOUNT_ID}, name: {selected_account.name}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.4. Select project" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "projects = client.projects.list({\"isDefault\": True, \"owner._id\": ACCOUNT_ID})\n", + "project_id = projects[0][\"_id\"]\n", + "print(f\"✅ Using project: {projects[0]['name']} ({project_id})\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Create material\n", + "\n", + "The workflow's IO unit fetches this material into the job, and the script builds its ASE molecule from it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.made.material import Material\n", + "from mat3ra.notebooks_utils.ipython.entity.material.visualize import visualize_materials as visualize\n", + "from mat3ra.notebooks_utils.material import load_material_from_folder\n", + "from mat3ra.standata.materials import Materials\n", + "\n", + "material = load_material_from_folder(FOLDER, MATERIAL_NAME) or Material.create(\n", + " Materials.get_by_name_first_match(MATERIAL_NAME))\n", + "\n", + "visualize(material)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3.2. Save material to the platform" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.core.entity.material.api import get_or_create_material\n", + "\n", + "saved_material_response = get_or_create_material(client, material, ACCOUNT_ID)\n", + "saved_material = Material.create(saved_material_response)\n", + "print(f\"Material ID: {saved_material.id} ({saved_material.name})\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Find the workflow on the platform\n", + "\n", + "Editing the workflow in the Workflow Designer changes what this notebook submits — no code change here." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.ipython.entity.workflow.visualize import visualize_workflow\n", + "from mat3ra.wode.workflows import Workflow\n", + "\n", + "matches = client.workflows.list({\"name\": WORKFLOW_NAME, \"owner._id\": ACCOUNT_ID})\n", + "assert matches, (\n", + " f\"No workflow named {WORKFLOW_NAME!r} under this account. \"\n", + " \"Check the name in the Workflows page, or that the account is the owner.\"\n", + ")\n", + "if len(matches) > 1:\n", + " print(f\"⚠️ {len(matches)} workflows share this name; using the most recently updated\")\n", + " matches = sorted(matches, key=lambda w: w.get(\"updatedAt\", \"\"), reverse=True)\n", + "\n", + "workflow_document = matches[0] # create_job needs the plain dict, not the wode object\n", + "saved_workflow = Workflow.create(workflow_document)\n", + "print(f\"Workflow ID: {saved_workflow.id} ({saved_workflow.name})\")\n", + "\n", + "visualize_workflow(saved_workflow)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 4.2. Check what the workflow says it publishes\n", + "\n", + "A file the script writes but the unit does not declare is never uploaded, and a declared name that does not\n", + "match a written file is silently dropped — so it is worth comparing the two before spending a job." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "declared = {\n", + " r[\"basename\"]: r.get(\"filetype\")\n", + " for sw in saved_workflow.to_dict()[\"subworkflows\"]\n", + " for u in sw[\"units\"]\n", + " for r in (u.get(\"results\") or [])\n", + " if r.get(\"name\") == \"file_content\" and \"basename\" in r\n", + "}\n", + "print(f\"declared by the workflow: {len(declared)}\")\n", + "for basename, filetype in RESULT_FILES.items():\n", + " mark = \"ok\" if declared.get(basename) == filetype else \"MISSING or wrong filetype\"\n", + " print(f\" {basename:28} {mark}\")\n", + "extra = set(declared) - set(RESULT_FILES)\n", + "if extra:\n", + " print(f\" declared but not expected: {sorted(extra)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Create the compute configuration" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "clusters = client.clusters.list()\n", + "print(f\"Available clusters: {[c['hostname'] for c in clusters]}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 5.2. Create compute configuration for the job" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.ide.compute import Compute\n", + "\n", + "if CLUSTER_NAME:\n", + " cluster = next((c for c in clusters if CLUSTER_NAME in c[\"hostname\"]), None)\n", + "else:\n", + " cluster = clusters[0]\n", + "\n", + "compute = Compute(cluster=cluster, queue=QUEUE_NAME, ppn=PPN)\n", + "print(f\"Using cluster: {compute.cluster.hostname}, queue: {QUEUE_NAME}, ppn: {PPN}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Create the job with material and workflow configuration" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.job import create_job\n", + "from mat3ra.notebooks_utils.ui import display_JSON\n", + "import copy\n", + "\n", + "from mat3ra.utils.namespace import dict_to_namespace_recursive\n", + "\n", + "job_name = f\"{WORKFLOW_NAME} {saved_material.name} {timestamp}\"\n", + "job_response = create_job(\n", + " api_client=client,\n", + " materials=[saved_material_response], # dicts, not entity objects: create_job subscripts them\n", + " workflow=copy.deepcopy(workflow_document), # create_job mutates it (pops _id)\n", + " project_id=project_id,\n", + " owner_id=ACCOUNT_ID,\n", + " prefix=job_name,\n", + " compute=compute.to_dict(),\n", + ")\n", + "\n", + "job = dict_to_namespace_recursive(job_response)\n", + "job_id = job._id\n", + "print(f\"✅ Job created: {job_id}\")\n", + "\n", + "# create_job embeds a copy of the workflow; if the wrong one is attached the job runs silently\n", + "# against the account default, which looks like a working job with unrelated results.\n", + "attached = job_response[\"workflow\"][\"name\"]\n", + "assert attached == WORKFLOW_NAME, f\"job got workflow {attached!r}, expected {WORKFLOW_NAME!r}\"\n", + "print(f\"✅ Workflow attached: {attached}\")\n", + "display_JSON(job_response)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Submit the job and monitor the status" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "client.jobs.submit(job_id)\n", + "print(f\"✅ Job {job_id} submitted successfully!\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.api.job import wait_for_jobs_to_finish_async\n", + "\n", + "await wait_for_jobs_to_finish_async(client.jobs, [job_id], poll_interval=POLL_INTERVAL)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 8. Results\n", + "\n", + "An empty list with a finished job means the files were never uploaded — a cluster problem rather than a\n", + "workflow one. The *Files* tab of the job distinguishes the two: if `script.py` is missing there too, nothing\n", + "about the workflow is being tested." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "files = {f[\"key\"].rsplit(\"/\", 1)[-1]: f[\"signedUrl\"] for f in client.jobs.list_files(job_id)}\n", + "print(f\"files uploaded: {len(files)}\")\n", + "for basename in RESULT_FILES:\n", + " print(f\" {basename:28} {'published' if basename in files else 'MISSING'}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 8.2. Which device did the work\n", + "\n", + "A fast job is not evidence of a GPU. These lines are." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.io import read_from_url\n", + "\n", + "stdout_key = next((k for k in files if k.endswith(\".out\")), None)\n", + "if stdout_key:\n", + " for line in (await read_from_url(files[stdout_key])).splitlines():\n", + " if line.startswith((\"PyTorch version\", \"CUDA compiled\", \"Is CUDA available\",\n", + " \"Running simulation on:\", \"Reactant Energy\",\n", + " \"Refined Activation Energy\", \"Found \")):\n", + " print(line)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 8.3. The figures and the numbers" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from IPython.display import Image, display\n", + "\n", + "for basename in (\"afir_energy_profile.png\", \"afir_bond_distances.png\"):\n", + " if basename in files:\n", + " display(Image(data=await read_from_url(files[basename], as_bytes=True)))\n", + "\n", + "if \"results.csv\" in files:\n", + " print(await read_from_url(files[\"results.csv\"]))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 9. Save the reaction path as materials\n", + "\n", + "The job writes `structures.json` — the input, every AFIR image and the refined transition state, each a\n", + "full Mat3ra material with lattice, basis, metadata and hashes. A python execution unit cannot register a\n", + "material itself, so that happens here; the result is a set of real materials, each usable as input to a\n", + "next job.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "\n", + "assert \"structures.json\" in files, \"the job did not publish structures.json\"\n", + "\n", + "path_structures = json.loads(await read_from_url(files[\"structures.json\"]))\n", + "print(f\"{len(path_structures)} structures along the path\\n\")\n", + "\n", + "saved_structures = []\n", + "for entry in path_structures:\n", + " mat = Material.create(entry[\"material\"])\n", + " response = get_or_create_material(client, mat, ACCOUNT_ID)\n", + " saved = Material.create(response)\n", + " saved_structures.append(saved)\n", + " energy = entry[\"energy_eV\"]\n", + " energy_text = f\"{energy:.3f} eV\" if energy is not None else \"-\"\n", + " print(f\" {entry['label']:38} {energy_text:>12} {saved.id}\")\n", + "\n", + "print(f\"\\n✅ {len(saved_structures)} materials saved under account {ACCOUNT_ID}\")\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11.10" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/other/materials_designer/workflows/scripts/afir_alchemi_requirements.txt b/other/materials_designer/workflows/scripts/afir_alchemi_requirements.txt new file mode 100644 index 00000000..21aa34b6 --- /dev/null +++ b/other/materials_designer/workflows/scripts/afir_alchemi_requirements.txt @@ -0,0 +1,28 @@ +# ----------------------------------------------------------------- # +# # +# Example Python package requirements for the Mat3ra platform # +# # +# Will be used as follows: # +# # +# 1. A runtime directory for this calculation is created # +# 2. This list is used to populate a Python virtual environment # +# 3. The virtual environment is activated # +# 4. The Python process running the script included within this # +# job is started # +# # +# For more information visit: # +# - https://pip.pypa.io/en/stable/reference/pip_install # +# - https://virtualenv.pypa.io/en/stable/ # +# # +# Adjust the list to include your preferred packages. # +# # +# ----------------------------------------------------------------- # + +--extra-index-url https://download.pytorch.org/whl/cu128 +torch>=2.10.0 +torchvision>=0.25.0 +torchaudio +nvalchemi-toolkit[mace] +matplotlib +numpy +mat3ra-made[tools] diff --git a/other/materials_designer/workflows/scripts/mace_afir_alchemi.py b/other/materials_designer/workflows/scripts/mace_afir_alchemi.py new file mode 100644 index 00000000..f5e8378c --- /dev/null +++ b/other/materials_designer/workflows/scripts/mace_afir_alchemi.py @@ -0,0 +1,387 @@ +# ---------------------------------------------------------------- # +# # +# AFIR reaction-path search with ALCHEMI (MACE), Mat3ra platform. # +# # +# Material taken from the Material IO +# # +# ---------------------------------------------------------------- # + +import csv +import time + + +import matplotlib + +matplotlib.use("Agg") # no display on a compute node + +import matplotlib.pyplot as plt +import numpy as np +import torch +from ase.calculators.calculator import Calculator, all_changes +from ase.constraints import ExternalForce +from ase.io import read, write +from ase.mep import DimerControl, MinModeAtoms, MinModeTranslate +from ase.optimize import BFGS +from ase.vibrations import Vibrations + +# NVIDIA ALCHEMI Imports +from nvalchemi.data import AtomicData, Batch +from nvalchemi.dynamics import DynamicsStage +from nvalchemi.hooks import NeighborListHook +from nvalchemi.models.base import ModelConfig +from nvalchemi.models.mace import MACEWrapper + + +print(f"PyTorch version: {torch.__version__}") +print(f"CUDA compiled with PyTorch: {torch.version.cuda}") +print(f"Is CUDA available: {torch.cuda.is_available()}") + + +start_time = time.time() + + +device = "cuda:0" if torch.cuda.is_available() else "cpu" +if device.startswith("cuda"): + torch.cuda.set_device(device) +print(f"Running simulation on: {device}") + +# ========================================== +# 1. Direct Structure Loading & Settings +# ========================================== +import json +from mat3ra.made.tools.convert import to_ase +material = json.loads(r"""{{ MATERIAL | default({}) | tojson }}""") +molecule = to_ase(material) + + +# Defined reaction pairs (atom indices in the structure file) +BOND_FORMING_PAIR = (4, 5) +BOND_BREAKING_PAIR = (0, 1) + +# AFIR and optimization parameters +AFIR_FORCE_RAMP = [1.0, 2.0, 3.0, 4.0] # eV/Å +RELAXATION_FMAX = 0.03 # eV/Å +AFIR_FMAX = 0.05 +SADDLE_FMAX = 0.02 + +# Resolved by mace-torch: a name from its catalogue, an https URL, or a file:// path. +# A name is downloaded once into ~/.cache/mace and needs no cluster-specific filesystem. +CHECKPOINT_PATH = "large-0b2" +MODEL_NAME = "mace-mp-0b2-large" +EV_TO_KCAL = 23.060548 + + +# ========================================== +# 2. ALCHEMI MACE ASE Calculator Adapter +# ========================================== +class AlchemiMaceCalculator(Calculator): + """ASE Calculator adapter wrapping NVIDIA ALCHEMI MACE and NeighborListHook.""" + + implemented_properties = ["energy", "forces"] + + def __init__(self, checkpoint_path, device="cuda:0", **kwargs): + super().__init__(**kwargs) + self.device = device + + # 1. Model output configuration + model_config = ModelConfig(outputs=frozenset({"energy", "forces"})) + + # 2. Initialize ALCHEMI MACE model + self.model = MACEWrapper.from_checkpoint( + checkpoint_path=checkpoint_path, + model_config=model_config, + device=self.device, + ).eval() + + # 3. Neighbor list hook using model's neighbor config + self.nl_hook = NeighborListHook( + config=self.model.model_config.neighbor_config, + stage=DynamicsStage.BEFORE_COMPUTE, + ) + + def calculate( + self, + atoms=None, + properties=["energy", "forces"], + system_changes=all_changes, + ): + super().calculate(atoms, properties, system_changes) + + # Convert ASE Atoms -> ALCHEMI AtomicData -> Batch + atomic_system = AtomicData.from_atoms(self.atoms, device=self.device) + batch = Batch.from_data_list([atomic_system]) + + # Generate neighbor list on the batch before model evaluation + try: + from nvalchemi.hooks import DynamicsContext + + context = DynamicsContext(batch=batch, step_count=0) + except (ImportError, TypeError): + from nvalchemi.hooks import HookContext + + context = HookContext(batch=batch) + + self.nl_hook(context, DynamicsStage.BEFORE_COMPUTE) + + # Evaluate model energy and forces + outputs = self.model(batch) + + # Extract energy scalar cleanly + raw_energy = outputs["energy"].detach().cpu().numpy() + self.results["energy"] = float( + raw_energy.item() if raw_energy.size == 1 else raw_energy.flat[0] + ) + + # Shape forces tensor strictly to (N_atoms, 3) expected by ASE + raw_forces = outputs["forces"].detach().cpu().numpy() + self.results["forces"] = raw_forces.reshape(-1, 3) + + +# Instantiate ALCHEMI Calculator +calculator = AlchemiMaceCalculator(checkpoint_path=CHECKPOINT_PATH, device=device) + +# ========================================== +# 3. Relax Reactant Minimum +# ========================================== +reactant = molecule.copy() +reactant.calc = calculator + +print("Relaxing initial reactant...") +opt = BFGS(reactant) +opt.run(fmax=RELAXATION_FMAX) +reactant_energy = reactant.get_potential_energy() +print(f"Reactant Energy: {reactant_energy:.3f} eV") + +# ========================================== +# 4. Apply Artificial Force (AFIR) +# ========================================== +structure = reactant.copy() +structure.calc = calculator +afir_trajectory = [structure.copy()] +applied_alpha = [0.0] + +print("\nRunning AFIR force ramp...") +for alpha in AFIR_FORCE_RAMP: + structure.set_constraint(ExternalForce(*BOND_FORMING_PAIR, -alpha)) + + opt = BFGS(structure, maxstep=0.1) + opt.run(fmax=AFIR_FMAX, steps=150) + + afir_trajectory.append(structure.copy()) + applied_alpha.append(alpha) + dist = structure.get_distance(*BOND_FORMING_PAIR) + print(f"Force α = {alpha:.1f} eV/Å | Distance: {dist:.2f} Å") + +structure.set_constraint() + +# ========================================== +# 5. Extract Transition State (TS) Guess +# ========================================== +# Re-evaluate path energies without the constraint bias +unbiased_energies = [] +forming_distances = [] +breaking_distances = [] +for img in afir_trajectory: + img.set_constraint() + img.calc = calculator + unbiased_energies.append(img.get_potential_energy()) + forming_distances.append(img.get_distance(*BOND_FORMING_PAIR)) + breaking_distances.append(img.get_distance(*BOND_BREAKING_PAIR)) + +ts_guess_idx = int(np.argmax(unbiased_energies)) +ts_guess = afir_trajectory[ts_guess_idx].copy() +ts_guess.calc = calculator + +barrier_guess = (unbiased_energies[ts_guess_idx] - reactant_energy) * EV_TO_KCAL +print( + f"\nTS Guess found at step {ts_guess_idx} with estimated barrier: {barrier_guess:.1f} kcal/mol" +) + +# ========================================== +# 6. Refine TS with Dimer Method +# ========================================== +direction = np.zeros_like(ts_guess.positions) +for pair, sign in [(BOND_FORMING_PAIR, 1.0), (BOND_BREAKING_PAIR, -1.0)]: + vec = ts_guess.positions[pair[1]] - ts_guess.positions[pair[0]] + direction[pair[0]] += sign * (vec / np.linalg.norm(vec)) + direction[pair[1]] -= sign * (vec / np.linalg.norm(vec)) +direction /= np.linalg.norm(direction) + +dimer_control = DimerControl( + initial_eigenmode_method="displacement", displacement_method="vector" +) +dimer = MinModeAtoms(ts_guess, dimer_control) +dimer.displace(displacement_vector=0.05 * direction) + +dimer_opt = MinModeTranslate(dimer) +dimer_opt.run(fmax=SADDLE_FMAX, steps=200) + +ts_energy = ts_guess.get_potential_energy() +activation_energy = (ts_energy - reactant_energy) * EV_TO_KCAL +print(f"Refined Activation Energy (Barrier): {activation_energy:.1f} kcal/mol") + +write("transition_state.xyz", ts_guess) + +ts_forming = ts_guess.get_distance(*BOND_FORMING_PAIR) +ts_breaking = ts_guess.get_distance(*BOND_BREAKING_PAIR) + +# ========================================== +# 7. Verify Imaginary Frequencies +# ========================================== +vib = Vibrations(ts_guess, name="ts_vib") +vib.run() +freqs = vib.get_frequencies() + +imaginary_freqs = [ + f.imag for f in freqs if np.iscomplex(f) and abs(f.imag) > 50 +] +print( + f"Found {len(imaginary_freqs)} imaginary frequency mode(s) > 50 cm^-1:" +) +for f in imaginary_freqs: + print(f" {f:.1f}i cm^-1") + +vib.clean() + +largest_imaginary = max((abs(f) for f in imaginary_freqs), default=0.0) +wall_time = time.time() - start_time + +# ========================================== +# 8. Publish results +# ========================================== +relative_kcal = [ + (e - reactant_energy) * EV_TO_KCAL for e in unbiased_energies +] +steps = list(range(len(afir_trajectory))) + +# --- afir_path.csv ------------------------------------------------ +with open("afir_path.csv", "w", newline="") as csvfile: + writer = csv.writer(csvfile) + writer.writerow( + [ + "Step", + "Applied force (eV/A)", + "Forming bond (A)", + "Breaking bond (A)", + "Energy (eV)", + "Relative energy (kcal/mol)", + ] + ) + for i in steps: + writer.writerow( + [ + i, + f"{applied_alpha[i]:.6g}", + f"{forming_distances[i]:.6g}", + f"{breaking_distances[i]:.6g}", + f"{unbiased_energies[i]:.6g}", + f"{relative_kcal[i]:.6g}", + ] + ) + +# --- afir_energy_profile.png ------------------------------------- +fig, ax = plt.subplots(figsize=(6, 4), dpi=150) +ax.plot(steps, relative_kcal, "o-", color="#1f77b4", label="unbiased energy") +ax.plot( + [ts_guess_idx], + [relative_kcal[ts_guess_idx]], + "o", + color="#d62728", + markersize=11, + label=f"TS guess (step {ts_guess_idx})", +) +ax.axhline( + activation_energy, + color="#2ca02c", + linestyle="--", + label=f"refined barrier {activation_energy:.1f} kcal/mol", +) +ax.set_xlabel("AFIR step") +ax.set_ylabel("Energy relative to reactant (kcal/mol)") +ax.set_title("AFIR reaction path") +ax.set_xticks(steps) +ax.legend(frameon=False, fontsize=8) +fig.tight_layout() +fig.savefig("afir_energy_profile.png") +plt.close(fig) + +# --- afir_bond_distances.png ------------------------------------- +fig, ax = plt.subplots(figsize=(6, 4), dpi=150) +ax.plot( + steps, + forming_distances, + "o-", + color="#1f77b4", + label=f"forming {BOND_FORMING_PAIR}", +) +ax.plot( + steps, + breaking_distances, + "s-", + color="#ff7f0e", + label=f"breaking {BOND_BREAKING_PAIR}", +) +ax.axvline(ts_guess_idx, color="#d62728", linestyle=":", label="TS guess") +ax.set_xlabel("AFIR step") +ax.set_ylabel("Bond distance (A)") +ax.set_title("Reaction coordinate") +ax.set_xticks(steps) +ax.legend(frameon=False, fontsize=8) +fig.tight_layout() +fig.savefig("afir_bond_distances.png") +plt.close(fig) + +# --- results.csv -------------------------------------------------- +results = { + "Device": device, + "Model": MODEL_NAME, + "Activation energy (kcal/mol)": activation_energy, + "Activation energy (eV)": ts_energy - reactant_energy, + "Reactant energy (eV)": reactant_energy, + "Transition state energy (eV)": ts_energy, + "AFIR TS-guess barrier (kcal/mol)": barrier_guess, + "TS guess step": ts_guess_idx, + "Forming bond, reactant (A)": forming_distances[0], + "Forming bond, TS (A)": ts_forming, + "Breaking bond, reactant (A)": breaking_distances[0], + "Breaking bond, TS (A)": ts_breaking, + "Imaginary modes > 50 cm^-1": len(imaginary_freqs), + "Largest imaginary frequency (cm^-1)": largest_imaginary, + "Wall time (s)": wall_time, +} + +# --- structures as Mat3ra materials -------------------------------- +# from_ase keeps lattice, basis, metadata and hashes; an .xyz keeps only +# positions and symbols, so a Material rebuilt from one loses its provenance. +from mat3ra.made.material import Material +from mat3ra.made.tools.convert import from_ase + +input_name = material.get("name", "molecule") +for atoms, label in ((molecule, "reactant"), (ts_guess, "transition_state")): + out_material = Material.create(from_ase(atoms)) + out_material.name = f"{input_name} - {label.replace('_', ' ')}" + with open(f"{label}.json", "w") as f: + f.write(out_material.to_json()) + print(f"wrote {label}.json ({out_material.name})") + +with open("results.csv", "w", newline="") as csvfile: + writer = csv.DictWriter(csvfile, fieldnames=["Property", "Value"]) + writer.writeheader() + for prop, val in results.items(): + print(f"{prop:<38} = {val}") + writer.writerow( + { + "Property": prop, + "Value": val if isinstance(val, str) else f"{val:.6g}", + } + ) + +PAYLOAD_KEYS = { + "basis", "lattice", "isNonPeriodic", "formula", "unitCellFormula", + "derivedProperties", "external", "src", "name", "description", "tags", "metadata", +} +ts_material = Material.create(from_ase(ts_guess)) +ts_material.name = f"{input_name} - transition state" +print("---MATERIAL---") +print(json.dumps({k: v for k, v in ts_material.to_dict().items() if k in PAYLOAD_KEYS})) +print("---END---") From 0157fcd51df8ebeea02e65f99c513a791325c0c4 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Mon, 31 Aug 2026 14:44:40 -0700 Subject: [PATCH 2/3] SOF-8039: pin the CUDA build of torch --extra-index-url adds an index, it does not prioritise one: pip picks the highest version across both, and PyPI's is higher than anything on the cu128 index. So torch>=2.10.0 resolved to 2.13.0, whose Linux wheel is a CUDA 13 build (cuda-toolkit==13.0.3, nvidia-cudnn-cu13). On a CUDA 12 driver that reports torch.cuda.is_available() == False with no error, which is indistinguishable from having no GPU. Only the local-version form pins the build. Verified all three exist as cp311 manylinux_2_28_x86_64 wheels on the cu128 index, and that they satisfy nvidia-physicsnemo's floors (torch>=2.10.0, torchvision>=0.25.0a0). Co-Authored-By: Claude Opus 5 (1M context) --- .../workflows/scripts/afir_alchemi_requirements.txt | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/other/materials_designer/workflows/scripts/afir_alchemi_requirements.txt b/other/materials_designer/workflows/scripts/afir_alchemi_requirements.txt index 21aa34b6..7eacec31 100644 --- a/other/materials_designer/workflows/scripts/afir_alchemi_requirements.txt +++ b/other/materials_designer/workflows/scripts/afir_alchemi_requirements.txt @@ -19,9 +19,9 @@ # ----------------------------------------------------------------- # --extra-index-url https://download.pytorch.org/whl/cu128 -torch>=2.10.0 -torchvision>=0.25.0 -torchaudio +torch==2.11.0+cu128 +torchvision==0.26.0+cu128 +torchaudio==2.11.0+cu128 nvalchemi-toolkit[mace] matplotlib numpy From 71b8fc34b43b39616f0551311d5289a1ad704f79 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Mon, 31 Aug 2026 18:21:08 -0700 Subject: [PATCH 3/3] SOF-8039: sync the ALCHEMI AFIR script with the platform workflow Renamed from mace_afir_alchemi.py: the calculator is ALCHEMI's MACEWrapper, so the mace_ prefix named the wrong engine. Brings the file level with workflow NS6wav4QtzMoCfvto, which now runs on the cluster GPU. Four changes, each fixing something that failed silently: - to_ase reads a basis as fractional, so a material stored in cartesian units had every coordinate multiplied by the lattice. Bond lengths came out near 84 A in a 15 A cell and the barrier read 0.006 kcal/mol while the job went green. to_crystal() before the conversion. - to_ase routes through a pymatgen Structure, which is always periodic, so isNonPeriodic never survived the round trip and an exported molecule came back a crystal. Restore pbc from the input material. - Relax the last image of the biased path without the bias, giving the product and the reaction energy alongside the barrier. - Emit the transition state and the product as one marker block for the io units to save, tagged jobId- to match what MaterialDAO gives structures the platform creates itself. Co-Authored-By: Claude Opus 5 (1M context) --- .../workflows/scripts/afir_alchemi.py | 445 ++++++++++++++++++ .../workflows/scripts/mace_afir_alchemi.py | 387 --------------- 2 files changed, 445 insertions(+), 387 deletions(-) create mode 100644 other/materials_designer/workflows/scripts/afir_alchemi.py delete mode 100644 other/materials_designer/workflows/scripts/mace_afir_alchemi.py diff --git a/other/materials_designer/workflows/scripts/afir_alchemi.py b/other/materials_designer/workflows/scripts/afir_alchemi.py new file mode 100644 index 00000000..38c40fb1 --- /dev/null +++ b/other/materials_designer/workflows/scripts/afir_alchemi.py @@ -0,0 +1,445 @@ +# ---------------------------------------------------------------- # +# # +# AFIR reaction-path search with ALCHEMI (MACE), Mat3ra platform. # +# # +# Material taken from the Material IO +# # +# ---------------------------------------------------------------- # + +import csv +import time + + +import matplotlib + +matplotlib.use("Agg") # no display on a compute node + +import matplotlib.pyplot as plt +import numpy as np +import torch +from ase.calculators.calculator import Calculator, all_changes +from ase.constraints import ExternalForce +from ase.io import read, write +from ase.mep import DimerControl, MinModeAtoms, MinModeTranslate +from ase.optimize import BFGS +from ase.vibrations import Vibrations + +# NVIDIA ALCHEMI Imports +from nvalchemi.data import AtomicData, Batch +from nvalchemi.dynamics import DynamicsStage +from nvalchemi.hooks import NeighborListHook +from nvalchemi.models.base import ModelConfig +from nvalchemi.models.mace import MACEWrapper + + +print(f"PyTorch version: {torch.__version__}") +print(f"CUDA compiled with PyTorch: {torch.version.cuda}") +print(f"Is CUDA available: {torch.cuda.is_available()}") + + +start_time = time.time() + + +device = "cuda:0" if torch.cuda.is_available() else "cpu" +if device.startswith("cuda"): + torch.cuda.set_device(device) +print(f"Running simulation on: {device}") + +# ========================================== +# 1. Direct Structure Loading & Settings +# ========================================== +import json +from mat3ra.made.material import Material +from mat3ra.made.tools.convert import to_ase + +material = json.loads(r"""{{ MATERIAL | default({}) | tojson }}""") + +# to_ase reads the basis as fractional coordinates. A material stored in cartesian +# units -- which is how the platform returns this one -- must be converted first, or +# every coordinate is multiplied by the lattice and the structure is torn apart. +input_material = Material.create(material) +input_material.to_crystal() +molecule = to_ase(input_material) + +# to_ase routes through a pymatgen Structure, which is always periodic, so the input's +# isNonPeriodic never survives the trip and every exported structure comes back a crystal. +molecule.pbc = not material.get("isNonPeriodic", False) + + +# Defined reaction pairs (atom indices in the structure file) +BOND_FORMING_PAIR = (4, 5) +BOND_BREAKING_PAIR = (0, 1) + +# AFIR and optimization parameters +AFIR_FORCE_STRENGTHS = [1.0, 2.0, 3.0, 4.0] # eV/Å +RELAXATION_FMAX = 0.03 # eV/Å +AFIR_FMAX = 0.05 +SADDLE_FMAX = 0.02 + +# Resolved by mace-torch: a name from its catalogue, an https URL, or a file:// path. +# A name is downloaded once into ~/.cache/mace and needs no cluster-specific filesystem. +CHECKPOINT_PATH = "large-0b2" +MODEL_NAME = "mace-mp-0b2-large" +EV_TO_KCAL_PER_MOL = 23.060548 + + +# ========================================== +# 2. ALCHEMI MACE ASE Calculator Adapter +# ========================================== +class AlchemiMaceCalculator(Calculator): + """ASE Calculator adapter wrapping NVIDIA ALCHEMI MACE and NeighborListHook.""" + + implemented_properties = ["energy", "forces"] + + def __init__(self, checkpoint_path, device="cuda:0", **kwargs): + super().__init__(**kwargs) + self.device = device + + # 1. Model output configuration + model_config = ModelConfig(outputs=frozenset({"energy", "forces"})) + + # 2. Initialize ALCHEMI MACE model + self.model = MACEWrapper.from_checkpoint( + checkpoint_path=checkpoint_path, + model_config=model_config, + device=self.device, + ).eval() + + # 3. Neighbor list hook using model's neighbor config + self.neighbor_list_hook = NeighborListHook( + config=self.model.model_config.neighbor_config, + stage=DynamicsStage.BEFORE_COMPUTE, + ) + + def calculate( + self, + atoms=None, + properties=["energy", "forces"], + system_changes=all_changes, + ): + super().calculate(atoms, properties, system_changes) + + # Convert ASE Atoms -> ALCHEMI AtomicData -> Batch + atomic_data = AtomicData.from_atoms(self.atoms, device=self.device) + batch = Batch.from_data_list([atomic_data]) + + # Generate neighbor list on the batch before model evaluation + try: + from nvalchemi.hooks import DynamicsContext + + context = DynamicsContext(batch=batch, step_count=0) + except (ImportError, TypeError): + from nvalchemi.hooks import HookContext + + context = HookContext(batch=batch) + + self.neighbor_list_hook(context, DynamicsStage.BEFORE_COMPUTE) + + # Evaluate model energy and forces + outputs = self.model(batch) + + # Extract energy scalar cleanly + energy_output = outputs["energy"].detach().cpu().numpy() + self.results["energy"] = float( + energy_output.item() if energy_output.size == 1 else energy_output.flat[0] + ) + + # Shape forces tensor strictly to (N_atoms, 3) expected by ASE + forces_output = outputs["forces"].detach().cpu().numpy() + self.results["forces"] = forces_output.reshape(-1, 3) + + +# Instantiate ALCHEMI Calculator +calculator = AlchemiMaceCalculator(checkpoint_path=CHECKPOINT_PATH, device=device) + +# ========================================== +# 3. Relax Reactant Minimum +# ========================================== +reactant = molecule.copy() +reactant.calc = calculator + +print("Relaxing initial reactant...") +optimizer = BFGS(reactant) +optimizer.run(fmax=RELAXATION_FMAX) +reactant_energy = reactant.get_potential_energy() +print(f"Reactant Energy: {reactant_energy:.3f} eV") + +# ========================================== +# 4. Apply Artificial Force (AFIR) +# ========================================== +structure = reactant.copy() +structure.calc = calculator +afir_trajectory = [structure.copy()] +applied_force_strengths = [0.0] + +print("\nRunning AFIR force ramp...") +for force_strength in AFIR_FORCE_STRENGTHS: + structure.set_constraint(ExternalForce(*BOND_FORMING_PAIR, -force_strength)) + + optimizer = BFGS(structure, maxstep=0.1) + optimizer.run(fmax=AFIR_FMAX, steps=150) + + afir_trajectory.append(structure.copy()) + applied_force_strengths.append(force_strength) + forming_distance = structure.get_distance(*BOND_FORMING_PAIR) + print(f"Force α = {force_strength:.1f} eV/Å | Distance: {forming_distance:.2f} Å") + +structure.set_constraint() + +# ========================================== +# 5. Extract Transition State (TS) Guess +# ========================================== +# Re-evaluate path energies without the constraint bias +unbiased_energies = [] +forming_distances = [] +breaking_distances = [] +for image in afir_trajectory: + image.set_constraint() + image.calc = calculator + unbiased_energies.append(image.get_potential_energy()) + forming_distances.append(image.get_distance(*BOND_FORMING_PAIR)) + breaking_distances.append(image.get_distance(*BOND_BREAKING_PAIR)) + +transition_state_index = int(np.argmax(unbiased_energies)) +transition_state = afir_trajectory[transition_state_index].copy() +transition_state.calc = calculator + +estimated_barrier = (unbiased_energies[transition_state_index] - reactant_energy) * EV_TO_KCAL_PER_MOL +print( + f"\nTS Guess found at step {transition_state_index} with estimated barrier: {estimated_barrier:.1f} kcal/mol" +) + +# ========================================== +# 6. Refine TS with Dimer Method +# ========================================== +direction = np.zeros_like(transition_state.positions) +for pair, sign in [(BOND_FORMING_PAIR, 1.0), (BOND_BREAKING_PAIR, -1.0)]: + bond_vector = transition_state.positions[pair[1]] - transition_state.positions[pair[0]] + direction[pair[0]] += sign * (bond_vector / np.linalg.norm(bond_vector)) + direction[pair[1]] -= sign * (bond_vector / np.linalg.norm(bond_vector)) +direction /= np.linalg.norm(direction) + +dimer_control = DimerControl( + initial_eigenmode_method="displacement", displacement_method="vector" +) +dimer = MinModeAtoms(transition_state, dimer_control) +dimer.displace(displacement_vector=0.05 * direction) + +dimer_opt = MinModeTranslate(dimer) +dimer_opt.run(fmax=SADDLE_FMAX, steps=200) + +transition_state_energy = transition_state.get_potential_energy() +activation_energy = (transition_state_energy - reactant_energy) * EV_TO_KCAL_PER_MOL +print(f"Refined Activation Energy (Barrier): {activation_energy:.1f} kcal/mol") + +write("transition_state.xyz", transition_state) + +transition_state_forming_distance = transition_state.get_distance(*BOND_FORMING_PAIR) +transition_state_breaking_distance = transition_state.get_distance(*BOND_BREAKING_PAIR) + +# ========================================== +# 7. Verify Imaginary Frequencies +# ========================================== +vibrations = Vibrations(transition_state, name="transition_state_vibrations") +vibrations.run() +frequencies = vibrations.get_frequencies() + +imaginary_frequencies = [ + frequency.imag + for frequency in frequencies + if np.iscomplex(frequency) and abs(frequency.imag) > 50 +] +print( + f"Found {len(imaginary_frequencies)} imaginary frequency mode(s) > 50 cm^-1:" +) +for frequency in imaginary_frequencies: + print(f" {frequency:.1f}i cm^-1") + +vibrations.clean() + +largest_imaginary_frequency = max( + (abs(frequency) for frequency in imaginary_frequencies), default=0.0 +) +# ========================================== +# 8. Relax the Product +# ========================================== +# The last image of the biased path, relaxed without the artificial force, falls into +# the product minimum the search reached. The constraint was already cleared in step 5. +product = afir_trajectory[-1].copy() +product.calc = calculator + +print("\nRelaxing the discovered product...") +BFGS(product).run(fmax=RELAXATION_FMAX) +product_energy = product.get_potential_energy() +reaction_energy = (product_energy - reactant_energy) * EV_TO_KCAL_PER_MOL +print(f"Reaction energy: {reaction_energy:.1f} kcal/mol relative to the reactant") + +wall_time = time.time() - start_time + +# ========================================== +# 9. Publish results +# ========================================== +relative_energies_kcal = [ + (e - reactant_energy) * EV_TO_KCAL_PER_MOL for e in unbiased_energies +] +step_indices = list(range(len(afir_trajectory))) + +# --- afir_path.csv ------------------------------------------------ +with open("afir_path.csv", "w", newline="") as csv_file: + writer = csv.writer(csv_file) + writer.writerow( + [ + "Step", + "Applied force (eV/A)", + "Forming bond (A)", + "Breaking bond (A)", + "Energy (eV)", + "Relative energy (kcal/mol)", + ] + ) + for i in step_indices: + writer.writerow( + [ + i, + f"{applied_force_strengths[i]:.6g}", + f"{forming_distances[i]:.6g}", + f"{breaking_distances[i]:.6g}", + f"{unbiased_energies[i]:.6g}", + f"{relative_energies_kcal[i]:.6g}", + ] + ) + +# --- afir_energy_profile.png ------------------------------------- +figure, axes = plt.subplots(figsize=(6, 4), dpi=150) +axes.plot(step_indices, relative_energies_kcal, "o-", color="#1f77b4", label="unbiased energy") +axes.plot( + [transition_state_index], + [relative_energies_kcal[transition_state_index]], + "o", + color="#d62728", + markersize=11, + label=f"TS guess (step {transition_state_index})", +) +axes.axhline( + activation_energy, + color="#2ca02c", + linestyle="--", + label=f"refined barrier {activation_energy:.1f} kcal/mol", +) +axes.set_xlabel("AFIR step") +axes.set_ylabel("Energy relative to reactant (kcal/mol)") +axes.set_title("AFIR reaction path") +axes.set_xticks(step_indices) +axes.legend(frameon=False, fontsize=8) +figure.tight_layout() +figure.savefig("afir_energy_profile.png") +plt.close(figure) + +# --- afir_bond_distances.png ------------------------------------- +figure, axes = plt.subplots(figsize=(6, 4), dpi=150) +axes.plot( + step_indices, + forming_distances, + "o-", + color="#1f77b4", + label=f"forming {BOND_FORMING_PAIR}", +) +axes.plot( + step_indices, + breaking_distances, + "s-", + color="#ff7f0e", + label=f"breaking {BOND_BREAKING_PAIR}", +) +axes.axvline(transition_state_index, color="#d62728", linestyle=":", label="TS guess") +axes.set_xlabel("AFIR step") +axes.set_ylabel("Bond distance (A)") +axes.set_title("Reaction coordinate") +axes.set_xticks(step_indices) +axes.legend(frameon=False, fontsize=8) +figure.tight_layout() +figure.savefig("afir_bond_distances.png") +plt.close(figure) + +# --- results.csv -------------------------------------------------- +results = { + "Device": device, + "Model": MODEL_NAME, + "Activation energy (kcal/mol)": activation_energy, + "Activation energy (eV)": transition_state_energy - reactant_energy, + "Reaction energy (kcal/mol)": reaction_energy, + "Reactant energy (eV)": reactant_energy, + "Transition state energy (eV)": transition_state_energy, + "AFIR TS-guess barrier (kcal/mol)": estimated_barrier, + "TS guess step": transition_state_index, + "Forming bond, reactant (A)": forming_distances[0], + "Forming bond, TS (A)": transition_state_forming_distance, + "Breaking bond, reactant (A)": breaking_distances[0], + "Breaking bond, TS (A)": transition_state_breaking_distance, + "Forming bond, product (A)": product.get_distance(*BOND_FORMING_PAIR), + "Breaking bond, product (A)": product.get_distance(*BOND_BREAKING_PAIR), + "Imaginary modes > 50 cm^-1": len(imaginary_frequencies), + "Largest imaginary frequency (cm^-1)": largest_imaginary_frequency, + "Wall time (s)": wall_time, +} + +# --- structures as Mat3ra materials -------------------------------- +# from_ase keeps lattice, basis, metadata and hashes; an .xyz keeps only +# positions and symbols, so a Material rebuilt from one loses its provenance. +from mat3ra.made.material import Material +from mat3ra.made.tools.convert import from_ase + +input_material_name = material.get("name", "molecule") +for atoms, label in ((molecule, "reactant"), (transition_state, "transition_state")): + output_material = Material.create(from_ase(atoms)) + output_material.name = f"{input_material_name} - {label.replace('_', ' ')}" + with open(f"{label}.json", "w") as material_file: + material_file.write(output_material.to_json()) + print(f"wrote {label}.json ({output_material.name})") + +with open("results.csv", "w", newline="") as csv_file: + writer = csv.DictWriter(csv_file, fieldnames=["Property", "Value"]) + writer.writeheader() + for property_name, value in results.items(): + print(f"{property_name:<38} = {value}") + writer.writerow( + { + "Property": property_name, + "Value": value if isinstance(value, str) else f"{value:.6g}", + } + ) + +MATERIAL_PAYLOAD_KEYS = { + "basis", "lattice", "isNonPeriodic", "formula", "unitCellFormula", + "derivedProperties", "external", "src", "name", "description", "tags", "metadata", +} +# MaterialDAO tags a job-produced structure "jobId-" (MaterialDAO.ts:650). Following +# that convention means these are found the same way as platform-created structures. +JOB_TAG = "jobId-{{ JOB_ID }}" + + +def build_material_payload(atoms, role): + output_material = Material.create(from_ase(atoms)) + output_material.name = f"{input_material_name} - {role}" + payload = { + key: value + for key, value in output_material.to_dict().items() + if key in MATERIAL_PAYLOAD_KEYS + } + payload["tags"] = [JOB_TAG] + # The REST layer reads the account from owner._id, and falls back to the calling + # user's default account when it is absent -- which fails for a workflow-issued + # request. The input material already carries the right owner, so reuse it. + if material.get("owner"): + payload["owner"] = material["owner"] + return payload + + +exported_structures = { + "transition_state": build_material_payload(transition_state, "transition state"), + "product": build_material_payload(product, "product"), +} + +print("---MATERIALS---") +print(json.dumps(exported_structures)) +print("---END---") diff --git a/other/materials_designer/workflows/scripts/mace_afir_alchemi.py b/other/materials_designer/workflows/scripts/mace_afir_alchemi.py deleted file mode 100644 index f5e8378c..00000000 --- a/other/materials_designer/workflows/scripts/mace_afir_alchemi.py +++ /dev/null @@ -1,387 +0,0 @@ -# ---------------------------------------------------------------- # -# # -# AFIR reaction-path search with ALCHEMI (MACE), Mat3ra platform. # -# # -# Material taken from the Material IO -# # -# ---------------------------------------------------------------- # - -import csv -import time - - -import matplotlib - -matplotlib.use("Agg") # no display on a compute node - -import matplotlib.pyplot as plt -import numpy as np -import torch -from ase.calculators.calculator import Calculator, all_changes -from ase.constraints import ExternalForce -from ase.io import read, write -from ase.mep import DimerControl, MinModeAtoms, MinModeTranslate -from ase.optimize import BFGS -from ase.vibrations import Vibrations - -# NVIDIA ALCHEMI Imports -from nvalchemi.data import AtomicData, Batch -from nvalchemi.dynamics import DynamicsStage -from nvalchemi.hooks import NeighborListHook -from nvalchemi.models.base import ModelConfig -from nvalchemi.models.mace import MACEWrapper - - -print(f"PyTorch version: {torch.__version__}") -print(f"CUDA compiled with PyTorch: {torch.version.cuda}") -print(f"Is CUDA available: {torch.cuda.is_available()}") - - -start_time = time.time() - - -device = "cuda:0" if torch.cuda.is_available() else "cpu" -if device.startswith("cuda"): - torch.cuda.set_device(device) -print(f"Running simulation on: {device}") - -# ========================================== -# 1. Direct Structure Loading & Settings -# ========================================== -import json -from mat3ra.made.tools.convert import to_ase -material = json.loads(r"""{{ MATERIAL | default({}) | tojson }}""") -molecule = to_ase(material) - - -# Defined reaction pairs (atom indices in the structure file) -BOND_FORMING_PAIR = (4, 5) -BOND_BREAKING_PAIR = (0, 1) - -# AFIR and optimization parameters -AFIR_FORCE_RAMP = [1.0, 2.0, 3.0, 4.0] # eV/Å -RELAXATION_FMAX = 0.03 # eV/Å -AFIR_FMAX = 0.05 -SADDLE_FMAX = 0.02 - -# Resolved by mace-torch: a name from its catalogue, an https URL, or a file:// path. -# A name is downloaded once into ~/.cache/mace and needs no cluster-specific filesystem. -CHECKPOINT_PATH = "large-0b2" -MODEL_NAME = "mace-mp-0b2-large" -EV_TO_KCAL = 23.060548 - - -# ========================================== -# 2. ALCHEMI MACE ASE Calculator Adapter -# ========================================== -class AlchemiMaceCalculator(Calculator): - """ASE Calculator adapter wrapping NVIDIA ALCHEMI MACE and NeighborListHook.""" - - implemented_properties = ["energy", "forces"] - - def __init__(self, checkpoint_path, device="cuda:0", **kwargs): - super().__init__(**kwargs) - self.device = device - - # 1. Model output configuration - model_config = ModelConfig(outputs=frozenset({"energy", "forces"})) - - # 2. Initialize ALCHEMI MACE model - self.model = MACEWrapper.from_checkpoint( - checkpoint_path=checkpoint_path, - model_config=model_config, - device=self.device, - ).eval() - - # 3. Neighbor list hook using model's neighbor config - self.nl_hook = NeighborListHook( - config=self.model.model_config.neighbor_config, - stage=DynamicsStage.BEFORE_COMPUTE, - ) - - def calculate( - self, - atoms=None, - properties=["energy", "forces"], - system_changes=all_changes, - ): - super().calculate(atoms, properties, system_changes) - - # Convert ASE Atoms -> ALCHEMI AtomicData -> Batch - atomic_system = AtomicData.from_atoms(self.atoms, device=self.device) - batch = Batch.from_data_list([atomic_system]) - - # Generate neighbor list on the batch before model evaluation - try: - from nvalchemi.hooks import DynamicsContext - - context = DynamicsContext(batch=batch, step_count=0) - except (ImportError, TypeError): - from nvalchemi.hooks import HookContext - - context = HookContext(batch=batch) - - self.nl_hook(context, DynamicsStage.BEFORE_COMPUTE) - - # Evaluate model energy and forces - outputs = self.model(batch) - - # Extract energy scalar cleanly - raw_energy = outputs["energy"].detach().cpu().numpy() - self.results["energy"] = float( - raw_energy.item() if raw_energy.size == 1 else raw_energy.flat[0] - ) - - # Shape forces tensor strictly to (N_atoms, 3) expected by ASE - raw_forces = outputs["forces"].detach().cpu().numpy() - self.results["forces"] = raw_forces.reshape(-1, 3) - - -# Instantiate ALCHEMI Calculator -calculator = AlchemiMaceCalculator(checkpoint_path=CHECKPOINT_PATH, device=device) - -# ========================================== -# 3. Relax Reactant Minimum -# ========================================== -reactant = molecule.copy() -reactant.calc = calculator - -print("Relaxing initial reactant...") -opt = BFGS(reactant) -opt.run(fmax=RELAXATION_FMAX) -reactant_energy = reactant.get_potential_energy() -print(f"Reactant Energy: {reactant_energy:.3f} eV") - -# ========================================== -# 4. Apply Artificial Force (AFIR) -# ========================================== -structure = reactant.copy() -structure.calc = calculator -afir_trajectory = [structure.copy()] -applied_alpha = [0.0] - -print("\nRunning AFIR force ramp...") -for alpha in AFIR_FORCE_RAMP: - structure.set_constraint(ExternalForce(*BOND_FORMING_PAIR, -alpha)) - - opt = BFGS(structure, maxstep=0.1) - opt.run(fmax=AFIR_FMAX, steps=150) - - afir_trajectory.append(structure.copy()) - applied_alpha.append(alpha) - dist = structure.get_distance(*BOND_FORMING_PAIR) - print(f"Force α = {alpha:.1f} eV/Å | Distance: {dist:.2f} Å") - -structure.set_constraint() - -# ========================================== -# 5. Extract Transition State (TS) Guess -# ========================================== -# Re-evaluate path energies without the constraint bias -unbiased_energies = [] -forming_distances = [] -breaking_distances = [] -for img in afir_trajectory: - img.set_constraint() - img.calc = calculator - unbiased_energies.append(img.get_potential_energy()) - forming_distances.append(img.get_distance(*BOND_FORMING_PAIR)) - breaking_distances.append(img.get_distance(*BOND_BREAKING_PAIR)) - -ts_guess_idx = int(np.argmax(unbiased_energies)) -ts_guess = afir_trajectory[ts_guess_idx].copy() -ts_guess.calc = calculator - -barrier_guess = (unbiased_energies[ts_guess_idx] - reactant_energy) * EV_TO_KCAL -print( - f"\nTS Guess found at step {ts_guess_idx} with estimated barrier: {barrier_guess:.1f} kcal/mol" -) - -# ========================================== -# 6. Refine TS with Dimer Method -# ========================================== -direction = np.zeros_like(ts_guess.positions) -for pair, sign in [(BOND_FORMING_PAIR, 1.0), (BOND_BREAKING_PAIR, -1.0)]: - vec = ts_guess.positions[pair[1]] - ts_guess.positions[pair[0]] - direction[pair[0]] += sign * (vec / np.linalg.norm(vec)) - direction[pair[1]] -= sign * (vec / np.linalg.norm(vec)) -direction /= np.linalg.norm(direction) - -dimer_control = DimerControl( - initial_eigenmode_method="displacement", displacement_method="vector" -) -dimer = MinModeAtoms(ts_guess, dimer_control) -dimer.displace(displacement_vector=0.05 * direction) - -dimer_opt = MinModeTranslate(dimer) -dimer_opt.run(fmax=SADDLE_FMAX, steps=200) - -ts_energy = ts_guess.get_potential_energy() -activation_energy = (ts_energy - reactant_energy) * EV_TO_KCAL -print(f"Refined Activation Energy (Barrier): {activation_energy:.1f} kcal/mol") - -write("transition_state.xyz", ts_guess) - -ts_forming = ts_guess.get_distance(*BOND_FORMING_PAIR) -ts_breaking = ts_guess.get_distance(*BOND_BREAKING_PAIR) - -# ========================================== -# 7. Verify Imaginary Frequencies -# ========================================== -vib = Vibrations(ts_guess, name="ts_vib") -vib.run() -freqs = vib.get_frequencies() - -imaginary_freqs = [ - f.imag for f in freqs if np.iscomplex(f) and abs(f.imag) > 50 -] -print( - f"Found {len(imaginary_freqs)} imaginary frequency mode(s) > 50 cm^-1:" -) -for f in imaginary_freqs: - print(f" {f:.1f}i cm^-1") - -vib.clean() - -largest_imaginary = max((abs(f) for f in imaginary_freqs), default=0.0) -wall_time = time.time() - start_time - -# ========================================== -# 8. Publish results -# ========================================== -relative_kcal = [ - (e - reactant_energy) * EV_TO_KCAL for e in unbiased_energies -] -steps = list(range(len(afir_trajectory))) - -# --- afir_path.csv ------------------------------------------------ -with open("afir_path.csv", "w", newline="") as csvfile: - writer = csv.writer(csvfile) - writer.writerow( - [ - "Step", - "Applied force (eV/A)", - "Forming bond (A)", - "Breaking bond (A)", - "Energy (eV)", - "Relative energy (kcal/mol)", - ] - ) - for i in steps: - writer.writerow( - [ - i, - f"{applied_alpha[i]:.6g}", - f"{forming_distances[i]:.6g}", - f"{breaking_distances[i]:.6g}", - f"{unbiased_energies[i]:.6g}", - f"{relative_kcal[i]:.6g}", - ] - ) - -# --- afir_energy_profile.png ------------------------------------- -fig, ax = plt.subplots(figsize=(6, 4), dpi=150) -ax.plot(steps, relative_kcal, "o-", color="#1f77b4", label="unbiased energy") -ax.plot( - [ts_guess_idx], - [relative_kcal[ts_guess_idx]], - "o", - color="#d62728", - markersize=11, - label=f"TS guess (step {ts_guess_idx})", -) -ax.axhline( - activation_energy, - color="#2ca02c", - linestyle="--", - label=f"refined barrier {activation_energy:.1f} kcal/mol", -) -ax.set_xlabel("AFIR step") -ax.set_ylabel("Energy relative to reactant (kcal/mol)") -ax.set_title("AFIR reaction path") -ax.set_xticks(steps) -ax.legend(frameon=False, fontsize=8) -fig.tight_layout() -fig.savefig("afir_energy_profile.png") -plt.close(fig) - -# --- afir_bond_distances.png ------------------------------------- -fig, ax = plt.subplots(figsize=(6, 4), dpi=150) -ax.plot( - steps, - forming_distances, - "o-", - color="#1f77b4", - label=f"forming {BOND_FORMING_PAIR}", -) -ax.plot( - steps, - breaking_distances, - "s-", - color="#ff7f0e", - label=f"breaking {BOND_BREAKING_PAIR}", -) -ax.axvline(ts_guess_idx, color="#d62728", linestyle=":", label="TS guess") -ax.set_xlabel("AFIR step") -ax.set_ylabel("Bond distance (A)") -ax.set_title("Reaction coordinate") -ax.set_xticks(steps) -ax.legend(frameon=False, fontsize=8) -fig.tight_layout() -fig.savefig("afir_bond_distances.png") -plt.close(fig) - -# --- results.csv -------------------------------------------------- -results = { - "Device": device, - "Model": MODEL_NAME, - "Activation energy (kcal/mol)": activation_energy, - "Activation energy (eV)": ts_energy - reactant_energy, - "Reactant energy (eV)": reactant_energy, - "Transition state energy (eV)": ts_energy, - "AFIR TS-guess barrier (kcal/mol)": barrier_guess, - "TS guess step": ts_guess_idx, - "Forming bond, reactant (A)": forming_distances[0], - "Forming bond, TS (A)": ts_forming, - "Breaking bond, reactant (A)": breaking_distances[0], - "Breaking bond, TS (A)": ts_breaking, - "Imaginary modes > 50 cm^-1": len(imaginary_freqs), - "Largest imaginary frequency (cm^-1)": largest_imaginary, - "Wall time (s)": wall_time, -} - -# --- structures as Mat3ra materials -------------------------------- -# from_ase keeps lattice, basis, metadata and hashes; an .xyz keeps only -# positions and symbols, so a Material rebuilt from one loses its provenance. -from mat3ra.made.material import Material -from mat3ra.made.tools.convert import from_ase - -input_name = material.get("name", "molecule") -for atoms, label in ((molecule, "reactant"), (ts_guess, "transition_state")): - out_material = Material.create(from_ase(atoms)) - out_material.name = f"{input_name} - {label.replace('_', ' ')}" - with open(f"{label}.json", "w") as f: - f.write(out_material.to_json()) - print(f"wrote {label}.json ({out_material.name})") - -with open("results.csv", "w", newline="") as csvfile: - writer = csv.DictWriter(csvfile, fieldnames=["Property", "Value"]) - writer.writeheader() - for prop, val in results.items(): - print(f"{prop:<38} = {val}") - writer.writerow( - { - "Property": prop, - "Value": val if isinstance(val, str) else f"{val:.6g}", - } - ) - -PAYLOAD_KEYS = { - "basis", "lattice", "isNonPeriodic", "formula", "unitCellFormula", - "derivedProperties", "external", "src", "name", "description", "tags", "metadata", -} -ts_material = Material.create(from_ase(ts_guess)) -ts_material.name = f"{input_name} - transition state" -print("---MATERIAL---") -print(json.dumps({k: v for k, v in ts_material.to_dict().items() if k in PAYLOAD_KEYS})) -print("---END---")