From b35a92252407f2c65e92b4f2657687f1b83dd73c Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Mon, 2 Nov 2026 14:50:25 -0800 Subject: [PATCH 01/31] afir_molecule_surface_alchemi.ipynb: random placements of a molecule over a slab, relaxed together under AFIR and then without it, ranked by energy One python step on the job's two materials (slab and the non-periodic molecule), ALCHEMI 0.2.0 batched FIRE2 with the slab frozen; publishes results.csv and final_structures.extxyz. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../afir_molecule_surface_alchemi.ipynb | 716 ++++++++++++++++++ 1 file changed, 716 insertions(+) create mode 100644 other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb new file mode 100644 index 000000000..dfaca87f0 --- /dev/null +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -0,0 +1,716 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# AFIR for a molecule on a surface with NVIDIA ALCHEMI, on a platform GPU\n", + "\n", + "Places a molecule at random over a slab many times and relaxes every placement with MACE through the\n", + "[NVIDIA ALCHEMI toolkit](https://github.com/NVIDIA/nvalchemi-toolkit), all placements batched on one GPU: first\n", + "under the artificial force (AFIR) between molecule and slab, then without it. The final structures are ranked by\n", + "energy.\n", + "\n", + "This notebook builds the slab, fetches the molecule, saves both as materials, creates the workflow, runs one job on\n", + "the two materials and shows the ranked configurations.\n", + "\n", + "Outside the platform's JupyterLite, use Python 3.10 or 3.11 and install the packages once with `pip install \"mat3ra-notebooks-utils[all]\"`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Set up the environment and parameters\n", + "### 1.1. Install packages (JupyterLite)" + ] + }, + { + "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", + "CRYSTAL_NAME = \"Cu\" # from Materials Standata\n", + "MILLER_INDICES = (1, 1, 1)\n", + "NUMBER_OF_LAYERS = 2\n", + "XY_SUPERCELL_MATRIX = [[4, 0], [0, 4]]\n", + "USE_CONVENTIONAL_CELL = False # the supercell repeats the primitive surface cell\n", + "VACUUM = 15.0 # in angstroms\n", + "MOLECULE_NAME = \"water\" # from PubChem\n", + "\n", + "# 4. Workflow parameters\n", + "N_CONFIGURATIONS = 100 # random placements of the molecule over the slab\n", + "GAMMA = 100 # kJ/mol, model collision energy of the AFIR force\n", + "SEED = 0 # of the random placements\n", + "PYTHON_VERSION = \"3.11.10\"\n", + "PYTHON_BUILD = \"Default\"\n", + "\n", + "# 5. Compute parameters\n", + "CLUSTER_NAME = \"cluster-003\"\n", + "QUEUE_NAME = QueueName.GOF # H100 nodes on cluster-003 (GOF on-demand, GSF spot)\n", + "NODES = 1\n", + "PPN = 1\n", + "TIME_LIMIT = \"01:00:00\"\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\n", + "### 2.1. Authenticate" + ] + }, + { + "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()" + ] + }, + { + "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 materials\n", + "### 3.1. Build the slab" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.made.material import Material\n", + "from mat3ra.made.tools.helpers import create_slab\n", + "from mat3ra.notebooks_utils.ipython.entity.material.visualize import visualize_materials as visualize\n", + "from mat3ra.standata.materials import Materials\n", + "\n", + "crystal = Material.create(Materials.get_by_name_first_match(CRYSTAL_NAME))\n", + "slab = create_slab(\n", + " crystal,\n", + " miller_indices=MILLER_INDICES,\n", + " number_of_layers=NUMBER_OF_LAYERS,\n", + " vacuum=VACUUM,\n", + " xy_supercell_matrix=XY_SUPERCELL_MATRIX,\n", + " use_conventional_cell=USE_CONVENTIONAL_CELL,\n", + ")\n", + "visualize(slab)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3.2. Fetch the molecule from PubChem" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import io\n", + "from urllib.parse import quote\n", + "\n", + "from ase.io import read\n", + "from mat3ra.made.tools.convert import from_ase\n", + "from mat3ra.notebooks_utils.io import read_from_url\n", + "\n", + "PUBCHEM_STRUCTURE_URL = \"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/{name}/SDF?record_type=3d\"\n", + "\n", + "molecule_sdf = await read_from_url(PUBCHEM_STRUCTURE_URL.format(name=quote(MOLECULE_NAME)))\n", + "molecule = Material.create(from_ase(read(io.StringIO(molecule_sdf), format=\"sdf\")))\n", + "molecule.name = MOLECULE_NAME\n", + "visualize(molecule)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3.3. Save materials 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_materials = [get_or_create_material(client, material, ACCOUNT_ID) for material in (slab, molecule)]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Create the workflow\n", + "### 4.1. Workflow script\n", + "\n", + "The script the workflow runs, and the packages installed for it on the compute node. It places the molecule\n", + "`N_CONFIGURATIONS` times over the slab, each at a random orientation and a random point of the surface cell, with\n", + "the closest molecule–slab pair at the sum of their covalent radii plus 0.5 Å. All placements are relaxed together\n", + "with the slab fixed, first under the AFIR force of strength `GAMMA` between molecule and slab, then without it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "AFIR_SCRIPT = r'''#!/usr/bin/env python3\n", + "\"\"\"Places the job's molecule at random over its slab, relaxes every placement under the AFIR force between\n", + "molecule and slab, relaxes the results again without it, and ranks the final structures by energy.\"\"\"\n", + "\n", + "import csv\n", + "import json\n", + "\n", + "import numpy as np\n", + "import torch\n", + "from ase.calculators.singlepoint import SinglePointCalculator\n", + "from ase.data import covalent_radii\n", + "from ase.geometry import get_distances\n", + "from ase.io import write\n", + "from loguru import logger\n", + "from mat3ra.made.material import Material\n", + "from mat3ra.made.tools.convert import to_ase\n", + "from mat3ra.made.tools.convert.interface_parts_enum import InterfacePartsEnum\n", + "from nvalchemi._typing import AtomCategory\n", + "from nvalchemi.data import AtomicData\n", + "from nvalchemi.dynamics import FIRE2, FusedStage, HostMemory, SizeAwareSampler\n", + "from nvalchemi.dynamics.base import ConvergenceHook, DynamicsStage\n", + "from nvalchemi.dynamics.hooks import FreezeAtomsHook\n", + "from nvalchemi.hooks import BiasedPotentialHook\n", + "from nvalchemi.models.mace import MACEWrapper\n", + "from scipy.spatial.transform import Rotation\n", + "\n", + "DEVICE = \"cuda:0\"\n", + "CHECKPOINT = \"medium-mpa-0\"\n", + "DTYPE = \"float32\"\n", + "FMAX = 0.03 # eV/A\n", + "MAX_STEPS = 500\n", + "TIME_STEP = 0.5\n", + "BATCH_SIZE = 10\n", + "N_CONFIGURATIONS = 100\n", + "GAMMA = 100 # kJ/mol, model collision energy of the AFIR force\n", + "SEED = 0\n", + "CONTACT_OFFSET = 0.5 # A, closest molecule-slab pair placed at R_i + R_j + CONTACT_OFFSET (GRRM17 RandDC)\n", + "AFIR_R0 = 3.8164 # A, GRRM17 Eq. (3)\n", + "AFIR_EPSILON = 1.0061 # kJ/mol, GRRM17 Eq. (3)\n", + "KJ_PER_MOL_PER_EV = 96.485\n", + "# The job's materials, written here by the platform when the job is created.\n", + "MATERIALS = json.loads(r\"\"\"{{ materials | dump | safe }}\"\"\")\n", + "\n", + "ALPHA = GAMMA / ((2 ** (-1 / 6) - (1 + (1 + GAMMA / AFIR_EPSILON) ** 0.5) ** (-1 / 6)) * AFIR_R0) / KJ_PER_MOL_PER_EV\n", + "CATEGORY_BY_PART = {InterfacePartsEnum.SUBSTRATE.value: AtomCategory.SURFACE.value,\n", + " InterfacePartsEnum.FILM.value: AtomCategory.GAS.value}\n", + "IN_PLANE = torch.tensor([1.0, 1.0, 0.0], device=DEVICE)\n", + "IN_PLANE_SHIFTS = torch.tensor([[a, b, 0.0] for a in (-1, 0, 1) for b in (-1, 0, 1)], device=DEVICE)\n", + "COVALENT_RADII = torch.tensor(covalent_radii, device=DEVICE, dtype=getattr(torch, DTYPE))\n", + "\n", + "\n", + "def material_to_atoms(document, part):\n", + " material = Material.create(document)\n", + " material.to_crystal()\n", + " atoms = to_ase(material)\n", + " atoms.set_tags(part.value)\n", + " return atoms\n", + "\n", + "\n", + "def place_molecule(slab, molecule, random_generator):\n", + " placed = molecule.copy()\n", + " rotation = Rotation.random(random_state=random_generator)\n", + " placed.positions = rotation.apply(placed.positions - placed.positions.mean(axis=0))\n", + " placed.positions += slab.cell.cartesian_positions([random_generator.random(), random_generator.random(), 0.0])\n", + " vectors, _ = get_distances(placed.positions, slab.positions, cell=slab.cell, pbc=[True, True, False])\n", + " contact = covalent_radii[placed.numbers][:, None] + covalent_radii[slab.numbers][None, :] + CONTACT_OFFSET\n", + " heights = slab.positions[None, :, 2] - placed.positions[:, None, 2] + np.sqrt(\n", + " np.clip(contact**2 - (vectors[..., :2] ** 2).sum(axis=-1), 0.0, None))\n", + " placed.positions[:, 2] += heights.max()\n", + " return slab + placed\n", + "\n", + "\n", + "def in_plane_minimum_image(vectors, cells):\n", + " fractional = torch.einsum(\"pi,pij->pj\", vectors, torch.linalg.inv(cells))\n", + " fractional = fractional - torch.round(fractional) * IN_PLANE\n", + " images = torch.einsum(\"psi,pij->psj\", fractional.unsqueeze(1) + IN_PLANE_SHIFTS, cells)\n", + " nearest = images.norm(dim=-1).argmin(dim=1)\n", + " return torch.take_along_dim(images, nearest[:, None, None], dim=1).squeeze(1)\n", + "\n", + "\n", + "def afir_energy(positions, batch):\n", + " molecule_atoms = torch.where(batch.atom_categories == AtomCategory.GAS.value)[0]\n", + " surface_atoms = torch.where(batch.atom_categories == AtomCategory.SURFACE.value)[0]\n", + " molecule_index, surface_index = torch.where(\n", + " batch.batch_idx[molecule_atoms][:, None] == batch.batch_idx[surface_atoms][None, :])\n", + " pair_molecule_atoms, pair_surface_atoms = molecule_atoms[molecule_index], surface_atoms[surface_index]\n", + " pair_structures = batch.batch_idx[pair_molecule_atoms]\n", + " distances = in_plane_minimum_image(positions[pair_molecule_atoms] - positions[pair_surface_atoms],\n", + " batch.cell[pair_structures]).norm(dim=-1)\n", + " radii = (COVALENT_RADII[batch.atomic_numbers[pair_molecule_atoms]]\n", + " + COVALENT_RADII[batch.atomic_numbers[pair_surface_atoms]])\n", + " weights = (radii / distances) ** 6\n", + " numerator = torch.zeros(batch.num_graphs, device=positions.device, dtype=positions.dtype)\n", + " denominator = torch.zeros_like(numerator)\n", + " return ALPHA * (numerator.index_add(0, pair_structures, weights * distances)\n", + " / denominator.index_add(0, pair_structures, weights))\n", + "\n", + "\n", + "def afir_bias(batch):\n", + " positions = batch.positions.clone().requires_grad_(True)\n", + " energy = afir_energy(positions, batch)\n", + " forces = -torch.autograd.grad(energy.sum(), positions)[0]\n", + " return energy.detach().unsqueeze(-1), forces\n", + "\n", + "\n", + "def atomic_data(structures, device, dtype):\n", + " data_list = []\n", + " for i, atoms in enumerate(structures):\n", + " data = AtomicData.from_atoms(atoms, device=device, dtype=dtype)\n", + " data.add_system_property(\"structure_id\", torch.tensor([[i]], device=device,\n", + " dtype=torch.long))\n", + " data[\"atom_categories\"] = torch.tensor([CATEGORY_BY_PART[tag] for tag in atoms.get_tags()],\n", + " device=device, dtype=torch.long)\n", + " data[\"forces\"] = torch.zeros_like(data.positions)\n", + " data[\"energy\"] = torch.zeros((1, 1), device=device, dtype=dtype)\n", + " data_list.append(data)\n", + " return data_list\n", + "\n", + "\n", + "def alchemi_optimizer(model):\n", + " optimizer = FIRE2(\n", + " model=model, dt=TIME_STEP, n_steps=MAX_STEPS,\n", + " convergence_hook=ConvergenceHook(criteria=[{\n", + " \"key\": \"forces\", \"threshold\": FMAX, \"reduce_op\": \"norm\", \"reduce_dims\": -1,\n", + " }]),\n", + " )\n", + " optimizer.register_hook(FreezeAtomsHook(freeze_category=AtomCategory.SURFACE.value))\n", + " return optimizer\n", + "\n", + "\n", + "def log_alchemi_completion(row, mode, steps):\n", + " maximum = float(torch.linalg.vector_norm(row.forces, dim=-1).max().item())\n", + " logger.info(\"{} structure_id={} energy={:+.6f} eV fmax={:.6f} eV/A \"\n", + " \"steps={} status={}\", mode, int(row.structure_id.item()),\n", + " float(row.energy.item()), maximum, steps,\n", + " \"converged\" if maximum < FMAX else \"max_steps\")\n", + "\n", + "\n", + "class StructureDataset:\n", + " def __init__(self, structures, device, dtype):\n", + " self.data_list = atomic_data(structures, device, dtype)\n", + "\n", + " def __len__(self):\n", + " return len(self.data_list)\n", + "\n", + " def __getitem__(self, index):\n", + " return self.data_list[index].clone(), {}\n", + "\n", + " def get_metadata(self, index):\n", + " return self.data_list[index].num_nodes, 0\n", + "\n", + "\n", + "class CompletionSink(HostMemory):\n", + " def __init__(self, capacity, mode):\n", + " super().__init__(capacity)\n", + " self.mode = mode\n", + "\n", + " def write(self, batch, mask=None):\n", + " previous = len(self._data_list)\n", + " super().write(batch, mask=mask)\n", + " for row in self._data_list[previous:]:\n", + " counter = getattr(row, \"n_steps_counter_0\", None)\n", + " steps = int(counter.item()) if counter is not None and counter.item() else MAX_STEPS\n", + " log_alchemi_completion(row, self.mode, steps)\n", + "\n", + "\n", + "def relax_inflight(structures, model, device, dtype, bias_hooks, mode):\n", + " sampler = SizeAwareSampler(StructureDataset(structures, device, dtype), max_batch_size=BATCH_SIZE)\n", + " sink = CompletionSink(len(structures), mode)\n", + " stage = FusedStage(sub_stages=[(0, alchemi_optimizer(model))], sampler=sampler, sinks=[sink])\n", + " for hook in model.make_neighbor_hooks():\n", + " stage.register_hook(hook, stage=DynamicsStage.BEFORE_COMPUTE)\n", + " for hook in bias_hooks:\n", + " stage.register_hook(hook, stage=DynamicsStage.AFTER_COMPUTE)\n", + " stage.run()\n", + " return sink.read()\n", + "\n", + "\n", + "def materialize_alchemi(batch, structures):\n", + " by_id = {int(row.structure_id.item()): row for row in batch.to_data_list()}\n", + " relaxed = []\n", + " for i, original in enumerate(structures):\n", + " atoms = original.copy()\n", + " atoms.positions = by_id[i].positions.cpu().numpy()\n", + " atoms.calc = SinglePointCalculator(atoms, energy=float(by_id[i].energy.item()))\n", + " relaxed.append(atoms)\n", + " return relaxed\n", + "\n", + "\n", + "def load_model(device, dtype):\n", + " return MACEWrapper.from_checkpoint(checkpoint_path=CHECKPOINT, device=device, dtype=dtype)\n", + "\n", + "\n", + "def write_results(afir_structures, final_structures):\n", + " order = np.argsort([atoms.get_potential_energy() for atoms in final_structures])\n", + " with open(\"results.csv\", \"w\", newline=\"\") as results_file:\n", + " writer = csv.writer(results_file)\n", + " writer.writerow([\"configuration\", \"rank\", \"afir_function_eV\", \"energy_eV\"])\n", + " for rank, index in enumerate(order, start=1):\n", + " writer.writerow([index, rank, afir_structures[index].get_potential_energy(),\n", + " final_structures[index].get_potential_energy()])\n", + " final_structures[index].info.update(configuration=int(index), rank=rank)\n", + " write(\"final_structures.extxyz\", [final_structures[index] for index in order])\n", + "\n", + "\n", + "def main():\n", + " device = torch.device(DEVICE)\n", + " dtype = getattr(torch, DTYPE)\n", + " logger.info(\"AFIR force α, γ = {} kJ/mol: {:.3f} eV/Å\", GAMMA, ALPHA)\n", + " slab = material_to_atoms(next(item for item in MATERIALS if not item.get(\"isNonPeriodic\")),\n", + " InterfacePartsEnum.SUBSTRATE)\n", + " molecule = material_to_atoms(next(item for item in MATERIALS if item.get(\"isNonPeriodic\")),\n", + " InterfacePartsEnum.FILM)\n", + " random_generator = np.random.default_rng(SEED)\n", + " configurations = [place_molecule(slab, molecule, random_generator) for _ in range(N_CONFIGURATIONS)]\n", + " model = load_model(device, dtype)\n", + " afir_structures = materialize_alchemi(relax_inflight(\n", + " configurations, model, device, dtype, [BiasedPotentialHook(bias_fn=afir_bias)], \"AFIR\"), configurations)\n", + " final_structures = materialize_alchemi(relax_inflight(\n", + " afir_structures, model, device, dtype, [], \"free\"), afir_structures)\n", + " write_results(afir_structures, final_structures)\n", + "\n", + "\n", + "if __name__ == \"__main__\":\n", + " main()\n", + "'''" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "REQUIREMENTS = \"\"\"--find-links https://download.pytorch.org/whl/cu128/torch/\n", + "--find-links https://download.pytorch.org/whl/cu128/torchvision/\n", + "--find-links https://download.pytorch.org/whl/cu128/torchaudio/\n", + "--find-links https://github.com/NVIDIA/warp/releases/expanded_assets/v1.18.0\n", + "torch==2.11.0+cu128\n", + "torchvision==0.26.0+cu128\n", + "torchaudio==2.11.0+cu128\n", + "warp-lang==1.18.0+cu12\n", + "nvalchemi-toolkit[mace,ase]==0.2.0\n", + "matplotlib\n", + "numpy\n", + "ase\n", + "mace-torch\n", + "loguru\n", + "mat3ra-made[tools]\n", + "\"\"\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 4.2. Create the workflow\n", + "\n", + "The workflow is one step, `afir`. When the job is created, the platform writes the job's materials into the script\n", + "(`materials`); the non-periodic one is the molecule. The step publishes `results.csv`, the configurations ranked by\n", + "final energy, and `final_structures.extxyz`, their structures in the same order." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import copy\n", + "import re\n", + "\n", + "from mat3ra.notebooks_utils.core.entity.workflow.api import set_execution_unit_input\n", + "from mat3ra.notebooks_utils.ipython.entity.workflow.visualize import visualize_workflow\n", + "from mat3ra.standata.workflows import WorkflowStandata\n", + "from mat3ra.wode.workflows import Workflow\n", + "\n", + "script = AFIR_SCRIPT\n", + "for name, value in {\"N_CONFIGURATIONS\": N_CONFIGURATIONS, \"GAMMA\": GAMMA, \"SEED\": SEED}.items():\n", + " script = re.sub(rf\"^({name} = )\\S+\", rf\"\\g<1>{value}\", script, flags=re.M)\n", + "\n", + "workflow_config = copy.deepcopy(\n", + " WorkflowStandata.filter_by_application(\"python\").get_by_name_first_match(\"custom_script.json\")\n", + ")\n", + "workflow_config.update(name=\"AFIR molecule on surface\", isMultiMaterial=True)\n", + "workflow_config[\"units\"][0][\"name\"] = workflow_config[\"name\"]\n", + "subworkflow = workflow_config[\"subworkflows\"][0]\n", + "afir_unit = subworkflow[\"units\"][-1]\n", + "afir_unit.update(\n", + " name=\"afir\",\n", + " head=True,\n", + " results=[\n", + " {\"name\": \"file_content\", \"filetype\": \"csv\", \"basename\": \"results.csv\"},\n", + " {\"name\": \"file_content\", \"filetype\": \"text\", \"basename\": \"final_structures.extxyz\"},\n", + " ],\n", + ")\n", + "for application in (subworkflow[\"application\"], afir_unit[\"application\"]):\n", + " application.update(version=PYTHON_VERSION, build=PYTHON_BUILD)\n", + "script_input = next(unit_input for unit_input in afir_unit[\"input\"] if unit_input[\"template\"][\"name\"] == \"script.py\")\n", + "script_input[\"template\"][\"content\"] = script\n", + "set_execution_unit_input(afir_unit, \"requirements.txt\", REQUIREMENTS)\n", + "subworkflow.update(name=workflow_config[\"name\"], units=[afir_unit])\n", + "workflow = Workflow.create(workflow_config)\n", + "visualize_workflow(workflow)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Create the compute configuration\n", + "### 5.1. Get list of clusters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "clusters = client.clusters.list()\n", + "print(f\"Available clusters: {[cluster['hostname'] for cluster 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", + "cluster = next(cluster for cluster in clusters if CLUSTER_NAME in cluster[\"hostname\"])\n", + "compute = Compute(cluster=cluster, queue=QUEUE_NAME, nodes=NODES, ppn=PPN, timeLimit=TIME_LIMIT)\n", + "print(f\"Using cluster: {compute.cluster.hostname}, queue: {QUEUE_NAME}, nodes: {NODES}, ppn: {PPN}, time limit: {TIME_LIMIT}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Create the job" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.job import create_job\n", + "\n", + "job = create_job(\n", + " api_client=client,\n", + " materials=saved_materials,\n", + " workflow=workflow,\n", + " project_id=project_id,\n", + " owner_id=ACCOUNT_ID,\n", + " prefix=f\"{workflow.name} {MOLECULE_NAME} {N_CONFIGURATIONS} {timestamp}\",\n", + " compute=compute.to_dict(),\n", + ")\n", + "job_id = job[\"_id\"]\n", + "print(f\"✅ Created job: {job_id}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Submit the job and monitor the status" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.api.job import submit_jobs\n", + "\n", + "submit_jobs(client.jobs, [job_id])" + ] + }, + { + "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", + "### 8.1. Configurations ranked by final energy\n", + "\n", + "`afir_function_eV` is the MACE energy plus the AFIR term at the end of the first relaxation, `energy_eV` the MACE\n", + "energy of the final structure." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "\n", + "files = {file[\"key\"].rsplit(\"/\", 1)[-1]: file[\"signedUrl\"] for file in client.jobs.list_files(job_id)}\n", + "pd.read_csv(io.StringIO(await read_from_url(files[\"results.csv\"])))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## References\n", + "\n", + "[1] S. Maeda, Y. Harabuchi, M. Takagi, K. Saita, K. Suzuki, T. Ichino, Y. Sumiya, K. Sugiyama and Y. Ono,\n", + "\"Implementation and performance of the artificial force induced reaction method in the GRRM17 program\",\n", + "J. Comput. Chem. 39, 233 (2018). https://doi.org/10.1002/jcc.25106\n", + "\n", + "[2] S. Maeda and K. Morokuma, \"Communications: A systematic method for locating transition structures of A + B → X\n", + "type reactions\", J. Chem. Phys. 132, 241102 (2010). https://doi.org/10.1063/1.3457903\n", + "\n", + "[3] K. Sugiyama, Y. Sumiya, M. Takagi, K. Saita and S. Maeda, \"Understanding CO oxidation on the Pt(111) surface\n", + "based on a reaction route network\", Phys. Chem. Chem. Phys. 21, 14366 (2019). https://doi.org/10.1039/c8cp06856a" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11.10" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} From 5778d09915bf6cf7329da93e937336f28544e6e1 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Mon, 2 Nov 2026 16:53:12 -0800 Subject: [PATCH 02/31] afir_molecule_surface_alchemi.ipynb: two steps, configurations with made then AFIR with ALCHEMI Configurations places the molecule over the slab with made rotate, set_lattice, translate_by_vector and merge, and prints CONFIGURATIONS to the job scope. AFIR (ALCHEMI) relaxes them as one batch, under AFIR then without it, with the slab frozen. Built from Standata custom_script and python_script, joined with add_subworkflow. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../afir_molecule_surface_alchemi.ipynb | 340 +++++++++--------- 1 file changed, 169 insertions(+), 171 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index dfaca87f0..bbf7389e5 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -11,8 +11,10 @@ "under the artificial force (AFIR) between molecule and slab, then without it. The final structures are ranked by\n", "energy.\n", "\n", - "This notebook builds the slab, fetches the molecule, saves both as materials, creates the workflow, runs one job on\n", - "the two materials and shows the ranked configurations.\n", + "The workflow has two steps, each in its own Python environment: *Configurations* builds the placements from the\n", + "job's two materials with made, and *AFIR (ALCHEMI)* relaxes them. This notebook builds the slab, fetches the\n", + "molecule, saves both as materials, creates the workflow, runs one job on the two materials and shows the ranked\n", + "configurations.\n", "\n", "Outside the platform's JupyterLite, use Python 3.10 or 3.11 and install the packages once with `pip install \"mat3ra-notebooks-utils[all]\"`." ] @@ -75,8 +77,6 @@ "# 5. Compute parameters\n", "CLUSTER_NAME = \"cluster-003\"\n", "QUEUE_NAME = QueueName.GOF # H100 nodes on cluster-003 (GOF on-demand, GSF spot)\n", - "NODES = 1\n", - "PPN = 1\n", "TIME_LIMIT = \"01:00:00\"\n", "\n", "# 6. Job parameters\n", @@ -244,12 +244,11 @@ "metadata": {}, "source": [ "## 4. Create the workflow\n", - "### 4.1. Workflow script\n", + "### 4.1. Step 1: configurations\n", "\n", - "The script the workflow runs, and the packages installed for it on the compute node. It places the molecule\n", - "`N_CONFIGURATIONS` times over the slab, each at a random orientation and a random point of the surface cell, with\n", - "the closest molecule–slab pair at the sum of their covalent radii plus 0.5 Å. All placements are relaxed together\n", - "with the slab fixed, first under the AFIR force of strength `GAMMA` between molecule and slab, then without it." + "Places the molecule `N_CONFIGURATIONS` times over the slab, each at a random orientation and a random point of the\n", + "surface cell, with its lowest atom above the top slab atom by the sum of their covalent radii plus 0.5 Å. Slab atoms\n", + "are labelled 0, molecule atoms 1." ] }, { @@ -258,9 +257,98 @@ "metadata": {}, "outputs": [], "source": [ - "AFIR_SCRIPT = r'''#!/usr/bin/env python3\n", - "\"\"\"Places the job's molecule at random over its slab, relaxes every placement under the AFIR force between\n", - "molecule and slab, relaxes the results again without it, and ranks the final structures by energy.\"\"\"\n", + "CONFIGURATIONS_SCRIPT = r'''\"\"\"Places the job's molecule over its slab N_CONFIGURATIONS times, each at a random orientation and a random point of\n", + "the surface cell, and passes the configurations to the next step as CONFIGURATIONS.\"\"\"\n", + "\n", + "import json\n", + "\n", + "import numpy as np\n", + "from ase.data import atomic_numbers, covalent_radii\n", + "from mat3ra.made.material import Material\n", + "from mat3ra.made.tools.convert.interface_parts_enum import InterfacePartsEnum\n", + "from mat3ra.made.tools.modify import translate_by_vector\n", + "from mat3ra.made.tools.operations.core.binary import merge\n", + "from mat3ra.made.tools.operations.core.unary import rotate\n", + "\n", + "N_CONFIGURATIONS = 100\n", + "SEED = 0\n", + "CONTACT_OFFSET = 0.5 # A, lowest molecule atom at R_molecule + R_slab + CONTACT_OFFSET above the top slab atom (GRRM17 RandDC)\n", + "# The job's materials, written here by the platform when the job is created.\n", + "MATERIALS = json.loads(r\"\"\"{{ materials | dump | safe }}\"\"\")\n", + "\n", + "\n", + "def material_in_crystal_units(document):\n", + " material = Material.create(document)\n", + " material.to_crystal()\n", + " return material\n", + "\n", + "\n", + "def extreme_atom_height_and_radius(material, choose):\n", + " cartesian = material.clone()\n", + " cartesian.to_cartesian()\n", + " height, element = choose(zip([coordinate[2] for coordinate in cartesian.basis.coordinates.values],\n", + " cartesian.basis.elements.values))\n", + " return height, covalent_radii[atomic_numbers[element]]\n", + "\n", + "\n", + "def create_configuration(slab, molecule, random_generator):\n", + " quaternion = random_generator.normal(size=4)\n", + " angle = 2 * np.degrees(np.arccos(quaternion[0] / np.linalg.norm(quaternion)))\n", + " rotated = rotate(molecule, axis=quaternion[1:].tolist(), angle=angle, wrap=False)\n", + " rotated.set_lattice(slab.lattice)\n", + " top_height, top_radius = extreme_atom_height_and_radius(slab, max)\n", + " lowest_height, lowest_radius = extreme_atom_height_and_radius(rotated, min)\n", + " point = slab.basis.cell.convert_point_to_cartesian([random_generator.random(), random_generator.random(), 0.0])\n", + " center = molecule.basis.cell.convert_point_to_cartesian([0.5, 0.5, 0.5])\n", + " height = top_height + top_radius + lowest_radius + CONTACT_OFFSET - lowest_height\n", + " placed = translate_by_vector(rotated, vector=[point[0] - center[0], point[1] - center[1], height],\n", + " use_cartesian_coordinates=True)\n", + " placed.set_labels_from_value(InterfacePartsEnum.FILM.value)\n", + " return merge([slab, placed])\n", + "\n", + "\n", + "def main():\n", + " slab = material_in_crystal_units(next(item for item in MATERIALS if not item.get(\"isNonPeriodic\")))\n", + " molecule = material_in_crystal_units(next(item for item in MATERIALS if item.get(\"isNonPeriodic\")))\n", + " slab.set_labels_from_value(InterfacePartsEnum.SUBSTRATE.value)\n", + " random_generator = np.random.default_rng(SEED)\n", + " configurations = [create_configuration(slab, molecule, random_generator).to_dict() for _ in range(N_CONFIGURATIONS)]\n", + " print(json.dumps({\"name\": \"CONFIGURATIONS\", \"value\": configurations, \"scope\": \"global\"}))\n", + "\n", + "\n", + "if __name__ == \"__main__\":\n", + " main()\n", + "'''" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "CONFIGURATIONS_REQUIREMENTS = \"mat3ra-made[tools]\\n\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 4.2. Step 2: AFIR with ALCHEMI\n", + "\n", + "Relaxes all configurations together with the slab fixed, first under the AFIR force of strength `GAMMA` between\n", + "molecule and slab, then without it, and writes `results.csv` and `final_structures.extxyz`. Another backend or\n", + "another ALCHEMI version replaces this script and its packages; step 1 stays." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "AFIR_SCRIPT = r'''\"\"\"Relaxes the configurations of the previous step with MACE, all batched with the slab fixed: first under the AFIR\n", + "force between molecule and slab, then without it. Ranks the final structures by energy.\"\"\"\n", "\n", "import csv\n", "import json\n", @@ -269,65 +357,37 @@ "import torch\n", "from ase.calculators.singlepoint import SinglePointCalculator\n", "from ase.data import covalent_radii\n", - "from ase.geometry import get_distances\n", "from ase.io import write\n", "from loguru import logger\n", "from mat3ra.made.material import Material\n", "from mat3ra.made.tools.convert import to_ase\n", "from mat3ra.made.tools.convert.interface_parts_enum import InterfacePartsEnum\n", "from nvalchemi._typing import AtomCategory\n", - "from nvalchemi.data import AtomicData\n", - "from nvalchemi.dynamics import FIRE2, FusedStage, HostMemory, SizeAwareSampler\n", + "from nvalchemi.data import AtomicData, Batch\n", + "from nvalchemi.dynamics import FIRE2, FusedStage\n", "from nvalchemi.dynamics.base import ConvergenceHook, DynamicsStage\n", "from nvalchemi.dynamics.hooks import FreezeAtomsHook\n", "from nvalchemi.hooks import BiasedPotentialHook\n", "from nvalchemi.models.mace import MACEWrapper\n", - "from scipy.spatial.transform import Rotation\n", "\n", "DEVICE = \"cuda:0\"\n", "CHECKPOINT = \"medium-mpa-0\"\n", - "DTYPE = \"float32\"\n", + "DTYPE = torch.float32\n", "FMAX = 0.03 # eV/A\n", "MAX_STEPS = 500\n", "TIME_STEP = 0.5\n", - "BATCH_SIZE = 10\n", - "N_CONFIGURATIONS = 100\n", "GAMMA = 100 # kJ/mol, model collision energy of the AFIR force\n", - "SEED = 0\n", - "CONTACT_OFFSET = 0.5 # A, closest molecule-slab pair placed at R_i + R_j + CONTACT_OFFSET (GRRM17 RandDC)\n", "AFIR_R0 = 3.8164 # A, GRRM17 Eq. (3)\n", "AFIR_EPSILON = 1.0061 # kJ/mol, GRRM17 Eq. (3)\n", "KJ_PER_MOL_PER_EV = 96.485\n", - "# The job's materials, written here by the platform when the job is created.\n", - "MATERIALS = json.loads(r\"\"\"{{ materials | dump | safe }}\"\"\")\n", + "CONFIGURATIONS = json.loads(r\"\"\"{{ CONFIGURATIONS | tojson }}\"\"\")\n", "\n", "ALPHA = GAMMA / ((2 ** (-1 / 6) - (1 + (1 + GAMMA / AFIR_EPSILON) ** 0.5) ** (-1 / 6)) * AFIR_R0) / KJ_PER_MOL_PER_EV\n", "CATEGORY_BY_PART = {InterfacePartsEnum.SUBSTRATE.value: AtomCategory.SURFACE.value,\n", " InterfacePartsEnum.FILM.value: AtomCategory.GAS.value}\n", "IN_PLANE = torch.tensor([1.0, 1.0, 0.0], device=DEVICE)\n", "IN_PLANE_SHIFTS = torch.tensor([[a, b, 0.0] for a in (-1, 0, 1) for b in (-1, 0, 1)], device=DEVICE)\n", - "COVALENT_RADII = torch.tensor(covalent_radii, device=DEVICE, dtype=getattr(torch, DTYPE))\n", - "\n", - "\n", - "def material_to_atoms(document, part):\n", - " material = Material.create(document)\n", - " material.to_crystal()\n", - " atoms = to_ase(material)\n", - " atoms.set_tags(part.value)\n", - " return atoms\n", - "\n", - "\n", - "def place_molecule(slab, molecule, random_generator):\n", - " placed = molecule.copy()\n", - " rotation = Rotation.random(random_state=random_generator)\n", - " placed.positions = rotation.apply(placed.positions - placed.positions.mean(axis=0))\n", - " placed.positions += slab.cell.cartesian_positions([random_generator.random(), random_generator.random(), 0.0])\n", - " vectors, _ = get_distances(placed.positions, slab.positions, cell=slab.cell, pbc=[True, True, False])\n", - " contact = covalent_radii[placed.numbers][:, None] + covalent_radii[slab.numbers][None, :] + CONTACT_OFFSET\n", - " heights = slab.positions[None, :, 2] - placed.positions[:, None, 2] + np.sqrt(\n", - " np.clip(contact**2 - (vectors[..., :2] ** 2).sum(axis=-1), 0.0, None))\n", - " placed.positions[:, 2] += heights.max()\n", - " return slab + placed\n", + "COVALENT_RADII = torch.tensor(covalent_radii, device=DEVICE, dtype=DTYPE)\n", "\n", "\n", "def in_plane_minimum_image(vectors, cells):\n", @@ -363,121 +423,57 @@ " return energy.detach().unsqueeze(-1), forces\n", "\n", "\n", - "def atomic_data(structures, device, dtype):\n", - " data_list = []\n", - " for i, atoms in enumerate(structures):\n", - " data = AtomicData.from_atoms(atoms, device=device, dtype=dtype)\n", - " data.add_system_property(\"structure_id\", torch.tensor([[i]], device=device,\n", - " dtype=torch.long))\n", - " data[\"atom_categories\"] = torch.tensor([CATEGORY_BY_PART[tag] for tag in atoms.get_tags()],\n", - " device=device, dtype=torch.long)\n", - " data[\"forces\"] = torch.zeros_like(data.positions)\n", - " data[\"energy\"] = torch.zeros((1, 1), device=device, dtype=dtype)\n", - " data_list.append(data)\n", - " return data_list\n", - "\n", - "\n", - "def alchemi_optimizer(model):\n", - " optimizer = FIRE2(\n", - " model=model, dt=TIME_STEP, n_steps=MAX_STEPS,\n", - " convergence_hook=ConvergenceHook(criteria=[{\n", - " \"key\": \"forces\", \"threshold\": FMAX, \"reduce_op\": \"norm\", \"reduce_dims\": -1,\n", - " }]),\n", - " )\n", - " optimizer.register_hook(FreezeAtomsHook(freeze_category=AtomCategory.SURFACE.value))\n", - " return optimizer\n", - "\n", - "\n", - "def log_alchemi_completion(row, mode, steps):\n", - " maximum = float(torch.linalg.vector_norm(row.forces, dim=-1).max().item())\n", - " logger.info(\"{} structure_id={} energy={:+.6f} eV fmax={:.6f} eV/A \"\n", - " \"steps={} status={}\", mode, int(row.structure_id.item()),\n", - " float(row.energy.item()), maximum, steps,\n", - " \"converged\" if maximum < FMAX else \"max_steps\")\n", - "\n", - "\n", - "class StructureDataset:\n", - " def __init__(self, structures, device, dtype):\n", - " self.data_list = atomic_data(structures, device, dtype)\n", - "\n", - " def __len__(self):\n", - " return len(self.data_list)\n", - "\n", - " def __getitem__(self, index):\n", - " return self.data_list[index].clone(), {}\n", - "\n", - " def get_metadata(self, index):\n", - " return self.data_list[index].num_nodes, 0\n", + "def atomic_data(atoms):\n", + " data = AtomicData.from_atoms(atoms, device=DEVICE, dtype=DTYPE)\n", + " data[\"atom_categories\"] = torch.tensor([CATEGORY_BY_PART[tag] for tag in atoms.get_tags()], device=DEVICE)\n", + " data[\"forces\"] = torch.zeros_like(data.positions)\n", + " data[\"energy\"] = torch.zeros((1, 1), device=DEVICE, dtype=DTYPE)\n", + " return data\n", "\n", "\n", - "class CompletionSink(HostMemory):\n", - " def __init__(self, capacity, mode):\n", - " super().__init__(capacity)\n", - " self.mode = mode\n", - "\n", - " def write(self, batch, mask=None):\n", - " previous = len(self._data_list)\n", - " super().write(batch, mask=mask)\n", - " for row in self._data_list[previous:]:\n", - " counter = getattr(row, \"n_steps_counter_0\", None)\n", - " steps = int(counter.item()) if counter is not None and counter.item() else MAX_STEPS\n", - " log_alchemi_completion(row, self.mode, steps)\n", + "def relaxed_atoms(original, row):\n", + " atoms = original.copy()\n", + " atoms.positions = row.positions.cpu().numpy()\n", + " atoms.calc = SinglePointCalculator(atoms, energy=float(row.energy.item()), forces=row.forces.cpu().numpy())\n", + " return atoms\n", "\n", "\n", - "def relax_inflight(structures, model, device, dtype, bias_hooks, mode):\n", - " sampler = SizeAwareSampler(StructureDataset(structures, device, dtype), max_batch_size=BATCH_SIZE)\n", - " sink = CompletionSink(len(structures), mode)\n", - " stage = FusedStage(sub_stages=[(0, alchemi_optimizer(model))], sampler=sampler, sinks=[sink])\n", + "def relax(structures, model, bias_hooks):\n", + " optimizer = FIRE2(model=model, dt=TIME_STEP, n_steps=MAX_STEPS, convergence_hook=ConvergenceHook(\n", + " criteria=[{\"key\": \"forces\", \"threshold\": FMAX, \"reduce_op\": \"norm\", \"reduce_dims\": -1}]))\n", + " optimizer.register_hook(FreezeAtomsHook(freeze_category=AtomCategory.SURFACE.value))\n", + " stage = FusedStage(sub_stages=[(0, optimizer)])\n", " for hook in model.make_neighbor_hooks():\n", " stage.register_hook(hook, stage=DynamicsStage.BEFORE_COMPUTE)\n", " for hook in bias_hooks:\n", " stage.register_hook(hook, stage=DynamicsStage.AFTER_COMPUTE)\n", - " stage.run()\n", - " return sink.read()\n", - "\n", - "\n", - "def materialize_alchemi(batch, structures):\n", - " by_id = {int(row.structure_id.item()): row for row in batch.to_data_list()}\n", - " relaxed = []\n", - " for i, original in enumerate(structures):\n", - " atoms = original.copy()\n", - " atoms.positions = by_id[i].positions.cpu().numpy()\n", - " atoms.calc = SinglePointCalculator(atoms, energy=float(by_id[i].energy.item()))\n", - " relaxed.append(atoms)\n", - " return relaxed\n", + " batch = stage.run(Batch.from_data_list([atomic_data(atoms) for atoms in structures]))\n", + " return [relaxed_atoms(atoms, row) for atoms, row in zip(structures, batch.to_data_list())]\n", "\n", "\n", - "def load_model(device, dtype):\n", - " return MACEWrapper.from_checkpoint(checkpoint_path=CHECKPOINT, device=device, dtype=dtype)\n", + "def energy_and_max_force(atoms):\n", + " return atoms.get_potential_energy(), np.linalg.norm(atoms.get_forces(), axis=1).max()\n", "\n", "\n", "def write_results(afir_structures, final_structures):\n", " order = np.argsort([atoms.get_potential_energy() for atoms in final_structures])\n", " with open(\"results.csv\", \"w\", newline=\"\") as results_file:\n", " writer = csv.writer(results_file)\n", - " writer.writerow([\"configuration\", \"rank\", \"afir_function_eV\", \"energy_eV\"])\n", + " writer.writerow([\"configuration\", \"rank\", \"afir_function_eV\", \"afir_max_force_eV_per_A\", \"energy_eV\",\n", + " \"max_force_eV_per_A\"])\n", " for rank, index in enumerate(order, start=1):\n", - " writer.writerow([index, rank, afir_structures[index].get_potential_energy(),\n", - " final_structures[index].get_potential_energy()])\n", + " writer.writerow([index, rank, *energy_and_max_force(afir_structures[index]),\n", + " *energy_and_max_force(final_structures[index])])\n", " final_structures[index].info.update(configuration=int(index), rank=rank)\n", " write(\"final_structures.extxyz\", [final_structures[index] for index in order])\n", "\n", "\n", "def main():\n", - " device = torch.device(DEVICE)\n", - " dtype = getattr(torch, DTYPE)\n", " logger.info(\"AFIR force α, γ = {} kJ/mol: {:.3f} eV/Å\", GAMMA, ALPHA)\n", - " slab = material_to_atoms(next(item for item in MATERIALS if not item.get(\"isNonPeriodic\")),\n", - " InterfacePartsEnum.SUBSTRATE)\n", - " molecule = material_to_atoms(next(item for item in MATERIALS if item.get(\"isNonPeriodic\")),\n", - " InterfacePartsEnum.FILM)\n", - " random_generator = np.random.default_rng(SEED)\n", - " configurations = [place_molecule(slab, molecule, random_generator) for _ in range(N_CONFIGURATIONS)]\n", - " model = load_model(device, dtype)\n", - " afir_structures = materialize_alchemi(relax_inflight(\n", - " configurations, model, device, dtype, [BiasedPotentialHook(bias_fn=afir_bias)], \"AFIR\"), configurations)\n", - " final_structures = materialize_alchemi(relax_inflight(\n", - " afir_structures, model, device, dtype, [], \"free\"), afir_structures)\n", + " structures = [to_ase(Material.create(configuration)) for configuration in CONFIGURATIONS]\n", + " model = MACEWrapper.from_checkpoint(checkpoint_path=CHECKPOINT, device=torch.device(DEVICE), dtype=DTYPE)\n", + " afir_structures = relax(structures, model, [BiasedPotentialHook(bias_fn=afir_bias)])\n", + " final_structures = relax(afir_structures, model, [])\n", " write_results(afir_structures, final_structures)\n", "\n", "\n", @@ -492,7 +488,7 @@ "metadata": {}, "outputs": [], "source": [ - "REQUIREMENTS = \"\"\"--find-links https://download.pytorch.org/whl/cu128/torch/\n", + "AFIR_REQUIREMENTS = \"\"\"--find-links https://download.pytorch.org/whl/cu128/torch/\n", "--find-links https://download.pytorch.org/whl/cu128/torchvision/\n", "--find-links https://download.pytorch.org/whl/cu128/torchaudio/\n", "--find-links https://github.com/NVIDIA/warp/releases/expanded_assets/v1.18.0\n", @@ -514,11 +510,11 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### 4.2. Create the workflow\n", + "### 4.3. Create the workflow\n", "\n", - "The workflow is one step, `afir`. When the job is created, the platform writes the job's materials into the script\n", - "(`materials`); the non-periodic one is the molecule. The step publishes `results.csv`, the configurations ranked by\n", - "final energy, and `final_structures.extxyz`, their structures in the same order." + "When the job is created, the platform writes the job's materials into step 1's script (`materials`; the\n", + "non-periodic one is the molecule). Step 1 prints `CONFIGURATIONS` to the job's scope, and the compute node writes\n", + "them into step 2's script before it runs." ] }, { @@ -532,35 +528,37 @@ "\n", "from mat3ra.notebooks_utils.core.entity.workflow.api import set_execution_unit_input\n", "from mat3ra.notebooks_utils.ipython.entity.workflow.visualize import visualize_workflow\n", + "from mat3ra.standata.subworkflows import SubworkflowStandata\n", "from mat3ra.standata.workflows import WorkflowStandata\n", - "from mat3ra.wode.workflows import Workflow\n", - "\n", - "script = AFIR_SCRIPT\n", - "for name, value in {\"N_CONFIGURATIONS\": N_CONFIGURATIONS, \"GAMMA\": GAMMA, \"SEED\": SEED}.items():\n", - " script = re.sub(rf\"^({name} = )\\S+\", rf\"\\g<1>{value}\", script, flags=re.M)\n", + "from mat3ra.wode import Subworkflow, Workflow\n", + "\n", + "workflow_config = copy.deepcopy(WorkflowStandata.filter_by_application(\"python\").get_by_name_first_match(\"custom_script.json\"))\n", + "workflow_config.update(name=\"AFIR, molecule on surface\", isMultiMaterial=True)\n", + "workflow_config[\"units\"][0][\"name\"] = \"Configurations\"\n", + "configurations_subworkflow = workflow_config[\"subworkflows\"][0]\n", + "afir_subworkflow = copy.deepcopy(SubworkflowStandata.filter_by_application(\"python\").get_by_name_first_match(\"python_script.json\"))\n", + "afir_results = [\n", + " {\"name\": \"file_content\", \"filetype\": \"csv\", \"basename\": \"results.csv\"},\n", + " {\"name\": \"file_content\", \"filetype\": \"text\", \"basename\": \"final_structures.extxyz\"},\n", + "]\n", + "for subworkflow, name, unit_name, requirements, results in (\n", + " (configurations_subworkflow, \"Configurations\", \"configurations\", CONFIGURATIONS_REQUIREMENTS, []),\n", + " (afir_subworkflow, \"AFIR (ALCHEMI)\", \"afir\", AFIR_REQUIREMENTS, afir_results),\n", + "):\n", + " unit = subworkflow[\"units\"][-1]\n", + " unit.update(name=unit_name, head=True, results=results)\n", + " for application in (subworkflow[\"application\"], unit[\"application\"]):\n", + " application.update(version=PYTHON_VERSION, build=PYTHON_BUILD)\n", + " set_execution_unit_input(unit, \"requirements.txt\", requirements)\n", + " subworkflow.update(name=name, units=[unit])\n", + "\n", + "configurations_script = re.sub(r\"^N_CONFIGURATIONS = \\S+\", f\"N_CONFIGURATIONS = {N_CONFIGURATIONS}\", CONFIGURATIONS_SCRIPT, flags=re.M)\n", + "script_input = next(unit_input for unit_input in configurations_subworkflow[\"units\"][0][\"input\"] if unit_input[\"template\"][\"name\"] == \"script.py\")\n", + "script_input[\"template\"][\"content\"] = re.sub(r\"^SEED = \\S+\", f\"SEED = {SEED}\", configurations_script, flags=re.M)\n", + "set_execution_unit_input(afir_subworkflow[\"units\"][0], \"script.py\", re.sub(r\"^GAMMA = \\S+\", f\"GAMMA = {GAMMA}\", AFIR_SCRIPT, flags=re.M))\n", "\n", - "workflow_config = copy.deepcopy(\n", - " WorkflowStandata.filter_by_application(\"python\").get_by_name_first_match(\"custom_script.json\")\n", - ")\n", - "workflow_config.update(name=\"AFIR molecule on surface\", isMultiMaterial=True)\n", - "workflow_config[\"units\"][0][\"name\"] = workflow_config[\"name\"]\n", - "subworkflow = workflow_config[\"subworkflows\"][0]\n", - "afir_unit = subworkflow[\"units\"][-1]\n", - "afir_unit.update(\n", - " name=\"afir\",\n", - " head=True,\n", - " results=[\n", - " {\"name\": \"file_content\", \"filetype\": \"csv\", \"basename\": \"results.csv\"},\n", - " {\"name\": \"file_content\", \"filetype\": \"text\", \"basename\": \"final_structures.extxyz\"},\n", - " ],\n", - ")\n", - "for application in (subworkflow[\"application\"], afir_unit[\"application\"]):\n", - " application.update(version=PYTHON_VERSION, build=PYTHON_BUILD)\n", - "script_input = next(unit_input for unit_input in afir_unit[\"input\"] if unit_input[\"template\"][\"name\"] == \"script.py\")\n", - "script_input[\"template\"][\"content\"] = script\n", - "set_execution_unit_input(afir_unit, \"requirements.txt\", REQUIREMENTS)\n", - "subworkflow.update(name=workflow_config[\"name\"], units=[afir_unit])\n", "workflow = Workflow.create(workflow_config)\n", + "workflow.add_subworkflow(Subworkflow(**afir_subworkflow))\n", "visualize_workflow(workflow)" ] }, @@ -598,8 +596,8 @@ "from mat3ra.ide.compute import Compute\n", "\n", "cluster = next(cluster for cluster in clusters if CLUSTER_NAME in cluster[\"hostname\"])\n", - "compute = Compute(cluster=cluster, queue=QUEUE_NAME, nodes=NODES, ppn=PPN, timeLimit=TIME_LIMIT)\n", - "print(f\"Using cluster: {compute.cluster.hostname}, queue: {QUEUE_NAME}, nodes: {NODES}, ppn: {PPN}, time limit: {TIME_LIMIT}\")" + "compute = Compute(cluster=cluster, queue=QUEUE_NAME, timeLimit=TIME_LIMIT)\n", + "print(f\"Using cluster: {compute.cluster.hostname}, queue: {QUEUE_NAME}, nodes: {compute.nodes}, ppn: {compute.ppn}, time limit: {TIME_LIMIT}\")" ] }, { @@ -667,7 +665,7 @@ "### 8.1. Configurations ranked by final energy\n", "\n", "`afir_function_eV` is the MACE energy plus the AFIR term at the end of the first relaxation, `energy_eV` the MACE\n", - "energy of the final structure." + "energy of the final structure; the `max_force` columns give the largest force on an atom at the end of each." ] }, { From 0eda0bf1c4fbb98de02948f4c8b1818c62c6debf Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Mon, 2 Nov 2026 16:56:15 -0800 Subject: [PATCH 03/31] afir_molecule_surface_alchemi.ipynb: take the AFIR subworkflow from the python_script workflow, whose application is complete The subworkflow standata copy carries only the application name; job creation returned 422 with shortName and summary REQUIRED. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../experiments/alchemi/afir_molecule_surface_alchemi.ipynb | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index bbf7389e5..1655a63de 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -528,15 +528,15 @@ "\n", "from mat3ra.notebooks_utils.core.entity.workflow.api import set_execution_unit_input\n", "from mat3ra.notebooks_utils.ipython.entity.workflow.visualize import visualize_workflow\n", - "from mat3ra.standata.subworkflows import SubworkflowStandata\n", "from mat3ra.standata.workflows import WorkflowStandata\n", "from mat3ra.wode import Subworkflow, Workflow\n", "\n", - "workflow_config = copy.deepcopy(WorkflowStandata.filter_by_application(\"python\").get_by_name_first_match(\"custom_script.json\"))\n", + "python_workflows = WorkflowStandata.filter_by_application(\"python\")\n", + "workflow_config = copy.deepcopy(python_workflows.get_by_name_first_match(\"custom_script.json\"))\n", "workflow_config.update(name=\"AFIR, molecule on surface\", isMultiMaterial=True)\n", "workflow_config[\"units\"][0][\"name\"] = \"Configurations\"\n", "configurations_subworkflow = workflow_config[\"subworkflows\"][0]\n", - "afir_subworkflow = copy.deepcopy(SubworkflowStandata.filter_by_application(\"python\").get_by_name_first_match(\"python_script.json\"))\n", + "afir_subworkflow = copy.deepcopy(python_workflows.get_by_name_first_match(\"python_script.json\")[\"subworkflows\"][0])\n", "afir_results = [\n", " {\"name\": \"file_content\", \"filetype\": \"csv\", \"basename\": \"results.csv\"},\n", " {\"name\": \"file_content\", \"filetype\": \"text\", \"basename\": \"final_structures.extxyz\"},\n", From b7a2f5b2f60c11d7136a006545d5129949140117 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Mon, 2 Nov 2026 17:57:21 -0800 Subject: [PATCH 04/31] SOF-8102: afir_molecule_surface_alchemi.ipynb: step 3 saves the lowest final structures; review round 1 Step 3 (made only) turns the lowest NUMBER_OF_FINAL_STRUCTURES into one POST structures request each, with its configuration as the initial structure. The notebook saves the workflow, shows the final structures, names each step's files after it, pins mat3ra-made and drops the packages nvalchemi-toolkit already brings. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../afir_molecule_surface_alchemi.ipynb | 279 ++++++++++++------ 1 file changed, 192 insertions(+), 87 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 1655a63de..fc7624d5e 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -9,12 +9,12 @@ "Places a molecule at random over a slab many times and relaxes every placement with MACE through the\n", "[NVIDIA ALCHEMI toolkit](https://github.com/NVIDIA/nvalchemi-toolkit), all placements batched on one GPU: first\n", "under the artificial force (AFIR) between molecule and slab, then without it. The final structures are ranked by\n", - "energy.\n", + "energy, and the lowest are saved as materials.\n", "\n", - "The workflow has two steps, each in its own Python environment: *Configurations* builds the placements from the\n", - "job's two materials with made, and *AFIR (ALCHEMI)* relaxes them. This notebook builds the slab, fetches the\n", - "molecule, saves both as materials, creates the workflow, runs one job on the two materials and shows the ranked\n", - "configurations.\n", + "The workflow has three steps, each in its own Python environment: *Configurations* builds the placements from the\n", + "job's two materials with made, *AFIR (ALCHEMI)* relaxes them, and *Final structures* saves the lowest as the job's\n", + "final structures. This notebook builds the slab, fetches the molecule, saves both as materials, creates the workflow,\n", + "runs one job on the two materials and shows the ranked configurations and the lowest final structures.\n", "\n", "Outside the platform's JupyterLite, use Python 3.10 or 3.11 and install the packages once with `pip install \"mat3ra-notebooks-utils[all]\"`." ] @@ -68,9 +68,10 @@ "MOLECULE_NAME = \"water\" # from PubChem\n", "\n", "# 4. Workflow parameters\n", - "N_CONFIGURATIONS = 100 # random placements of the molecule over the slab\n", + "NUMBER_OF_CONFIGURATIONS = 100 # random placements of the molecule over the slab\n", "GAMMA = 100 # kJ/mol, model collision energy of the AFIR force\n", - "SEED = 0 # of the random placements\n", + "RANDOM_SEED = 0 # of the random placements\n", + "NUMBER_OF_FINAL_STRUCTURES = 5 # the lowest-energy final structures saved as materials\n", "PYTHON_VERSION = \"3.11.10\"\n", "PYTHON_BUILD = \"Default\"\n", "\n", @@ -246,9 +247,9 @@ "## 4. Create the workflow\n", "### 4.1. Step 1: configurations\n", "\n", - "Places the molecule `N_CONFIGURATIONS` times over the slab, each at a random orientation and a random point of the\n", - "surface cell, with its lowest atom above the top slab atom by the sum of their covalent radii plus 0.5 Å. Slab atoms\n", - "are labelled 0, molecule atoms 1." + "Places the molecule `NUMBER_OF_CONFIGURATIONS` times over the slab, each at a random orientation and a random point\n", + "of the surface cell, with its lowest atom above the top slab atom by the sum of their covalent radii plus 0.5 Å, as\n", + "GRRM's random initial structures [1]. Slab atoms are labelled 0, molecule atoms 1." ] }, { @@ -257,8 +258,8 @@ "metadata": {}, "outputs": [], "source": [ - "CONFIGURATIONS_SCRIPT = r'''\"\"\"Places the job's molecule over its slab N_CONFIGURATIONS times, each at a random orientation and a random point of\n", - "the surface cell, and passes the configurations to the next step as CONFIGURATIONS.\"\"\"\n", + "CONFIGURATIONS_SCRIPT = r'''\"\"\"Places the job's molecule over its slab NUMBER_OF_CONFIGURATIONS times, each at a random orientation and a random\n", + "point of the surface cell, and passes the configurations to the next step as CONFIGURATIONS.\"\"\"\n", "\n", "import json\n", "\n", @@ -270,49 +271,48 @@ "from mat3ra.made.tools.operations.core.binary import merge\n", "from mat3ra.made.tools.operations.core.unary import rotate\n", "\n", - "N_CONFIGURATIONS = 100\n", - "SEED = 0\n", + "NUMBER_OF_CONFIGURATIONS = 100\n", + "RANDOM_SEED = 0\n", "CONTACT_OFFSET = 0.5 # A, lowest molecule atom at R_molecule + R_slab + CONTACT_OFFSET above the top slab atom (GRRM17 RandDC)\n", "# The job's materials, written here by the platform when the job is created.\n", "MATERIALS = json.loads(r\"\"\"{{ materials | dump | safe }}\"\"\")\n", "\n", "\n", - "def material_in_crystal_units(document):\n", - " material = Material.create(document)\n", - " material.to_crystal()\n", - " return material\n", - "\n", - "\n", - "def extreme_atom_height_and_radius(material, choose):\n", + "def extreme_atom_height_and_radius(material, extremum):\n", " cartesian = material.clone()\n", " cartesian.to_cartesian()\n", - " height, element = choose(zip([coordinate[2] for coordinate in cartesian.basis.coordinates.values],\n", - " cartesian.basis.elements.values))\n", + " height, element = extremum(zip([coordinate[2] for coordinate in cartesian.basis.coordinates.values],\n", + " cartesian.basis.elements.values))\n", " return height, covalent_radii[atomic_numbers[element]]\n", "\n", "\n", - "def create_configuration(slab, molecule, random_generator):\n", + "def create_configuration(slab, molecule, top_atom, random_generator, index):\n", + " \"\"\"Random orientation and point over the surface cell; lowest molecule atom R + R + CONTACT_OFFSET above the slab.\"\"\"\n", " quaternion = random_generator.normal(size=4)\n", " angle = 2 * np.degrees(np.arccos(quaternion[0] / np.linalg.norm(quaternion)))\n", " rotated = rotate(molecule, axis=quaternion[1:].tolist(), angle=angle, wrap=False)\n", " rotated.set_lattice(slab.lattice)\n", - " top_height, top_radius = extreme_atom_height_and_radius(slab, max)\n", - " lowest_height, lowest_radius = extreme_atom_height_and_radius(rotated, min)\n", - " point = slab.basis.cell.convert_point_to_cartesian([random_generator.random(), random_generator.random(), 0.0])\n", - " center = molecule.basis.cell.convert_point_to_cartesian([0.5, 0.5, 0.5])\n", - " height = top_height + top_radius + lowest_radius + CONTACT_OFFSET - lowest_height\n", - " placed = translate_by_vector(rotated, vector=[point[0] - center[0], point[1] - center[1], height],\n", + " (top_height, top_radius), (lowest_height, lowest_radius) = top_atom, extreme_atom_height_and_radius(rotated, min)\n", + " surface_point = slab.basis.cell.convert_point_to_cartesian([random_generator.random(), random_generator.random(), 0])\n", + " molecule_cell_center = molecule.basis.cell.convert_point_to_cartesian([0.5, 0.5, 0.5])\n", + " vertical_shift = top_height + top_radius + lowest_radius + CONTACT_OFFSET - lowest_height\n", + " placed = translate_by_vector(rotated, vector=[surface_point[0] - molecule_cell_center[0],\n", + " surface_point[1] - molecule_cell_center[1], vertical_shift],\n", " use_cartesian_coordinates=True)\n", " placed.set_labels_from_value(InterfacePartsEnum.FILM.value)\n", - " return merge([slab, placed])\n", + " return merge([slab, placed], material_name=f\"{molecule.name} on {slab.name}, configuration {index}\")\n", "\n", "\n", "def main():\n", - " slab = material_in_crystal_units(next(item for item in MATERIALS if not item.get(\"isNonPeriodic\")))\n", - " molecule = material_in_crystal_units(next(item for item in MATERIALS if item.get(\"isNonPeriodic\")))\n", + " slab = Material.create(next(config for config in MATERIALS if not config.get(\"isNonPeriodic\")))\n", + " molecule = Material.create(next(config for config in MATERIALS if config.get(\"isNonPeriodic\")))\n", + " for material in (slab, molecule):\n", + " material.to_crystal()\n", " slab.set_labels_from_value(InterfacePartsEnum.SUBSTRATE.value)\n", - " random_generator = np.random.default_rng(SEED)\n", - " configurations = [create_configuration(slab, molecule, random_generator).to_dict() for _ in range(N_CONFIGURATIONS)]\n", + " top_atom = extreme_atom_height_and_radius(slab, max)\n", + " random_generator = np.random.default_rng(RANDOM_SEED)\n", + " configurations = [create_configuration(slab, molecule, top_atom, random_generator, index).to_dict()\n", + " for index in range(NUMBER_OF_CONFIGURATIONS)]\n", " print(json.dumps({\"name\": \"CONFIGURATIONS\", \"value\": configurations, \"scope\": \"global\"}))\n", "\n", "\n", @@ -321,24 +321,16 @@ "'''" ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "CONFIGURATIONS_REQUIREMENTS = \"mat3ra-made[tools]\\n\"" - ] - }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 4.2. Step 2: AFIR with ALCHEMI\n", "\n", - "Relaxes all configurations together with the slab fixed, first under the AFIR force of strength `GAMMA` between\n", - "molecule and slab, then without it, and writes `results.csv` and `final_structures.extxyz`. Another backend or\n", - "another ALCHEMI version replaces this script and its packages; step 1 stays." + "Relaxes all configurations together with the slab fixed [3], first under the AFIR force [1] of strength `GAMMA`\n", + "(100 kJ/mol with random orientations [2]) between molecule and slab, then without it, and writes `results.csv` and\n", + "`final_structures.extxyz`. Another backend or another ALCHEMI version replaces this script and its packages; steps 1\n", + "and 3 stay." ] }, { @@ -358,6 +350,7 @@ "from ase.calculators.singlepoint import SinglePointCalculator\n", "from ase.data import covalent_radii\n", "from ase.io import write\n", + "from ase.units import kJ, mol\n", "from loguru import logger\n", "from mat3ra.made.material import Material\n", "from mat3ra.made.tools.convert import to_ase\n", @@ -375,30 +368,31 @@ "DTYPE = torch.float32\n", "FMAX = 0.03 # eV/A\n", "MAX_STEPS = 500\n", - "TIME_STEP = 0.5\n", + "TIME_STEP = 0.5 # initial FIRE2 step, in ALCHEMI's internal time unit (10.18 fs)\n", "GAMMA = 100 # kJ/mol, model collision energy of the AFIR force\n", "AFIR_R0 = 3.8164 # A, GRRM17 Eq. (3)\n", "AFIR_EPSILON = 1.0061 # kJ/mol, GRRM17 Eq. (3)\n", - "KJ_PER_MOL_PER_EV = 96.485\n", "CONFIGURATIONS = json.loads(r\"\"\"{{ CONFIGURATIONS | tojson }}\"\"\")\n", "\n", - "ALPHA = GAMMA / ((2 ** (-1 / 6) - (1 + (1 + GAMMA / AFIR_EPSILON) ** 0.5) ** (-1 / 6)) * AFIR_R0) / KJ_PER_MOL_PER_EV\n", + "ALPHA = GAMMA * kJ / mol / ((2 ** (-1 / 6) - (1 + (1 + GAMMA / AFIR_EPSILON) ** 0.5) ** (-1 / 6)) * AFIR_R0) # Eq. (3)\n", "CATEGORY_BY_PART = {InterfacePartsEnum.SUBSTRATE.value: AtomCategory.SURFACE.value,\n", " InterfacePartsEnum.FILM.value: AtomCategory.GAS.value}\n", - "IN_PLANE = torch.tensor([1.0, 1.0, 0.0], device=DEVICE)\n", + "IN_PLANE_MASK = torch.tensor([1.0, 1.0, 0.0], device=DEVICE)\n", "IN_PLANE_SHIFTS = torch.tensor([[a, b, 0.0] for a in (-1, 0, 1) for b in (-1, 0, 1)], device=DEVICE)\n", "COVALENT_RADII = torch.tensor(covalent_radii, device=DEVICE, dtype=DTYPE)\n", "\n", "\n", "def in_plane_minimum_image(vectors, cells):\n", + " \"\"\"Shortest periodic image of each vector in the surface plane (a, b); no image along c.\"\"\"\n", " fractional = torch.einsum(\"pi,pij->pj\", vectors, torch.linalg.inv(cells))\n", - " fractional = fractional - torch.round(fractional) * IN_PLANE\n", + " fractional = fractional - torch.round(fractional) * IN_PLANE_MASK\n", " images = torch.einsum(\"psi,pij->psj\", fractional.unsqueeze(1) + IN_PLANE_SHIFTS, cells)\n", " nearest = images.norm(dim=-1).argmin(dim=1)\n", " return torch.take_along_dim(images, nearest[:, None, None], dim=1).squeeze(1)\n", "\n", "\n", "def afir_energy(positions, batch):\n", + " \"\"\"AFIR term of GRRM17 Eq. (1), α Σ ω_ij r_ij / Σ ω_ij, ω_ij = ((R_i + R_j) / r_ij)^6 (Eq. 2), per structure.\"\"\"\n", " molecule_atoms = torch.where(batch.atom_categories == AtomCategory.GAS.value)[0]\n", " surface_atoms = torch.where(batch.atom_categories == AtomCategory.SURFACE.value)[0]\n", " molecule_index, surface_index = torch.where(\n", @@ -431,14 +425,15 @@ " return data\n", "\n", "\n", - "def relaxed_atoms(original, row):\n", + "def relaxed_atoms(original, relaxed_data):\n", " atoms = original.copy()\n", - " atoms.positions = row.positions.cpu().numpy()\n", - " atoms.calc = SinglePointCalculator(atoms, energy=float(row.energy.item()), forces=row.forces.cpu().numpy())\n", + " atoms.positions = relaxed_data.positions.cpu().numpy()\n", + " atoms.calc = SinglePointCalculator(atoms, energy=float(relaxed_data.energy.item()),\n", + " forces=relaxed_data.forces.cpu().numpy())\n", " return atoms\n", "\n", "\n", - "def relax(structures, model, bias_hooks):\n", + "def relax(ase_structures, model, bias_hooks):\n", " optimizer = FIRE2(model=model, dt=TIME_STEP, n_steps=MAX_STEPS, convergence_hook=ConvergenceHook(\n", " criteria=[{\"key\": \"forces\", \"threshold\": FMAX, \"reduce_op\": \"norm\", \"reduce_dims\": -1}]))\n", " optimizer.register_hook(FreezeAtomsHook(freeze_category=AtomCategory.SURFACE.value))\n", @@ -447,8 +442,8 @@ " stage.register_hook(hook, stage=DynamicsStage.BEFORE_COMPUTE)\n", " for hook in bias_hooks:\n", " stage.register_hook(hook, stage=DynamicsStage.AFTER_COMPUTE)\n", - " batch = stage.run(Batch.from_data_list([atomic_data(atoms) for atoms in structures]))\n", - " return [relaxed_atoms(atoms, row) for atoms, row in zip(structures, batch.to_data_list())]\n", + " batch = stage.run(Batch.from_data_list([atomic_data(atoms) for atoms in ase_structures]))\n", + " return [relaxed_atoms(atoms, relaxed_data) for atoms, relaxed_data in zip(ase_structures, batch.to_data_list())]\n", "\n", "\n", "def energy_and_max_force(atoms):\n", @@ -470,9 +465,9 @@ "\n", "def main():\n", " logger.info(\"AFIR force α, γ = {} kJ/mol: {:.3f} eV/Å\", GAMMA, ALPHA)\n", - " structures = [to_ase(Material.create(configuration)) for configuration in CONFIGURATIONS]\n", + " ase_structures = [to_ase(Material.create(configuration)) for configuration in CONFIGURATIONS]\n", " model = MACEWrapper.from_checkpoint(checkpoint_path=CHECKPOINT, device=torch.device(DEVICE), dtype=DTYPE)\n", - " afir_structures = relax(structures, model, [BiasedPotentialHook(bias_fn=afir_bias)])\n", + " afir_structures = relax(ase_structures, model, [BiasedPotentialHook(bias_fn=afir_bias)])\n", " final_structures = relax(afir_structures, model, [])\n", " write_results(afir_structures, final_structures)\n", "\n", @@ -497,12 +492,7 @@ "torchaudio==2.11.0+cu128\n", "warp-lang==1.18.0+cu12\n", "nvalchemi-toolkit[mace,ase]==0.2.0\n", - "matplotlib\n", - "numpy\n", - "ase\n", - "mace-torch\n", - "loguru\n", - "mat3ra-made[tools]\n", + "mat3ra-made[tools]==2026.10.10.post0\n", "\"\"\"" ] }, @@ -510,11 +500,55 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### 4.3. Create the workflow\n", + "### 4.3. Step 3: final structures\n", + "\n", + "Takes the `NUMBER_OF_FINAL_STRUCTURES` lowest final structures, each with the configuration it started from, and\n", + "saves them as the job's final structures: each is a material, shown in the Results tab of the AFIR step." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "FINAL_STRUCTURES_SCRIPT = r'''\"\"\"Takes the NUMBER_OF_FINAL_STRUCTURES lowest-energy structures of the previous step, each with the configuration it\n", + "started from, and passes one final-structure request body per structure to the next units.\"\"\"\n", + "\n", + "import json\n", + "\n", + "from ase.io import read\n", + "from mat3ra.made.material import Material\n", + "from mat3ra.made.tools.convert import from_ase\n", + "\n", + "NUMBER_OF_FINAL_STRUCTURES = 5\n", + "AFIR_UNIT_FLOWCHART_ID = \"\"\n", + "CONFIGURATIONS = json.loads(r\"\"\"{{ CONFIGURATIONS | tojson }}\"\"\")\n", + "\n", + "\n", + "def request_body(rank_index, atoms):\n", + " # The platform keeps one final_structure per unit and repetition, so each structure gets its own repetition.\n", + " final = {**Material.create(from_ase(atoms)).to_dict(), \"repetition\": rank_index}\n", + " initial = CONFIGURATIONS[atoms.info[\"configuration\"]]\n", + " body = {\"jobId\": \"{{ JOB_ID }}\", \"unitId\": AFIR_UNIT_FLOWCHART_ID, \"structures\": [{\"initial\": initial, \"final\": final}]}\n", + " return json.dumps(json.dumps(body))[1:-1]\n", + "\n", + "\n", + "final_structures = read(\"{{ JOB_WORK_DIR }}/final_structures.extxyz\", index=f\":{NUMBER_OF_FINAL_STRUCTURES}\")\n", + "requests = [request_body(rank_index, atoms) for rank_index, atoms in enumerate(final_structures)]\n", + "print(json.dumps({\"name\": \"FINAL_STRUCTURE_REQUESTS\", \"value\": requests, \"scope\": \"global\"}))\n", + "'''" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 4.4. Configure the steps\n", "\n", "When the job is created, the platform writes the job's materials into step 1's script (`materials`; the\n", "non-periodic one is the molecule). Step 1 prints `CONFIGURATIONS` to the job's scope, and the compute node writes\n", - "them into step 2's script before it runs." + "them into the scripts of steps 2 and 3 before they run. Each step's files are named after it." ] }, { @@ -527,39 +561,91 @@ "import re\n", "\n", "from mat3ra.notebooks_utils.core.entity.workflow.api import set_execution_unit_input\n", - "from mat3ra.notebooks_utils.ipython.entity.workflow.visualize import visualize_workflow\n", "from mat3ra.standata.workflows import WorkflowStandata\n", - "from mat3ra.wode import Subworkflow, Workflow\n", + "\n", + "MADE_REQUIREMENTS = \"mat3ra-made[tools]==2026.10.10.post0\\n\"\n", + "\n", + "\n", + "def set_constants(script, constants):\n", + " for name, value in constants.items():\n", + " script = re.sub(rf\"^({name} = )\\S+\", rf\"\\g<1>{value}\", script, flags=re.M)\n", + " return script\n", + "\n", "\n", "python_workflows = WorkflowStandata.filter_by_application(\"python\")\n", "workflow_config = copy.deepcopy(python_workflows.get_by_name_first_match(\"custom_script.json\"))\n", "workflow_config.update(name=\"AFIR, molecule on surface\", isMultiMaterial=True)\n", "workflow_config[\"units\"][0][\"name\"] = \"Configurations\"\n", - "configurations_subworkflow = workflow_config[\"subworkflows\"][0]\n", - "afir_subworkflow = copy.deepcopy(python_workflows.get_by_name_first_match(\"python_script.json\")[\"subworkflows\"][0])\n", + "python_script_subworkflow = python_workflows.get_by_name_first_match(\"python_script.json\")[\"subworkflows\"][0]\n", + "afir_subworkflow_config = copy.deepcopy(python_script_subworkflow)\n", + "final_structures_subworkflow_config = copy.deepcopy(python_script_subworkflow)\n", + "final_structures_subworkflow_config.pop(\"_id\")\n", + "next(unit for unit in final_structures_subworkflow_config[\"units\"] if unit[\"name\"] == \"python\").pop(\"flowchartId\")\n", + "afir_unit_flowchart_id = next(unit for unit in afir_subworkflow_config[\"units\"] if unit[\"name\"] == \"python\")[\"flowchartId\"]\n", "afir_results = [\n", " {\"name\": \"file_content\", \"filetype\": \"csv\", \"basename\": \"results.csv\"},\n", " {\"name\": \"file_content\", \"filetype\": \"text\", \"basename\": \"final_structures.extxyz\"},\n", "]\n", - "for subworkflow, name, unit_name, requirements, results in (\n", - " (configurations_subworkflow, \"Configurations\", \"configurations\", CONFIGURATIONS_REQUIREMENTS, []),\n", - " (afir_subworkflow, \"AFIR (ALCHEMI)\", \"afir\", AFIR_REQUIREMENTS, afir_results),\n", - "):\n", - " unit = subworkflow[\"units\"][-1]\n", + "steps = [\n", + " (workflow_config[\"subworkflows\"][0], \"custom_script\", \"Configurations\", \"configurations\", MADE_REQUIREMENTS, [], False,\n", + " set_constants(CONFIGURATIONS_SCRIPT, {\"NUMBER_OF_CONFIGURATIONS\": NUMBER_OF_CONFIGURATIONS, \"RANDOM_SEED\": RANDOM_SEED})),\n", + " (afir_subworkflow_config, \"python\", \"AFIR (ALCHEMI)\", \"afir\", AFIR_REQUIREMENTS, afir_results, True,\n", + " set_constants(AFIR_SCRIPT, {\"GAMMA\": GAMMA})),\n", + " (final_structures_subworkflow_config, \"python\", \"Final structures\", \"final_structures\", MADE_REQUIREMENTS, [], True,\n", + " set_constants(FINAL_STRUCTURES_SCRIPT, {\"NUMBER_OF_FINAL_STRUCTURES\": NUMBER_OF_FINAL_STRUCTURES,\n", + " \"AFIR_UNIT_FLOWCHART_ID\": f'\"{afir_unit_flowchart_id}\"'})),\n", + "]\n", + "for subworkflow, template_unit_name, name, unit_name, requirements, results, is_rendered_at_run_time, script in steps:\n", + " unit = next(unit for unit in subworkflow[\"units\"] if unit[\"name\"] == template_unit_name)\n", " unit.update(name=unit_name, head=True, results=results)\n", " for application in (subworkflow[\"application\"], unit[\"application\"]):\n", " application.update(version=PYTHON_VERSION, build=PYTHON_BUILD)\n", + " set_execution_unit_input(unit, \"script.py\", script)\n", " set_execution_unit_input(unit, \"requirements.txt\", requirements)\n", - " subworkflow.update(name=name, units=[unit])\n", + " inputs = {unit_input[\"template\"][\"name\"]: unit_input for unit_input in unit[\"input\"]}\n", + " inputs[\"script.py\"][\"isManuallyChanged\"] = is_rendered_at_run_time\n", + " inputs[\"script.py\"][\"template\"][\"name\"] = f\"{unit_name}.py\"\n", + " inputs[\"requirements.txt\"][\"template\"][\"name\"] = f\"{unit_name}_requirements.txt\"\n", + " subworkflow.update(name=name, units=[unit])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 4.5. Create and save the workflow\n", "\n", - "configurations_script = re.sub(r\"^N_CONFIGURATIONS = \\S+\", f\"N_CONFIGURATIONS = {N_CONFIGURATIONS}\", CONFIGURATIONS_SCRIPT, flags=re.M)\n", - "script_input = next(unit_input for unit_input in configurations_subworkflow[\"units\"][0][\"input\"] if unit_input[\"template\"][\"name\"] == \"script.py\")\n", - "script_input[\"template\"][\"content\"] = re.sub(r\"^SEED = \\S+\", f\"SEED = {SEED}\", configurations_script, flags=re.M)\n", - "set_execution_unit_input(afir_subworkflow[\"units\"][0], \"script.py\", re.sub(r\"^GAMMA = \\S+\", f\"GAMMA = {GAMMA}\", AFIR_SCRIPT, flags=re.M))\n", + "Step 3 ends with one request per final structure, each saving it with the starting configuration as its initial\n", + "structure." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.notebooks_utils.ipython.entity.workflow.visualize import visualize_workflow\n", + "from mat3ra.wode import Subworkflow, Workflow\n", + "from mat3ra.wode.units import Unit\n", "\n", "workflow = Workflow.create(workflow_config)\n", - "workflow.add_subworkflow(Subworkflow(**afir_subworkflow))\n", - "visualize_workflow(workflow)" + "workflow.add_subworkflow(Subworkflow(**afir_subworkflow_config))\n", + "final_structures_subworkflow = Subworkflow(**final_structures_subworkflow_config)\n", + "for index in range(NUMBER_OF_FINAL_STRUCTURES):\n", + " request = {\"method\": \"POST\", \"data\": f\"{{{{ FINAL_STRUCTURE_REQUESTS[{index}] }}}}\"}\n", + " final_structures_subworkflow.add_unit(Unit(\n", + " type=\"io\", subtype=\"output\", source=\"api\", enableRender=True, name=f\"save-final-structure-{index + 1}\",\n", + " input=[{\"type\": \"api\", \"endpoint\": \"structures\", \"endpoint_options\": request, \"name\": \"FINAL_STRUCTURE\"}],\n", + " ))\n", + "final_structures_subworkflow.add_unit(\n", + " Unit(type=\"assignment\", name=\"final-structures-saved\", operand=\"FINAL_STRUCTURES_SAVED\", value=\"True\", input=[])\n", + ")\n", + "workflow.add_subworkflow(final_structures_subworkflow)\n", + "visualize_workflow(workflow)\n", + "\n", + "saved_workflow = client.workflows.create(workflow.to_dict_without_special_keys(), owner_id=ACCOUNT_ID)\n", + "print(f\"✅ Workflow saved to the collection: {saved_workflow['_id']}\")" ] }, { @@ -621,7 +707,7 @@ " workflow=workflow,\n", " project_id=project_id,\n", " owner_id=ACCOUNT_ID,\n", - " prefix=f\"{workflow.name} {MOLECULE_NAME} {N_CONFIGURATIONS} {timestamp}\",\n", + " prefix=f\"{workflow.name}: {MOLECULE_NAME} on {slab.name}, {NUMBER_OF_CONFIGURATIONS} configurations {timestamp}\",\n", " compute=compute.to_dict(),\n", ")\n", "job_id = job[\"_id\"]\n", @@ -676,8 +762,27 @@ "source": [ "import pandas as pd\n", "\n", - "files = {file[\"key\"].rsplit(\"/\", 1)[-1]: file[\"signedUrl\"] for file in client.jobs.list_files(job_id)}\n", - "pd.read_csv(io.StringIO(await read_from_url(files[\"results.csv\"])))" + "file_url_by_name = {file[\"key\"].rsplit(\"/\", 1)[-1]: file[\"signedUrl\"] for file in client.jobs.list_files(job_id)}\n", + "pd.read_csv(io.StringIO(await read_from_url(file_url_by_name[\"results.csv\"])))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 8.2. The lowest final structures" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.prode import PropertyName\n", + "\n", + "final_structures = client.properties.get_for_job(job_id, property_name=PropertyName.non_scalar.final_structure.value)\n", + "visualize([Material.create(client.materials.get(final_structure[\"materialId\"])) for final_structure in final_structures])" ] }, { From 635bcddfcae35a493bb6998c5626265ff734f9c8 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Mon, 2 Nov 2026 18:31:28 -0800 Subject: [PATCH 05/31] SOF-8102: afir_molecule_surface_alchemi.ipynb: final-structure requests send Content-Type JSON; script headers and docstrings rupy's io unit sends endpoint_options.headers plus its auth headers only; without Content-Type the server parsed no body and answered 422 FORMAT_ERROR (jobId, unitId, structures REQUIRED). rupy's own publish_structures sets it. Each step's script gets the Standata header box and made-style docstrings. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../afir_molecule_surface_alchemi.ipynb | 186 ++++++++++++++++-- 1 file changed, 175 insertions(+), 11 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index fc7624d5e..3ce928493 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -258,8 +258,25 @@ "metadata": {}, "outputs": [], "source": [ - "CONFIGURATIONS_SCRIPT = r'''\"\"\"Places the job's molecule over its slab NUMBER_OF_CONFIGURATIONS times, each at a random orientation and a random\n", - "point of the surface cell, and passes the configurations to the next step as CONFIGURATIONS.\"\"\"\n", + "CONFIGURATIONS_SCRIPT = r'''# ---------------------------------------------------------------- #\n", + "# #\n", + "# Configurations step of the AFIR molecule-on-surface workflow. #\n", + "# #\n", + "# Places the job's molecule over its slab #\n", + "# NUMBER_OF_CONFIGURATIONS times with made, each at a random #\n", + "# orientation and a random point of the surface cell, as #\n", + "# GRRM17's random initial structures [1]. #\n", + "# #\n", + "# Reads: the job's two materials, written into this script by #\n", + "# the platform when the job is created; the non-periodic #\n", + "# one is the molecule. #\n", + "# Writes: CONFIGURATIONS to the job's global scope: the placed #\n", + "# structures as materials, slab atoms labelled 0 and #\n", + "# molecule atoms 1. It is the only output of the script. #\n", + "# #\n", + "# [1] S. Maeda et al., J. Comput. Chem. 39, 233 (2018). #\n", + "# #\n", + "# ---------------------------------------------------------------- #\n", "\n", "import json\n", "\n", @@ -279,6 +296,16 @@ "\n", "\n", "def extreme_atom_height_and_radius(material, extremum):\n", + " \"\"\"\n", + " Returns the height and the covalent radius of the highest or the lowest atom of a material.\n", + "\n", + " Args:\n", + " material (Material): The material to search.\n", + " extremum (Callable): max for the highest atom, min for the lowest.\n", + "\n", + " Returns:\n", + " Tuple[float, float]: The atom's cartesian z and its covalent radius, both in Å.\n", + " \"\"\"\n", " cartesian = material.clone()\n", " cartesian.to_cartesian()\n", " height, element = extremum(zip([coordinate[2] for coordinate in cartesian.basis.coordinates.values],\n", @@ -287,7 +314,23 @@ "\n", "\n", "def create_configuration(slab, molecule, top_atom, random_generator, index):\n", - " \"\"\"Random orientation and point over the surface cell; lowest molecule atom R + R + CONTACT_OFFSET above the slab.\"\"\"\n", + " \"\"\"\n", + " Places the molecule once over the slab, as a random initial structure of GRRM17's MC-AFIR.\n", + "\n", + " The orientation is uniformly random (a normalised Gaussian quaternion), the molecule's center goes to a random\n", + " point of the surface cell, and its lowest atom sits R_molecule + R_slab + CONTACT_OFFSET above the top slab atom,\n", + " the RandDC rule of GRRM17 (S. Maeda et al., J. Comput. Chem. 39, 233 (2018)).\n", + "\n", + " Args:\n", + " slab (Material): The slab, in crystal units, labelled 0.\n", + " molecule (Material): The molecule, in crystal units.\n", + " top_atom (Tuple[float, float]): Height and covalent radius of the top slab atom, in Å.\n", + " random_generator (numpy.random.Generator): The seeded generator of all placements.\n", + " index (int): The configuration's index, used in its name.\n", + "\n", + " Returns:\n", + " MaterialWithBuildMetadata: The slab and the placed molecule, molecule atoms labelled 1.\n", + " \"\"\"\n", " quaternion = random_generator.normal(size=4)\n", " angle = 2 * np.degrees(np.arccos(quaternion[0] / np.linalg.norm(quaternion)))\n", " rotated = rotate(molecule, axis=quaternion[1:].tolist(), angle=angle, wrap=False)\n", @@ -304,6 +347,9 @@ "\n", "\n", "def main():\n", + " \"\"\"\n", + " Builds NUMBER_OF_CONFIGURATIONS configurations from the job's materials and prints them as CONFIGURATIONS.\n", + " \"\"\"\n", " slab = Material.create(next(config for config in MATERIALS if not config.get(\"isNonPeriodic\")))\n", " molecule = Material.create(next(config for config in MATERIALS if config.get(\"isNonPeriodic\")))\n", " for material in (slab, molecule):\n", @@ -339,8 +385,25 @@ "metadata": {}, "outputs": [], "source": [ - "AFIR_SCRIPT = r'''\"\"\"Relaxes the configurations of the previous step with MACE, all batched with the slab fixed: first under the AFIR\n", - "force between molecule and slab, then without it. Ranks the final structures by energy.\"\"\"\n", + "AFIR_SCRIPT = r'''# ---------------------------------------------------------------- #\n", + "# #\n", + "# AFIR step of the AFIR molecule-on-surface workflow, with MACE #\n", + "# through NVIDIA ALCHEMI on one GPU. #\n", + "# #\n", + "# Relaxes all configurations as one batch with the slab fixed, #\n", + "# first under the AFIR force between molecule and slab (GRRM17 #\n", + "# Eqs. 1-3 [1], strength GAMMA), then without it, and ranks the #\n", + "# final structures by energy. #\n", + "# #\n", + "# Reads: CONFIGURATIONS from the job's global scope, written #\n", + "# into this script by the compute node before it runs. #\n", + "# Writes: results.csv, the configurations ranked by final energy, #\n", + "# and final_structures.extxyz, their structures in that #\n", + "# order; both are results of this unit. #\n", + "# #\n", + "# [1] S. Maeda et al., J. Comput. Chem. 39, 233 (2018). #\n", + "# #\n", + "# ---------------------------------------------------------------- #\n", "\n", "import csv\n", "import json\n", @@ -383,7 +446,16 @@ "\n", "\n", "def in_plane_minimum_image(vectors, cells):\n", - " \"\"\"Shortest periodic image of each vector in the surface plane (a, b); no image along c.\"\"\"\n", + " \"\"\"\n", + " Returns the shortest periodic image of each vector in the surface plane, a and b; there is no image along c.\n", + "\n", + " Args:\n", + " vectors (torch.Tensor): Cartesian vectors, [P, 3], in Å.\n", + " cells (torch.Tensor): The lattice vectors of each vector's structure, [P, 3, 3], rows a, b, c.\n", + "\n", + " Returns:\n", + " torch.Tensor: The minimum-image vectors, [P, 3], in Å.\n", + " \"\"\"\n", " fractional = torch.einsum(\"pi,pij->pj\", vectors, torch.linalg.inv(cells))\n", " fractional = fractional - torch.round(fractional) * IN_PLANE_MASK\n", " images = torch.einsum(\"psi,pij->psj\", fractional.unsqueeze(1) + IN_PLANE_SHIFTS, cells)\n", @@ -392,7 +464,20 @@ "\n", "\n", "def afir_energy(positions, batch):\n", - " \"\"\"AFIR term of GRRM17 Eq. (1), α Σ ω_ij r_ij / Σ ω_ij, ω_ij = ((R_i + R_j) / r_ij)^6 (Eq. 2), per structure.\"\"\"\n", + " \"\"\"\n", + " Returns the AFIR term of every structure in the batch, between its molecule and its slab atoms.\n", + "\n", + " GRRM17 Eq. (1), α Σ ω_ij r_ij / Σ ω_ij over molecule atoms i and slab atoms j, with ω_ij = ((R_i + R_j) / r_ij)^6\n", + " (Eq. 2), R the covalent radii and r_ij at the minimum image in the surface plane; α is ALPHA, Eq. (3)\n", + " (S. Maeda et al., J. Comput. Chem. 39, 233 (2018)).\n", + "\n", + " Args:\n", + " positions (torch.Tensor): Atomic positions of the batch, [V, 3], in Å.\n", + " batch (Batch): The batch, with atom_categories GAS for the molecule and SURFACE for the slab.\n", + "\n", + " Returns:\n", + " torch.Tensor: The AFIR term of each structure, [B], in eV.\n", + " \"\"\"\n", " molecule_atoms = torch.where(batch.atom_categories == AtomCategory.GAS.value)[0]\n", " surface_atoms = torch.where(batch.atom_categories == AtomCategory.SURFACE.value)[0]\n", " molecule_index, surface_index = torch.where(\n", @@ -411,6 +496,16 @@ "\n", "\n", "def afir_bias(batch):\n", + " \"\"\"\n", + " Returns the AFIR energy and forces of the batch, as BiasedPotentialHook expects them.\n", + "\n", + " Args:\n", + " batch (Batch): The batch after the model's forward pass.\n", + "\n", + " Returns:\n", + " Tuple[torch.Tensor, torch.Tensor]: Energy per structure, [B, 1], in eV, and force per atom, [V, 3],\n", + " in eV/Å.\n", + " \"\"\"\n", " positions = batch.positions.clone().requires_grad_(True)\n", " energy = afir_energy(positions, batch)\n", " forces = -torch.autograd.grad(energy.sum(), positions)[0]\n", @@ -418,6 +513,15 @@ "\n", "\n", "def atomic_data(atoms):\n", + " \"\"\"\n", + " Converts a configuration to ALCHEMI's AtomicData, its labels to atom categories: 0 to SURFACE, 1 to GAS.\n", + "\n", + " Args:\n", + " atoms (ase.Atoms): The configuration, slab atoms tagged 0 and molecule atoms 1.\n", + "\n", + " Returns:\n", + " AtomicData: The configuration on DEVICE, with zeroed forces and energy.\n", + " \"\"\"\n", " data = AtomicData.from_atoms(atoms, device=DEVICE, dtype=DTYPE)\n", " data[\"atom_categories\"] = torch.tensor([CATEGORY_BY_PART[tag] for tag in atoms.get_tags()], device=DEVICE)\n", " data[\"forces\"] = torch.zeros_like(data.positions)\n", @@ -426,6 +530,16 @@ "\n", "\n", "def relaxed_atoms(original, relaxed_data):\n", + " \"\"\"\n", + " Returns a copy of a structure at its relaxed positions, with its energy and forces.\n", + "\n", + " Args:\n", + " original (ase.Atoms): The structure before the relaxation.\n", + " relaxed_data (AtomicData): The same structure after it.\n", + "\n", + " Returns:\n", + " ase.Atoms: The relaxed structure, energy and forces on a SinglePointCalculator.\n", + " \"\"\"\n", " atoms = original.copy()\n", " atoms.positions = relaxed_data.positions.cpu().numpy()\n", " atoms.calc = SinglePointCalculator(atoms, energy=float(relaxed_data.energy.item()),\n", @@ -434,6 +548,17 @@ "\n", "\n", "def relax(ase_structures, model, bias_hooks):\n", + " \"\"\"\n", + " Relaxes all structures as one batch with fixed-cell FIRE2, slab atoms frozen.\n", + "\n", + " Args:\n", + " ase_structures (List[ase.Atoms]): The structures, slab atoms tagged 0 and molecule atoms 1.\n", + " model (MACEWrapper): The MACE model.\n", + " bias_hooks (List[BiasedPotentialHook]): Bias potentials added after each forward pass; [] for none.\n", + "\n", + " Returns:\n", + " List[ase.Atoms]: The relaxed structures, in the input order.\n", + " \"\"\"\n", " optimizer = FIRE2(model=model, dt=TIME_STEP, n_steps=MAX_STEPS, convergence_hook=ConvergenceHook(\n", " criteria=[{\"key\": \"forces\", \"threshold\": FMAX, \"reduce_op\": \"norm\", \"reduce_dims\": -1}]))\n", " optimizer.register_hook(FreezeAtomsHook(freeze_category=AtomCategory.SURFACE.value))\n", @@ -447,10 +572,20 @@ "\n", "\n", "def energy_and_max_force(atoms):\n", + " \"\"\"\n", + " Returns the energy of a structure, in eV, and the largest force on any of its atoms, in eV/Å.\n", + " \"\"\"\n", " return atoms.get_potential_energy(), np.linalg.norm(atoms.get_forces(), axis=1).max()\n", "\n", "\n", "def write_results(afir_structures, final_structures):\n", + " \"\"\"\n", + " Ranks the configurations by final energy and writes results.csv and final_structures.extxyz in that order.\n", + "\n", + " Args:\n", + " afir_structures (List[ase.Atoms]): The structures after the relaxation under the AFIR force.\n", + " final_structures (List[ase.Atoms]): The structures after the relaxation without it.\n", + " \"\"\"\n", " order = np.argsort([atoms.get_potential_energy() for atoms in final_structures])\n", " with open(\"results.csv\", \"w\", newline=\"\") as results_file:\n", " writer = csv.writer(results_file)\n", @@ -464,6 +599,9 @@ "\n", "\n", "def main():\n", + " \"\"\"\n", + " Relaxes CONFIGURATIONS under the AFIR force, then without it, and writes the ranked results.\n", + " \"\"\"\n", " logger.info(\"AFIR force α, γ = {} kJ/mol: {:.3f} eV/Å\", GAMMA, ALPHA)\n", " ase_structures = [to_ase(Material.create(configuration)) for configuration in CONFIGURATIONS]\n", " model = MACEWrapper.from_checkpoint(checkpoint_path=CHECKPOINT, device=torch.device(DEVICE), dtype=DTYPE)\n", @@ -512,8 +650,22 @@ "metadata": {}, "outputs": [], "source": [ - "FINAL_STRUCTURES_SCRIPT = r'''\"\"\"Takes the NUMBER_OF_FINAL_STRUCTURES lowest-energy structures of the previous step, each with the configuration it\n", - "started from, and passes one final-structure request body per structure to the next units.\"\"\"\n", + "FINAL_STRUCTURES_SCRIPT = r'''# ---------------------------------------------------------------- #\n", + "# #\n", + "# Final structures step of the AFIR molecule-on-surface workflow. #\n", + "# #\n", + "# Takes the NUMBER_OF_FINAL_STRUCTURES lowest final structures, #\n", + "# each with the configuration it started from, and makes one #\n", + "# request body per structure for the platform's final-structure #\n", + "# endpoint, which the io units after this one send. #\n", + "# #\n", + "# Reads: final_structures.extxyz in the job's work folder, and #\n", + "# CONFIGURATIONS from the job's global scope, written #\n", + "# into this script by the compute node before it runs. #\n", + "# Writes: FINAL_STRUCTURE_REQUESTS to the job's global scope, as #\n", + "# the only output of the script. #\n", + "# #\n", + "# ---------------------------------------------------------------- #\n", "\n", "import json\n", "\n", @@ -527,7 +679,18 @@ "\n", "\n", "def request_body(rank_index, atoms):\n", - " # The platform keeps one final_structure per unit and repetition, so each structure gets its own repetition.\n", + " \"\"\"\n", + " Returns the body of one final-structure request, escaped to sit inside the next io unit's JSON string.\n", + "\n", + " The platform keeps one final_structure per unit and repetition, so each structure gets its own repetition.\n", + "\n", + " Args:\n", + " rank_index (int): The structure's rank by final energy, from 0.\n", + " atoms (ase.Atoms): The final structure, with its configuration's index in atoms.info.\n", + "\n", + " Returns:\n", + " str: {\"jobId\", \"unitId\", \"structures\": [{\"initial\", \"final\"}]} as JSON, escaped for a JSON string.\n", + " \"\"\"\n", " final = {**Material.create(from_ase(atoms)).to_dict(), \"repetition\": rank_index}\n", " initial = CONFIGURATIONS[atoms.info[\"configuration\"]]\n", " body = {\"jobId\": \"{{ JOB_ID }}\", \"unitId\": AFIR_UNIT_FLOWCHART_ID, \"structures\": [{\"initial\": initial, \"final\": final}]}\n", @@ -633,7 +796,8 @@ "workflow.add_subworkflow(Subworkflow(**afir_subworkflow_config))\n", "final_structures_subworkflow = Subworkflow(**final_structures_subworkflow_config)\n", "for index in range(NUMBER_OF_FINAL_STRUCTURES):\n", - " request = {\"method\": \"POST\", \"data\": f\"{{{{ FINAL_STRUCTURE_REQUESTS[{index}] }}}}\"}\n", + " request = {\"method\": \"POST\", \"headers\": {\"Content-Type\": \"application/json\"},\n", + " \"data\": f\"{{{{ FINAL_STRUCTURE_REQUESTS[{index}] }}}}\"}\n", " final_structures_subworkflow.add_unit(Unit(\n", " type=\"io\", subtype=\"output\", source=\"api\", enableRender=True, name=f\"save-final-structure-{index + 1}\",\n", " input=[{\"type\": \"api\", \"endpoint\": \"structures\", \"endpoint_options\": request, \"name\": \"FINAL_STRUCTURE\"}],\n", From bbbe44c978776d0f8e966387c64b6a9af0369051 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Mon, 2 Nov 2026 19:06:37 -0800 Subject: [PATCH 06/31] SOF-8102: afir_molecule_surface_alchemi.ipynb: AFIR unit declares final_structure; review round 2 The AFIR unit lists final_structure among its results, so the Results tab shows the saved structures there. The number saved is the smaller of NUMBER_OF_FINAL_STRUCTURES and NUMBER_OF_CONFIGURATIONS. Each step is configured by one function with keyword arguments, the made pin is written once, and the viewer titles each structure by its rank. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../afir_molecule_surface_alchemi.ipynb | 138 +++++++++++------- 1 file changed, 86 insertions(+), 52 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 3ce928493..41d3c3e0f 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -290,7 +290,7 @@ "\n", "NUMBER_OF_CONFIGURATIONS = 100\n", "RANDOM_SEED = 0\n", - "CONTACT_OFFSET = 0.5 # A, lowest molecule atom at R_molecule + R_slab + CONTACT_OFFSET above the top slab atom (GRRM17 RandDC)\n", + "CONTACT_OFFSET = 0.5 # Å, lowest molecule atom at R_molecule + R_slab + CONTACT_OFFSET above the top slab atom\n", "# The job's materials, written here by the platform when the job is created.\n", "MATERIALS = json.loads(r\"\"\"{{ materials | dump | safe }}\"\"\")\n", "\n", @@ -313,7 +313,7 @@ " return height, covalent_radii[atomic_numbers[element]]\n", "\n", "\n", - "def create_configuration(slab, molecule, top_atom, random_generator, index):\n", + "def create_configuration(slab, molecule, top_atom_height_and_radius, random_generator, index):\n", " \"\"\"\n", " Places the molecule once over the slab, as a random initial structure of GRRM17's MC-AFIR.\n", "\n", @@ -324,7 +324,7 @@ " Args:\n", " slab (Material): The slab, in crystal units, labelled 0.\n", " molecule (Material): The molecule, in crystal units.\n", - " top_atom (Tuple[float, float]): Height and covalent radius of the top slab atom, in Å.\n", + " top_atom_height_and_radius (Tuple[float, float]): Height and covalent radius of the top slab atom, in Å.\n", " random_generator (numpy.random.Generator): The seeded generator of all placements.\n", " index (int): The configuration's index, used in its name.\n", "\n", @@ -335,8 +335,10 @@ " angle = 2 * np.degrees(np.arccos(quaternion[0] / np.linalg.norm(quaternion)))\n", " rotated = rotate(molecule, axis=quaternion[1:].tolist(), angle=angle, wrap=False)\n", " rotated.set_lattice(slab.lattice)\n", - " (top_height, top_radius), (lowest_height, lowest_radius) = top_atom, extreme_atom_height_and_radius(rotated, min)\n", - " surface_point = slab.basis.cell.convert_point_to_cartesian([random_generator.random(), random_generator.random(), 0])\n", + " top_height, top_radius = top_atom_height_and_radius\n", + " lowest_height, lowest_radius = extreme_atom_height_and_radius(rotated, min)\n", + " surface_fractions = [random_generator.random(), random_generator.random(), 0]\n", + " surface_point = slab.basis.cell.convert_point_to_cartesian(surface_fractions)\n", " molecule_cell_center = molecule.basis.cell.convert_point_to_cartesian([0.5, 0.5, 0.5])\n", " vertical_shift = top_height + top_radius + lowest_radius + CONTACT_OFFSET - lowest_height\n", " placed = translate_by_vector(rotated, vector=[surface_point[0] - molecule_cell_center[0],\n", @@ -350,15 +352,19 @@ " \"\"\"\n", " Builds NUMBER_OF_CONFIGURATIONS configurations from the job's materials and prints them as CONFIGURATIONS.\n", " \"\"\"\n", - " slab = Material.create(next(config for config in MATERIALS if not config.get(\"isNonPeriodic\")))\n", - " molecule = Material.create(next(config for config in MATERIALS if config.get(\"isNonPeriodic\")))\n", + " slab = Material.create(next(material_config for material_config in MATERIALS\n", + " if not material_config.get(\"isNonPeriodic\")))\n", + " molecule = Material.create(next(material_config for material_config in MATERIALS\n", + " if material_config.get(\"isNonPeriodic\")))\n", " for material in (slab, molecule):\n", " material.to_crystal()\n", " slab.set_labels_from_value(InterfacePartsEnum.SUBSTRATE.value)\n", - " top_atom = extreme_atom_height_and_radius(slab, max)\n", + " top_atom_height_and_radius = extreme_atom_height_and_radius(slab, max)\n", " random_generator = np.random.default_rng(RANDOM_SEED)\n", - " configurations = [create_configuration(slab, molecule, top_atom, random_generator, index).to_dict()\n", - " for index in range(NUMBER_OF_CONFIGURATIONS)]\n", + " configurations = [\n", + " create_configuration(slab, molecule, top_atom_height_and_radius, random_generator, index).to_dict()\n", + " for index in range(NUMBER_OF_CONFIGURATIONS)\n", + " ]\n", " print(json.dumps({\"name\": \"CONFIGURATIONS\", \"value\": configurations, \"scope\": \"global\"}))\n", "\n", "\n", @@ -427,10 +433,10 @@ "from nvalchemi.models.mace import MACEWrapper\n", "\n", "DEVICE = \"cuda:0\"\n", - "CHECKPOINT = \"medium-mpa-0\"\n", + "CHECKPOINT = \"medium-mpa-0\" # MACE-MP-0 model, medium MPA-0\n", "DTYPE = torch.float32\n", - "FMAX = 0.03 # eV/A\n", - "MAX_STEPS = 500\n", + "FMAX = 0.03 # eV/Å, largest atomic force at convergence\n", + "MAX_STEPS = 500 # FIRE2 steps per relaxation at most\n", "TIME_STEP = 0.5 # initial FIRE2 step, in ALCHEMI's internal time unit (10.18 fs)\n", "GAMMA = 100 # kJ/mol, model collision energy of the AFIR force\n", "AFIR_R0 = 3.8164 # A, GRRM17 Eq. (3)\n", @@ -621,6 +627,7 @@ "metadata": {}, "outputs": [], "source": [ + "MADE_REQUIREMENTS = \"mat3ra-made[tools]==2026.10.10.post0\\n\"\n", "AFIR_REQUIREMENTS = \"\"\"--find-links https://download.pytorch.org/whl/cu128/torch/\n", "--find-links https://download.pytorch.org/whl/cu128/torchvision/\n", "--find-links https://download.pytorch.org/whl/cu128/torchaudio/\n", @@ -630,8 +637,7 @@ "torchaudio==2.11.0+cu128\n", "warp-lang==1.18.0+cu12\n", "nvalchemi-toolkit[mace,ase]==0.2.0\n", - "mat3ra-made[tools]==2026.10.10.post0\n", - "\"\"\"" + "\"\"\" + MADE_REQUIREMENTS" ] }, { @@ -693,13 +699,23 @@ " \"\"\"\n", " final = {**Material.create(from_ase(atoms)).to_dict(), \"repetition\": rank_index}\n", " initial = CONFIGURATIONS[atoms.info[\"configuration\"]]\n", - " body = {\"jobId\": \"{{ JOB_ID }}\", \"unitId\": AFIR_UNIT_FLOWCHART_ID, \"structures\": [{\"initial\": initial, \"final\": final}]}\n", + " body = {\"jobId\": \"{{ JOB_ID }}\", \"unitId\": AFIR_UNIT_FLOWCHART_ID,\n", + " \"structures\": [{\"initial\": initial, \"final\": final}]}\n", " return json.dumps(json.dumps(body))[1:-1]\n", "\n", "\n", - "final_structures = read(\"{{ JOB_WORK_DIR }}/final_structures.extxyz\", index=f\":{NUMBER_OF_FINAL_STRUCTURES}\")\n", - "requests = [request_body(rank_index, atoms) for rank_index, atoms in enumerate(final_structures)]\n", - "print(json.dumps({\"name\": \"FINAL_STRUCTURE_REQUESTS\", \"value\": requests, \"scope\": \"global\"}))\n", + "def main():\n", + " \"\"\"\n", + " Reads the lowest NUMBER_OF_FINAL_STRUCTURES final structures and prints their request bodies as\n", + " FINAL_STRUCTURE_REQUESTS.\n", + " \"\"\"\n", + " final_structures = read(\"{{ JOB_WORK_DIR }}/final_structures.extxyz\", index=f\":{NUMBER_OF_FINAL_STRUCTURES}\")\n", + " request_bodies = [request_body(rank_index, atoms) for rank_index, atoms in enumerate(final_structures)]\n", + " print(json.dumps({\"name\": \"FINAL_STRUCTURE_REQUESTS\", \"value\": request_bodies, \"scope\": \"global\"}))\n", + "\n", + "\n", + "if __name__ == \"__main__\":\n", + " main()\n", "'''" ] }, @@ -726,39 +742,17 @@ "from mat3ra.notebooks_utils.core.entity.workflow.api import set_execution_unit_input\n", "from mat3ra.standata.workflows import WorkflowStandata\n", "\n", - "MADE_REQUIREMENTS = \"mat3ra-made[tools]==2026.10.10.post0\\n\"\n", - "\n", "\n", "def set_constants(script, constants):\n", + " \"\"\"Sets each named constant of a script to its value.\"\"\"\n", " for name, value in constants.items():\n", " script = re.sub(rf\"^({name} = )\\S+\", rf\"\\g<1>{value}\", script, flags=re.M)\n", " return script\n", "\n", "\n", - "python_workflows = WorkflowStandata.filter_by_application(\"python\")\n", - "workflow_config = copy.deepcopy(python_workflows.get_by_name_first_match(\"custom_script.json\"))\n", - "workflow_config.update(name=\"AFIR, molecule on surface\", isMultiMaterial=True)\n", - "workflow_config[\"units\"][0][\"name\"] = \"Configurations\"\n", - "python_script_subworkflow = python_workflows.get_by_name_first_match(\"python_script.json\")[\"subworkflows\"][0]\n", - "afir_subworkflow_config = copy.deepcopy(python_script_subworkflow)\n", - "final_structures_subworkflow_config = copy.deepcopy(python_script_subworkflow)\n", - "final_structures_subworkflow_config.pop(\"_id\")\n", - "next(unit for unit in final_structures_subworkflow_config[\"units\"] if unit[\"name\"] == \"python\").pop(\"flowchartId\")\n", - "afir_unit_flowchart_id = next(unit for unit in afir_subworkflow_config[\"units\"] if unit[\"name\"] == \"python\")[\"flowchartId\"]\n", - "afir_results = [\n", - " {\"name\": \"file_content\", \"filetype\": \"csv\", \"basename\": \"results.csv\"},\n", - " {\"name\": \"file_content\", \"filetype\": \"text\", \"basename\": \"final_structures.extxyz\"},\n", - "]\n", - "steps = [\n", - " (workflow_config[\"subworkflows\"][0], \"custom_script\", \"Configurations\", \"configurations\", MADE_REQUIREMENTS, [], False,\n", - " set_constants(CONFIGURATIONS_SCRIPT, {\"NUMBER_OF_CONFIGURATIONS\": NUMBER_OF_CONFIGURATIONS, \"RANDOM_SEED\": RANDOM_SEED})),\n", - " (afir_subworkflow_config, \"python\", \"AFIR (ALCHEMI)\", \"afir\", AFIR_REQUIREMENTS, afir_results, True,\n", - " set_constants(AFIR_SCRIPT, {\"GAMMA\": GAMMA})),\n", - " (final_structures_subworkflow_config, \"python\", \"Final structures\", \"final_structures\", MADE_REQUIREMENTS, [], True,\n", - " set_constants(FINAL_STRUCTURES_SCRIPT, {\"NUMBER_OF_FINAL_STRUCTURES\": NUMBER_OF_FINAL_STRUCTURES,\n", - " \"AFIR_UNIT_FLOWCHART_ID\": f'\"{afir_unit_flowchart_id}\"'})),\n", - "]\n", - "for subworkflow, template_unit_name, name, unit_name, requirements, results, is_rendered_at_run_time, script in steps:\n", + "def configure_python_step(subworkflow, template_unit_name, name, unit_name, script, requirements, results,\n", + " is_rendered_at_run_time):\n", + " \"\"\"Makes a Standata python subworkflow one step: its unit, Python version, files named after it, and results.\"\"\"\n", " unit = next(unit for unit in subworkflow[\"units\"] if unit[\"name\"] == template_unit_name)\n", " unit.update(name=unit_name, head=True, results=results)\n", " for application in (subworkflow[\"application\"], unit[\"application\"]):\n", @@ -769,7 +763,42 @@ " inputs[\"script.py\"][\"isManuallyChanged\"] = is_rendered_at_run_time\n", " inputs[\"script.py\"][\"template\"][\"name\"] = f\"{unit_name}.py\"\n", " inputs[\"requirements.txt\"][\"template\"][\"name\"] = f\"{unit_name}_requirements.txt\"\n", - " subworkflow.update(name=name, units=[unit])" + " subworkflow.update(name=name, units=[unit])\n", + "\n", + "\n", + "python_workflows = WorkflowStandata.filter_by_application(\"python\")\n", + "workflow_config = copy.deepcopy(python_workflows.get_by_name_first_match(\"custom_script.json\"))\n", + "workflow_config.update(name=\"AFIR, molecule on surface\", isMultiMaterial=True)\n", + "workflow_config[\"units\"][0][\"name\"] = \"Configurations\"\n", + "python_script_subworkflow = python_workflows.get_by_name_first_match(\"python_script.json\")[\"subworkflows\"][0]\n", + "afir_subworkflow_config = copy.deepcopy(python_script_subworkflow)\n", + "final_structures_subworkflow_config = copy.deepcopy(python_script_subworkflow)\n", + "final_structures_subworkflow_config.pop(\"_id\")\n", + "next(unit for unit in final_structures_subworkflow_config[\"units\"] if unit[\"name\"] == \"python\").pop(\"flowchartId\")\n", + "afir_unit_flowchart_id = next(unit for unit in afir_subworkflow_config[\"units\"] if unit[\"name\"] == \"python\")[\"flowchartId\"]\n", + "number_of_saved_structures = min(NUMBER_OF_FINAL_STRUCTURES, NUMBER_OF_CONFIGURATIONS)\n", + "\n", + "configure_python_step(\n", + " workflow_config[\"subworkflows\"][0], template_unit_name=\"custom_script\", name=\"Configurations\",\n", + " unit_name=\"configurations\", requirements=MADE_REQUIREMENTS, results=[], is_rendered_at_run_time=False,\n", + " script=set_constants(CONFIGURATIONS_SCRIPT, {\"NUMBER_OF_CONFIGURATIONS\": NUMBER_OF_CONFIGURATIONS,\n", + " \"RANDOM_SEED\": RANDOM_SEED}),\n", + ")\n", + "configure_python_step(\n", + " afir_subworkflow_config, template_unit_name=\"python\", name=\"AFIR (ALCHEMI)\", unit_name=\"afir\",\n", + " requirements=AFIR_REQUIREMENTS, is_rendered_at_run_time=True, script=set_constants(AFIR_SCRIPT, {\"GAMMA\": GAMMA}),\n", + " results=[\n", + " {\"name\": \"file_content\", \"filetype\": \"csv\", \"basename\": \"results.csv\"},\n", + " {\"name\": \"final_structure\"},\n", + " {\"name\": \"file_content\", \"filetype\": \"text\", \"basename\": \"final_structures.extxyz\"},\n", + " ],\n", + ")\n", + "configure_python_step(\n", + " final_structures_subworkflow_config, template_unit_name=\"python\", name=\"Final structures\",\n", + " unit_name=\"final_structures\", requirements=MADE_REQUIREMENTS, results=[], is_rendered_at_run_time=True,\n", + " script=set_constants(FINAL_STRUCTURES_SCRIPT, {\"NUMBER_OF_FINAL_STRUCTURES\": number_of_saved_structures,\n", + " \"AFIR_UNIT_FLOWCHART_ID\": f'\"{afir_unit_flowchart_id}\"'}),\n", + ")" ] }, { @@ -795,12 +824,12 @@ "workflow = Workflow.create(workflow_config)\n", "workflow.add_subworkflow(Subworkflow(**afir_subworkflow_config))\n", "final_structures_subworkflow = Subworkflow(**final_structures_subworkflow_config)\n", - "for index in range(NUMBER_OF_FINAL_STRUCTURES):\n", - " request = {\"method\": \"POST\", \"headers\": {\"Content-Type\": \"application/json\"},\n", - " \"data\": f\"{{{{ FINAL_STRUCTURE_REQUESTS[{index}] }}}}\"}\n", + "for index in range(number_of_saved_structures):\n", + " endpoint_options = {\"method\": \"POST\", \"headers\": {\"Content-Type\": \"application/json\"},\n", + " \"data\": \"{{ FINAL_STRUCTURE_REQUESTS[\" + str(index) + \"] }}\"}\n", " final_structures_subworkflow.add_unit(Unit(\n", " type=\"io\", subtype=\"output\", source=\"api\", enableRender=True, name=f\"save-final-structure-{index + 1}\",\n", - " input=[{\"type\": \"api\", \"endpoint\": \"structures\", \"endpoint_options\": request, \"name\": \"FINAL_STRUCTURE\"}],\n", + " input=[{\"type\": \"api\", \"endpoint\": \"structures\", \"endpoint_options\": endpoint_options, \"name\": \"FINAL_STRUCTURE\"}],\n", " ))\n", "final_structures_subworkflow.add_unit(\n", " Unit(type=\"assignment\", name=\"final-structures-saved\", operand=\"FINAL_STRUCTURES_SAVED\", value=\"True\", input=[])\n", @@ -945,8 +974,13 @@ "source": [ "from mat3ra.prode import PropertyName\n", "\n", - "final_structures = client.properties.get_for_job(job_id, property_name=PropertyName.non_scalar.final_structure.value)\n", - "visualize([Material.create(client.materials.get(final_structure[\"materialId\"])) for final_structure in final_structures])" + "final_structure_properties = client.properties.get_for_job(\n", + " job_id, property_name=PropertyName.non_scalar.final_structure.value\n", + ")\n", + "visualize([\n", + " {\"material\": Material.create(client.materials.get(final_structure_property[\"materialId\"])), \"title\": f\"Rank {rank}\"}\n", + " for rank, final_structure_property in enumerate(final_structure_properties, start=1)\n", + "])" ] }, { From 2c50c8c473dfb6da672957a5f5150cd48e86a2d9 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Mon, 2 Nov 2026 19:31:30 -0800 Subject: [PATCH 07/31] SOF-8102: afir_molecule_surface_alchemi.ipynb: review round 3 nits CHECKPOINT names MACE-MPA-0 and AFIR_R0 its unit; configure_python_step takes subworkflow_name, has a full docstring, and its three calls pass keywords in the signature's order; two lines wrapped under 120 characters. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../afir_molecule_surface_alchemi.ipynb | 46 +++++++++++++------ 1 file changed, 33 insertions(+), 13 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 41d3c3e0f..52a1f6046 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -433,13 +433,13 @@ "from nvalchemi.models.mace import MACEWrapper\n", "\n", "DEVICE = \"cuda:0\"\n", - "CHECKPOINT = \"medium-mpa-0\" # MACE-MP-0 model, medium MPA-0\n", + "CHECKPOINT = \"medium-mpa-0\" # MACE-MPA-0, medium\n", "DTYPE = torch.float32\n", "FMAX = 0.03 # eV/Å, largest atomic force at convergence\n", "MAX_STEPS = 500 # FIRE2 steps per relaxation at most\n", "TIME_STEP = 0.5 # initial FIRE2 step, in ALCHEMI's internal time unit (10.18 fs)\n", "GAMMA = 100 # kJ/mol, model collision energy of the AFIR force\n", - "AFIR_R0 = 3.8164 # A, GRRM17 Eq. (3)\n", + "AFIR_R0 = 3.8164 # Å, GRRM17 Eq. (3)\n", "AFIR_EPSILON = 1.0061 # kJ/mol, GRRM17 Eq. (3)\n", "CONFIGURATIONS = json.loads(r\"\"\"{{ CONFIGURATIONS | tojson }}\"\"\")\n", "\n", @@ -750,9 +750,23 @@ " return script\n", "\n", "\n", - "def configure_python_step(subworkflow, template_unit_name, name, unit_name, script, requirements, results,\n", + "def configure_python_step(subworkflow, template_unit_name, subworkflow_name, unit_name, script, requirements, results,\n", " is_rendered_at_run_time):\n", - " \"\"\"Makes a Standata python subworkflow one step: its unit, Python version, files named after it, and results.\"\"\"\n", + " \"\"\"\n", + " Makes a Standata python subworkflow one step of the workflow.\n", + "\n", + " Args:\n", + " subworkflow (dict): The subworkflow config, edited in place.\n", + " template_unit_name (str): The name of the template's python unit.\n", + " subworkflow_name (str): The step's name in the workflow.\n", + " unit_name (str): The python unit's name; its script and requirements files are named after it.\n", + " script (str): The unit's script.\n", + " requirements (str): The packages installed for the unit, as a requirements.txt.\n", + " results (list): The results the unit declares.\n", + " is_rendered_at_run_time (bool): True to leave the script's placeholders to the compute node, which fills them\n", + " from the job's scope when the step runs; False to let the platform fill them, with the job's materials,\n", + " when the job is created.\n", + " \"\"\"\n", " unit = next(unit for unit in subworkflow[\"units\"] if unit[\"name\"] == template_unit_name)\n", " unit.update(name=unit_name, head=True, results=results)\n", " for application in (subworkflow[\"application\"], unit[\"application\"]):\n", @@ -763,7 +777,7 @@ " inputs[\"script.py\"][\"isManuallyChanged\"] = is_rendered_at_run_time\n", " inputs[\"script.py\"][\"template\"][\"name\"] = f\"{unit_name}.py\"\n", " inputs[\"requirements.txt\"][\"template\"][\"name\"] = f\"{unit_name}_requirements.txt\"\n", - " subworkflow.update(name=name, units=[unit])\n", + " subworkflow.update(name=subworkflow_name, units=[unit])\n", "\n", "\n", "python_workflows = WorkflowStandata.filter_by_application(\"python\")\n", @@ -775,29 +789,34 @@ "final_structures_subworkflow_config = copy.deepcopy(python_script_subworkflow)\n", "final_structures_subworkflow_config.pop(\"_id\")\n", "next(unit for unit in final_structures_subworkflow_config[\"units\"] if unit[\"name\"] == \"python\").pop(\"flowchartId\")\n", - "afir_unit_flowchart_id = next(unit for unit in afir_subworkflow_config[\"units\"] if unit[\"name\"] == \"python\")[\"flowchartId\"]\n", + "afir_unit = next(unit for unit in afir_subworkflow_config[\"units\"] if unit[\"name\"] == \"python\")\n", + "afir_unit_flowchart_id = afir_unit[\"flowchartId\"]\n", "number_of_saved_structures = min(NUMBER_OF_FINAL_STRUCTURES, NUMBER_OF_CONFIGURATIONS)\n", "\n", "configure_python_step(\n", - " workflow_config[\"subworkflows\"][0], template_unit_name=\"custom_script\", name=\"Configurations\",\n", - " unit_name=\"configurations\", requirements=MADE_REQUIREMENTS, results=[], is_rendered_at_run_time=False,\n", + " workflow_config[\"subworkflows\"][0], template_unit_name=\"custom_script\", subworkflow_name=\"Configurations\",\n", + " unit_name=\"configurations\",\n", " script=set_constants(CONFIGURATIONS_SCRIPT, {\"NUMBER_OF_CONFIGURATIONS\": NUMBER_OF_CONFIGURATIONS,\n", " \"RANDOM_SEED\": RANDOM_SEED}),\n", + " requirements=MADE_REQUIREMENTS, results=[], is_rendered_at_run_time=False,\n", ")\n", "configure_python_step(\n", - " afir_subworkflow_config, template_unit_name=\"python\", name=\"AFIR (ALCHEMI)\", unit_name=\"afir\",\n", - " requirements=AFIR_REQUIREMENTS, is_rendered_at_run_time=True, script=set_constants(AFIR_SCRIPT, {\"GAMMA\": GAMMA}),\n", + " afir_subworkflow_config, template_unit_name=\"python\", subworkflow_name=\"AFIR (ALCHEMI)\", unit_name=\"afir\",\n", + " script=set_constants(AFIR_SCRIPT, {\"GAMMA\": GAMMA}),\n", + " requirements=AFIR_REQUIREMENTS,\n", " results=[\n", " {\"name\": \"file_content\", \"filetype\": \"csv\", \"basename\": \"results.csv\"},\n", " {\"name\": \"final_structure\"},\n", " {\"name\": \"file_content\", \"filetype\": \"text\", \"basename\": \"final_structures.extxyz\"},\n", " ],\n", + " is_rendered_at_run_time=True,\n", ")\n", "configure_python_step(\n", - " final_structures_subworkflow_config, template_unit_name=\"python\", name=\"Final structures\",\n", - " unit_name=\"final_structures\", requirements=MADE_REQUIREMENTS, results=[], is_rendered_at_run_time=True,\n", + " final_structures_subworkflow_config, template_unit_name=\"python\", subworkflow_name=\"Final structures\",\n", + " unit_name=\"final_structures\",\n", " script=set_constants(FINAL_STRUCTURES_SCRIPT, {\"NUMBER_OF_FINAL_STRUCTURES\": number_of_saved_structures,\n", " \"AFIR_UNIT_FLOWCHART_ID\": f'\"{afir_unit_flowchart_id}\"'}),\n", + " requirements=MADE_REQUIREMENTS, results=[], is_rendered_at_run_time=True,\n", ")" ] }, @@ -829,7 +848,8 @@ " \"data\": \"{{ FINAL_STRUCTURE_REQUESTS[\" + str(index) + \"] }}\"}\n", " final_structures_subworkflow.add_unit(Unit(\n", " type=\"io\", subtype=\"output\", source=\"api\", enableRender=True, name=f\"save-final-structure-{index + 1}\",\n", - " input=[{\"type\": \"api\", \"endpoint\": \"structures\", \"endpoint_options\": endpoint_options, \"name\": \"FINAL_STRUCTURE\"}],\n", + " input=[{\"type\": \"api\", \"endpoint\": \"structures\", \"endpoint_options\": endpoint_options,\n", + " \"name\": \"FINAL_STRUCTURE\"}],\n", " ))\n", "final_structures_subworkflow.add_unit(\n", " Unit(type=\"assignment\", name=\"final-structures-saved\", operand=\"FINAL_STRUCTURES_SAVED\", value=\"True\", input=[])\n", From af6d0027c8fff178fc8104350db98b7bd3aaf27c Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Mon, 2 Nov 2026 23:44:38 -0800 Subject: [PATCH 08/31] SOF-8102: afir_molecule_surface_alchemi.ipynb: save the final structures in one loop; keep the lowest; energy plot; orientation and fragments per configuration Step 3 saves its structures with one io unit in a loop (index, increment, condition), as Standata's k-point convergence loop does. The AFIR step keeps only the NUMBER_OF_FINAL_STRUCTURES lowest in final_structures.extxyz, plots every configuration's final energy to energies.png, and writes each configuration's starting rotation and position (recorded by step 1 in its metadata) and its molecule's fragment count to results.csv. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../afir_molecule_surface_alchemi.ipynb | 177 +++++++++++++----- 1 file changed, 132 insertions(+), 45 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 52a1f6046..128c239f2 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -71,7 +71,7 @@ "NUMBER_OF_CONFIGURATIONS = 100 # random placements of the molecule over the slab\n", "GAMMA = 100 # kJ/mol, model collision energy of the AFIR force\n", "RANDOM_SEED = 0 # of the random placements\n", - "NUMBER_OF_FINAL_STRUCTURES = 5 # the lowest-energy final structures saved as materials\n", + "NUMBER_OF_FINAL_STRUCTURES = 5 # the lowest-energy final structures kept and saved as materials\n", "PYTHON_VERSION = \"3.11.10\"\n", "PYTHON_BUILD = \"Default\"\n", "\n", @@ -249,7 +249,8 @@ "\n", "Places the molecule `NUMBER_OF_CONFIGURATIONS` times over the slab, each at a random orientation and a random point\n", "of the surface cell, with its lowest atom above the top slab atom by the sum of their covalent radii plus 0.5 Å, as\n", - "GRRM's random initial structures [1]. Slab atoms are labelled 0, molecule atoms 1." + "GRRM's random initial structures [1]. Slab atoms are labelled 0, molecule atoms 1, and each configuration\n", + "keeps its rotation and position in its metadata." ] }, { @@ -272,7 +273,9 @@ "# one is the molecule. #\n", "# Writes: CONFIGURATIONS to the job's global scope: the placed #\n", "# structures as materials, slab atoms labelled 0 and #\n", - "# molecule atoms 1. It is the only output of the script. #\n", + "# molecule atoms 1, each with its rotation and fractional #\n", + "# position in metadata.placement. It is the only output #\n", + "# of the script. #\n", "# #\n", "# [1] S. Maeda et al., J. Comput. Chem. 39, 233 (2018). #\n", "# #\n", @@ -329,11 +332,13 @@ " index (int): The configuration's index, used in its name.\n", "\n", " Returns:\n", - " MaterialWithBuildMetadata: The slab and the placed molecule, molecule atoms labelled 1.\n", + " MaterialWithBuildMetadata: The slab and the placed molecule, molecule atoms labelled 1, with the rotation and\n", + " the fractional position in metadata.placement.\n", " \"\"\"\n", " quaternion = random_generator.normal(size=4)\n", - " angle = 2 * np.degrees(np.arccos(quaternion[0] / np.linalg.norm(quaternion)))\n", - " rotated = rotate(molecule, axis=quaternion[1:].tolist(), angle=angle, wrap=False)\n", + " axis = (quaternion[1:] / np.linalg.norm(quaternion[1:])).tolist()\n", + " angle = float(2 * np.degrees(np.arccos(quaternion[0] / np.linalg.norm(quaternion))))\n", + " rotated = rotate(molecule, axis=axis, angle=angle, wrap=False)\n", " rotated.set_lattice(slab.lattice)\n", " top_height, top_radius = top_atom_height_and_radius\n", " lowest_height, lowest_radius = extreme_atom_height_and_radius(rotated, min)\n", @@ -345,7 +350,10 @@ " surface_point[1] - molecule_cell_center[1], vertical_shift],\n", " use_cartesian_coordinates=True)\n", " placed.set_labels_from_value(InterfacePartsEnum.FILM.value)\n", - " return merge([slab, placed], material_name=f\"{molecule.name} on {slab.name}, configuration {index}\")\n", + " configuration = merge([slab, placed], material_name=f\"{molecule.name} on {slab.name}, configuration {index}\")\n", + " configuration.metadata.placement = {\"rotation_axis\": axis, \"rotation_angle_deg\": angle,\n", + " \"position_u\": surface_fractions[0], \"position_v\": surface_fractions[1]}\n", + " return configuration\n", "\n", "\n", "def main():\n", @@ -380,9 +388,10 @@ "### 4.2. Step 2: AFIR with ALCHEMI\n", "\n", "Relaxes all configurations together with the slab fixed [3], first under the AFIR force [1] of strength `GAMMA`\n", - "(100 kJ/mol with random orientations [2]) between molecule and slab, then without it, and writes `results.csv` and\n", - "`final_structures.extxyz`. Another backend or another ALCHEMI version replaces this script and its packages; steps 1\n", - "and 3 stay." + "(100 kJ/mol with random orientations [2]) between molecule and slab, then without it. It writes `results.csv` with\n", + "every configuration ranked by final energy, its starting orientation and the fragments its molecule ends in;\n", + "`final_structures.extxyz` with the `NUMBER_OF_FINAL_STRUCTURES` lowest; and `energies.png`, the final energy of each\n", + "configuration. Another backend or another ALCHEMI version replaces this script and its packages; steps 1 and 3 stay." ] }, { @@ -403,9 +412,13 @@ "# #\n", "# Reads: CONFIGURATIONS from the job's global scope, written #\n", "# into this script by the compute node before it runs. #\n", - "# Writes: results.csv, the configurations ranked by final energy, #\n", - "# and final_structures.extxyz, their structures in that #\n", - "# order; both are results of this unit. #\n", + "# Writes: results.csv, every configuration ranked by final #\n", + "# energy, with its starting rotation and position and the #\n", + "# number of fragments its molecule ends in; #\n", + "# final_structures.extxyz, the NUMBER_OF_FINAL_STRUCTURES #\n", + "# lowest in that order; and energies.png, final energy #\n", + "# against configuration index. All three are results of #\n", + "# this unit. #\n", "# #\n", "# [1] S. Maeda et al., J. Comput. Chem. 39, 233 (2018). #\n", "# #\n", @@ -414,8 +427,10 @@ "import csv\n", "import json\n", "\n", + "import matplotlib\n", "import numpy as np\n", "import torch\n", + "from ase.build import separate\n", "from ase.calculators.singlepoint import SinglePointCalculator\n", "from ase.data import covalent_radii\n", "from ase.io import write\n", @@ -432,6 +447,9 @@ "from nvalchemi.hooks import BiasedPotentialHook\n", "from nvalchemi.models.mace import MACEWrapper\n", "\n", + "matplotlib.use(\"Agg\")\n", + "import matplotlib.pyplot as plt\n", + "\n", "DEVICE = \"cuda:0\"\n", "CHECKPOINT = \"medium-mpa-0\" # MACE-MPA-0, medium\n", "DTYPE = torch.float32\n", @@ -439,6 +457,8 @@ "MAX_STEPS = 500 # FIRE2 steps per relaxation at most\n", "TIME_STEP = 0.5 # initial FIRE2 step, in ALCHEMI's internal time unit (10.18 fs)\n", "GAMMA = 100 # kJ/mol, model collision energy of the AFIR force\n", + "NUMBER_OF_FINAL_STRUCTURES = 5 # the lowest final structures kept in final_structures.extxyz\n", + "BOND_TOLERANCE = 1.2 # two molecule atoms are bonded within BOND_TOLERANCE * (R_i + R_j)\n", "AFIR_R0 = 3.8164 # Å, GRRM17 Eq. (3)\n", "AFIR_EPSILON = 1.0061 # kJ/mol, GRRM17 Eq. (3)\n", "CONFIGURATIONS = json.loads(r\"\"\"{{ CONFIGURATIONS | tojson }}\"\"\")\n", @@ -577,6 +597,19 @@ " return [relaxed_atoms(atoms, relaxed_data) for atoms, relaxed_data in zip(ase_structures, batch.to_data_list())]\n", "\n", "\n", + "def molecule_fragments(atoms):\n", + " \"\"\"\n", + " Returns the number of fragments the molecule of a structure is in: 1 if intact, more if it split.\n", + "\n", + " Args:\n", + " atoms (ase.Atoms): The structure, molecule atoms tagged 1.\n", + "\n", + " Returns:\n", + " int: The number of connected groups of molecule atoms, bonded within BOND_TOLERANCE times their covalent radii.\n", + " \"\"\"\n", + " return len(separate(atoms[atoms.get_tags() == InterfacePartsEnum.FILM.value], scale=BOND_TOLERANCE))\n", + "\n", + "\n", "def energy_and_max_force(atoms):\n", " \"\"\"\n", " Returns the energy of a structure, in eV, and the largest force on any of its atoms, in eV/Å.\n", @@ -584,9 +617,30 @@ " return atoms.get_potential_energy(), np.linalg.norm(atoms.get_forces(), axis=1).max()\n", "\n", "\n", + "def plot_energies(final_structures, lowest_indices):\n", + " \"\"\"\n", + " Plots the final energy of every configuration against its index, the kept lowest marked, to energies.png.\n", + "\n", + " Args:\n", + " final_structures (List[ase.Atoms]): The structures after the relaxation without the AFIR force.\n", + " lowest_indices (List[int]): The indices of the kept lowest structures.\n", + " \"\"\"\n", + " energies = [atoms.get_potential_energy() for atoms in final_structures]\n", + " figure, axes = plt.subplots()\n", + " axes.scatter(range(len(energies)), energies, s=12, label=\"all configurations\")\n", + " axes.scatter(lowest_indices, [energies[index] for index in lowest_indices], s=40, color=\"tab:red\",\n", + " label=f\"{len(lowest_indices)} lowest, kept\")\n", + " axes.set(xlabel=\"Configuration index\", ylabel=\"Final energy (eV)\")\n", + " axes.ticklabel_format(axis=\"y\", useOffset=False)\n", + " axes.legend()\n", + " figure.savefig(\"energies.png\", dpi=150, bbox_inches=\"tight\")\n", + "\n", + "\n", "def write_results(afir_structures, final_structures):\n", " \"\"\"\n", - " Ranks the configurations by final energy and writes results.csv and final_structures.extxyz in that order.\n", + " Ranks the configurations by final energy, writes all of them to results.csv with the placement each started from\n", + " (atoms.info[\"metadata\"][\"placement\"]) and its molecule's fragments, the NUMBER_OF_FINAL_STRUCTURES lowest to\n", + " final_structures.extxyz in that order, and the energies to energies.png.\n", "\n", " Args:\n", " afir_structures (List[ase.Atoms]): The structures after the relaxation under the AFIR force.\n", @@ -596,12 +650,18 @@ " with open(\"results.csv\", \"w\", newline=\"\") as results_file:\n", " writer = csv.writer(results_file)\n", " writer.writerow([\"configuration\", \"rank\", \"afir_function_eV\", \"afir_max_force_eV_per_A\", \"energy_eV\",\n", - " \"max_force_eV_per_A\"])\n", + " \"max_force_eV_per_A\", \"rotation_axis_x\", \"rotation_axis_y\", \"rotation_axis_z\",\n", + " \"rotation_angle_deg\", \"position_u\", \"position_v\", \"molecule_fragments\"])\n", " for rank, index in enumerate(order, start=1):\n", + " placement = final_structures[index].info[\"metadata\"][\"placement\"]\n", " writer.writerow([index, rank, *energy_and_max_force(afir_structures[index]),\n", - " *energy_and_max_force(final_structures[index])])\n", + " *energy_and_max_force(final_structures[index]), *placement[\"rotation_axis\"],\n", + " placement[\"rotation_angle_deg\"], placement[\"position_u\"], placement[\"position_v\"],\n", + " molecule_fragments(final_structures[index])])\n", " final_structures[index].info.update(configuration=int(index), rank=rank)\n", - " write(\"final_structures.extxyz\", [final_structures[index] for index in order])\n", + " lowest_indices = order[:NUMBER_OF_FINAL_STRUCTURES].tolist()\n", + " write(\"final_structures.extxyz\", [final_structures[index] for index in lowest_indices])\n", + " plot_energies(final_structures, lowest_indices)\n", "\n", "\n", "def main():\n", @@ -637,6 +697,7 @@ "torchaudio==2.11.0+cu128\n", "warp-lang==1.18.0+cu12\n", "nvalchemi-toolkit[mace,ase]==0.2.0\n", + "matplotlib\n", "\"\"\" + MADE_REQUIREMENTS" ] }, @@ -646,8 +707,8 @@ "source": [ "### 4.3. Step 3: final structures\n", "\n", - "Takes the `NUMBER_OF_FINAL_STRUCTURES` lowest final structures, each with the configuration it started from, and\n", - "saves them as the job's final structures: each is a material, shown in the Results tab of the AFIR step." + "Takes the lowest final structures, each with the configuration it started from, and saves them as the job's final\n", + "structures: each is a material, shown in the Results tab of the AFIR step." ] }, { @@ -660,10 +721,10 @@ "# #\n", "# Final structures step of the AFIR molecule-on-surface workflow. #\n", "# #\n", - "# Takes the NUMBER_OF_FINAL_STRUCTURES lowest final structures, #\n", - "# each with the configuration it started from, and makes one #\n", - "# request body per structure for the platform's final-structure #\n", - "# endpoint, which the io units after this one send. #\n", + "# Takes the lowest final structures of the AFIR step, each with #\n", + "# the configuration it started from, and makes one request body #\n", + "# per structure for the platform's final-structure endpoint, #\n", + "# which the io unit after this one sends in a loop. #\n", "# #\n", "# Reads: final_structures.extxyz in the job's work folder, and #\n", "# CONFIGURATIONS from the job's global scope, written #\n", @@ -679,7 +740,6 @@ "from mat3ra.made.material import Material\n", "from mat3ra.made.tools.convert import from_ase\n", "\n", - "NUMBER_OF_FINAL_STRUCTURES = 5\n", "AFIR_UNIT_FLOWCHART_ID = \"\"\n", "CONFIGURATIONS = json.loads(r\"\"\"{{ CONFIGURATIONS | tojson }}\"\"\")\n", "\n", @@ -706,10 +766,9 @@ "\n", "def main():\n", " \"\"\"\n", - " Reads the lowest NUMBER_OF_FINAL_STRUCTURES final structures and prints their request bodies as\n", - " FINAL_STRUCTURE_REQUESTS.\n", + " Reads the lowest final structures and prints their request bodies as FINAL_STRUCTURE_REQUESTS.\n", " \"\"\"\n", - " final_structures = read(\"{{ JOB_WORK_DIR }}/final_structures.extxyz\", index=f\":{NUMBER_OF_FINAL_STRUCTURES}\")\n", + " final_structures = read(\"{{ JOB_WORK_DIR }}/final_structures.extxyz\", index=\":\")\n", " request_bodies = [request_body(rank_index, atoms) for rank_index, atoms in enumerate(final_structures)]\n", " print(json.dumps({\"name\": \"FINAL_STRUCTURE_REQUESTS\", \"value\": request_bodies, \"scope\": \"global\"}))\n", "\n", @@ -802,20 +861,20 @@ ")\n", "configure_python_step(\n", " afir_subworkflow_config, template_unit_name=\"python\", subworkflow_name=\"AFIR (ALCHEMI)\", unit_name=\"afir\",\n", - " script=set_constants(AFIR_SCRIPT, {\"GAMMA\": GAMMA}),\n", + " script=set_constants(AFIR_SCRIPT, {\"GAMMA\": GAMMA, \"NUMBER_OF_FINAL_STRUCTURES\": number_of_saved_structures}),\n", " requirements=AFIR_REQUIREMENTS,\n", " results=[\n", " {\"name\": \"file_content\", \"filetype\": \"csv\", \"basename\": \"results.csv\"},\n", " {\"name\": \"final_structure\"},\n", " {\"name\": \"file_content\", \"filetype\": \"text\", \"basename\": \"final_structures.extxyz\"},\n", + " {\"name\": \"file_content\", \"filetype\": \"image\", \"basename\": \"energies.png\"},\n", " ],\n", " is_rendered_at_run_time=True,\n", ")\n", "configure_python_step(\n", " final_structures_subworkflow_config, template_unit_name=\"python\", subworkflow_name=\"Final structures\",\n", " unit_name=\"final_structures\",\n", - " script=set_constants(FINAL_STRUCTURES_SCRIPT, {\"NUMBER_OF_FINAL_STRUCTURES\": number_of_saved_structures,\n", - " \"AFIR_UNIT_FLOWCHART_ID\": f'\"{afir_unit_flowchart_id}\"'}),\n", + " script=set_constants(FINAL_STRUCTURES_SCRIPT, {\"AFIR_UNIT_FLOWCHART_ID\": f'\"{afir_unit_flowchart_id}\"'}),\n", " requirements=MADE_REQUIREMENTS, results=[], is_rendered_at_run_time=True,\n", ")" ] @@ -826,7 +885,7 @@ "source": [ "### 4.5. Create and save the workflow\n", "\n", - "Step 3 ends with one request per final structure, each saving it with the starting configuration as its initial\n", + "Step 3 ends with a loop that saves one final structure per pass, with its starting configuration as the initial\n", "structure." ] }, @@ -840,20 +899,27 @@ "from mat3ra.wode import Subworkflow, Workflow\n", "from mat3ra.wode.units import Unit\n", "\n", + "endpoint_options = {\"method\": \"POST\", \"headers\": {\"Content-Type\": \"application/json\"},\n", + " \"data\": \"{{ FINAL_STRUCTURE_REQUESTS[INDEX] }}\"}\n", + "save_unit = Unit(type=\"io\", subtype=\"output\", source=\"api\", enableRender=True, name=\"save-final-structure\",\n", + " input=[{\"type\": \"api\", \"endpoint\": \"structures\", \"endpoint_options\": endpoint_options,\n", + " \"name\": \"FINAL_STRUCTURE\"}])\n", + "exit_unit = Unit(type=\"assignment\", name=\"final-structures-saved\", operand=\"NUMBER_OF_SAVED_STRUCTURES\", value=\"INDEX\",\n", + " input=[])\n", + "loop_units = [\n", + " Unit(type=\"assignment\", name=\"init-index\", operand=\"INDEX\", value=0, input=[]),\n", + " save_unit,\n", + " Unit(type=\"assignment\", name=\"increment-index\", operand=\"INDEX\", value=\"INDEX + 1\", input=[]),\n", + " Unit(**{\"type\": \"condition\", \"name\": \"check-index\", \"statement\": f\"INDEX >= {number_of_saved_structures}\",\n", + " \"then\": exit_unit.flowchartId, \"else\": save_unit.flowchartId,\n", + " \"maxOccurrences\": number_of_saved_structures, \"input\": [{\"name\": \"INDEX\", \"scope\": \"global\"}]}),\n", + " exit_unit,\n", + "]\n", + "final_structures_subworkflow = Subworkflow(**final_structures_subworkflow_config)\n", + "for unit in loop_units:\n", + " final_structures_subworkflow.add_unit(unit)\n", "workflow = Workflow.create(workflow_config)\n", "workflow.add_subworkflow(Subworkflow(**afir_subworkflow_config))\n", - "final_structures_subworkflow = Subworkflow(**final_structures_subworkflow_config)\n", - "for index in range(number_of_saved_structures):\n", - " endpoint_options = {\"method\": \"POST\", \"headers\": {\"Content-Type\": \"application/json\"},\n", - " \"data\": \"{{ FINAL_STRUCTURE_REQUESTS[\" + str(index) + \"] }}\"}\n", - " final_structures_subworkflow.add_unit(Unit(\n", - " type=\"io\", subtype=\"output\", source=\"api\", enableRender=True, name=f\"save-final-structure-{index + 1}\",\n", - " input=[{\"type\": \"api\", \"endpoint\": \"structures\", \"endpoint_options\": endpoint_options,\n", - " \"name\": \"FINAL_STRUCTURE\"}],\n", - " ))\n", - "final_structures_subworkflow.add_unit(\n", - " Unit(type=\"assignment\", name=\"final-structures-saved\", operand=\"FINAL_STRUCTURES_SAVED\", value=\"True\", input=[])\n", - ")\n", "workflow.add_subworkflow(final_structures_subworkflow)\n", "visualize_workflow(workflow)\n", "\n", @@ -964,7 +1030,10 @@ "### 8.1. Configurations ranked by final energy\n", "\n", "`afir_function_eV` is the MACE energy plus the AFIR term at the end of the first relaxation, `energy_eV` the MACE\n", - "energy of the final structure; the `max_force` columns give the largest force on an atom at the end of each." + "energy of the final structure; the `max_force` columns give the largest force on an atom at the end of each. The\n", + "`rotation_*` columns give the rotation applied to the molecule about its center, `position_u` and `position_v` the\n", + "fractional point of the surface cell it was placed over, and `molecule_fragments` how many pieces the molecule\n", + "ends in: 1 if intact, more if it split." ] }, { @@ -983,7 +1052,25 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### 8.2. The lowest final structures" + "### 8.2. Final energy of each configuration" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from IPython.display import Image\n", + "\n", + "Image(data=await read_from_url(file_url_by_name[\"energies.png\"], as_bytes=True))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 8.3. The lowest final structures" ] }, { From e39e9ad97df361d32e516c09f1412e482a828a6b Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Tue, 3 Nov 2026 09:01:48 -0800 Subject: [PATCH 09/31] SOF-8102: afir_molecule_surface_alchemi.ipynb: pin matplotlib==3.11.2 in the ALCHEMI requirements Co-Authored-By: Claude Opus 5.5 (1M context) --- other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 128c239f2..70a760131 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -697,7 +697,7 @@ "torchaudio==2.11.0+cu128\n", "warp-lang==1.18.0+cu12\n", "nvalchemi-toolkit[mace,ase]==0.2.0\n", - "matplotlib\n", + "matplotlib==3.11.2\n", "\"\"\" + MADE_REQUIREMENTS" ] }, From 1735d52d9436405f2faa39f760953636b69059e7 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Tue, 3 Nov 2026 11:25:07 -0800 Subject: [PATCH 10/31] SOF-8102: afir_molecule_surface_alchemi.ipynb: drop the matplotlib line from the ALCHEMI requirements matplotlib arrives through ase, as in the environment rounds 1 and 2 ran; the explicit line only forced new environment builds, which cluster-003 cannot complete now. Co-Authored-By: Claude Opus 5.5 (1M context) --- other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb | 1 - 1 file changed, 1 deletion(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 70a760131..790fe58c1 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -697,7 +697,6 @@ "torchaudio==2.11.0+cu128\n", "warp-lang==1.18.0+cu12\n", "nvalchemi-toolkit[mace,ase]==0.2.0\n", - "matplotlib==3.11.2\n", "\"\"\" + MADE_REQUIREMENTS" ] }, From 66889d7290429a806c460f202a07f4ed08868b3e Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Tue, 3 Nov 2026 13:00:39 -0800 Subject: [PATCH 11/31] SOF-8102: afir_molecule_surface_alchemi.ipynb: save the lowest final structure only; review round 4 suggestions NUMBER_OF_FINAL_STRUCTURES defaults to 1, and final_structures.extxyz stays in the job's files but is no longer a unit result. The loop index is FINAL_STRUCTURE_INDEX, the placement position is one pair, and the final energies are computed once for the ranking and the plot. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../afir_molecule_surface_alchemi.ipynb | 39 ++++++++++--------- 1 file changed, 20 insertions(+), 19 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 790fe58c1..2ecbc2266 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -71,7 +71,7 @@ "NUMBER_OF_CONFIGURATIONS = 100 # random placements of the molecule over the slab\n", "GAMMA = 100 # kJ/mol, model collision energy of the AFIR force\n", "RANDOM_SEED = 0 # of the random placements\n", - "NUMBER_OF_FINAL_STRUCTURES = 5 # the lowest-energy final structures kept and saved as materials\n", + "NUMBER_OF_FINAL_STRUCTURES = 1 # the lowest-energy final structures kept and saved as materials\n", "PYTHON_VERSION = \"3.11.10\"\n", "PYTHON_BUILD = \"Default\"\n", "\n", @@ -352,7 +352,7 @@ " placed.set_labels_from_value(InterfacePartsEnum.FILM.value)\n", " configuration = merge([slab, placed], material_name=f\"{molecule.name} on {slab.name}, configuration {index}\")\n", " configuration.metadata.placement = {\"rotation_axis\": axis, \"rotation_angle_deg\": angle,\n", - " \"position_u\": surface_fractions[0], \"position_v\": surface_fractions[1]}\n", + " \"position\": surface_fractions[:2]}\n", " return configuration\n", "\n", "\n", @@ -417,8 +417,8 @@ "# number of fragments its molecule ends in; #\n", "# final_structures.extxyz, the NUMBER_OF_FINAL_STRUCTURES #\n", "# lowest in that order; and energies.png, final energy #\n", - "# against configuration index. All three are results of #\n", - "# this unit. #\n", + "# against configuration index. results.csv and #\n", + "# energies.png are results of this unit. #\n", "# #\n", "# [1] S. Maeda et al., J. Comput. Chem. 39, 233 (2018). #\n", "# #\n", @@ -617,15 +617,14 @@ " return atoms.get_potential_energy(), np.linalg.norm(atoms.get_forces(), axis=1).max()\n", "\n", "\n", - "def plot_energies(final_structures, lowest_indices):\n", + "def plot_energies(energies, lowest_indices):\n", " \"\"\"\n", " Plots the final energy of every configuration against its index, the kept lowest marked, to energies.png.\n", "\n", " Args:\n", - " final_structures (List[ase.Atoms]): The structures after the relaxation without the AFIR force.\n", + " energies (List[float]): The final energy of each configuration, in eV.\n", " lowest_indices (List[int]): The indices of the kept lowest structures.\n", " \"\"\"\n", - " energies = [atoms.get_potential_energy() for atoms in final_structures]\n", " figure, axes = plt.subplots()\n", " axes.scatter(range(len(energies)), energies, s=12, label=\"all configurations\")\n", " axes.scatter(lowest_indices, [energies[index] for index in lowest_indices], s=40, color=\"tab:red\",\n", @@ -646,7 +645,8 @@ " afir_structures (List[ase.Atoms]): The structures after the relaxation under the AFIR force.\n", " final_structures (List[ase.Atoms]): The structures after the relaxation without it.\n", " \"\"\"\n", - " order = np.argsort([atoms.get_potential_energy() for atoms in final_structures])\n", + " energies = [atoms.get_potential_energy() for atoms in final_structures]\n", + " order = np.argsort(energies)\n", " with open(\"results.csv\", \"w\", newline=\"\") as results_file:\n", " writer = csv.writer(results_file)\n", " writer.writerow([\"configuration\", \"rank\", \"afir_function_eV\", \"afir_max_force_eV_per_A\", \"energy_eV\",\n", @@ -656,12 +656,12 @@ " placement = final_structures[index].info[\"metadata\"][\"placement\"]\n", " writer.writerow([index, rank, *energy_and_max_force(afir_structures[index]),\n", " *energy_and_max_force(final_structures[index]), *placement[\"rotation_axis\"],\n", - " placement[\"rotation_angle_deg\"], placement[\"position_u\"], placement[\"position_v\"],\n", + " placement[\"rotation_angle_deg\"], *placement[\"position\"],\n", " molecule_fragments(final_structures[index])])\n", " final_structures[index].info.update(configuration=int(index), rank=rank)\n", " lowest_indices = order[:NUMBER_OF_FINAL_STRUCTURES].tolist()\n", " write(\"final_structures.extxyz\", [final_structures[index] for index in lowest_indices])\n", - " plot_energies(final_structures, lowest_indices)\n", + " plot_energies(energies, lowest_indices)\n", "\n", "\n", "def main():\n", @@ -865,7 +865,6 @@ " results=[\n", " {\"name\": \"file_content\", \"filetype\": \"csv\", \"basename\": \"results.csv\"},\n", " {\"name\": \"final_structure\"},\n", - " {\"name\": \"file_content\", \"filetype\": \"text\", \"basename\": \"final_structures.extxyz\"},\n", " {\"name\": \"file_content\", \"filetype\": \"image\", \"basename\": \"energies.png\"},\n", " ],\n", " is_rendered_at_run_time=True,\n", @@ -899,19 +898,21 @@ "from mat3ra.wode.units import Unit\n", "\n", "endpoint_options = {\"method\": \"POST\", \"headers\": {\"Content-Type\": \"application/json\"},\n", - " \"data\": \"{{ FINAL_STRUCTURE_REQUESTS[INDEX] }}\"}\n", + " \"data\": \"{{ FINAL_STRUCTURE_REQUESTS[FINAL_STRUCTURE_INDEX] }}\"}\n", "save_unit = Unit(type=\"io\", subtype=\"output\", source=\"api\", enableRender=True, name=\"save-final-structure\",\n", " input=[{\"type\": \"api\", \"endpoint\": \"structures\", \"endpoint_options\": endpoint_options,\n", " \"name\": \"FINAL_STRUCTURE\"}])\n", - "exit_unit = Unit(type=\"assignment\", name=\"final-structures-saved\", operand=\"NUMBER_OF_SAVED_STRUCTURES\", value=\"INDEX\",\n", - " input=[])\n", + "exit_unit = Unit(type=\"assignment\", name=\"final-structures-saved\", operand=\"NUMBER_OF_SAVED_STRUCTURES\",\n", + " value=\"FINAL_STRUCTURE_INDEX\", input=[])\n", "loop_units = [\n", - " Unit(type=\"assignment\", name=\"init-index\", operand=\"INDEX\", value=0, input=[]),\n", + " Unit(type=\"assignment\", name=\"init-index\", operand=\"FINAL_STRUCTURE_INDEX\", value=0, input=[]),\n", " save_unit,\n", - " Unit(type=\"assignment\", name=\"increment-index\", operand=\"INDEX\", value=\"INDEX + 1\", input=[]),\n", - " Unit(**{\"type\": \"condition\", \"name\": \"check-index\", \"statement\": f\"INDEX >= {number_of_saved_structures}\",\n", - " \"then\": exit_unit.flowchartId, \"else\": save_unit.flowchartId,\n", - " \"maxOccurrences\": number_of_saved_structures, \"input\": [{\"name\": \"INDEX\", \"scope\": \"global\"}]}),\n", + " Unit(type=\"assignment\", name=\"increment-index\", operand=\"FINAL_STRUCTURE_INDEX\", value=\"FINAL_STRUCTURE_INDEX + 1\",\n", + " input=[]),\n", + " Unit(**{\"type\": \"condition\", \"name\": \"check-index\",\n", + " \"statement\": f\"FINAL_STRUCTURE_INDEX >= {number_of_saved_structures}\",\n", + " \"then\": exit_unit.flowchartId, \"else\": save_unit.flowchartId, \"maxOccurrences\": number_of_saved_structures,\n", + " \"input\": [{\"name\": \"FINAL_STRUCTURE_INDEX\", \"scope\": \"global\"}]}),\n", " exit_unit,\n", "]\n", "final_structures_subworkflow = Subworkflow(**final_structures_subworkflow_config)\n", From 88d60782215d53e2428d2d5e668312335def7c40 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Tue, 3 Nov 2026 13:42:55 -0800 Subject: [PATCH 12/31] SOF-8102: AFIR script default keeps one final structure Co-Authored-By: Claude Opus 5.5 (1M context) --- other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 2ecbc2266..f9b038f2c 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -457,7 +457,7 @@ "MAX_STEPS = 500 # FIRE2 steps per relaxation at most\n", "TIME_STEP = 0.5 # initial FIRE2 step, in ALCHEMI's internal time unit (10.18 fs)\n", "GAMMA = 100 # kJ/mol, model collision energy of the AFIR force\n", - "NUMBER_OF_FINAL_STRUCTURES = 5 # the lowest final structures kept in final_structures.extxyz\n", + "NUMBER_OF_FINAL_STRUCTURES = 1 # the lowest final structures kept in final_structures.extxyz\n", "BOND_TOLERANCE = 1.2 # two molecule atoms are bonded within BOND_TOLERANCE * (R_i + R_j)\n", "AFIR_R0 = 3.8164 # Å, GRRM17 Eq. (3)\n", "AFIR_EPSILON = 1.0061 # kJ/mol, GRRM17 Eq. (3)\n", From d4ae051a79c1ad1a3497eebc2dd6d41009d7a976 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Tue, 3 Nov 2026 17:06:37 -0800 Subject: [PATCH 13/31] SOF-8102: step 3 saves the final structure under its own unit; AFIR step relaxes in batches of BATCH_SIZE Co-Authored-By: Claude Opus 5.5 (1M context) --- .../afir_molecule_surface_alchemi.ipynb | 65 ++++++++++++------- 1 file changed, 43 insertions(+), 22 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index f9b038f2c..1a8e7707a 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -72,6 +72,7 @@ "GAMMA = 100 # kJ/mol, model collision energy of the AFIR force\n", "RANDOM_SEED = 0 # of the random placements\n", "NUMBER_OF_FINAL_STRUCTURES = 1 # the lowest-energy final structures kept and saved as materials\n", + "BATCH_SIZE = 100 # configurations relaxed together on the GPU\n", "PYTHON_VERSION = \"3.11.10\"\n", "PYTHON_BUILD = \"Default\"\n", "\n", @@ -387,11 +388,12 @@ "source": [ "### 4.2. Step 2: AFIR with ALCHEMI\n", "\n", - "Relaxes all configurations together with the slab fixed [3], first under the AFIR force [1] of strength `GAMMA`\n", - "(100 kJ/mol with random orientations [2]) between molecule and slab, then without it. It writes `results.csv` with\n", - "every configuration ranked by final energy, its starting orientation and the fragments its molecule ends in;\n", - "`final_structures.extxyz` with the `NUMBER_OF_FINAL_STRUCTURES` lowest; and `energies.png`, the final energy of each\n", - "configuration. Another backend or another ALCHEMI version replaces this script and its packages; steps 1 and 3 stay." + "Relaxes the configurations in batches of `BATCH_SIZE` with the slab fixed [3], first under the AFIR force [1] of\n", + "strength `GAMMA` (100 kJ/mol with random orientations [2]) between molecule and slab, then without it. It writes\n", + "`results.csv` with every configuration ranked by final energy, its starting orientation and the fragments its molecule\n", + "ends in; `final_structures.extxyz` with the `NUMBER_OF_FINAL_STRUCTURES` lowest; and `energies.png`, the final energy of\n", + "each configuration. Another backend or another ALCHEMI version replaces this script and its packages; steps 1 and 3\n", + "stay." ] }, { @@ -405,10 +407,10 @@ "# AFIR step of the AFIR molecule-on-surface workflow, with MACE #\n", "# through NVIDIA ALCHEMI on one GPU. #\n", "# #\n", - "# Relaxes all configurations as one batch with the slab fixed, #\n", - "# first under the AFIR force between molecule and slab (GRRM17 #\n", - "# Eqs. 1-3 [1], strength GAMMA), then without it, and ranks the #\n", - "# final structures by energy. #\n", + "# Relaxes the configurations in batches of BATCH_SIZE with the #\n", + "# slab fixed, first under the AFIR force between molecule and #\n", + "# slab (GRRM17 Eqs. 1-3 [1], strength GAMMA), then without it, #\n", + "# and ranks the final structures by energy. #\n", "# #\n", "# Reads: CONFIGURATIONS from the job's global scope, written #\n", "# into this script by the compute node before it runs. #\n", @@ -455,6 +457,7 @@ "DTYPE = torch.float32\n", "FMAX = 0.03 # eV/Å, largest atomic force at convergence\n", "MAX_STEPS = 500 # FIRE2 steps per relaxation at most\n", + "BATCH_SIZE = 100 # structures relaxed together on the GPU\n", "TIME_STEP = 0.5 # initial FIRE2 step, in ALCHEMI's internal time unit (10.18 fs)\n", "GAMMA = 100 # kJ/mol, model collision energy of the AFIR force\n", "NUMBER_OF_FINAL_STRUCTURES = 1 # the lowest final structures kept in final_structures.extxyz\n", @@ -573,9 +576,9 @@ " return atoms\n", "\n", "\n", - "def relax(ase_structures, model, bias_hooks):\n", + "def relax_batch(ase_structures, model, bias_hooks):\n", " \"\"\"\n", - " Relaxes all structures as one batch with fixed-cell FIRE2, slab atoms frozen.\n", + " Relaxes the structures as one batch with fixed-cell FIRE2, slab atoms frozen.\n", "\n", " Args:\n", " ase_structures (List[ase.Atoms]): The structures, slab atoms tagged 0 and molecule atoms 1.\n", @@ -597,6 +600,26 @@ " return [relaxed_atoms(atoms, relaxed_data) for atoms, relaxed_data in zip(ase_structures, batch.to_data_list())]\n", "\n", "\n", + "def relax(ase_structures, model, bias_hooks):\n", + " \"\"\"\n", + " Relaxes all structures in batches of BATCH_SIZE, one after another, so the GPU memory a batch needs stays bounded.\n", + "\n", + " Args:\n", + " ase_structures (List[ase.Atoms]): The structures, slab atoms tagged 0 and molecule atoms 1.\n", + " model (MACEWrapper): The MACE model.\n", + " bias_hooks (List[BiasedPotentialHook]): Bias potentials added after each forward pass; [] for none.\n", + "\n", + " Returns:\n", + " List[ase.Atoms]: The relaxed structures, in the input order.\n", + " \"\"\"\n", + " relaxed = []\n", + " for start in range(0, len(ase_structures), BATCH_SIZE):\n", + " end = min(start + BATCH_SIZE, len(ase_structures))\n", + " logger.info(\"Relaxing structures {}–{} of {}\", start + 1, end, len(ase_structures))\n", + " relaxed += relax_batch(ase_structures[start:end], model, bias_hooks)\n", + " return relaxed\n", + "\n", + "\n", "def molecule_fragments(atoms):\n", " \"\"\"\n", " Returns the number of fragments the molecule of a structure is in: 1 if intact, more if it split.\n", @@ -707,7 +730,7 @@ "### 4.3. Step 3: final structures\n", "\n", "Takes the lowest final structures, each with the configuration it started from, and saves them as the job's final\n", - "structures: each is a material, shown in the Results tab of the AFIR step." + "structures: each is a material, shown in this step's Results." ] }, { @@ -723,7 +746,9 @@ "# Takes the lowest final structures of the AFIR step, each with #\n", "# the configuration it started from, and makes one request body #\n", "# per structure for the platform's final-structure endpoint, #\n", - "# which the io unit after this one sends in a loop. #\n", + "# which the io unit after this one sends in a loop. Each request #\n", + "# names this unit (UNIT_ID), so the structures are this step's #\n", + "# results. #\n", "# #\n", "# Reads: final_structures.extxyz in the job's work folder, and #\n", "# CONFIGURATIONS from the job's global scope, written #\n", @@ -739,7 +764,6 @@ "from mat3ra.made.material import Material\n", "from mat3ra.made.tools.convert import from_ase\n", "\n", - "AFIR_UNIT_FLOWCHART_ID = \"\"\n", "CONFIGURATIONS = json.loads(r\"\"\"{{ CONFIGURATIONS | tojson }}\"\"\")\n", "\n", "\n", @@ -758,8 +782,7 @@ " \"\"\"\n", " final = {**Material.create(from_ase(atoms)).to_dict(), \"repetition\": rank_index}\n", " initial = CONFIGURATIONS[atoms.info[\"configuration\"]]\n", - " body = {\"jobId\": \"{{ JOB_ID }}\", \"unitId\": AFIR_UNIT_FLOWCHART_ID,\n", - " \"structures\": [{\"initial\": initial, \"final\": final}]}\n", + " body = {\"jobId\": \"{{ JOB_ID }}\", \"unitId\": \"{{ UNIT_ID }}\", \"structures\": [{\"initial\": initial, \"final\": final}]}\n", " return json.dumps(json.dumps(body))[1:-1]\n", "\n", "\n", @@ -847,8 +870,6 @@ "final_structures_subworkflow_config = copy.deepcopy(python_script_subworkflow)\n", "final_structures_subworkflow_config.pop(\"_id\")\n", "next(unit for unit in final_structures_subworkflow_config[\"units\"] if unit[\"name\"] == \"python\").pop(\"flowchartId\")\n", - "afir_unit = next(unit for unit in afir_subworkflow_config[\"units\"] if unit[\"name\"] == \"python\")\n", - "afir_unit_flowchart_id = afir_unit[\"flowchartId\"]\n", "number_of_saved_structures = min(NUMBER_OF_FINAL_STRUCTURES, NUMBER_OF_CONFIGURATIONS)\n", "\n", "configure_python_step(\n", @@ -860,11 +881,11 @@ ")\n", "configure_python_step(\n", " afir_subworkflow_config, template_unit_name=\"python\", subworkflow_name=\"AFIR (ALCHEMI)\", unit_name=\"afir\",\n", - " script=set_constants(AFIR_SCRIPT, {\"GAMMA\": GAMMA, \"NUMBER_OF_FINAL_STRUCTURES\": number_of_saved_structures}),\n", + " script=set_constants(AFIR_SCRIPT, {\"GAMMA\": GAMMA, \"NUMBER_OF_FINAL_STRUCTURES\": number_of_saved_structures,\n", + " \"BATCH_SIZE\": BATCH_SIZE}),\n", " requirements=AFIR_REQUIREMENTS,\n", " results=[\n", " {\"name\": \"file_content\", \"filetype\": \"csv\", \"basename\": \"results.csv\"},\n", - " {\"name\": \"final_structure\"},\n", " {\"name\": \"file_content\", \"filetype\": \"image\", \"basename\": \"energies.png\"},\n", " ],\n", " is_rendered_at_run_time=True,\n", @@ -872,8 +893,8 @@ "configure_python_step(\n", " final_structures_subworkflow_config, template_unit_name=\"python\", subworkflow_name=\"Final structures\",\n", " unit_name=\"final_structures\",\n", - " script=set_constants(FINAL_STRUCTURES_SCRIPT, {\"AFIR_UNIT_FLOWCHART_ID\": f'\"{afir_unit_flowchart_id}\"'}),\n", - " requirements=MADE_REQUIREMENTS, results=[], is_rendered_at_run_time=True,\n", + " script=FINAL_STRUCTURES_SCRIPT, requirements=MADE_REQUIREMENTS, results=[{\"name\": \"final_structure\"}],\n", + " is_rendered_at_run_time=True,\n", ")" ] }, From f351219260394226daf62c485a311c8cbc151c61 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Tue, 3 Nov 2026 17:42:45 -0800 Subject: [PATCH 14/31] SOF-8102: final-structure loop runs until every request is sent; cell 0 says batches Co-Authored-By: Claude Opus 5.5 (1M context) --- .../alchemi/afir_molecule_surface_alchemi.ipynb | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 1a8e7707a..d9e8978df 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -7,7 +7,7 @@ "# AFIR for a molecule on a surface with NVIDIA ALCHEMI, on a platform GPU\n", "\n", "Places a molecule at random over a slab many times and relaxes every placement with MACE through the\n", - "[NVIDIA ALCHEMI toolkit](https://github.com/NVIDIA/nvalchemi-toolkit), all placements batched on one GPU: first\n", + "[NVIDIA ALCHEMI toolkit](https://github.com/NVIDIA/nvalchemi-toolkit), in batches on one GPU: first\n", "under the artificial force (AFIR) between molecule and slab, then without it. The final structures are ranked by\n", "energy, and the lowest are saved as materials.\n", "\n", @@ -870,7 +870,6 @@ "final_structures_subworkflow_config = copy.deepcopy(python_script_subworkflow)\n", "final_structures_subworkflow_config.pop(\"_id\")\n", "next(unit for unit in final_structures_subworkflow_config[\"units\"] if unit[\"name\"] == \"python\").pop(\"flowchartId\")\n", - "number_of_saved_structures = min(NUMBER_OF_FINAL_STRUCTURES, NUMBER_OF_CONFIGURATIONS)\n", "\n", "configure_python_step(\n", " workflow_config[\"subworkflows\"][0], template_unit_name=\"custom_script\", subworkflow_name=\"Configurations\",\n", @@ -881,7 +880,7 @@ ")\n", "configure_python_step(\n", " afir_subworkflow_config, template_unit_name=\"python\", subworkflow_name=\"AFIR (ALCHEMI)\", unit_name=\"afir\",\n", - " script=set_constants(AFIR_SCRIPT, {\"GAMMA\": GAMMA, \"NUMBER_OF_FINAL_STRUCTURES\": number_of_saved_structures,\n", + " script=set_constants(AFIR_SCRIPT, {\"GAMMA\": GAMMA, \"NUMBER_OF_FINAL_STRUCTURES\": NUMBER_OF_FINAL_STRUCTURES,\n", " \"BATCH_SIZE\": BATCH_SIZE}),\n", " requirements=AFIR_REQUIREMENTS,\n", " results=[\n", @@ -931,9 +930,10 @@ " Unit(type=\"assignment\", name=\"increment-index\", operand=\"FINAL_STRUCTURE_INDEX\", value=\"FINAL_STRUCTURE_INDEX + 1\",\n", " input=[]),\n", " Unit(**{\"type\": \"condition\", \"name\": \"check-index\",\n", - " \"statement\": f\"FINAL_STRUCTURE_INDEX >= {number_of_saved_structures}\",\n", - " \"then\": exit_unit.flowchartId, \"else\": save_unit.flowchartId, \"maxOccurrences\": number_of_saved_structures,\n", - " \"input\": [{\"name\": \"FINAL_STRUCTURE_INDEX\", \"scope\": \"global\"}]}),\n", + " \"statement\": \"FINAL_STRUCTURE_INDEX >= len(FINAL_STRUCTURE_REQUESTS)\",\n", + " \"then\": exit_unit.flowchartId, \"else\": save_unit.flowchartId, \"maxOccurrences\": NUMBER_OF_CONFIGURATIONS,\n", + " \"input\": [{\"name\": \"FINAL_STRUCTURE_INDEX\", \"scope\": \"global\"},\n", + " {\"name\": \"FINAL_STRUCTURE_REQUESTS\", \"scope\": \"global\"}]}),\n", " exit_unit,\n", "]\n", "final_structures_subworkflow = Subworkflow(**final_structures_subworkflow_config)\n", From ebc0a6327a44c405bbd479ed18f2b2b124f583d0 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 16:51:47 -0800 Subject: [PATCH 15/31] SOF-8102: AFIR step runs NVIDIA's ALCHEMI pre-release with cuEquivariance; steps share configurations.extxyz Co-Authored-By: Claude Opus 5.5 --- .../afir_molecule_surface_alchemi.ipynb | 111 ++++++++++-------- 1 file changed, 62 insertions(+), 49 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index d9e8978df..3f2f8ed32 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -250,8 +250,8 @@ "\n", "Places the molecule `NUMBER_OF_CONFIGURATIONS` times over the slab, each at a random orientation and a random point\n", "of the surface cell, with its lowest atom above the top slab atom by the sum of their covalent radii plus 0.5 Å, as\n", - "GRRM's random initial structures [1]. Slab atoms are labelled 0, molecule atoms 1, and each configuration\n", - "keeps its rotation and position in its metadata." + "GRRM's random initial structures [1]. It writes them to `configurations.extxyz`, slab atoms tagged 0 and molecule\n", + "atoms 1, each with its index, rotation and position." ] }, { @@ -272,11 +272,11 @@ "# Reads: the job's two materials, written into this script by #\n", "# the platform when the job is created; the non-periodic #\n", "# one is the molecule. #\n", - "# Writes: CONFIGURATIONS to the job's global scope: the placed #\n", - "# structures as materials, slab atoms labelled 0 and #\n", - "# molecule atoms 1, each with its rotation and fractional #\n", - "# position in metadata.placement. It is the only output #\n", - "# of the script. #\n", + "# Writes: configurations.extxyz in the job's work folder: the #\n", + "# placed structures, slab atoms tagged 0 and molecule #\n", + "# atoms 1, each with its index in configuration and its #\n", + "# rotation and fractional position in metadata.placement. #\n", + "# It is the only output of the script. #\n", "# #\n", "# [1] S. Maeda et al., J. Comput. Chem. 39, 233 (2018). #\n", "# #\n", @@ -286,7 +286,9 @@ "\n", "import numpy as np\n", "from ase.data import atomic_numbers, covalent_radii\n", + "from ase.io import write\n", "from mat3ra.made.material import Material\n", + "from mat3ra.made.tools.convert import to_ase\n", "from mat3ra.made.tools.convert.interface_parts_enum import InterfacePartsEnum\n", "from mat3ra.made.tools.modify import translate_by_vector\n", "from mat3ra.made.tools.operations.core.binary import merge\n", @@ -359,7 +361,7 @@ "\n", "def main():\n", " \"\"\"\n", - " Builds NUMBER_OF_CONFIGURATIONS configurations from the job's materials and prints them as CONFIGURATIONS.\n", + " Builds NUMBER_OF_CONFIGURATIONS configurations from the job's materials and writes them to configurations.extxyz.\n", " \"\"\"\n", " slab = Material.create(next(material_config for material_config in MATERIALS\n", " if not material_config.get(\"isNonPeriodic\")))\n", @@ -370,11 +372,12 @@ " slab.set_labels_from_value(InterfacePartsEnum.SUBSTRATE.value)\n", " top_atom_height_and_radius = extreme_atom_height_and_radius(slab, max)\n", " random_generator = np.random.default_rng(RANDOM_SEED)\n", - " configurations = [\n", - " create_configuration(slab, molecule, top_atom_height_and_radius, random_generator, index).to_dict()\n", - " for index in range(NUMBER_OF_CONFIGURATIONS)\n", - " ]\n", - " print(json.dumps({\"name\": \"CONFIGURATIONS\", \"value\": configurations, \"scope\": \"global\"}))\n", + " configurations = []\n", + " for index in range(NUMBER_OF_CONFIGURATIONS):\n", + " atoms = to_ase(create_configuration(slab, molecule, top_atom_height_and_radius, random_generator, index))\n", + " atoms.info[\"configuration\"] = index\n", + " configurations.append(atoms)\n", + " write(\"configurations.extxyz\", configurations)\n", "\n", "\n", "if __name__ == \"__main__\":\n", @@ -388,12 +391,12 @@ "source": [ "### 4.2. Step 2: AFIR with ALCHEMI\n", "\n", - "Relaxes the configurations in batches of `BATCH_SIZE` with the slab fixed [3], first under the AFIR force [1] of\n", - "strength `GAMMA` (100 kJ/mol with random orientations [2]) between molecule and slab, then without it. It writes\n", - "`results.csv` with every configuration ranked by final energy, its starting orientation and the fragments its molecule\n", - "ends in; `final_structures.extxyz` with the `NUMBER_OF_FINAL_STRUCTURES` lowest; and `energies.png`, the final energy of\n", - "each configuration. Another backend or another ALCHEMI version replaces this script and its packages; steps 1 and 3\n", - "stay." + "Relaxes the configurations of `configurations.extxyz` in batches of `BATCH_SIZE` with the slab fixed [3], first under\n", + "the AFIR force [1] of strength `GAMMA` (100 kJ/mol with random orientations [2]) between molecule and slab, then\n", + "without it. It writes `results.csv` with every configuration ranked by final energy, its starting orientation and the\n", + "fragments its molecule ends in; `final_structures.extxyz` with the `NUMBER_OF_FINAL_STRUCTURES` lowest; and\n", + "`energies.png`, the final energy of each configuration. Another backend or another ALCHEMI version replaces this\n", + "script and its packages; steps 1 and 3 stay." ] }, { @@ -412,8 +415,9 @@ "# slab (GRRM17 Eqs. 1-3 [1], strength GAMMA), then without it, #\n", "# and ranks the final structures by energy. #\n", "# #\n", - "# Reads: CONFIGURATIONS from the job's global scope, written #\n", - "# into this script by the compute node before it runs. #\n", + "# Reads: configurations.extxyz, written by the configurations #\n", + "# step in the job's work folder; slab atoms tagged 0, #\n", + "# molecule atoms 1. #\n", "# Writes: results.csv, every configuration ranked by final #\n", "# energy, with its starting rotation and position and the #\n", "# number of fragments its molecule ends in; #\n", @@ -427,7 +431,7 @@ "# ---------------------------------------------------------------- #\n", "\n", "import csv\n", - "import json\n", + "from importlib.metadata import version\n", "\n", "import matplotlib\n", "import numpy as np\n", @@ -435,12 +439,9 @@ "from ase.build import separate\n", "from ase.calculators.singlepoint import SinglePointCalculator\n", "from ase.data import covalent_radii\n", - "from ase.io import write\n", + "from ase.io import read, write\n", "from ase.units import kJ, mol\n", "from loguru import logger\n", - "from mat3ra.made.material import Material\n", - "from mat3ra.made.tools.convert import to_ase\n", - "from mat3ra.made.tools.convert.interface_parts_enum import InterfacePartsEnum\n", "from nvalchemi._typing import AtomCategory\n", "from nvalchemi.data import AtomicData, Batch\n", "from nvalchemi.dynamics import FIRE2, FusedStage\n", @@ -462,13 +463,13 @@ "GAMMA = 100 # kJ/mol, model collision energy of the AFIR force\n", "NUMBER_OF_FINAL_STRUCTURES = 1 # the lowest final structures kept in final_structures.extxyz\n", "BOND_TOLERANCE = 1.2 # two molecule atoms are bonded within BOND_TOLERANCE * (R_i + R_j)\n", + "SLAB_TAG = 0 # atom tags in configurations.extxyz\n", + "MOLECULE_TAG = 1\n", "AFIR_R0 = 3.8164 # Å, GRRM17 Eq. (3)\n", "AFIR_EPSILON = 1.0061 # kJ/mol, GRRM17 Eq. (3)\n", - "CONFIGURATIONS = json.loads(r\"\"\"{{ CONFIGURATIONS | tojson }}\"\"\")\n", "\n", "ALPHA = GAMMA * kJ / mol / ((2 ** (-1 / 6) - (1 + (1 + GAMMA / AFIR_EPSILON) ** 0.5) ** (-1 / 6)) * AFIR_R0) # Eq. (3)\n", - "CATEGORY_BY_PART = {InterfacePartsEnum.SUBSTRATE.value: AtomCategory.SURFACE.value,\n", - " InterfacePartsEnum.FILM.value: AtomCategory.GAS.value}\n", + "CATEGORY_BY_TAG = {SLAB_TAG: AtomCategory.SURFACE.value, MOLECULE_TAG: AtomCategory.GAS.value}\n", "IN_PLANE_MASK = torch.tensor([1.0, 1.0, 0.0], device=DEVICE)\n", "IN_PLANE_SHIFTS = torch.tensor([[a, b, 0.0] for a in (-1, 0, 1) for b in (-1, 0, 1)], device=DEVICE)\n", "COVALENT_RADII = torch.tensor(covalent_radii, device=DEVICE, dtype=DTYPE)\n", @@ -543,7 +544,8 @@ "\n", "def atomic_data(atoms):\n", " \"\"\"\n", - " Converts a configuration to ALCHEMI's AtomicData, its labels to atom categories: 0 to SURFACE, 1 to GAS.\n", + " Converts a configuration to ALCHEMI's AtomicData, its tags to atom categories: SLAB_TAG to SURFACE, MOLECULE_TAG\n", + " to GAS.\n", "\n", " Args:\n", " atoms (ase.Atoms): The configuration, slab atoms tagged 0 and molecule atoms 1.\n", @@ -552,7 +554,7 @@ " AtomicData: The configuration on DEVICE, with zeroed forces and energy.\n", " \"\"\"\n", " data = AtomicData.from_atoms(atoms, device=DEVICE, dtype=DTYPE)\n", - " data[\"atom_categories\"] = torch.tensor([CATEGORY_BY_PART[tag] for tag in atoms.get_tags()], device=DEVICE)\n", + " data[\"atom_categories\"] = torch.tensor([CATEGORY_BY_TAG[tag] for tag in atoms.get_tags()], device=DEVICE)\n", " data[\"forces\"] = torch.zeros_like(data.positions)\n", " data[\"energy\"] = torch.zeros((1, 1), device=DEVICE, dtype=DTYPE)\n", " return data\n", @@ -630,7 +632,7 @@ " Returns:\n", " int: The number of connected groups of molecule atoms, bonded within BOND_TOLERANCE times their covalent radii.\n", " \"\"\"\n", - " return len(separate(atoms[atoms.get_tags() == InterfacePartsEnum.FILM.value], scale=BOND_TOLERANCE))\n", + " return len(separate(atoms[atoms.get_tags() == MOLECULE_TAG], scale=BOND_TOLERANCE))\n", "\n", "\n", "def energy_and_max_force(atoms):\n", @@ -681,7 +683,7 @@ " *energy_and_max_force(final_structures[index]), *placement[\"rotation_axis\"],\n", " placement[\"rotation_angle_deg\"], *placement[\"position\"],\n", " molecule_fragments(final_structures[index])])\n", - " final_structures[index].info.update(configuration=int(index), rank=rank)\n", + " final_structures[index].info.update(rank=rank)\n", " lowest_indices = order[:NUMBER_OF_FINAL_STRUCTURES].tolist()\n", " write(\"final_structures.extxyz\", [final_structures[index] for index in lowest_indices])\n", " plot_energies(energies, lowest_indices)\n", @@ -689,11 +691,14 @@ "\n", "def main():\n", " \"\"\"\n", - " Relaxes CONFIGURATIONS under the AFIR force, then without it, and writes the ranked results.\n", + " Relaxes the configurations of configurations.extxyz under the AFIR force, then without it, and writes the ranked\n", + " results.\n", " \"\"\"\n", " logger.info(\"AFIR force α, γ = {} kJ/mol: {:.3f} eV/Å\", GAMMA, ALPHA)\n", - " ase_structures = [to_ase(Material.create(configuration)) for configuration in CONFIGURATIONS]\n", - " model = MACEWrapper.from_checkpoint(checkpoint_path=CHECKPOINT, device=torch.device(DEVICE), dtype=DTYPE)\n", + " ase_structures = read(\"{{ JOB_WORK_DIR }}/configurations.extxyz\", index=\":\")\n", + " model = MACEWrapper.from_checkpoint(checkpoint_path=CHECKPOINT, device=torch.device(DEVICE), dtype=DTYPE,\n", + " enable_cueq=True)\n", + " logger.info(\"nvalchemi-toolkit {}, MACE with cuEquivariance\", version(\"nvalchemi-toolkit\"))\n", " afir_structures = relax(ase_structures, model, [BiasedPotentialHook(bias_fn=afir_bias)])\n", " final_structures = relax(afir_structures, model, [])\n", " write_results(afir_structures, final_structures)\n", @@ -715,12 +720,19 @@ "--find-links https://download.pytorch.org/whl/cu128/torchvision/\n", "--find-links https://download.pytorch.org/whl/cu128/torchaudio/\n", "--find-links https://github.com/NVIDIA/warp/releases/expanded_assets/v1.18.0\n", + "#--extra-index-url https://download.pytorch.org/whl/cu128\n", "torch==2.11.0+cu128\n", "torchvision==0.26.0+cu128\n", "torchaudio==2.11.0+cu128\n", "warp-lang==1.18.0+cu12\n", - "nvalchemi-toolkit[mace,ase]==0.2.0\n", - "\"\"\" + MADE_REQUIREMENTS" + "nvalchemi-toolkit[cu12,mace,ase] @ git+https://github.com/NVIDIA/nvalchemi-toolkit.git@618916442783c152cd53b919164f07f80286f3da\n", + "matplotlib\n", + "numpy\n", + "ase\n", + "mace-torch\n", + "loguru\n", + "mat3ra-made[tools]\n", + "\"\"\"" ] }, { @@ -750,9 +762,8 @@ "# names this unit (UNIT_ID), so the structures are this step's #\n", "# results. #\n", "# #\n", - "# Reads: final_structures.extxyz in the job's work folder, and #\n", - "# CONFIGURATIONS from the job's global scope, written #\n", - "# into this script by the compute node before it runs. #\n", + "# Reads: final_structures.extxyz and configurations.extxyz in #\n", + "# the job's work folder. #\n", "# Writes: FINAL_STRUCTURE_REQUESTS to the job's global scope, as #\n", "# the only output of the script. #\n", "# #\n", @@ -764,10 +775,8 @@ "from mat3ra.made.material import Material\n", "from mat3ra.made.tools.convert import from_ase\n", "\n", - "CONFIGURATIONS = json.loads(r\"\"\"{{ CONFIGURATIONS | tojson }}\"\"\")\n", - "\n", "\n", - "def request_body(rank_index, atoms):\n", + "def request_body(rank_index, atoms, configurations):\n", " \"\"\"\n", " Returns the body of one final-structure request, escaped to sit inside the next io unit's JSON string.\n", "\n", @@ -776,22 +785,26 @@ " Args:\n", " rank_index (int): The structure's rank by final energy, from 0.\n", " atoms (ase.Atoms): The final structure, with its configuration's index in atoms.info.\n", + " configurations (List[ase.Atoms]): All configurations, in the order of their indices.\n", "\n", " Returns:\n", " str: {\"jobId\", \"unitId\", \"structures\": [{\"initial\", \"final\"}]} as JSON, escaped for a JSON string.\n", " \"\"\"\n", " final = {**Material.create(from_ase(atoms)).to_dict(), \"repetition\": rank_index}\n", - " initial = CONFIGURATIONS[atoms.info[\"configuration\"]]\n", + " initial = Material.create(from_ase(configurations[atoms.info[\"configuration\"]])).to_dict()\n", " body = {\"jobId\": \"{{ JOB_ID }}\", \"unitId\": \"{{ UNIT_ID }}\", \"structures\": [{\"initial\": initial, \"final\": final}]}\n", " return json.dumps(json.dumps(body))[1:-1]\n", "\n", "\n", "def main():\n", " \"\"\"\n", - " Reads the lowest final structures and prints their request bodies as FINAL_STRUCTURE_REQUESTS.\n", + " Reads the lowest final structures and the configurations, and prints the request bodies as\n", + " FINAL_STRUCTURE_REQUESTS.\n", " \"\"\"\n", " final_structures = read(\"{{ JOB_WORK_DIR }}/final_structures.extxyz\", index=\":\")\n", - " request_bodies = [request_body(rank_index, atoms) for rank_index, atoms in enumerate(final_structures)]\n", + " configurations = read(\"{{ JOB_WORK_DIR }}/configurations.extxyz\", index=\":\")\n", + " request_bodies = [request_body(rank_index, atoms, configurations)\n", + " for rank_index, atoms in enumerate(final_structures)]\n", " print(json.dumps({\"name\": \"FINAL_STRUCTURE_REQUESTS\", \"value\": request_bodies, \"scope\": \"global\"}))\n", "\n", "\n", @@ -807,8 +820,8 @@ "### 4.4. Configure the steps\n", "\n", "When the job is created, the platform writes the job's materials into step 1's script (`materials`; the\n", - "non-periodic one is the molecule). Step 1 prints `CONFIGURATIONS` to the job's scope, and the compute node writes\n", - "them into the scripts of steps 2 and 3 before they run. Each step's files are named after it." + "non-periodic one is the molecule). Step 1 writes the configurations to `configurations.extxyz` in the\n", + "job's work folder, where steps 2 and 3 read them. Each step's files are named after it." ] }, { From 13b49969a1984f5153a14a66349985a0100dc0de Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 18:04:11 -0800 Subject: [PATCH 16/31] SOF-8102: workflow, steps, units, script and result files renamed; AFIR spelled out Co-Authored-By: Claude Opus 5.5 --- .../afir_molecule_surface_alchemi.ipynb | 152 ++++++++++-------- 1 file changed, 81 insertions(+), 71 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 3f2f8ed32..7a79e87bf 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -8,13 +8,14 @@ "\n", "Places a molecule at random over a slab many times and relaxes every placement with MACE through the\n", "[NVIDIA ALCHEMI toolkit](https://github.com/NVIDIA/nvalchemi-toolkit), in batches on one GPU: first\n", - "under the artificial force (AFIR) between molecule and slab, then without it. The final structures are ranked by\n", - "energy, and the lowest are saved as materials.\n", + "under the artificial force of the artificial force induced reaction (AFIR) method between molecule and slab, then\n", + "without it. The final structures are ranked by energy, and the lowest are saved as materials.\n", "\n", - "The workflow has three steps, each in its own Python environment: *Configurations* builds the placements from the\n", - "job's two materials with made, *AFIR (ALCHEMI)* relaxes them, and *Final structures* saves the lowest as the job's\n", - "final structures. This notebook builds the slab, fetches the molecule, saves both as materials, creates the workflow,\n", - "runs one job on the two materials and shows the ranked configurations and the lowest final structures.\n", + "The workflow has three steps, each in its own Python environment: *Create Initial Configurations* builds the\n", + "placements from the job's two materials with made, *Relax under AFIR Force (ALCHEMI)* relaxes them, and *Save Lowest\n", + "Final Structures* saves the lowest as the job's final structures. This notebook builds the slab, fetches the\n", + "molecule, saves both as materials, creates the workflow, runs one job on the two materials and shows the ranked\n", + "configurations and the lowest final structures.\n", "\n", "Outside the platform's JupyterLite, use Python 3.10 or 3.11 and install the packages once with `pip install \"mat3ra-notebooks-utils[all]\"`." ] @@ -246,7 +247,7 @@ "metadata": {}, "source": [ "## 4. Create the workflow\n", - "### 4.1. Step 1: configurations\n", + "### 4.1. Step 1: Create Initial Configurations\n", "\n", "Places the molecule `NUMBER_OF_CONFIGURATIONS` times over the slab, each at a random orientation and a random point\n", "of the surface cell, with its lowest atom above the top slab atom by the sum of their covalent radii plus 0.5 Å, as\n", @@ -262,7 +263,9 @@ "source": [ "CONFIGURATIONS_SCRIPT = r'''# ---------------------------------------------------------------- #\n", "# #\n", - "# Configurations step of the AFIR molecule-on-surface workflow. #\n", + "# Step 1, \"Create Initial Configurations\", of the workflow #\n", + "# \"Molecule on Surface: AFIR Search (ALCHEMI)\", a search with the #\n", + "# artificial force induced reaction (AFIR) method. #\n", "# #\n", "# Places the job's molecule over its slab #\n", "# NUMBER_OF_CONFIGURATIONS times with made, each at a random #\n", @@ -274,9 +277,10 @@ "# one is the molecule. #\n", "# Writes: configurations.extxyz in the job's work folder: the #\n", "# placed structures, slab atoms tagged 0 and molecule #\n", - "# atoms 1, each with its index in configuration and its #\n", - "# rotation and fractional position in metadata.placement. #\n", - "# It is the only output of the script. #\n", + "# atoms 1, each with its index in the info key #\n", + "# configuration and its rotation and fractional position #\n", + "# in metadata.placement. It is the only output of the #\n", + "# script. #\n", "# #\n", "# [1] S. Maeda et al., J. Comput. Chem. 39, 233 (2018). #\n", "# #\n", @@ -389,14 +393,14 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### 4.2. Step 2: AFIR with ALCHEMI\n", + "### 4.2. Step 2: Relax under AFIR Force (ALCHEMI)\n", "\n", "Relaxes the configurations of `configurations.extxyz` in batches of `BATCH_SIZE` with the slab fixed [3], first under\n", "the AFIR force [1] of strength `GAMMA` (100 kJ/mol with random orientations [2]) between molecule and slab, then\n", - "without it. It writes `results.csv` with every configuration ranked by final energy, its starting orientation and the\n", - "fragments its molecule ends in; `final_structures.extxyz` with the `NUMBER_OF_FINAL_STRUCTURES` lowest; and\n", - "`energies.png`, the final energy of each configuration. Another backend or another ALCHEMI version replaces this\n", - "script and its packages; steps 1 and 3 stay." + "without it. It writes `ranked_configurations.csv` with every configuration ranked by final energy, its starting\n", + "orientation and the fragments its molecule ends in; `final_structures.extxyz` with the `NUMBER_OF_FINAL_STRUCTURES`\n", + "lowest; and `final_energies.png`, the final energy of each configuration. Another backend or another ALCHEMI version\n", + "replaces this script and its packages; steps 1 and 3 stay." ] }, { @@ -407,24 +411,27 @@ "source": [ "AFIR_SCRIPT = r'''# ---------------------------------------------------------------- #\n", "# #\n", - "# AFIR step of the AFIR molecule-on-surface workflow, with MACE #\n", - "# through NVIDIA ALCHEMI on one GPU. #\n", + "# Step 2, \"Relax under AFIR Force (ALCHEMI)\", of the workflow #\n", + "# \"Molecule on Surface: AFIR Search (ALCHEMI)\", with MACE through #\n", + "# NVIDIA ALCHEMI on one GPU. #\n", "# #\n", "# Relaxes the configurations in batches of BATCH_SIZE with the #\n", - "# slab fixed, first under the AFIR force between molecule and #\n", - "# slab (GRRM17 Eqs. 1-3 [1], strength GAMMA), then without it, #\n", - "# and ranks the final structures by energy. #\n", + "# slab fixed, first under the artificial force of the artificial #\n", + "# force induced reaction (AFIR) method between molecule and slab #\n", + "# (GRRM17 Eqs. 1-3 [1], strength GAMMA), then without it, and #\n", + "# ranks the final structures by energy. #\n", "# #\n", - "# Reads: configurations.extxyz, written by the configurations #\n", - "# step in the job's work folder; slab atoms tagged 0, #\n", - "# molecule atoms 1. #\n", - "# Writes: results.csv, every configuration ranked by final #\n", - "# energy, with its starting rotation and position and the #\n", - "# number of fragments its molecule ends in; #\n", - "# final_structures.extxyz, the NUMBER_OF_FINAL_STRUCTURES #\n", - "# lowest in that order; and energies.png, final energy #\n", - "# against configuration index. results.csv and #\n", - "# energies.png are results of this unit. #\n", + "# Reads: configurations.extxyz, written by \"Create Initial #\n", + "# Configurations\" in the job's work folder; slab atoms #\n", + "# tagged 0, molecule atoms 1. #\n", + "# Writes: ranked_configurations.csv, every configuration ranked #\n", + "# by final energy, with its starting rotation and #\n", + "# position and the number of fragments its molecule ends #\n", + "# in; final_structures.extxyz, the #\n", + "# NUMBER_OF_FINAL_STRUCTURES lowest in that order; and #\n", + "# final_energies.png, final energy against configuration #\n", + "# index. ranked_configurations.csv and final_energies.png #\n", + "# are results of this unit. #\n", "# #\n", "# [1] S. Maeda et al., J. Comput. Chem. 39, 233 (2018). #\n", "# #\n", @@ -644,7 +651,7 @@ "\n", "def plot_energies(energies, lowest_indices):\n", " \"\"\"\n", - " Plots the final energy of every configuration against its index, the kept lowest marked, to energies.png.\n", + " Plots the final energy of every configuration against its index, the kept lowest marked, to final_energies.png.\n", "\n", " Args:\n", " energies (List[float]): The final energy of each configuration, in eV.\n", @@ -657,14 +664,14 @@ " axes.set(xlabel=\"Configuration index\", ylabel=\"Final energy (eV)\")\n", " axes.ticklabel_format(axis=\"y\", useOffset=False)\n", " axes.legend()\n", - " figure.savefig(\"energies.png\", dpi=150, bbox_inches=\"tight\")\n", + " figure.savefig(\"final_energies.png\", dpi=150, bbox_inches=\"tight\")\n", "\n", "\n", "def write_results(afir_structures, final_structures):\n", " \"\"\"\n", - " Ranks the configurations by final energy, writes all of them to results.csv with the placement each started from\n", - " (atoms.info[\"metadata\"][\"placement\"]) and its molecule's fragments, the NUMBER_OF_FINAL_STRUCTURES lowest to\n", - " final_structures.extxyz in that order, and the energies to energies.png.\n", + " Ranks the configurations by final energy, writes all of them to ranked_configurations.csv with the placement each\n", + " started from (atoms.info[\"metadata\"][\"placement\"]) and its molecule's fragments, the NUMBER_OF_FINAL_STRUCTURES\n", + " lowest to final_structures.extxyz in that order, and the energies to final_energies.png.\n", "\n", " Args:\n", " afir_structures (List[ase.Atoms]): The structures after the relaxation under the AFIR force.\n", @@ -672,7 +679,7 @@ " \"\"\"\n", " energies = [atoms.get_potential_energy() for atoms in final_structures]\n", " order = np.argsort(energies)\n", - " with open(\"results.csv\", \"w\", newline=\"\") as results_file:\n", + " with open(\"ranked_configurations.csv\", \"w\", newline=\"\") as results_file:\n", " writer = csv.writer(results_file)\n", " writer.writerow([\"configuration\", \"rank\", \"afir_function_eV\", \"afir_max_force_eV_per_A\", \"energy_eV\",\n", " \"max_force_eV_per_A\", \"rotation_axis_x\", \"rotation_axis_y\", \"rotation_axis_z\",\n", @@ -739,7 +746,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### 4.3. Step 3: final structures\n", + "### 4.3. Step 3: Save Lowest Final Structures\n", "\n", "Takes the lowest final structures, each with the configuration it started from, and saves them as the job's final\n", "structures: each is a material, shown in this step's Results." @@ -753,14 +760,16 @@ "source": [ "FINAL_STRUCTURES_SCRIPT = r'''# ---------------------------------------------------------------- #\n", "# #\n", - "# Final structures step of the AFIR molecule-on-surface workflow. #\n", + "# Step 3, \"Save Lowest Final Structures\", of the workflow #\n", + "# \"Molecule on Surface: AFIR Search (ALCHEMI)\", a search with the #\n", + "# artificial force induced reaction (AFIR) method. #\n", "# #\n", - "# Takes the lowest final structures of the AFIR step, each with #\n", - "# the configuration it started from, and makes one request body #\n", - "# per structure for the platform's final-structure endpoint, #\n", - "# which the io unit after this one sends in a loop. Each request #\n", - "# names this unit (UNIT_ID), so the structures are this step's #\n", - "# results. #\n", + "# Takes the lowest final structures of \"Relax under AFIR Force #\n", + "# (ALCHEMI)\", each with the configuration it started from, and #\n", + "# makes one request body per structure for the platform's #\n", + "# final-structure endpoint, which the io unit after this one #\n", + "# sends in a loop. Each request names this unit (UNIT_ID), so the #\n", + "# structures are this step's results. #\n", "# #\n", "# Reads: final_structures.extxyz and configurations.extxyz in #\n", "# the job's work folder. #\n", @@ -844,7 +853,7 @@ " return script\n", "\n", "\n", - "def configure_python_step(subworkflow, template_unit_name, subworkflow_name, unit_name, script, requirements, results,\n", + "def configure_python_step(subworkflow, template_unit_name, subworkflow_name, file_name, script, requirements, results,\n", " is_rendered_at_run_time):\n", " \"\"\"\n", " Makes a Standata python subworkflow one step of the workflow.\n", @@ -852,32 +861,32 @@ " Args:\n", " subworkflow (dict): The subworkflow config, edited in place.\n", " template_unit_name (str): The name of the template's python unit.\n", - " subworkflow_name (str): The step's name in the workflow.\n", - " unit_name (str): The python unit's name; its script and requirements files are named after it.\n", + " subworkflow_name (str): The step's name in the workflow, also the name of its python unit.\n", + " file_name (str): The base name of the unit's script and requirements files.\n", " script (str): The unit's script.\n", " requirements (str): The packages installed for the unit, as a requirements.txt.\n", " results (list): The results the unit declares.\n", - " is_rendered_at_run_time (bool): True to leave the script's placeholders to the compute node, which fills them\n", - " from the job's scope when the step runs; False to let the platform fill them, with the job's materials,\n", - " when the job is created.\n", + " is_rendered_at_run_time (bool): True to leave the script's placeholders (JOB_WORK_DIR, JOB_ID, UNIT_ID) to the\n", + " compute node, which fills them when the step runs; False to let the platform fill them, with the job's\n", + " materials, when the job is created.\n", " \"\"\"\n", " unit = next(unit for unit in subworkflow[\"units\"] if unit[\"name\"] == template_unit_name)\n", - " unit.update(name=unit_name, head=True, results=results)\n", + " unit.update(name=subworkflow_name, head=True, results=results)\n", " for application in (subworkflow[\"application\"], unit[\"application\"]):\n", " application.update(version=PYTHON_VERSION, build=PYTHON_BUILD)\n", " set_execution_unit_input(unit, \"script.py\", script)\n", " set_execution_unit_input(unit, \"requirements.txt\", requirements)\n", " inputs = {unit_input[\"template\"][\"name\"]: unit_input for unit_input in unit[\"input\"]}\n", " inputs[\"script.py\"][\"isManuallyChanged\"] = is_rendered_at_run_time\n", - " inputs[\"script.py\"][\"template\"][\"name\"] = f\"{unit_name}.py\"\n", - " inputs[\"requirements.txt\"][\"template\"][\"name\"] = f\"{unit_name}_requirements.txt\"\n", + " inputs[\"script.py\"][\"template\"][\"name\"] = f\"{file_name}.py\"\n", + " inputs[\"requirements.txt\"][\"template\"][\"name\"] = f\"{file_name}_requirements.txt\"\n", " subworkflow.update(name=subworkflow_name, units=[unit])\n", "\n", "\n", "python_workflows = WorkflowStandata.filter_by_application(\"python\")\n", "workflow_config = copy.deepcopy(python_workflows.get_by_name_first_match(\"custom_script.json\"))\n", - "workflow_config.update(name=\"AFIR, molecule on surface\", isMultiMaterial=True)\n", - "workflow_config[\"units\"][0][\"name\"] = \"Configurations\"\n", + "workflow_config.update(name=\"Molecule on Surface: AFIR Search (ALCHEMI)\", isMultiMaterial=True)\n", + "workflow_config[\"units\"][0][\"name\"] = \"Create Initial Configurations\"\n", "python_script_subworkflow = python_workflows.get_by_name_first_match(\"python_script.json\")[\"subworkflows\"][0]\n", "afir_subworkflow_config = copy.deepcopy(python_script_subworkflow)\n", "final_structures_subworkflow_config = copy.deepcopy(python_script_subworkflow)\n", @@ -885,26 +894,27 @@ "next(unit for unit in final_structures_subworkflow_config[\"units\"] if unit[\"name\"] == \"python\").pop(\"flowchartId\")\n", "\n", "configure_python_step(\n", - " workflow_config[\"subworkflows\"][0], template_unit_name=\"custom_script\", subworkflow_name=\"Configurations\",\n", - " unit_name=\"configurations\",\n", + " workflow_config[\"subworkflows\"][0], template_unit_name=\"custom_script\",\n", + " subworkflow_name=\"Create Initial Configurations\", file_name=\"create_initial_configurations\",\n", " script=set_constants(CONFIGURATIONS_SCRIPT, {\"NUMBER_OF_CONFIGURATIONS\": NUMBER_OF_CONFIGURATIONS,\n", " \"RANDOM_SEED\": RANDOM_SEED}),\n", " requirements=MADE_REQUIREMENTS, results=[], is_rendered_at_run_time=False,\n", ")\n", "configure_python_step(\n", - " afir_subworkflow_config, template_unit_name=\"python\", subworkflow_name=\"AFIR (ALCHEMI)\", unit_name=\"afir\",\n", + " afir_subworkflow_config, template_unit_name=\"python\", subworkflow_name=\"Relax under AFIR Force (ALCHEMI)\",\n", + " file_name=\"relax_under_afir_force\",\n", " script=set_constants(AFIR_SCRIPT, {\"GAMMA\": GAMMA, \"NUMBER_OF_FINAL_STRUCTURES\": NUMBER_OF_FINAL_STRUCTURES,\n", " \"BATCH_SIZE\": BATCH_SIZE}),\n", " requirements=AFIR_REQUIREMENTS,\n", " results=[\n", - " {\"name\": \"file_content\", \"filetype\": \"csv\", \"basename\": \"results.csv\"},\n", - " {\"name\": \"file_content\", \"filetype\": \"image\", \"basename\": \"energies.png\"},\n", + " {\"name\": \"file_content\", \"filetype\": \"csv\", \"basename\": \"ranked_configurations.csv\"},\n", + " {\"name\": \"file_content\", \"filetype\": \"image\", \"basename\": \"final_energies.png\"},\n", " ],\n", " is_rendered_at_run_time=True,\n", ")\n", "configure_python_step(\n", - " final_structures_subworkflow_config, template_unit_name=\"python\", subworkflow_name=\"Final structures\",\n", - " unit_name=\"final_structures\",\n", + " final_structures_subworkflow_config, template_unit_name=\"python\", subworkflow_name=\"Save Lowest Final Structures\",\n", + " file_name=\"save_lowest_final_structures\",\n", " script=FINAL_STRUCTURES_SCRIPT, requirements=MADE_REQUIREMENTS, results=[{\"name\": \"final_structure\"}],\n", " is_rendered_at_run_time=True,\n", ")" @@ -932,17 +942,17 @@ "\n", "endpoint_options = {\"method\": \"POST\", \"headers\": {\"Content-Type\": \"application/json\"},\n", " \"data\": \"{{ FINAL_STRUCTURE_REQUESTS[FINAL_STRUCTURE_INDEX] }}\"}\n", - "save_unit = Unit(type=\"io\", subtype=\"output\", source=\"api\", enableRender=True, name=\"save-final-structure\",\n", + "save_unit = Unit(type=\"io\", subtype=\"output\", source=\"api\", enableRender=True, name=\"Save final structure\",\n", " input=[{\"type\": \"api\", \"endpoint\": \"structures\", \"endpoint_options\": endpoint_options,\n", " \"name\": \"FINAL_STRUCTURE\"}])\n", - "exit_unit = Unit(type=\"assignment\", name=\"final-structures-saved\", operand=\"NUMBER_OF_SAVED_STRUCTURES\",\n", + "exit_unit = Unit(type=\"assignment\", name=\"Count saved structures\", operand=\"NUMBER_OF_SAVED_STRUCTURES\",\n", " value=\"FINAL_STRUCTURE_INDEX\", input=[])\n", "loop_units = [\n", - " Unit(type=\"assignment\", name=\"init-index\", operand=\"FINAL_STRUCTURE_INDEX\", value=0, input=[]),\n", + " Unit(type=\"assignment\", name=\"Init structure index\", operand=\"FINAL_STRUCTURE_INDEX\", value=0, input=[]),\n", " save_unit,\n", - " Unit(type=\"assignment\", name=\"increment-index\", operand=\"FINAL_STRUCTURE_INDEX\", value=\"FINAL_STRUCTURE_INDEX + 1\",\n", - " input=[]),\n", - " Unit(**{\"type\": \"condition\", \"name\": \"check-index\",\n", + " Unit(type=\"assignment\", name=\"Increment structure index\", operand=\"FINAL_STRUCTURE_INDEX\",\n", + " value=\"FINAL_STRUCTURE_INDEX + 1\", input=[]),\n", + " Unit(**{\"type\": \"condition\", \"name\": \"Check all structures saved\",\n", " \"statement\": \"FINAL_STRUCTURE_INDEX >= len(FINAL_STRUCTURE_REQUESTS)\",\n", " \"then\": exit_unit.flowchartId, \"else\": save_unit.flowchartId, \"maxOccurrences\": NUMBER_OF_CONFIGURATIONS,\n", " \"input\": [{\"name\": \"FINAL_STRUCTURE_INDEX\", \"scope\": \"global\"},\n", @@ -1079,7 +1089,7 @@ "import pandas as pd\n", "\n", "file_url_by_name = {file[\"key\"].rsplit(\"/\", 1)[-1]: file[\"signedUrl\"] for file in client.jobs.list_files(job_id)}\n", - "pd.read_csv(io.StringIO(await read_from_url(file_url_by_name[\"results.csv\"])))" + "pd.read_csv(io.StringIO(await read_from_url(file_url_by_name[\"ranked_configurations.csv\"])))" ] }, { @@ -1097,7 +1107,7 @@ "source": [ "from IPython.display import Image\n", "\n", - "Image(data=await read_from_url(file_url_by_name[\"energies.png\"], as_bytes=True))" + "Image(data=await read_from_url(file_url_by_name[\"final_energies.png\"], as_bytes=True))" ] }, { From 15044dc03ba6544f1fd8834728135268b5aeb97d Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 18:27:40 -0800 Subject: [PATCH 17/31] SOF-8102: AFIR wording in the introduction and step 2's header; job name without a second colon Co-Authored-By: Claude Opus 5.5 --- .../afir_molecule_surface_alchemi.ipynb | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 7a79e87bf..5052ec4b9 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -8,14 +8,14 @@ "\n", "Places a molecule at random over a slab many times and relaxes every placement with MACE through the\n", "[NVIDIA ALCHEMI toolkit](https://github.com/NVIDIA/nvalchemi-toolkit), in batches on one GPU: first\n", - "under the artificial force of the artificial force induced reaction (AFIR) method between molecule and slab, then\n", - "without it. The final structures are ranked by energy, and the lowest are saved as materials.\n", + "under an artificial force between molecule and slab, as in the artificial force induced reaction (AFIR) method,\n", + "then without it. The final structures are ranked by energy, and the lowest are saved as materials.\n", "\n", - "The workflow has three steps, each in its own Python environment: *Create Initial Configurations* builds the\n", - "placements from the job's two materials with made, *Relax under AFIR Force (ALCHEMI)* relaxes them, and *Save Lowest\n", - "Final Structures* saves the lowest as the job's final structures. This notebook builds the slab, fetches the\n", - "molecule, saves both as materials, creates the workflow, runs one job on the two materials and shows the ranked\n", - "configurations and the lowest final structures.\n", + "The workflow has three steps, each in its own Python environment: *Create Initial Configurations* places the\n", + "molecule over the slab with made, from the job's two materials, *Relax under AFIR Force (ALCHEMI)* relaxes the\n", + "configurations, and *Save Lowest Final Structures* saves the lowest as the job's final structures. This notebook\n", + "builds the slab, fetches the molecule, saves both as materials, creates the workflow, runs one job on the two\n", + "materials and shows the ranked configurations and the lowest final structures.\n", "\n", "Outside the platform's JupyterLite, use Python 3.10 or 3.11 and install the packages once with `pip install \"mat3ra-notebooks-utils[all]\"`." ] @@ -416,10 +416,10 @@ "# NVIDIA ALCHEMI on one GPU. #\n", "# #\n", "# Relaxes the configurations in batches of BATCH_SIZE with the #\n", - "# slab fixed, first under the artificial force of the artificial #\n", - "# force induced reaction (AFIR) method between molecule and slab #\n", - "# (GRRM17 Eqs. 1-3 [1], strength GAMMA), then without it, and #\n", - "# ranks the final structures by energy. #\n", + "# slab fixed, first under an artificial force between molecule #\n", + "# and slab, as in the artificial force induced reaction (AFIR) #\n", + "# method (GRRM17 Eqs. 1-3 [1], strength GAMMA), then without it, #\n", + "# and ranks the final structures by energy. #\n", "# #\n", "# Reads: configurations.extxyz, written by \"Create Initial #\n", "# Configurations\" in the job's work folder; slab atoms #\n", @@ -1030,7 +1030,7 @@ " workflow=workflow,\n", " project_id=project_id,\n", " owner_id=ACCOUNT_ID,\n", - " prefix=f\"{workflow.name}: {MOLECULE_NAME} on {slab.name}, {NUMBER_OF_CONFIGURATIONS} configurations {timestamp}\",\n", + " prefix=f\"{workflow.name} {MOLECULE_NAME} on {slab.name}, {NUMBER_OF_CONFIGURATIONS} configurations {timestamp}\",\n", " compute=compute.to_dict(),\n", ")\n", "job_id = job[\"_id\"]\n", From 17ce4f270c11a4b71406b2b1403bf84b912a03c2 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 18:28:36 -0800 Subject: [PATCH 18/31] SOF-8102: step 3 named Save Final Structures; script files back to short names Co-Authored-By: Claude Opus 5.5 --- .../alchemi/afir_molecule_surface_alchemi.ipynb | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 5052ec4b9..cb1ee7bce 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -13,7 +13,7 @@ "\n", "The workflow has three steps, each in its own Python environment: *Create Initial Configurations* places the\n", "molecule over the slab with made, from the job's two materials, *Relax under AFIR Force (ALCHEMI)* relaxes the\n", - "configurations, and *Save Lowest Final Structures* saves the lowest as the job's final structures. This notebook\n", + "configurations, and *Save Final Structures* saves the lowest as the job's final structures. This notebook\n", "builds the slab, fetches the molecule, saves both as materials, creates the workflow, runs one job on the two\n", "materials and shows the ranked configurations and the lowest final structures.\n", "\n", @@ -746,7 +746,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### 4.3. Step 3: Save Lowest Final Structures\n", + "### 4.3. Step 3: Save Final Structures\n", "\n", "Takes the lowest final structures, each with the configuration it started from, and saves them as the job's final\n", "structures: each is a material, shown in this step's Results." @@ -760,7 +760,7 @@ "source": [ "FINAL_STRUCTURES_SCRIPT = r'''# ---------------------------------------------------------------- #\n", "# #\n", - "# Step 3, \"Save Lowest Final Structures\", of the workflow #\n", + "# Step 3, \"Save Final Structures\", of the workflow #\n", "# \"Molecule on Surface: AFIR Search (ALCHEMI)\", a search with the #\n", "# artificial force induced reaction (AFIR) method. #\n", "# #\n", @@ -895,14 +895,14 @@ "\n", "configure_python_step(\n", " workflow_config[\"subworkflows\"][0], template_unit_name=\"custom_script\",\n", - " subworkflow_name=\"Create Initial Configurations\", file_name=\"create_initial_configurations\",\n", + " subworkflow_name=\"Create Initial Configurations\", file_name=\"configurations\",\n", " script=set_constants(CONFIGURATIONS_SCRIPT, {\"NUMBER_OF_CONFIGURATIONS\": NUMBER_OF_CONFIGURATIONS,\n", " \"RANDOM_SEED\": RANDOM_SEED}),\n", " requirements=MADE_REQUIREMENTS, results=[], is_rendered_at_run_time=False,\n", ")\n", "configure_python_step(\n", " afir_subworkflow_config, template_unit_name=\"python\", subworkflow_name=\"Relax under AFIR Force (ALCHEMI)\",\n", - " file_name=\"relax_under_afir_force\",\n", + " file_name=\"afir\",\n", " script=set_constants(AFIR_SCRIPT, {\"GAMMA\": GAMMA, \"NUMBER_OF_FINAL_STRUCTURES\": NUMBER_OF_FINAL_STRUCTURES,\n", " \"BATCH_SIZE\": BATCH_SIZE}),\n", " requirements=AFIR_REQUIREMENTS,\n", @@ -913,8 +913,8 @@ " is_rendered_at_run_time=True,\n", ")\n", "configure_python_step(\n", - " final_structures_subworkflow_config, template_unit_name=\"python\", subworkflow_name=\"Save Lowest Final Structures\",\n", - " file_name=\"save_lowest_final_structures\",\n", + " final_structures_subworkflow_config, template_unit_name=\"python\", subworkflow_name=\"Save Final Structures\",\n", + " file_name=\"final_structures\",\n", " script=FINAL_STRUCTURES_SCRIPT, requirements=MADE_REQUIREMENTS, results=[{\"name\": \"final_structure\"}],\n", " is_rendered_at_run_time=True,\n", ")" From bad13ddc102f3ff0314377bc25fc8bcb8a3e2c25 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 20:07:15 -0800 Subject: [PATCH 19/31] SOF-8102: introduction notes the ALCHEMI 0.3.0 pre-release from a GitHub commit Co-Authored-By: Claude Opus 5.5 --- other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index cb1ee7bce..7157ed78a 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -17,6 +17,10 @@ "builds the slab, fetches the molecule, saves both as materials, creates the workflow, runs one job on the two\n", "materials and shows the ranked configurations and the lowest final structures.\n", "\n", + "**Note:** *Relax under AFIR Force (ALCHEMI)* uses the ALCHEMI 0.3.0 pre-release (nvalchemi-toolkit), installed from\n", + "GitHub commit [6189164](https://github.com/NVIDIA/nvalchemi-toolkit/commit/618916442783c152cd53b919164f07f80286f3da),\n", + "not from a PyPI release.\n", + "\n", "Outside the platform's JupyterLite, use Python 3.10 or 3.11 and install the packages once with `pip install \"mat3ra-notebooks-utils[all]\"`." ] }, From 80059872bdcddaf1eac2a4cc36504fc6caa4dc2e Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 20:08:09 -0800 Subject: [PATCH 20/31] SOF-8102: requirements explained above their cell, ALCHEMI pre-release noted there instead of the introduction Co-Authored-By: Claude Opus 5.5 --- .../afir_molecule_surface_alchemi.ipynb | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 7157ed78a..67dfd0a76 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -17,10 +17,6 @@ "builds the slab, fetches the molecule, saves both as materials, creates the workflow, runs one job on the two\n", "materials and shows the ranked configurations and the lowest final structures.\n", "\n", - "**Note:** *Relax under AFIR Force (ALCHEMI)* uses the ALCHEMI 0.3.0 pre-release (nvalchemi-toolkit), installed from\n", - "GitHub commit [6189164](https://github.com/NVIDIA/nvalchemi-toolkit/commit/618916442783c152cd53b919164f07f80286f3da),\n", - "not from a PyPI release.\n", - "\n", "Outside the platform's JupyterLite, use Python 3.10 or 3.11 and install the packages once with `pip install \"mat3ra-notebooks-utils[all]\"`." ] }, @@ -720,6 +716,22 @@ "'''" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Each step runs in its own Python environment on the cluster, installed from its list the first time and reused\n", + "afterwards.\n", + "\n", + "- `MADE_REQUIREMENTS`: steps 1 and 3 (*Create Initial Configurations*, *Save Final Structures*), the made library,\n", + " which builds the configurations and converts the structures.\n", + "- `AFIR_REQUIREMENTS`: step 2 (*Relax under AFIR Force (ALCHEMI)*), NVIDIA ALCHEMI with MACE on the GPU. **ALCHEMI here\n", + " is a pre-release: version 0.3.0 installed from GitHub commit\n", + " [6189164](https://github.com/NVIDIA/nvalchemi-toolkit/commit/618916442783c152cd53b919164f07f80286f3da), not a PyPI\n", + " release.** PyTorch is the CUDA 12.8 build and Warp the CUDA 12 build (the `--find-links` lines). Building this\n", + " environment the first time takes about 10 minutes (measured on cluster-003)." + ] + }, { "cell_type": "code", "execution_count": null, From 4a41222a7134cc788a2ee9b2d59821a023092cde Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 20:09:07 -0800 Subject: [PATCH 21/31] SOF-8102: AFIR step's requirements without the made and commented index lines Co-Authored-By: Claude Opus 5.5 --- other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb | 2 -- 1 file changed, 2 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 67dfd0a76..12d08dcb6 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -743,7 +743,6 @@ "--find-links https://download.pytorch.org/whl/cu128/torchvision/\n", "--find-links https://download.pytorch.org/whl/cu128/torchaudio/\n", "--find-links https://github.com/NVIDIA/warp/releases/expanded_assets/v1.18.0\n", - "#--extra-index-url https://download.pytorch.org/whl/cu128\n", "torch==2.11.0+cu128\n", "torchvision==0.26.0+cu128\n", "torchaudio==2.11.0+cu128\n", @@ -754,7 +753,6 @@ "ase\n", "mace-torch\n", "loguru\n", - "mat3ra-made[tools]\n", "\"\"\"" ] }, From 10fbd8681f530e451704cc6db6b7dbbfc9318a0b Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 20:27:09 -0800 Subject: [PATCH 22/31] SOF-8102: step scripts and requirements as files beside the notebook, read by one cell Co-Authored-By: Claude Opus 5.5 --- .../afir_molecule_surface_alchemi.ipynb | 569 +----------------- .../afir_molecule_surface_alchemi/afir.py | 304 ++++++++++ .../afir_requirements.txt | 14 + .../configurations.py | 125 ++++ .../configurations_requirements.txt | 1 + .../final_structures.py | 61 ++ .../final_structures_requirements.txt | 1 + 7 files changed, 530 insertions(+), 545 deletions(-) create mode 100644 other/experiments/alchemi/afir_molecule_surface_alchemi/afir.py create mode 100644 other/experiments/alchemi/afir_molecule_surface_alchemi/afir_requirements.txt create mode 100644 other/experiments/alchemi/afir_molecule_surface_alchemi/configurations.py create mode 100644 other/experiments/alchemi/afir_molecule_surface_alchemi/configurations_requirements.txt create mode 100644 other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures.py create mode 100644 other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures_requirements.txt diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 12d08dcb6..822f97455 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -255,140 +255,6 @@ "atoms 1, each with its index, rotation and position." ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "CONFIGURATIONS_SCRIPT = r'''# ---------------------------------------------------------------- #\n", - "# #\n", - "# Step 1, \"Create Initial Configurations\", of the workflow #\n", - "# \"Molecule on Surface: AFIR Search (ALCHEMI)\", a search with the #\n", - "# artificial force induced reaction (AFIR) method. #\n", - "# #\n", - "# Places the job's molecule over its slab #\n", - "# NUMBER_OF_CONFIGURATIONS times with made, each at a random #\n", - "# orientation and a random point of the surface cell, as #\n", - "# GRRM17's random initial structures [1]. #\n", - "# #\n", - "# Reads: the job's two materials, written into this script by #\n", - "# the platform when the job is created; the non-periodic #\n", - "# one is the molecule. #\n", - "# Writes: configurations.extxyz in the job's work folder: the #\n", - "# placed structures, slab atoms tagged 0 and molecule #\n", - "# atoms 1, each with its index in the info key #\n", - "# configuration and its rotation and fractional position #\n", - "# in metadata.placement. It is the only output of the #\n", - "# script. #\n", - "# #\n", - "# [1] S. Maeda et al., J. Comput. Chem. 39, 233 (2018). #\n", - "# #\n", - "# ---------------------------------------------------------------- #\n", - "\n", - "import json\n", - "\n", - "import numpy as np\n", - "from ase.data import atomic_numbers, covalent_radii\n", - "from ase.io import write\n", - "from mat3ra.made.material import Material\n", - "from mat3ra.made.tools.convert import to_ase\n", - "from mat3ra.made.tools.convert.interface_parts_enum import InterfacePartsEnum\n", - "from mat3ra.made.tools.modify import translate_by_vector\n", - "from mat3ra.made.tools.operations.core.binary import merge\n", - "from mat3ra.made.tools.operations.core.unary import rotate\n", - "\n", - "NUMBER_OF_CONFIGURATIONS = 100\n", - "RANDOM_SEED = 0\n", - "CONTACT_OFFSET = 0.5 # Å, lowest molecule atom at R_molecule + R_slab + CONTACT_OFFSET above the top slab atom\n", - "# The job's materials, written here by the platform when the job is created.\n", - "MATERIALS = json.loads(r\"\"\"{{ materials | dump | safe }}\"\"\")\n", - "\n", - "\n", - "def extreme_atom_height_and_radius(material, extremum):\n", - " \"\"\"\n", - " Returns the height and the covalent radius of the highest or the lowest atom of a material.\n", - "\n", - " Args:\n", - " material (Material): The material to search.\n", - " extremum (Callable): max for the highest atom, min for the lowest.\n", - "\n", - " Returns:\n", - " Tuple[float, float]: The atom's cartesian z and its covalent radius, both in Å.\n", - " \"\"\"\n", - " cartesian = material.clone()\n", - " cartesian.to_cartesian()\n", - " height, element = extremum(zip([coordinate[2] for coordinate in cartesian.basis.coordinates.values],\n", - " cartesian.basis.elements.values))\n", - " return height, covalent_radii[atomic_numbers[element]]\n", - "\n", - "\n", - "def create_configuration(slab, molecule, top_atom_height_and_radius, random_generator, index):\n", - " \"\"\"\n", - " Places the molecule once over the slab, as a random initial structure of GRRM17's MC-AFIR.\n", - "\n", - " The orientation is uniformly random (a normalised Gaussian quaternion), the molecule's center goes to a random\n", - " point of the surface cell, and its lowest atom sits R_molecule + R_slab + CONTACT_OFFSET above the top slab atom,\n", - " the RandDC rule of GRRM17 (S. Maeda et al., J. Comput. Chem. 39, 233 (2018)).\n", - "\n", - " Args:\n", - " slab (Material): The slab, in crystal units, labelled 0.\n", - " molecule (Material): The molecule, in crystal units.\n", - " top_atom_height_and_radius (Tuple[float, float]): Height and covalent radius of the top slab atom, in Å.\n", - " random_generator (numpy.random.Generator): The seeded generator of all placements.\n", - " index (int): The configuration's index, used in its name.\n", - "\n", - " Returns:\n", - " MaterialWithBuildMetadata: The slab and the placed molecule, molecule atoms labelled 1, with the rotation and\n", - " the fractional position in metadata.placement.\n", - " \"\"\"\n", - " quaternion = random_generator.normal(size=4)\n", - " axis = (quaternion[1:] / np.linalg.norm(quaternion[1:])).tolist()\n", - " angle = float(2 * np.degrees(np.arccos(quaternion[0] / np.linalg.norm(quaternion))))\n", - " rotated = rotate(molecule, axis=axis, angle=angle, wrap=False)\n", - " rotated.set_lattice(slab.lattice)\n", - " top_height, top_radius = top_atom_height_and_radius\n", - " lowest_height, lowest_radius = extreme_atom_height_and_radius(rotated, min)\n", - " surface_fractions = [random_generator.random(), random_generator.random(), 0]\n", - " surface_point = slab.basis.cell.convert_point_to_cartesian(surface_fractions)\n", - " molecule_cell_center = molecule.basis.cell.convert_point_to_cartesian([0.5, 0.5, 0.5])\n", - " vertical_shift = top_height + top_radius + lowest_radius + CONTACT_OFFSET - lowest_height\n", - " placed = translate_by_vector(rotated, vector=[surface_point[0] - molecule_cell_center[0],\n", - " surface_point[1] - molecule_cell_center[1], vertical_shift],\n", - " use_cartesian_coordinates=True)\n", - " placed.set_labels_from_value(InterfacePartsEnum.FILM.value)\n", - " configuration = merge([slab, placed], material_name=f\"{molecule.name} on {slab.name}, configuration {index}\")\n", - " configuration.metadata.placement = {\"rotation_axis\": axis, \"rotation_angle_deg\": angle,\n", - " \"position\": surface_fractions[:2]}\n", - " return configuration\n", - "\n", - "\n", - "def main():\n", - " \"\"\"\n", - " Builds NUMBER_OF_CONFIGURATIONS configurations from the job's materials and writes them to configurations.extxyz.\n", - " \"\"\"\n", - " slab = Material.create(next(material_config for material_config in MATERIALS\n", - " if not material_config.get(\"isNonPeriodic\")))\n", - " molecule = Material.create(next(material_config for material_config in MATERIALS\n", - " if material_config.get(\"isNonPeriodic\")))\n", - " for material in (slab, molecule):\n", - " material.to_crystal()\n", - " slab.set_labels_from_value(InterfacePartsEnum.SUBSTRATE.value)\n", - " top_atom_height_and_radius = extreme_atom_height_and_radius(slab, max)\n", - " random_generator = np.random.default_rng(RANDOM_SEED)\n", - " configurations = []\n", - " for index in range(NUMBER_OF_CONFIGURATIONS):\n", - " atoms = to_ase(create_configuration(slab, molecule, top_atom_height_and_radius, random_generator, index))\n", - " atoms.info[\"configuration\"] = index\n", - " configurations.append(atoms)\n", - " write(\"configurations.extxyz\", configurations)\n", - "\n", - "\n", - "if __name__ == \"__main__\":\n", - " main()\n", - "'''" - ] - }, { "cell_type": "markdown", "metadata": {}, @@ -400,333 +266,32 @@ "without it. It writes `ranked_configurations.csv` with every configuration ranked by final energy, its starting\n", "orientation and the fragments its molecule ends in; `final_structures.extxyz` with the `NUMBER_OF_FINAL_STRUCTURES`\n", "lowest; and `final_energies.png`, the final energy of each configuration. Another backend or another ALCHEMI version\n", - "replaces this script and its packages; steps 1 and 3 stay." + "replaces `afir.py` and `afir_requirements.txt`; steps 1 and 3 stay." ] }, { - "cell_type": "code", - "execution_count": null, + "cell_type": "markdown", "metadata": {}, - "outputs": [], "source": [ - "AFIR_SCRIPT = r'''# ---------------------------------------------------------------- #\n", - "# #\n", - "# Step 2, \"Relax under AFIR Force (ALCHEMI)\", of the workflow #\n", - "# \"Molecule on Surface: AFIR Search (ALCHEMI)\", with MACE through #\n", - "# NVIDIA ALCHEMI on one GPU. #\n", - "# #\n", - "# Relaxes the configurations in batches of BATCH_SIZE with the #\n", - "# slab fixed, first under an artificial force between molecule #\n", - "# and slab, as in the artificial force induced reaction (AFIR) #\n", - "# method (GRRM17 Eqs. 1-3 [1], strength GAMMA), then without it, #\n", - "# and ranks the final structures by energy. #\n", - "# #\n", - "# Reads: configurations.extxyz, written by \"Create Initial #\n", - "# Configurations\" in the job's work folder; slab atoms #\n", - "# tagged 0, molecule atoms 1. #\n", - "# Writes: ranked_configurations.csv, every configuration ranked #\n", - "# by final energy, with its starting rotation and #\n", - "# position and the number of fragments its molecule ends #\n", - "# in; final_structures.extxyz, the #\n", - "# NUMBER_OF_FINAL_STRUCTURES lowest in that order; and #\n", - "# final_energies.png, final energy against configuration #\n", - "# index. ranked_configurations.csv and final_energies.png #\n", - "# are results of this unit. #\n", - "# #\n", - "# [1] S. Maeda et al., J. Comput. Chem. 39, 233 (2018). #\n", - "# #\n", - "# ---------------------------------------------------------------- #\n", - "\n", - "import csv\n", - "from importlib.metadata import version\n", - "\n", - "import matplotlib\n", - "import numpy as np\n", - "import torch\n", - "from ase.build import separate\n", - "from ase.calculators.singlepoint import SinglePointCalculator\n", - "from ase.data import covalent_radii\n", - "from ase.io import read, write\n", - "from ase.units import kJ, mol\n", - "from loguru import logger\n", - "from nvalchemi._typing import AtomCategory\n", - "from nvalchemi.data import AtomicData, Batch\n", - "from nvalchemi.dynamics import FIRE2, FusedStage\n", - "from nvalchemi.dynamics.base import ConvergenceHook, DynamicsStage\n", - "from nvalchemi.dynamics.hooks import FreezeAtomsHook\n", - "from nvalchemi.hooks import BiasedPotentialHook\n", - "from nvalchemi.models.mace import MACEWrapper\n", - "\n", - "matplotlib.use(\"Agg\")\n", - "import matplotlib.pyplot as plt\n", - "\n", - "DEVICE = \"cuda:0\"\n", - "CHECKPOINT = \"medium-mpa-0\" # MACE-MPA-0, medium\n", - "DTYPE = torch.float32\n", - "FMAX = 0.03 # eV/Å, largest atomic force at convergence\n", - "MAX_STEPS = 500 # FIRE2 steps per relaxation at most\n", - "BATCH_SIZE = 100 # structures relaxed together on the GPU\n", - "TIME_STEP = 0.5 # initial FIRE2 step, in ALCHEMI's internal time unit (10.18 fs)\n", - "GAMMA = 100 # kJ/mol, model collision energy of the AFIR force\n", - "NUMBER_OF_FINAL_STRUCTURES = 1 # the lowest final structures kept in final_structures.extxyz\n", - "BOND_TOLERANCE = 1.2 # two molecule atoms are bonded within BOND_TOLERANCE * (R_i + R_j)\n", - "SLAB_TAG = 0 # atom tags in configurations.extxyz\n", - "MOLECULE_TAG = 1\n", - "AFIR_R0 = 3.8164 # Å, GRRM17 Eq. (3)\n", - "AFIR_EPSILON = 1.0061 # kJ/mol, GRRM17 Eq. (3)\n", - "\n", - "ALPHA = GAMMA * kJ / mol / ((2 ** (-1 / 6) - (1 + (1 + GAMMA / AFIR_EPSILON) ** 0.5) ** (-1 / 6)) * AFIR_R0) # Eq. (3)\n", - "CATEGORY_BY_TAG = {SLAB_TAG: AtomCategory.SURFACE.value, MOLECULE_TAG: AtomCategory.GAS.value}\n", - "IN_PLANE_MASK = torch.tensor([1.0, 1.0, 0.0], device=DEVICE)\n", - "IN_PLANE_SHIFTS = torch.tensor([[a, b, 0.0] for a in (-1, 0, 1) for b in (-1, 0, 1)], device=DEVICE)\n", - "COVALENT_RADII = torch.tensor(covalent_radii, device=DEVICE, dtype=DTYPE)\n", - "\n", - "\n", - "def in_plane_minimum_image(vectors, cells):\n", - " \"\"\"\n", - " Returns the shortest periodic image of each vector in the surface plane, a and b; there is no image along c.\n", - "\n", - " Args:\n", - " vectors (torch.Tensor): Cartesian vectors, [P, 3], in Å.\n", - " cells (torch.Tensor): The lattice vectors of each vector's structure, [P, 3, 3], rows a, b, c.\n", - "\n", - " Returns:\n", - " torch.Tensor: The minimum-image vectors, [P, 3], in Å.\n", - " \"\"\"\n", - " fractional = torch.einsum(\"pi,pij->pj\", vectors, torch.linalg.inv(cells))\n", - " fractional = fractional - torch.round(fractional) * IN_PLANE_MASK\n", - " images = torch.einsum(\"psi,pij->psj\", fractional.unsqueeze(1) + IN_PLANE_SHIFTS, cells)\n", - " nearest = images.norm(dim=-1).argmin(dim=1)\n", - " return torch.take_along_dim(images, nearest[:, None, None], dim=1).squeeze(1)\n", - "\n", - "\n", - "def afir_energy(positions, batch):\n", - " \"\"\"\n", - " Returns the AFIR term of every structure in the batch, between its molecule and its slab atoms.\n", - "\n", - " GRRM17 Eq. (1), α Σ ω_ij r_ij / Σ ω_ij over molecule atoms i and slab atoms j, with ω_ij = ((R_i + R_j) / r_ij)^6\n", - " (Eq. 2), R the covalent radii and r_ij at the minimum image in the surface plane; α is ALPHA, Eq. (3)\n", - " (S. Maeda et al., J. Comput. Chem. 39, 233 (2018)).\n", - "\n", - " Args:\n", - " positions (torch.Tensor): Atomic positions of the batch, [V, 3], in Å.\n", - " batch (Batch): The batch, with atom_categories GAS for the molecule and SURFACE for the slab.\n", - "\n", - " Returns:\n", - " torch.Tensor: The AFIR term of each structure, [B], in eV.\n", - " \"\"\"\n", - " molecule_atoms = torch.where(batch.atom_categories == AtomCategory.GAS.value)[0]\n", - " surface_atoms = torch.where(batch.atom_categories == AtomCategory.SURFACE.value)[0]\n", - " molecule_index, surface_index = torch.where(\n", - " batch.batch_idx[molecule_atoms][:, None] == batch.batch_idx[surface_atoms][None, :])\n", - " pair_molecule_atoms, pair_surface_atoms = molecule_atoms[molecule_index], surface_atoms[surface_index]\n", - " pair_structures = batch.batch_idx[pair_molecule_atoms]\n", - " distances = in_plane_minimum_image(positions[pair_molecule_atoms] - positions[pair_surface_atoms],\n", - " batch.cell[pair_structures]).norm(dim=-1)\n", - " radii = (COVALENT_RADII[batch.atomic_numbers[pair_molecule_atoms]]\n", - " + COVALENT_RADII[batch.atomic_numbers[pair_surface_atoms]])\n", - " weights = (radii / distances) ** 6\n", - " numerator = torch.zeros(batch.num_graphs, device=positions.device, dtype=positions.dtype)\n", - " denominator = torch.zeros_like(numerator)\n", - " return ALPHA * (numerator.index_add(0, pair_structures, weights * distances)\n", - " / denominator.index_add(0, pair_structures, weights))\n", - "\n", - "\n", - "def afir_bias(batch):\n", - " \"\"\"\n", - " Returns the AFIR energy and forces of the batch, as BiasedPotentialHook expects them.\n", - "\n", - " Args:\n", - " batch (Batch): The batch after the model's forward pass.\n", - "\n", - " Returns:\n", - " Tuple[torch.Tensor, torch.Tensor]: Energy per structure, [B, 1], in eV, and force per atom, [V, 3],\n", - " in eV/Å.\n", - " \"\"\"\n", - " positions = batch.positions.clone().requires_grad_(True)\n", - " energy = afir_energy(positions, batch)\n", - " forces = -torch.autograd.grad(energy.sum(), positions)[0]\n", - " return energy.detach().unsqueeze(-1), forces\n", - "\n", - "\n", - "def atomic_data(atoms):\n", - " \"\"\"\n", - " Converts a configuration to ALCHEMI's AtomicData, its tags to atom categories: SLAB_TAG to SURFACE, MOLECULE_TAG\n", - " to GAS.\n", - "\n", - " Args:\n", - " atoms (ase.Atoms): The configuration, slab atoms tagged 0 and molecule atoms 1.\n", - "\n", - " Returns:\n", - " AtomicData: The configuration on DEVICE, with zeroed forces and energy.\n", - " \"\"\"\n", - " data = AtomicData.from_atoms(atoms, device=DEVICE, dtype=DTYPE)\n", - " data[\"atom_categories\"] = torch.tensor([CATEGORY_BY_TAG[tag] for tag in atoms.get_tags()], device=DEVICE)\n", - " data[\"forces\"] = torch.zeros_like(data.positions)\n", - " data[\"energy\"] = torch.zeros((1, 1), device=DEVICE, dtype=DTYPE)\n", - " return data\n", - "\n", - "\n", - "def relaxed_atoms(original, relaxed_data):\n", - " \"\"\"\n", - " Returns a copy of a structure at its relaxed positions, with its energy and forces.\n", - "\n", - " Args:\n", - " original (ase.Atoms): The structure before the relaxation.\n", - " relaxed_data (AtomicData): The same structure after it.\n", - "\n", - " Returns:\n", - " ase.Atoms: The relaxed structure, energy and forces on a SinglePointCalculator.\n", - " \"\"\"\n", - " atoms = original.copy()\n", - " atoms.positions = relaxed_data.positions.cpu().numpy()\n", - " atoms.calc = SinglePointCalculator(atoms, energy=float(relaxed_data.energy.item()),\n", - " forces=relaxed_data.forces.cpu().numpy())\n", - " return atoms\n", - "\n", - "\n", - "def relax_batch(ase_structures, model, bias_hooks):\n", - " \"\"\"\n", - " Relaxes the structures as one batch with fixed-cell FIRE2, slab atoms frozen.\n", - "\n", - " Args:\n", - " ase_structures (List[ase.Atoms]): The structures, slab atoms tagged 0 and molecule atoms 1.\n", - " model (MACEWrapper): The MACE model.\n", - " bias_hooks (List[BiasedPotentialHook]): Bias potentials added after each forward pass; [] for none.\n", - "\n", - " Returns:\n", - " List[ase.Atoms]: The relaxed structures, in the input order.\n", - " \"\"\"\n", - " optimizer = FIRE2(model=model, dt=TIME_STEP, n_steps=MAX_STEPS, convergence_hook=ConvergenceHook(\n", - " criteria=[{\"key\": \"forces\", \"threshold\": FMAX, \"reduce_op\": \"norm\", \"reduce_dims\": -1}]))\n", - " optimizer.register_hook(FreezeAtomsHook(freeze_category=AtomCategory.SURFACE.value))\n", - " stage = FusedStage(sub_stages=[(0, optimizer)])\n", - " for hook in model.make_neighbor_hooks():\n", - " stage.register_hook(hook, stage=DynamicsStage.BEFORE_COMPUTE)\n", - " for hook in bias_hooks:\n", - " stage.register_hook(hook, stage=DynamicsStage.AFTER_COMPUTE)\n", - " batch = stage.run(Batch.from_data_list([atomic_data(atoms) for atoms in ase_structures]))\n", - " return [relaxed_atoms(atoms, relaxed_data) for atoms, relaxed_data in zip(ase_structures, batch.to_data_list())]\n", - "\n", - "\n", - "def relax(ase_structures, model, bias_hooks):\n", - " \"\"\"\n", - " Relaxes all structures in batches of BATCH_SIZE, one after another, so the GPU memory a batch needs stays bounded.\n", - "\n", - " Args:\n", - " ase_structures (List[ase.Atoms]): The structures, slab atoms tagged 0 and molecule atoms 1.\n", - " model (MACEWrapper): The MACE model.\n", - " bias_hooks (List[BiasedPotentialHook]): Bias potentials added after each forward pass; [] for none.\n", - "\n", - " Returns:\n", - " List[ase.Atoms]: The relaxed structures, in the input order.\n", - " \"\"\"\n", - " relaxed = []\n", - " for start in range(0, len(ase_structures), BATCH_SIZE):\n", - " end = min(start + BATCH_SIZE, len(ase_structures))\n", - " logger.info(\"Relaxing structures {}–{} of {}\", start + 1, end, len(ase_structures))\n", - " relaxed += relax_batch(ase_structures[start:end], model, bias_hooks)\n", - " return relaxed\n", - "\n", - "\n", - "def molecule_fragments(atoms):\n", - " \"\"\"\n", - " Returns the number of fragments the molecule of a structure is in: 1 if intact, more if it split.\n", - "\n", - " Args:\n", - " atoms (ase.Atoms): The structure, molecule atoms tagged 1.\n", - "\n", - " Returns:\n", - " int: The number of connected groups of molecule atoms, bonded within BOND_TOLERANCE times their covalent radii.\n", - " \"\"\"\n", - " return len(separate(atoms[atoms.get_tags() == MOLECULE_TAG], scale=BOND_TOLERANCE))\n", - "\n", - "\n", - "def energy_and_max_force(atoms):\n", - " \"\"\"\n", - " Returns the energy of a structure, in eV, and the largest force on any of its atoms, in eV/Å.\n", - " \"\"\"\n", - " return atoms.get_potential_energy(), np.linalg.norm(atoms.get_forces(), axis=1).max()\n", - "\n", - "\n", - "def plot_energies(energies, lowest_indices):\n", - " \"\"\"\n", - " Plots the final energy of every configuration against its index, the kept lowest marked, to final_energies.png.\n", - "\n", - " Args:\n", - " energies (List[float]): The final energy of each configuration, in eV.\n", - " lowest_indices (List[int]): The indices of the kept lowest structures.\n", - " \"\"\"\n", - " figure, axes = plt.subplots()\n", - " axes.scatter(range(len(energies)), energies, s=12, label=\"all configurations\")\n", - " axes.scatter(lowest_indices, [energies[index] for index in lowest_indices], s=40, color=\"tab:red\",\n", - " label=f\"{len(lowest_indices)} lowest, kept\")\n", - " axes.set(xlabel=\"Configuration index\", ylabel=\"Final energy (eV)\")\n", - " axes.ticklabel_format(axis=\"y\", useOffset=False)\n", - " axes.legend()\n", - " figure.savefig(\"final_energies.png\", dpi=150, bbox_inches=\"tight\")\n", - "\n", - "\n", - "def write_results(afir_structures, final_structures):\n", - " \"\"\"\n", - " Ranks the configurations by final energy, writes all of them to ranked_configurations.csv with the placement each\n", - " started from (atoms.info[\"metadata\"][\"placement\"]) and its molecule's fragments, the NUMBER_OF_FINAL_STRUCTURES\n", - " lowest to final_structures.extxyz in that order, and the energies to final_energies.png.\n", - "\n", - " Args:\n", - " afir_structures (List[ase.Atoms]): The structures after the relaxation under the AFIR force.\n", - " final_structures (List[ase.Atoms]): The structures after the relaxation without it.\n", - " \"\"\"\n", - " energies = [atoms.get_potential_energy() for atoms in final_structures]\n", - " order = np.argsort(energies)\n", - " with open(\"ranked_configurations.csv\", \"w\", newline=\"\") as results_file:\n", - " writer = csv.writer(results_file)\n", - " writer.writerow([\"configuration\", \"rank\", \"afir_function_eV\", \"afir_max_force_eV_per_A\", \"energy_eV\",\n", - " \"max_force_eV_per_A\", \"rotation_axis_x\", \"rotation_axis_y\", \"rotation_axis_z\",\n", - " \"rotation_angle_deg\", \"position_u\", \"position_v\", \"molecule_fragments\"])\n", - " for rank, index in enumerate(order, start=1):\n", - " placement = final_structures[index].info[\"metadata\"][\"placement\"]\n", - " writer.writerow([index, rank, *energy_and_max_force(afir_structures[index]),\n", - " *energy_and_max_force(final_structures[index]), *placement[\"rotation_axis\"],\n", - " placement[\"rotation_angle_deg\"], *placement[\"position\"],\n", - " molecule_fragments(final_structures[index])])\n", - " final_structures[index].info.update(rank=rank)\n", - " lowest_indices = order[:NUMBER_OF_FINAL_STRUCTURES].tolist()\n", - " write(\"final_structures.extxyz\", [final_structures[index] for index in lowest_indices])\n", - " plot_energies(energies, lowest_indices)\n", - "\n", - "\n", - "def main():\n", - " \"\"\"\n", - " Relaxes the configurations of configurations.extxyz under the AFIR force, then without it, and writes the ranked\n", - " results.\n", - " \"\"\"\n", - " logger.info(\"AFIR force α, γ = {} kJ/mol: {:.3f} eV/Å\", GAMMA, ALPHA)\n", - " ase_structures = read(\"{{ JOB_WORK_DIR }}/configurations.extxyz\", index=\":\")\n", - " model = MACEWrapper.from_checkpoint(checkpoint_path=CHECKPOINT, device=torch.device(DEVICE), dtype=DTYPE,\n", - " enable_cueq=True)\n", - " logger.info(\"nvalchemi-toolkit {}, MACE with cuEquivariance\", version(\"nvalchemi-toolkit\"))\n", - " afir_structures = relax(ase_structures, model, [BiasedPotentialHook(bias_fn=afir_bias)])\n", - " final_structures = relax(afir_structures, model, [])\n", - " write_results(afir_structures, final_structures)\n", - "\n", + "### 4.3. Step 3: Save Final Structures\n", "\n", - "if __name__ == \"__main__\":\n", - " main()\n", - "'''" + "Takes the lowest final structures, each with the configuration it started from, and saves them as the job's final\n", + "structures: each is a material, shown in this step's Results." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "Each step runs in its own Python environment on the cluster, installed from its list the first time and reused\n", - "afterwards.\n", + "Each step's script and requirements are files in `afir_molecule_surface_alchemi/` beside this notebook, under the\n", + "names the workflow gives them on the platform; the cell below reads them. Each step runs in its own Python environment\n", + "on the cluster, installed from its requirements the first time and reused afterwards.\n", "\n", - "- `MADE_REQUIREMENTS`: steps 1 and 3 (*Create Initial Configurations*, *Save Final Structures*), the made library,\n", - " which builds the configurations and converts the structures.\n", - "- `AFIR_REQUIREMENTS`: step 2 (*Relax under AFIR Force (ALCHEMI)*), NVIDIA ALCHEMI with MACE on the GPU. **ALCHEMI here\n", - " is a pre-release: version 0.3.0 installed from GitHub commit\n", + "- `configurations_requirements.txt` and `final_structures_requirements.txt`: steps 1 and 3 (*Create Initial\n", + " Configurations*, *Save Final Structures*), the made library, which builds the configurations and converts the\n", + " structures.\n", + "- `afir_requirements.txt`: step 2 (*Relax under AFIR Force (ALCHEMI)*), NVIDIA ALCHEMI with MACE on the GPU. **ALCHEMI\n", + " here is a pre-release: version 0.3.0 installed from GitHub commit\n", " [6189164](https://github.com/NVIDIA/nvalchemi-toolkit/commit/618916442783c152cd53b919164f07f80286f3da), not a PyPI\n", " release.** PyTorch is the CUDA 12.8 build and Warp the CUDA 12 build (the `--find-links` lines). Building this\n", " environment the first time takes about 10 minutes (measured on cluster-003)." @@ -738,102 +303,15 @@ "metadata": {}, "outputs": [], "source": [ - "MADE_REQUIREMENTS = \"mat3ra-made[tools]==2026.10.10.post0\\n\"\n", - "AFIR_REQUIREMENTS = \"\"\"--find-links https://download.pytorch.org/whl/cu128/torch/\n", - "--find-links https://download.pytorch.org/whl/cu128/torchvision/\n", - "--find-links https://download.pytorch.org/whl/cu128/torchaudio/\n", - "--find-links https://github.com/NVIDIA/warp/releases/expanded_assets/v1.18.0\n", - "torch==2.11.0+cu128\n", - "torchvision==0.26.0+cu128\n", - "torchaudio==2.11.0+cu128\n", - "warp-lang==1.18.0+cu12\n", - "nvalchemi-toolkit[cu12,mace,ase] @ git+https://github.com/NVIDIA/nvalchemi-toolkit.git@618916442783c152cd53b919164f07f80286f3da\n", - "matplotlib\n", - "numpy\n", - "ase\n", - "mace-torch\n", - "loguru\n", - "\"\"\"" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 4.3. Step 3: Save Final Structures\n", - "\n", - "Takes the lowest final structures, each with the configuration it started from, and saves them as the job's final\n", - "structures: each is a material, shown in this step's Results." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "FINAL_STRUCTURES_SCRIPT = r'''# ---------------------------------------------------------------- #\n", - "# #\n", - "# Step 3, \"Save Final Structures\", of the workflow #\n", - "# \"Molecule on Surface: AFIR Search (ALCHEMI)\", a search with the #\n", - "# artificial force induced reaction (AFIR) method. #\n", - "# #\n", - "# Takes the lowest final structures of \"Relax under AFIR Force #\n", - "# (ALCHEMI)\", each with the configuration it started from, and #\n", - "# makes one request body per structure for the platform's #\n", - "# final-structure endpoint, which the io unit after this one #\n", - "# sends in a loop. Each request names this unit (UNIT_ID), so the #\n", - "# structures are this step's results. #\n", - "# #\n", - "# Reads: final_structures.extxyz and configurations.extxyz in #\n", - "# the job's work folder. #\n", - "# Writes: FINAL_STRUCTURE_REQUESTS to the job's global scope, as #\n", - "# the only output of the script. #\n", - "# #\n", - "# ---------------------------------------------------------------- #\n", - "\n", - "import json\n", - "\n", - "from ase.io import read\n", - "from mat3ra.made.material import Material\n", - "from mat3ra.made.tools.convert import from_ase\n", - "\n", - "\n", - "def request_body(rank_index, atoms, configurations):\n", - " \"\"\"\n", - " Returns the body of one final-structure request, escaped to sit inside the next io unit's JSON string.\n", - "\n", - " The platform keeps one final_structure per unit and repetition, so each structure gets its own repetition.\n", - "\n", - " Args:\n", - " rank_index (int): The structure's rank by final energy, from 0.\n", - " atoms (ase.Atoms): The final structure, with its configuration's index in atoms.info.\n", - " configurations (List[ase.Atoms]): All configurations, in the order of their indices.\n", - "\n", - " Returns:\n", - " str: {\"jobId\", \"unitId\", \"structures\": [{\"initial\", \"final\"}]} as JSON, escaped for a JSON string.\n", - " \"\"\"\n", - " final = {**Material.create(from_ase(atoms)).to_dict(), \"repetition\": rank_index}\n", - " initial = Material.create(from_ase(configurations[atoms.info[\"configuration\"]])).to_dict()\n", - " body = {\"jobId\": \"{{ JOB_ID }}\", \"unitId\": \"{{ UNIT_ID }}\", \"structures\": [{\"initial\": initial, \"final\": final}]}\n", - " return json.dumps(json.dumps(body))[1:-1]\n", - "\n", - "\n", - "def main():\n", - " \"\"\"\n", - " Reads the lowest final structures and the configurations, and prints the request bodies as\n", - " FINAL_STRUCTURE_REQUESTS.\n", - " \"\"\"\n", - " final_structures = read(\"{{ JOB_WORK_DIR }}/final_structures.extxyz\", index=\":\")\n", - " configurations = read(\"{{ JOB_WORK_DIR }}/configurations.extxyz\", index=\":\")\n", - " request_bodies = [request_body(rank_index, atoms, configurations)\n", - " for rank_index, atoms in enumerate(final_structures)]\n", - " print(json.dumps({\"name\": \"FINAL_STRUCTURE_REQUESTS\", \"value\": request_bodies, \"scope\": \"global\"}))\n", - "\n", + "from pathlib import Path\n", "\n", - "if __name__ == \"__main__\":\n", - " main()\n", - "'''" + "STEP_FILES_DIRECTORY = Path(\"afir_molecule_surface_alchemi\")\n", + "CONFIGURATIONS_SCRIPT = (STEP_FILES_DIRECTORY / \"configurations.py\").read_text()\n", + "CONFIGURATIONS_REQUIREMENTS = (STEP_FILES_DIRECTORY / \"configurations_requirements.txt\").read_text()\n", + "AFIR_SCRIPT = (STEP_FILES_DIRECTORY / \"afir.py\").read_text()\n", + "AFIR_REQUIREMENTS = (STEP_FILES_DIRECTORY / \"afir_requirements.txt\").read_text()\n", + "FINAL_STRUCTURES_SCRIPT = (STEP_FILES_DIRECTORY / \"final_structures.py\").read_text()\n", + "FINAL_STRUCTURES_REQUIREMENTS = (STEP_FILES_DIRECTORY / \"final_structures_requirements.txt\").read_text()" ] }, { @@ -912,7 +390,7 @@ " subworkflow_name=\"Create Initial Configurations\", file_name=\"configurations\",\n", " script=set_constants(CONFIGURATIONS_SCRIPT, {\"NUMBER_OF_CONFIGURATIONS\": NUMBER_OF_CONFIGURATIONS,\n", " \"RANDOM_SEED\": RANDOM_SEED}),\n", - " requirements=MADE_REQUIREMENTS, results=[], is_rendered_at_run_time=False,\n", + " requirements=CONFIGURATIONS_REQUIREMENTS, results=[], is_rendered_at_run_time=False,\n", ")\n", "configure_python_step(\n", " afir_subworkflow_config, template_unit_name=\"python\", subworkflow_name=\"Relax under AFIR Force (ALCHEMI)\",\n", @@ -929,7 +407,8 @@ "configure_python_step(\n", " final_structures_subworkflow_config, template_unit_name=\"python\", subworkflow_name=\"Save Final Structures\",\n", " file_name=\"final_structures\",\n", - " script=FINAL_STRUCTURES_SCRIPT, requirements=MADE_REQUIREMENTS, results=[{\"name\": \"final_structure\"}],\n", + " script=FINAL_STRUCTURES_SCRIPT, requirements=FINAL_STRUCTURES_REQUIREMENTS,\n", + " results=[{\"name\": \"final_structure\"}],\n", " is_rendered_at_run_time=True,\n", ")" ] diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/afir.py b/other/experiments/alchemi/afir_molecule_surface_alchemi/afir.py new file mode 100644 index 000000000..675d5cbf8 --- /dev/null +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi/afir.py @@ -0,0 +1,304 @@ +# ---------------------------------------------------------------- # +# # +# Step 2, "Relax under AFIR Force (ALCHEMI)", of the workflow # +# "Molecule on Surface: AFIR Search (ALCHEMI)", with MACE through # +# NVIDIA ALCHEMI on one GPU. # +# # +# Relaxes the configurations in batches of BATCH_SIZE with the # +# slab fixed, first under an artificial force between molecule # +# and slab, as in the artificial force induced reaction (AFIR) # +# method (GRRM17 Eqs. 1-3 [1], strength GAMMA), then without it, # +# and ranks the final structures by energy. # +# # +# Reads: configurations.extxyz, written by "Create Initial # +# Configurations" in the job's work folder; slab atoms # +# tagged 0, molecule atoms 1. # +# Writes: ranked_configurations.csv, every configuration ranked # +# by final energy, with its starting rotation and # +# position and the number of fragments its molecule ends # +# in; final_structures.extxyz, the # +# NUMBER_OF_FINAL_STRUCTURES lowest in that order; and # +# final_energies.png, final energy against configuration # +# index. ranked_configurations.csv and final_energies.png # +# are results of this unit. # +# # +# [1] S. Maeda et al., J. Comput. Chem. 39, 233 (2018). # +# # +# ---------------------------------------------------------------- # + +import csv +from importlib.metadata import version + +import matplotlib +import numpy as np +import torch +from ase.build import separate +from ase.calculators.singlepoint import SinglePointCalculator +from ase.data import covalent_radii +from ase.io import read, write +from ase.units import kJ, mol +from loguru import logger +from nvalchemi._typing import AtomCategory +from nvalchemi.data import AtomicData, Batch +from nvalchemi.dynamics import FIRE2, FusedStage +from nvalchemi.dynamics.base import ConvergenceHook, DynamicsStage +from nvalchemi.dynamics.hooks import FreezeAtomsHook +from nvalchemi.hooks import BiasedPotentialHook +from nvalchemi.models.mace import MACEWrapper + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +DEVICE = "cuda:0" +CHECKPOINT = "medium-mpa-0" # MACE-MPA-0, medium +DTYPE = torch.float32 +FMAX = 0.03 # eV/Å, largest atomic force at convergence +MAX_STEPS = 500 # FIRE2 steps per relaxation at most +BATCH_SIZE = 100 # structures relaxed together on the GPU +TIME_STEP = 0.5 # initial FIRE2 step, in ALCHEMI's internal time unit (10.18 fs) +GAMMA = 100 # kJ/mol, model collision energy of the AFIR force +NUMBER_OF_FINAL_STRUCTURES = 1 # the lowest final structures kept in final_structures.extxyz +BOND_TOLERANCE = 1.2 # two molecule atoms are bonded within BOND_TOLERANCE * (R_i + R_j) +SLAB_TAG = 0 # atom tags in configurations.extxyz +MOLECULE_TAG = 1 +AFIR_R0 = 3.8164 # Å, GRRM17 Eq. (3) +AFIR_EPSILON = 1.0061 # kJ/mol, GRRM17 Eq. (3) + +ALPHA = GAMMA * kJ / mol / ((2 ** (-1 / 6) - (1 + (1 + GAMMA / AFIR_EPSILON) ** 0.5) ** (-1 / 6)) * AFIR_R0) # Eq. (3) +CATEGORY_BY_TAG = {SLAB_TAG: AtomCategory.SURFACE.value, MOLECULE_TAG: AtomCategory.GAS.value} +IN_PLANE_MASK = torch.tensor([1.0, 1.0, 0.0], device=DEVICE) +IN_PLANE_SHIFTS = torch.tensor([[a, b, 0.0] for a in (-1, 0, 1) for b in (-1, 0, 1)], device=DEVICE) +COVALENT_RADII = torch.tensor(covalent_radii, device=DEVICE, dtype=DTYPE) + + +def in_plane_minimum_image(vectors, cells): + """ + Returns the shortest periodic image of each vector in the surface plane, a and b; there is no image along c. + + Args: + vectors (torch.Tensor): Cartesian vectors, [P, 3], in Å. + cells (torch.Tensor): The lattice vectors of each vector's structure, [P, 3, 3], rows a, b, c. + + Returns: + torch.Tensor: The minimum-image vectors, [P, 3], in Å. + """ + fractional = torch.einsum("pi,pij->pj", vectors, torch.linalg.inv(cells)) + fractional = fractional - torch.round(fractional) * IN_PLANE_MASK + images = torch.einsum("psi,pij->psj", fractional.unsqueeze(1) + IN_PLANE_SHIFTS, cells) + nearest = images.norm(dim=-1).argmin(dim=1) + return torch.take_along_dim(images, nearest[:, None, None], dim=1).squeeze(1) + + +def afir_energy(positions, batch): + """ + Returns the AFIR term of every structure in the batch, between its molecule and its slab atoms. + + GRRM17 Eq. (1), α Σ ω_ij r_ij / Σ ω_ij over molecule atoms i and slab atoms j, with ω_ij = ((R_i + R_j) / r_ij)^6 + (Eq. 2), R the covalent radii and r_ij at the minimum image in the surface plane; α is ALPHA, Eq. (3) + (S. Maeda et al., J. Comput. Chem. 39, 233 (2018)). + + Args: + positions (torch.Tensor): Atomic positions of the batch, [V, 3], in Å. + batch (Batch): The batch, with atom_categories GAS for the molecule and SURFACE for the slab. + + Returns: + torch.Tensor: The AFIR term of each structure, [B], in eV. + """ + molecule_atoms = torch.where(batch.atom_categories == AtomCategory.GAS.value)[0] + surface_atoms = torch.where(batch.atom_categories == AtomCategory.SURFACE.value)[0] + molecule_index, surface_index = torch.where( + batch.batch_idx[molecule_atoms][:, None] == batch.batch_idx[surface_atoms][None, :]) + pair_molecule_atoms, pair_surface_atoms = molecule_atoms[molecule_index], surface_atoms[surface_index] + pair_structures = batch.batch_idx[pair_molecule_atoms] + distances = in_plane_minimum_image(positions[pair_molecule_atoms] - positions[pair_surface_atoms], + batch.cell[pair_structures]).norm(dim=-1) + radii = (COVALENT_RADII[batch.atomic_numbers[pair_molecule_atoms]] + + COVALENT_RADII[batch.atomic_numbers[pair_surface_atoms]]) + weights = (radii / distances) ** 6 + numerator = torch.zeros(batch.num_graphs, device=positions.device, dtype=positions.dtype) + denominator = torch.zeros_like(numerator) + return ALPHA * (numerator.index_add(0, pair_structures, weights * distances) + / denominator.index_add(0, pair_structures, weights)) + + +def afir_bias(batch): + """ + Returns the AFIR energy and forces of the batch, as BiasedPotentialHook expects them. + + Args: + batch (Batch): The batch after the model's forward pass. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Energy per structure, [B, 1], in eV, and force per atom, [V, 3], + in eV/Å. + """ + positions = batch.positions.clone().requires_grad_(True) + energy = afir_energy(positions, batch) + forces = -torch.autograd.grad(energy.sum(), positions)[0] + return energy.detach().unsqueeze(-1), forces + + +def atomic_data(atoms): + """ + Converts a configuration to ALCHEMI's AtomicData, its tags to atom categories: SLAB_TAG to SURFACE, MOLECULE_TAG + to GAS. + + Args: + atoms (ase.Atoms): The configuration, slab atoms tagged 0 and molecule atoms 1. + + Returns: + AtomicData: The configuration on DEVICE, with zeroed forces and energy. + """ + data = AtomicData.from_atoms(atoms, device=DEVICE, dtype=DTYPE) + data["atom_categories"] = torch.tensor([CATEGORY_BY_TAG[tag] for tag in atoms.get_tags()], device=DEVICE) + data["forces"] = torch.zeros_like(data.positions) + data["energy"] = torch.zeros((1, 1), device=DEVICE, dtype=DTYPE) + return data + + +def relaxed_atoms(original, relaxed_data): + """ + Returns a copy of a structure at its relaxed positions, with its energy and forces. + + Args: + original (ase.Atoms): The structure before the relaxation. + relaxed_data (AtomicData): The same structure after it. + + Returns: + ase.Atoms: The relaxed structure, energy and forces on a SinglePointCalculator. + """ + atoms = original.copy() + atoms.positions = relaxed_data.positions.cpu().numpy() + atoms.calc = SinglePointCalculator(atoms, energy=float(relaxed_data.energy.item()), + forces=relaxed_data.forces.cpu().numpy()) + return atoms + + +def relax_batch(ase_structures, model, bias_hooks): + """ + Relaxes the structures as one batch with fixed-cell FIRE2, slab atoms frozen. + + Args: + ase_structures (List[ase.Atoms]): The structures, slab atoms tagged 0 and molecule atoms 1. + model (MACEWrapper): The MACE model. + bias_hooks (List[BiasedPotentialHook]): Bias potentials added after each forward pass; [] for none. + + Returns: + List[ase.Atoms]: The relaxed structures, in the input order. + """ + optimizer = FIRE2(model=model, dt=TIME_STEP, n_steps=MAX_STEPS, convergence_hook=ConvergenceHook( + criteria=[{"key": "forces", "threshold": FMAX, "reduce_op": "norm", "reduce_dims": -1}])) + optimizer.register_hook(FreezeAtomsHook(freeze_category=AtomCategory.SURFACE.value)) + stage = FusedStage(sub_stages=[(0, optimizer)]) + for hook in model.make_neighbor_hooks(): + stage.register_hook(hook, stage=DynamicsStage.BEFORE_COMPUTE) + for hook in bias_hooks: + stage.register_hook(hook, stage=DynamicsStage.AFTER_COMPUTE) + batch = stage.run(Batch.from_data_list([atomic_data(atoms) for atoms in ase_structures])) + return [relaxed_atoms(atoms, relaxed_data) for atoms, relaxed_data in zip(ase_structures, batch.to_data_list())] + + +def relax(ase_structures, model, bias_hooks): + """ + Relaxes all structures in batches of BATCH_SIZE, one after another, so the GPU memory a batch needs stays bounded. + + Args: + ase_structures (List[ase.Atoms]): The structures, slab atoms tagged 0 and molecule atoms 1. + model (MACEWrapper): The MACE model. + bias_hooks (List[BiasedPotentialHook]): Bias potentials added after each forward pass; [] for none. + + Returns: + List[ase.Atoms]: The relaxed structures, in the input order. + """ + relaxed = [] + for start in range(0, len(ase_structures), BATCH_SIZE): + end = min(start + BATCH_SIZE, len(ase_structures)) + logger.info("Relaxing structures {}–{} of {}", start + 1, end, len(ase_structures)) + relaxed += relax_batch(ase_structures[start:end], model, bias_hooks) + return relaxed + + +def molecule_fragments(atoms): + """ + Returns the number of fragments the molecule of a structure is in: 1 if intact, more if it split. + + Args: + atoms (ase.Atoms): The structure, molecule atoms tagged 1. + + Returns: + int: The number of connected groups of molecule atoms, bonded within BOND_TOLERANCE times their covalent radii. + """ + return len(separate(atoms[atoms.get_tags() == MOLECULE_TAG], scale=BOND_TOLERANCE)) + + +def energy_and_max_force(atoms): + """ + Returns the energy of a structure, in eV, and the largest force on any of its atoms, in eV/Å. + """ + return atoms.get_potential_energy(), np.linalg.norm(atoms.get_forces(), axis=1).max() + + +def plot_energies(energies, lowest_indices): + """ + Plots the final energy of every configuration against its index, the kept lowest marked, to final_energies.png. + + Args: + energies (List[float]): The final energy of each configuration, in eV. + lowest_indices (List[int]): The indices of the kept lowest structures. + """ + figure, axes = plt.subplots() + axes.scatter(range(len(energies)), energies, s=12, label="all configurations") + axes.scatter(lowest_indices, [energies[index] for index in lowest_indices], s=40, color="tab:red", + label=f"{len(lowest_indices)} lowest, kept") + axes.set(xlabel="Configuration index", ylabel="Final energy (eV)") + axes.ticklabel_format(axis="y", useOffset=False) + axes.legend() + figure.savefig("final_energies.png", dpi=150, bbox_inches="tight") + + +def write_results(afir_structures, final_structures): + """ + Ranks the configurations by final energy, writes all of them to ranked_configurations.csv with the placement each + started from (atoms.info["metadata"]["placement"]) and its molecule's fragments, the NUMBER_OF_FINAL_STRUCTURES + lowest to final_structures.extxyz in that order, and the energies to final_energies.png. + + Args: + afir_structures (List[ase.Atoms]): The structures after the relaxation under the AFIR force. + final_structures (List[ase.Atoms]): The structures after the relaxation without it. + """ + energies = [atoms.get_potential_energy() for atoms in final_structures] + order = np.argsort(energies) + with open("ranked_configurations.csv", "w", newline="") as results_file: + writer = csv.writer(results_file) + writer.writerow(["configuration", "rank", "afir_function_eV", "afir_max_force_eV_per_A", "energy_eV", + "max_force_eV_per_A", "rotation_axis_x", "rotation_axis_y", "rotation_axis_z", + "rotation_angle_deg", "position_u", "position_v", "molecule_fragments"]) + for rank, index in enumerate(order, start=1): + placement = final_structures[index].info["metadata"]["placement"] + writer.writerow([index, rank, *energy_and_max_force(afir_structures[index]), + *energy_and_max_force(final_structures[index]), *placement["rotation_axis"], + placement["rotation_angle_deg"], *placement["position"], + molecule_fragments(final_structures[index])]) + final_structures[index].info.update(rank=rank) + lowest_indices = order[:NUMBER_OF_FINAL_STRUCTURES].tolist() + write("final_structures.extxyz", [final_structures[index] for index in lowest_indices]) + plot_energies(energies, lowest_indices) + + +def main(): + """ + Relaxes the configurations of configurations.extxyz under the AFIR force, then without it, and writes the ranked + results. + """ + logger.info("AFIR force α, γ = {} kJ/mol: {:.3f} eV/Å", GAMMA, ALPHA) + ase_structures = read("{{ JOB_WORK_DIR }}/configurations.extxyz", index=":") + model = MACEWrapper.from_checkpoint(checkpoint_path=CHECKPOINT, device=torch.device(DEVICE), dtype=DTYPE, + enable_cueq=True) + logger.info("nvalchemi-toolkit {}, MACE with cuEquivariance", version("nvalchemi-toolkit")) + afir_structures = relax(ase_structures, model, [BiasedPotentialHook(bias_fn=afir_bias)]) + final_structures = relax(afir_structures, model, []) + write_results(afir_structures, final_structures) + + +if __name__ == "__main__": + main() diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/afir_requirements.txt b/other/experiments/alchemi/afir_molecule_surface_alchemi/afir_requirements.txt new file mode 100644 index 000000000..b0edfa40d --- /dev/null +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi/afir_requirements.txt @@ -0,0 +1,14 @@ +--find-links https://download.pytorch.org/whl/cu128/torch/ +--find-links https://download.pytorch.org/whl/cu128/torchvision/ +--find-links https://download.pytorch.org/whl/cu128/torchaudio/ +--find-links https://github.com/NVIDIA/warp/releases/expanded_assets/v1.18.0 +torch==2.11.0+cu128 +torchvision==0.26.0+cu128 +torchaudio==2.11.0+cu128 +warp-lang==1.18.0+cu12 +nvalchemi-toolkit[cu12,mace,ase] @ git+https://github.com/NVIDIA/nvalchemi-toolkit.git@618916442783c152cd53b919164f07f80286f3da +matplotlib +numpy +ase +mace-torch +loguru diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/configurations.py b/other/experiments/alchemi/afir_molecule_surface_alchemi/configurations.py new file mode 100644 index 000000000..dba63dbfb --- /dev/null +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi/configurations.py @@ -0,0 +1,125 @@ +# ---------------------------------------------------------------- # +# # +# Step 1, "Create Initial Configurations", of the workflow # +# "Molecule on Surface: AFIR Search (ALCHEMI)", a search with the # +# artificial force induced reaction (AFIR) method. # +# # +# Places the job's molecule over its slab # +# NUMBER_OF_CONFIGURATIONS times with made, each at a random # +# orientation and a random point of the surface cell, as # +# GRRM17's random initial structures [1]. # +# # +# Reads: the job's two materials, written into this script by # +# the platform when the job is created; the non-periodic # +# one is the molecule. # +# Writes: configurations.extxyz in the job's work folder: the # +# placed structures, slab atoms tagged 0 and molecule # +# atoms 1, each with its index in the info key # +# configuration and its rotation and fractional position # +# in metadata.placement. It is the only output of the # +# script. # +# # +# [1] S. Maeda et al., J. Comput. Chem. 39, 233 (2018). # +# # +# ---------------------------------------------------------------- # + +import json + +import numpy as np +from ase.data import atomic_numbers, covalent_radii +from ase.io import write +from mat3ra.made.material import Material +from mat3ra.made.tools.convert import to_ase +from mat3ra.made.tools.convert.interface_parts_enum import InterfacePartsEnum +from mat3ra.made.tools.modify import translate_by_vector +from mat3ra.made.tools.operations.core.binary import merge +from mat3ra.made.tools.operations.core.unary import rotate + +NUMBER_OF_CONFIGURATIONS = 100 +RANDOM_SEED = 0 +CONTACT_OFFSET = 0.5 # Å, lowest molecule atom at R_molecule + R_slab + CONTACT_OFFSET above the top slab atom +# The job's materials, written here by the platform when the job is created. +MATERIALS = json.loads(r"""{{ materials | dump | safe }}""") + + +def extreme_atom_height_and_radius(material, extremum): + """ + Returns the height and the covalent radius of the highest or the lowest atom of a material. + + Args: + material (Material): The material to search. + extremum (Callable): max for the highest atom, min for the lowest. + + Returns: + Tuple[float, float]: The atom's cartesian z and its covalent radius, both in Å. + """ + cartesian = material.clone() + cartesian.to_cartesian() + height, element = extremum(zip([coordinate[2] for coordinate in cartesian.basis.coordinates.values], + cartesian.basis.elements.values)) + return height, covalent_radii[atomic_numbers[element]] + + +def create_configuration(slab, molecule, top_atom_height_and_radius, random_generator, index): + """ + Places the molecule once over the slab, as a random initial structure of GRRM17's MC-AFIR. + + The orientation is uniformly random (a normalised Gaussian quaternion), the molecule's center goes to a random + point of the surface cell, and its lowest atom sits R_molecule + R_slab + CONTACT_OFFSET above the top slab atom, + the RandDC rule of GRRM17 (S. Maeda et al., J. Comput. Chem. 39, 233 (2018)). + + Args: + slab (Material): The slab, in crystal units, labelled 0. + molecule (Material): The molecule, in crystal units. + top_atom_height_and_radius (Tuple[float, float]): Height and covalent radius of the top slab atom, in Å. + random_generator (numpy.random.Generator): The seeded generator of all placements. + index (int): The configuration's index, used in its name. + + Returns: + MaterialWithBuildMetadata: The slab and the placed molecule, molecule atoms labelled 1, with the rotation and + the fractional position in metadata.placement. + """ + quaternion = random_generator.normal(size=4) + axis = (quaternion[1:] / np.linalg.norm(quaternion[1:])).tolist() + angle = float(2 * np.degrees(np.arccos(quaternion[0] / np.linalg.norm(quaternion)))) + rotated = rotate(molecule, axis=axis, angle=angle, wrap=False) + rotated.set_lattice(slab.lattice) + top_height, top_radius = top_atom_height_and_radius + lowest_height, lowest_radius = extreme_atom_height_and_radius(rotated, min) + surface_fractions = [random_generator.random(), random_generator.random(), 0] + surface_point = slab.basis.cell.convert_point_to_cartesian(surface_fractions) + molecule_cell_center = molecule.basis.cell.convert_point_to_cartesian([0.5, 0.5, 0.5]) + vertical_shift = top_height + top_radius + lowest_radius + CONTACT_OFFSET - lowest_height + placed = translate_by_vector(rotated, vector=[surface_point[0] - molecule_cell_center[0], + surface_point[1] - molecule_cell_center[1], vertical_shift], + use_cartesian_coordinates=True) + placed.set_labels_from_value(InterfacePartsEnum.FILM.value) + configuration = merge([slab, placed], material_name=f"{molecule.name} on {slab.name}, configuration {index}") + configuration.metadata.placement = {"rotation_axis": axis, "rotation_angle_deg": angle, + "position": surface_fractions[:2]} + return configuration + + +def main(): + """ + Builds NUMBER_OF_CONFIGURATIONS configurations from the job's materials and writes them to configurations.extxyz. + """ + slab = Material.create(next(material_config for material_config in MATERIALS + if not material_config.get("isNonPeriodic"))) + molecule = Material.create(next(material_config for material_config in MATERIALS + if material_config.get("isNonPeriodic"))) + for material in (slab, molecule): + material.to_crystal() + slab.set_labels_from_value(InterfacePartsEnum.SUBSTRATE.value) + top_atom_height_and_radius = extreme_atom_height_and_radius(slab, max) + random_generator = np.random.default_rng(RANDOM_SEED) + configurations = [] + for index in range(NUMBER_OF_CONFIGURATIONS): + atoms = to_ase(create_configuration(slab, molecule, top_atom_height_and_radius, random_generator, index)) + atoms.info["configuration"] = index + configurations.append(atoms) + write("configurations.extxyz", configurations) + + +if __name__ == "__main__": + main() diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/configurations_requirements.txt b/other/experiments/alchemi/afir_molecule_surface_alchemi/configurations_requirements.txt new file mode 100644 index 000000000..9aeb82595 --- /dev/null +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi/configurations_requirements.txt @@ -0,0 +1 @@ +mat3ra-made[tools]==2026.10.10.post0 diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures.py b/other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures.py new file mode 100644 index 000000000..41cc7a351 --- /dev/null +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures.py @@ -0,0 +1,61 @@ +# ---------------------------------------------------------------- # +# # +# Step 3, "Save Final Structures", of the workflow # +# "Molecule on Surface: AFIR Search (ALCHEMI)", a search with the # +# artificial force induced reaction (AFIR) method. # +# # +# Takes the lowest final structures of "Relax under AFIR Force # +# (ALCHEMI)", each with the configuration it started from, and # +# makes one request body per structure for the platform's # +# final-structure endpoint, which the io unit after this one # +# sends in a loop. Each request names this unit (UNIT_ID), so the # +# structures are this step's results. # +# # +# Reads: final_structures.extxyz and configurations.extxyz in # +# the job's work folder. # +# Writes: FINAL_STRUCTURE_REQUESTS to the job's global scope, as # +# the only output of the script. # +# # +# ---------------------------------------------------------------- # + +import json + +from ase.io import read +from mat3ra.made.material import Material +from mat3ra.made.tools.convert import from_ase + + +def request_body(rank_index, atoms, configurations): + """ + Returns the body of one final-structure request, escaped to sit inside the next io unit's JSON string. + + The platform keeps one final_structure per unit and repetition, so each structure gets its own repetition. + + Args: + rank_index (int): The structure's rank by final energy, from 0. + atoms (ase.Atoms): The final structure, with its configuration's index in atoms.info. + configurations (List[ase.Atoms]): All configurations, in the order of their indices. + + Returns: + str: {"jobId", "unitId", "structures": [{"initial", "final"}]} as JSON, escaped for a JSON string. + """ + final = {**Material.create(from_ase(atoms)).to_dict(), "repetition": rank_index} + initial = Material.create(from_ase(configurations[atoms.info["configuration"]])).to_dict() + body = {"jobId": "{{ JOB_ID }}", "unitId": "{{ UNIT_ID }}", "structures": [{"initial": initial, "final": final}]} + return json.dumps(json.dumps(body))[1:-1] + + +def main(): + """ + Reads the lowest final structures and the configurations, and prints the request bodies as + FINAL_STRUCTURE_REQUESTS. + """ + final_structures = read("{{ JOB_WORK_DIR }}/final_structures.extxyz", index=":") + configurations = read("{{ JOB_WORK_DIR }}/configurations.extxyz", index=":") + request_bodies = [request_body(rank_index, atoms, configurations) + for rank_index, atoms in enumerate(final_structures)] + print(json.dumps({"name": "FINAL_STRUCTURE_REQUESTS", "value": request_bodies, "scope": "global"})) + + +if __name__ == "__main__": + main() diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures_requirements.txt b/other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures_requirements.txt new file mode 100644 index 000000000..9aeb82595 --- /dev/null +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures_requirements.txt @@ -0,0 +1 @@ +mat3ra-made[tools]==2026.10.10.post0 From b30fda1fe92aaab4edd5c559494a9f45daf22d96 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 20:59:46 -0800 Subject: [PATCH 23/31] SOF-8102: AFIR step's ALCHEMI from the PyPI 0.3.0 release instead of a git commit Co-Authored-By: Claude Opus 5.5 --- .../experiments/alchemi/afir_molecule_surface_alchemi.ipynb | 6 ++---- .../afir_molecule_surface_alchemi/afir_requirements.txt | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 822f97455..f58473149 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -290,10 +290,8 @@ "- `configurations_requirements.txt` and `final_structures_requirements.txt`: steps 1 and 3 (*Create Initial\n", " Configurations*, *Save Final Structures*), the made library, which builds the configurations and converts the\n", " structures.\n", - "- `afir_requirements.txt`: step 2 (*Relax under AFIR Force (ALCHEMI)*), NVIDIA ALCHEMI with MACE on the GPU. **ALCHEMI\n", - " here is a pre-release: version 0.3.0 installed from GitHub commit\n", - " [6189164](https://github.com/NVIDIA/nvalchemi-toolkit/commit/618916442783c152cd53b919164f07f80286f3da), not a PyPI\n", - " release.** PyTorch is the CUDA 12.8 build and Warp the CUDA 12 build (the `--find-links` lines). Building this\n", + "- `afir_requirements.txt`: step 2 (*Relax under AFIR Force (ALCHEMI)*), NVIDIA ALCHEMI 0.3.0 (PyPI) with MACE on the\n", + " GPU. PyTorch is the CUDA 12.8 build and Warp the CUDA 12 build (the `--find-links` lines). Building this\n", " environment the first time takes about 10 minutes (measured on cluster-003)." ] }, diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/afir_requirements.txt b/other/experiments/alchemi/afir_molecule_surface_alchemi/afir_requirements.txt index b0edfa40d..0c4b31ef5 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi/afir_requirements.txt +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi/afir_requirements.txt @@ -6,7 +6,7 @@ torch==2.11.0+cu128 torchvision==0.26.0+cu128 torchaudio==2.11.0+cu128 warp-lang==1.18.0+cu12 -nvalchemi-toolkit[cu12,mace,ase] @ git+https://github.com/NVIDIA/nvalchemi-toolkit.git@618916442783c152cd53b919164f07f80286f3da +nvalchemi-toolkit[cu12,mace,ase]==0.3.0 matplotlib numpy ase From 0f5cea70d0450cc9aec4c26b434a5ab5fa348567 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 20:59:47 -0800 Subject: [PATCH 24/31] SOF-8102: drop the measured build time from the requirements note Co-Authored-By: Claude Opus 5.5 --- other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index f58473149..444146a9d 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -291,8 +291,7 @@ " Configurations*, *Save Final Structures*), the made library, which builds the configurations and converts the\n", " structures.\n", "- `afir_requirements.txt`: step 2 (*Relax under AFIR Force (ALCHEMI)*), NVIDIA ALCHEMI 0.3.0 (PyPI) with MACE on the\n", - " GPU. PyTorch is the CUDA 12.8 build and Warp the CUDA 12 build (the `--find-links` lines). Building this\n", - " environment the first time takes about 10 minutes (measured on cluster-003)." + " GPU. PyTorch is the CUDA 12.8 build and Warp the CUDA 12 build (the `--find-links` lines)." ] }, { From f840f6970d28d4b8c228d9f381e9cb6d9e05f698 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 20:59:53 -0800 Subject: [PATCH 25/31] SOF-8102: the step files get their own heading, 4.4 Co-Authored-By: Claude Opus 5.5 --- .../experiments/alchemi/afir_molecule_surface_alchemi.ipynb | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 444146a9d..007603078 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -283,6 +283,8 @@ "cell_type": "markdown", "metadata": {}, "source": [ + "### 4.4. Read the step scripts and requirements\n", + "\n", "Each step's script and requirements are files in `afir_molecule_surface_alchemi/` beside this notebook, under the\n", "names the workflow gives them on the platform; the cell below reads them. Each step runs in its own Python environment\n", "on the cluster, installed from its requirements the first time and reused afterwards.\n", @@ -315,7 +317,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### 4.4. Configure the steps\n", + "### 4.5. Configure the steps\n", "\n", "When the job is created, the platform writes the job's materials into step 1's script (`materials`; the\n", "non-periodic one is the molecule). Step 1 writes the configurations to `configurations.extxyz` in the\n", @@ -414,7 +416,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### 4.5. Create and save the workflow\n", + "### 4.6. Create and save the workflow\n", "\n", "Step 3 ends with a loop that saves one final structure per pass, with its starting configuration as the initial\n", "structure." From d2308ab2d1f1f328377483d46d631bdccdfe1c39 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 20:59:54 -0800 Subject: [PATCH 26/31] SOF-8102: steps 1 and 3 share a Python environment, step 2 has its own Co-Authored-By: Claude Opus 5.5 --- .../experiments/alchemi/afir_molecule_surface_alchemi.ipynb | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 007603078..3cbe1fe2b 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -286,8 +286,9 @@ "### 4.4. Read the step scripts and requirements\n", "\n", "Each step's script and requirements are files in `afir_molecule_surface_alchemi/` beside this notebook, under the\n", - "names the workflow gives them on the platform; the cell below reads them. Each step runs in its own Python environment\n", - "on the cluster, installed from its requirements the first time and reused afterwards.\n", + "names the workflow gives them on the platform; the cell below reads them. Steps 1 and 3 list the same packages\n", + "and share one Python environment on the cluster, step 2 has its own; each is installed from its requirements the first\n", + "time and reused afterwards.\n", "\n", "- `configurations_requirements.txt` and `final_structures_requirements.txt`: steps 1 and 3 (*Create Initial\n", " Configurations*, *Save Final Structures*), the made library, which builds the configurations and converts the\n", From b026b8ad3aa6e8b21ee36311cac3cafb9e87d400 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 21:22:26 -0800 Subject: [PATCH 27/31] SOF-8102: the introduction says steps 1 and 3 share a Python environment Co-Authored-By: Claude Opus 5.5 --- other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 3cbe1fe2b..a57d2702e 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -11,7 +11,7 @@ "under an artificial force between molecule and slab, as in the artificial force induced reaction (AFIR) method,\n", "then without it. The final structures are ranked by energy, and the lowest are saved as materials.\n", "\n", - "The workflow has three steps, each in its own Python environment: *Create Initial Configurations* places the\n", + "The workflow has three steps, in two Python environments (steps 1 and 3 share one, step 2 has its own): *Create Initial Configurations* places the\n", "molecule over the slab with made, from the job's two materials, *Relax under AFIR Force (ALCHEMI)* relaxes the\n", "configurations, and *Save Final Structures* saves the lowest as the job's final structures. This notebook\n", "builds the slab, fetches the molecule, saves both as materials, creates the workflow, runs one job on the two\n", From 25142101929333fa29574c7094feacac1e2fe41e Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 21:26:04 -0800 Subject: [PATCH 28/31] SOF-8102: step files in one folder per step, each with its requirements.txt Co-Authored-By: Claude Opus 5.5 --- .../afir_molecule_surface_alchemi.ipynb | 24 ++++++++++--------- .../{ => afir}/afir.py | 0 .../requirements.txt} | 0 .../{ => configurations}/configurations.py | 0 .../requirements.txt} | 0 .../final_structures.py | 0 .../requirements.txt} | 0 7 files changed, 13 insertions(+), 11 deletions(-) rename other/experiments/alchemi/afir_molecule_surface_alchemi/{ => afir}/afir.py (100%) rename other/experiments/alchemi/afir_molecule_surface_alchemi/{afir_requirements.txt => afir/requirements.txt} (100%) rename other/experiments/alchemi/afir_molecule_surface_alchemi/{ => configurations}/configurations.py (100%) rename other/experiments/alchemi/afir_molecule_surface_alchemi/{configurations_requirements.txt => configurations/requirements.txt} (100%) rename other/experiments/alchemi/afir_molecule_surface_alchemi/{ => final_structures}/final_structures.py (100%) rename other/experiments/alchemi/afir_molecule_surface_alchemi/{final_structures_requirements.txt => final_structures/requirements.txt} (100%) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index a57d2702e..b7d088029 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -266,7 +266,7 @@ "without it. It writes `ranked_configurations.csv` with every configuration ranked by final energy, its starting\n", "orientation and the fragments its molecule ends in; `final_structures.extxyz` with the `NUMBER_OF_FINAL_STRUCTURES`\n", "lowest; and `final_energies.png`, the final energy of each configuration. Another backend or another ALCHEMI version\n", - "replaces `afir.py` and `afir_requirements.txt`; steps 1 and 3 stay." + "replaces `afir/afir.py` and `afir/requirements.txt`; steps 1 and 3 stay." ] }, { @@ -285,15 +285,17 @@ "source": [ "### 4.4. Read the step scripts and requirements\n", "\n", - "Each step's script and requirements are files in `afir_molecule_surface_alchemi/` beside this notebook, under the\n", - "names the workflow gives them on the platform; the cell below reads them. Steps 1 and 3 list the same packages\n", + "Each step's script and requirements are files in `afir_molecule_surface_alchemi/` beside this notebook, in one folder\n", + "per step: `configurations/`, `afir/` and `final_structures/`, each holding the step's script and its\n", + "`requirements.txt`. The workflow gives them the names `.py` and `_requirements.txt` on the platform, and\n", + "the cell below reads them. Steps 1 and 3 list the same packages\n", "and share one Python environment on the cluster, step 2 has its own; each is installed from its requirements the first\n", "time and reused afterwards.\n", "\n", - "- `configurations_requirements.txt` and `final_structures_requirements.txt`: steps 1 and 3 (*Create Initial\n", + "- `configurations/requirements.txt` and `final_structures/requirements.txt`: steps 1 and 3 (*Create Initial\n", " Configurations*, *Save Final Structures*), the made library, which builds the configurations and converts the\n", " structures.\n", - "- `afir_requirements.txt`: step 2 (*Relax under AFIR Force (ALCHEMI)*), NVIDIA ALCHEMI 0.3.0 (PyPI) with MACE on the\n", + "- `afir/requirements.txt`: step 2 (*Relax under AFIR Force (ALCHEMI)*), NVIDIA ALCHEMI 0.3.0 (PyPI) with MACE on the\n", " GPU. PyTorch is the CUDA 12.8 build and Warp the CUDA 12 build (the `--find-links` lines)." ] }, @@ -306,12 +308,12 @@ "from pathlib import Path\n", "\n", "STEP_FILES_DIRECTORY = Path(\"afir_molecule_surface_alchemi\")\n", - "CONFIGURATIONS_SCRIPT = (STEP_FILES_DIRECTORY / \"configurations.py\").read_text()\n", - "CONFIGURATIONS_REQUIREMENTS = (STEP_FILES_DIRECTORY / \"configurations_requirements.txt\").read_text()\n", - "AFIR_SCRIPT = (STEP_FILES_DIRECTORY / \"afir.py\").read_text()\n", - "AFIR_REQUIREMENTS = (STEP_FILES_DIRECTORY / \"afir_requirements.txt\").read_text()\n", - "FINAL_STRUCTURES_SCRIPT = (STEP_FILES_DIRECTORY / \"final_structures.py\").read_text()\n", - "FINAL_STRUCTURES_REQUIREMENTS = (STEP_FILES_DIRECTORY / \"final_structures_requirements.txt\").read_text()" + "CONFIGURATIONS_SCRIPT = (STEP_FILES_DIRECTORY / \"configurations\" / \"configurations.py\").read_text()\n", + "CONFIGURATIONS_REQUIREMENTS = (STEP_FILES_DIRECTORY / \"configurations\" / \"requirements.txt\").read_text()\n", + "AFIR_SCRIPT = (STEP_FILES_DIRECTORY / \"afir\" / \"afir.py\").read_text()\n", + "AFIR_REQUIREMENTS = (STEP_FILES_DIRECTORY / \"afir\" / \"requirements.txt\").read_text()\n", + "FINAL_STRUCTURES_SCRIPT = (STEP_FILES_DIRECTORY / \"final_structures\" / \"final_structures.py\").read_text()\n", + "FINAL_STRUCTURES_REQUIREMENTS = (STEP_FILES_DIRECTORY / \"final_structures\" / \"requirements.txt\").read_text()" ] }, { diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/afir.py b/other/experiments/alchemi/afir_molecule_surface_alchemi/afir/afir.py similarity index 100% rename from other/experiments/alchemi/afir_molecule_surface_alchemi/afir.py rename to other/experiments/alchemi/afir_molecule_surface_alchemi/afir/afir.py diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/afir_requirements.txt b/other/experiments/alchemi/afir_molecule_surface_alchemi/afir/requirements.txt similarity index 100% rename from other/experiments/alchemi/afir_molecule_surface_alchemi/afir_requirements.txt rename to other/experiments/alchemi/afir_molecule_surface_alchemi/afir/requirements.txt diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/configurations.py b/other/experiments/alchemi/afir_molecule_surface_alchemi/configurations/configurations.py similarity index 100% rename from other/experiments/alchemi/afir_molecule_surface_alchemi/configurations.py rename to other/experiments/alchemi/afir_molecule_surface_alchemi/configurations/configurations.py diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/configurations_requirements.txt b/other/experiments/alchemi/afir_molecule_surface_alchemi/configurations/requirements.txt similarity index 100% rename from other/experiments/alchemi/afir_molecule_surface_alchemi/configurations_requirements.txt rename to other/experiments/alchemi/afir_molecule_surface_alchemi/configurations/requirements.txt diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures.py b/other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures/final_structures.py similarity index 100% rename from other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures.py rename to other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures/final_structures.py diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures_requirements.txt b/other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures/requirements.txt similarity index 100% rename from other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures_requirements.txt rename to other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures/requirements.txt From 705f86cf27410b794a901ad0a904f01b4565a09b Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 21:27:05 -0800 Subject: [PATCH 29/31] SOF-8102: step folders named after the workflow's subworkflows Co-Authored-By: Claude Opus 5.5 --- .../afir_molecule_surface_alchemi.ipynb | 31 ++++++++++--------- .../configurations.py | 0 .../requirements.txt | 0 .../afir.py | 0 .../requirements.txt | 0 .../final_structures.py | 0 .../requirements.txt | 0 7 files changed, 16 insertions(+), 15 deletions(-) rename other/experiments/alchemi/afir_molecule_surface_alchemi/{configurations => create_initial_configurations}/configurations.py (100%) rename other/experiments/alchemi/afir_molecule_surface_alchemi/{configurations => create_initial_configurations}/requirements.txt (100%) rename other/experiments/alchemi/afir_molecule_surface_alchemi/{afir => relax_under_afir_force_alchemi}/afir.py (100%) rename other/experiments/alchemi/afir_molecule_surface_alchemi/{afir => relax_under_afir_force_alchemi}/requirements.txt (100%) rename other/experiments/alchemi/afir_molecule_surface_alchemi/{final_structures => save_final_structures}/final_structures.py (100%) rename other/experiments/alchemi/afir_molecule_surface_alchemi/{final_structures => save_final_structures}/requirements.txt (100%) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index b7d088029..8ef4aac47 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -266,7 +266,7 @@ "without it. It writes `ranked_configurations.csv` with every configuration ranked by final energy, its starting\n", "orientation and the fragments its molecule ends in; `final_structures.extxyz` with the `NUMBER_OF_FINAL_STRUCTURES`\n", "lowest; and `final_energies.png`, the final energy of each configuration. Another backend or another ALCHEMI version\n", - "replaces `afir/afir.py` and `afir/requirements.txt`; steps 1 and 3 stay." + "replaces `relax_under_afir_force_alchemi/afir.py` and `relax_under_afir_force_alchemi/requirements.txt`; steps 1 and 3 stay." ] }, { @@ -286,17 +286,18 @@ "### 4.4. Read the step scripts and requirements\n", "\n", "Each step's script and requirements are files in `afir_molecule_surface_alchemi/` beside this notebook, in one folder\n", - "per step: `configurations/`, `afir/` and `final_structures/`, each holding the step's script and its\n", - "`requirements.txt`. The workflow gives them the names `.py` and `_requirements.txt` on the platform, and\n", - "the cell below reads them. Steps 1 and 3 list the same packages\n", + "per step, named after the step: `create_initial_configurations/`, `relax_under_afir_force_alchemi/` and\n", + "`save_final_structures/`, each holding the step's script and its `requirements.txt`. The workflow gives them the\n", + "names `.py` and `_requirements.txt` on the platform, and the cell below reads them. Steps 1 and 3 list\n", + "the same packages\n", "and share one Python environment on the cluster, step 2 has its own; each is installed from its requirements the first\n", "time and reused afterwards.\n", "\n", - "- `configurations/requirements.txt` and `final_structures/requirements.txt`: steps 1 and 3 (*Create Initial\n", - " Configurations*, *Save Final Structures*), the made library, which builds the configurations and converts the\n", - " structures.\n", - "- `afir/requirements.txt`: step 2 (*Relax under AFIR Force (ALCHEMI)*), NVIDIA ALCHEMI 0.3.0 (PyPI) with MACE on the\n", - " GPU. PyTorch is the CUDA 12.8 build and Warp the CUDA 12 build (the `--find-links` lines)." + "- `create_initial_configurations/requirements.txt` and `save_final_structures/requirements.txt`: steps 1 and 3\n", + " (*Create Initial Configurations*, *Save Final Structures*), the made library, which builds the configurations and\n", + " converts the structures.\n", + "- `relax_under_afir_force_alchemi/requirements.txt`: step 2 (*Relax under AFIR Force (ALCHEMI)*), NVIDIA ALCHEMI\n", + " 0.3.0 (PyPI) with MACE on the GPU. PyTorch is the CUDA 12.8 build and Warp the CUDA 12 build (the `--find-links` lines)." ] }, { @@ -308,12 +309,12 @@ "from pathlib import Path\n", "\n", "STEP_FILES_DIRECTORY = Path(\"afir_molecule_surface_alchemi\")\n", - "CONFIGURATIONS_SCRIPT = (STEP_FILES_DIRECTORY / \"configurations\" / \"configurations.py\").read_text()\n", - "CONFIGURATIONS_REQUIREMENTS = (STEP_FILES_DIRECTORY / \"configurations\" / \"requirements.txt\").read_text()\n", - "AFIR_SCRIPT = (STEP_FILES_DIRECTORY / \"afir\" / \"afir.py\").read_text()\n", - "AFIR_REQUIREMENTS = (STEP_FILES_DIRECTORY / \"afir\" / \"requirements.txt\").read_text()\n", - "FINAL_STRUCTURES_SCRIPT = (STEP_FILES_DIRECTORY / \"final_structures\" / \"final_structures.py\").read_text()\n", - "FINAL_STRUCTURES_REQUIREMENTS = (STEP_FILES_DIRECTORY / \"final_structures\" / \"requirements.txt\").read_text()" + "CONFIGURATIONS_SCRIPT = (STEP_FILES_DIRECTORY / \"create_initial_configurations\" / \"configurations.py\").read_text()\n", + "CONFIGURATIONS_REQUIREMENTS = (STEP_FILES_DIRECTORY / \"create_initial_configurations\" / \"requirements.txt\").read_text()\n", + "AFIR_SCRIPT = (STEP_FILES_DIRECTORY / \"relax_under_afir_force_alchemi\" / \"afir.py\").read_text()\n", + "AFIR_REQUIREMENTS = (STEP_FILES_DIRECTORY / \"relax_under_afir_force_alchemi\" / \"requirements.txt\").read_text()\n", + "FINAL_STRUCTURES_SCRIPT = (STEP_FILES_DIRECTORY / \"save_final_structures\" / \"final_structures.py\").read_text()\n", + "FINAL_STRUCTURES_REQUIREMENTS = (STEP_FILES_DIRECTORY / \"save_final_structures\" / \"requirements.txt\").read_text()" ] }, { diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/configurations/configurations.py b/other/experiments/alchemi/afir_molecule_surface_alchemi/create_initial_configurations/configurations.py similarity index 100% rename from other/experiments/alchemi/afir_molecule_surface_alchemi/configurations/configurations.py rename to other/experiments/alchemi/afir_molecule_surface_alchemi/create_initial_configurations/configurations.py diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/configurations/requirements.txt b/other/experiments/alchemi/afir_molecule_surface_alchemi/create_initial_configurations/requirements.txt similarity index 100% rename from other/experiments/alchemi/afir_molecule_surface_alchemi/configurations/requirements.txt rename to other/experiments/alchemi/afir_molecule_surface_alchemi/create_initial_configurations/requirements.txt diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/afir/afir.py b/other/experiments/alchemi/afir_molecule_surface_alchemi/relax_under_afir_force_alchemi/afir.py similarity index 100% rename from other/experiments/alchemi/afir_molecule_surface_alchemi/afir/afir.py rename to other/experiments/alchemi/afir_molecule_surface_alchemi/relax_under_afir_force_alchemi/afir.py diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/afir/requirements.txt b/other/experiments/alchemi/afir_molecule_surface_alchemi/relax_under_afir_force_alchemi/requirements.txt similarity index 100% rename from other/experiments/alchemi/afir_molecule_surface_alchemi/afir/requirements.txt rename to other/experiments/alchemi/afir_molecule_surface_alchemi/relax_under_afir_force_alchemi/requirements.txt diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures/final_structures.py b/other/experiments/alchemi/afir_molecule_surface_alchemi/save_final_structures/final_structures.py similarity index 100% rename from other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures/final_structures.py rename to other/experiments/alchemi/afir_molecule_surface_alchemi/save_final_structures/final_structures.py diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures/requirements.txt b/other/experiments/alchemi/afir_molecule_surface_alchemi/save_final_structures/requirements.txt similarity index 100% rename from other/experiments/alchemi/afir_molecule_surface_alchemi/final_structures/requirements.txt rename to other/experiments/alchemi/afir_molecule_surface_alchemi/save_final_structures/requirements.txt From 057ba634d2d5d19d98362057782819c5f2007f8a Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 21:30:33 -0800 Subject: [PATCH 30/31] SOF-8102: introduction and step-files note reworded for the folder layout Co-Authored-By: Claude Opus 5.5 --- .../alchemi/afir_molecule_surface_alchemi.ipynb | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index 8ef4aac47..fdd630ae0 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -11,11 +11,11 @@ "under an artificial force between molecule and slab, as in the artificial force induced reaction (AFIR) method,\n", "then without it. The final structures are ranked by energy, and the lowest are saved as materials.\n", "\n", - "The workflow has three steps, in two Python environments (steps 1 and 3 share one, step 2 has its own): *Create Initial Configurations* places the\n", - "molecule over the slab with made, from the job's two materials, *Relax under AFIR Force (ALCHEMI)* relaxes the\n", - "configurations, and *Save Final Structures* saves the lowest as the job's final structures. This notebook\n", - "builds the slab, fetches the molecule, saves both as materials, creates the workflow, runs one job on the two\n", - "materials and shows the ranked configurations and the lowest final structures.\n", + "The workflow has three steps: *Create Initial Configurations* places the molecule over the slab with made, from the\n", + "job's two materials; *Relax under AFIR Force (ALCHEMI)* relaxes the configurations; and *Save Final Structures* saves\n", + "the lowest as the job's final structures. Steps 1 and 3 share one Python environment and step 2 has its own. This\n", + "notebook builds the slab, fetches the molecule, saves both as materials, creates the workflow, runs one job on the\n", + "two materials and shows the ranked configurations and the lowest final structures.\n", "\n", "Outside the platform's JupyterLite, use Python 3.10 or 3.11 and install the packages once with `pip install \"mat3ra-notebooks-utils[all]\"`." ] @@ -287,10 +287,9 @@ "\n", "Each step's script and requirements are files in `afir_molecule_surface_alchemi/` beside this notebook, in one folder\n", "per step, named after the step: `create_initial_configurations/`, `relax_under_afir_force_alchemi/` and\n", - "`save_final_structures/`, each holding the step's script and its `requirements.txt`. The workflow gives them the\n", - "names `.py` and `_requirements.txt` on the platform, and the cell below reads them. Steps 1 and 3 list\n", - "the same packages\n", - "and share one Python environment on the cluster, step 2 has its own; each is installed from its requirements the first\n", + "`save_final_structures/`, each holding the step's script and its `requirements.txt`; the cell below reads them. On the\n", + "platform the requirements of a step are named `_requirements.txt`. Steps 1 and 3 list the same packages and\n", + "share one Python environment on the cluster and step 2 has its own; each is installed from its requirements the first\n", "time and reused afterwards.\n", "\n", "- `create_initial_configurations/requirements.txt` and `save_final_structures/requirements.txt`: steps 1 and 3\n", From f3ad03893d01c5a169998b271c1bbcc9e27ddc19 Mon Sep 17 00:00:00 2001 From: VsevolodX Date: Sun, 8 Nov 2026 21:47:31 -0800 Subject: [PATCH 31/31] SOF-8102: requirements file names stated as the platform gives them; long markdown lines wrapped Co-Authored-By: Claude Opus 5.5 --- .../alchemi/afir_molecule_surface_alchemi.ipynb | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb index fdd630ae0..0b9792eb1 100644 --- a/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -266,7 +266,8 @@ "without it. It writes `ranked_configurations.csv` with every configuration ranked by final energy, its starting\n", "orientation and the fragments its molecule ends in; `final_structures.extxyz` with the `NUMBER_OF_FINAL_STRUCTURES`\n", "lowest; and `final_energies.png`, the final energy of each configuration. Another backend or another ALCHEMI version\n", - "replaces `relax_under_afir_force_alchemi/afir.py` and `relax_under_afir_force_alchemi/requirements.txt`; steps 1 and 3 stay." + "replaces `relax_under_afir_force_alchemi/afir.py` and `relax_under_afir_force_alchemi/requirements.txt`; steps 1\n", + "and 3 stay." ] }, { @@ -288,15 +289,16 @@ "Each step's script and requirements are files in `afir_molecule_surface_alchemi/` beside this notebook, in one folder\n", "per step, named after the step: `create_initial_configurations/`, `relax_under_afir_force_alchemi/` and\n", "`save_final_structures/`, each holding the step's script and its `requirements.txt`; the cell below reads them. On the\n", - "platform the requirements of a step are named `_requirements.txt`. Steps 1 and 3 list the same packages and\n", - "share one Python environment on the cluster and step 2 has its own; each is installed from its requirements the first\n", - "time and reused afterwards.\n", + "platform each requirements file takes its script's name: `configurations_requirements.txt`, `afir_requirements.txt`\n", + "and `final_structures_requirements.txt`. Steps 1 and 3 list the same packages and share one Python environment on the\n", + "cluster and step 2 has its own; each is installed from its requirements the first time and reused afterwards.\n", "\n", "- `create_initial_configurations/requirements.txt` and `save_final_structures/requirements.txt`: steps 1 and 3\n", " (*Create Initial Configurations*, *Save Final Structures*), the made library, which builds the configurations and\n", " converts the structures.\n", "- `relax_under_afir_force_alchemi/requirements.txt`: step 2 (*Relax under AFIR Force (ALCHEMI)*), NVIDIA ALCHEMI\n", - " 0.3.0 (PyPI) with MACE on the GPU. PyTorch is the CUDA 12.8 build and Warp the CUDA 12 build (the `--find-links` lines)." + " 0.3.0 (PyPI) with MACE on the GPU. PyTorch is the CUDA 12.8 build and Warp the CUDA 12 build (the `--find-links`\n", + " lines)." ] }, { @@ -324,7 +326,7 @@ "\n", "When the job is created, the platform writes the job's materials into step 1's script (`materials`; the\n", "non-periodic one is the molecule). Step 1 writes the configurations to `configurations.extxyz` in the\n", - "job's work folder, where steps 2 and 3 read them. Each step's files are named after it." + "job's work folder, where steps 2 and 3 read them." ] }, {