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..0b9792eb1 --- /dev/null +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi.ipynb @@ -0,0 +1,665 @@ +{ + "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), in batches on one GPU: first\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: *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]\"`." + ] + }, + { + "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", + "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 = 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", + "# 5. Compute parameters\n", + "CLUSTER_NAME = \"cluster-003\"\n", + "QUEUE_NAME = QueueName.GOF # H100 nodes on cluster-003 (GOF on-demand, GSF spot)\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. 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", + "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." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 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 `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\n", + "and 3 stay." + ] + }, + { + "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": "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, 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 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`\n", + " lines)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "STEP_FILES_DIRECTORY = Path(\"afir_molecule_surface_alchemi\")\n", + "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()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 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", + "job's work folder, where steps 2 and 3 read them." + ] + }, + { + "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.standata.workflows import WorkflowStandata\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", + "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", + "\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, 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 (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=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\"{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=\"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", + "final_structures_subworkflow_config.pop(\"_id\")\n", + "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\",\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=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", + " 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", + " results=[\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=\"Save Final Structures\",\n", + " file_name=\"final_structures\",\n", + " script=FINAL_STRUCTURES_SCRIPT, requirements=FINAL_STRUCTURES_REQUIREMENTS,\n", + " results=[{\"name\": \"final_structure\"}],\n", + " is_rendered_at_run_time=True,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 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." + ] + }, + { + "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", + "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", + " input=[{\"type\": \"api\", \"endpoint\": \"structures\", \"endpoint_options\": endpoint_options,\n", + " \"name\": \"FINAL_STRUCTURE\"}])\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 structure index\", operand=\"FINAL_STRUCTURE_INDEX\", value=0, input=[]),\n", + " save_unit,\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", + " {\"name\": \"FINAL_STRUCTURE_REQUESTS\", \"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", + "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']}\")" + ] + }, + { + "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, timeLimit=TIME_LIMIT)\n", + "print(f\"Using cluster: {compute.cluster.hostname}, queue: {QUEUE_NAME}, nodes: {compute.nodes}, ppn: {compute.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} on {slab.name}, {NUMBER_OF_CONFIGURATIONS} 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; 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." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "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[\"ranked_configurations.csv\"])))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 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[\"final_energies.png\"], as_bytes=True))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 8.3. The lowest final structures" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from mat3ra.prode import PropertyName\n", + "\n", + "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", + "])" + ] + }, + { + "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 +} diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/create_initial_configurations/configurations.py b/other/experiments/alchemi/afir_molecule_surface_alchemi/create_initial_configurations/configurations.py new file mode 100644 index 000000000..dba63dbfb --- /dev/null +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi/create_initial_configurations/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/create_initial_configurations/requirements.txt b/other/experiments/alchemi/afir_molecule_surface_alchemi/create_initial_configurations/requirements.txt new file mode 100644 index 000000000..9aeb82595 --- /dev/null +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi/create_initial_configurations/requirements.txt @@ -0,0 +1 @@ +mat3ra-made[tools]==2026.10.10.post0 diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/relax_under_afir_force_alchemi/afir.py b/other/experiments/alchemi/afir_molecule_surface_alchemi/relax_under_afir_force_alchemi/afir.py new file mode 100644 index 000000000..675d5cbf8 --- /dev/null +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi/relax_under_afir_force_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/relax_under_afir_force_alchemi/requirements.txt b/other/experiments/alchemi/afir_molecule_surface_alchemi/relax_under_afir_force_alchemi/requirements.txt new file mode 100644 index 000000000..0c4b31ef5 --- /dev/null +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi/relax_under_afir_force_alchemi/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]==0.3.0 +matplotlib +numpy +ase +mace-torch +loguru diff --git a/other/experiments/alchemi/afir_molecule_surface_alchemi/save_final_structures/final_structures.py b/other/experiments/alchemi/afir_molecule_surface_alchemi/save_final_structures/final_structures.py new file mode 100644 index 000000000..41cc7a351 --- /dev/null +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi/save_final_structures/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/save_final_structures/requirements.txt b/other/experiments/alchemi/afir_molecule_surface_alchemi/save_final_structures/requirements.txt new file mode 100644 index 000000000..9aeb82595 --- /dev/null +++ b/other/experiments/alchemi/afir_molecule_surface_alchemi/save_final_structures/requirements.txt @@ -0,0 +1 @@ +mat3ra-made[tools]==2026.10.10.post0