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.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/afir_alchemi_requirements.txt b/other/materials_designer/workflows/scripts/afir_alchemi_requirements.txt new file mode 100644 index 00000000..7eacec31 --- /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.11.0+cu128 +torchvision==0.26.0+cu128 +torchaudio==2.11.0+cu128 +nvalchemi-toolkit[mace] +matplotlib +numpy +mat3ra-made[tools]