submit_reactions.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Submit nudget elastic band or replica chain calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%javascript\n",
    "IPython.OutputArea.prototype._should_scroll = function(lines) {\n",
    "    return false;\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# General imports.\n",
    "import ipywidgets as ipw\n",
    "from IPython.display import clear_output\n",
    "\n",
    "# AiiDA imports.\n",
    "%load_ext aiida\n",
    "%aiida\n",
    "\n",
    "# AiiDAlab imports.\n",
    "import aiidalab_widgets_base as awb\n",
    "from aiida import orm, plugins\n",
    "\n",
    "from surfaces_tools.utils import wfn\n",
    "\n",
    "# Custom imports.\n",
    "from surfaces_tools.widgets import cdxml, computational_resources, editors, inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Cp2kNebWorkChain = plugins.WorkflowFactory(\"nanotech_empa.cp2k.neb\")\n",
    "Cp2kReplicaWorkChain = plugins.WorkflowFactory(\"nanotech_empa.cp2k.replica\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Structure selector.\n",
    "build_slab = editors.BuildSlab(title=\"Build slab\")\n",
    "input_details = inputs.InputDetails()\n",
    "\n",
    "\n",
    "structure_selector = awb.StructureManagerWidget(\n",
    "    importers=[\n",
    "        awb.StructureUploadWidget(title=\"Import from computer\"),\n",
    "        awb.StructureBrowserWidget(title=\"AiiDA database\"),\n",
    "        awb.SmilesWidget(title=\"From SMILES\"),\n",
    "        cdxml.CdxmlUpload2GnrWidget(title=\"CDXML\"),\n",
    "    ],\n",
    "    editors=[awb.BasicStructureEditor(title=\"Edit structure\"), build_slab],\n",
    "    storable=False,\n",
    "    node_class=\"StructureData\",\n",
    ")\n",
    "\n",
    "ipw.dlink((structure_selector, \"structure\"), (build_slab, \"molecule\"))\n",
    "ipw.dlink((structure_selector, \"structure\"), (input_details, \"structure\"))\n",
    "ipw.dlink((input_details, \"details\"), (build_slab, \"details\"))\n",
    "\n",
    "\n",
    "display(structure_selector)\n",
    "\n",
    "replica_or_neb = ipw.ToggleButtons(\n",
    "    options=[\"Replica\", \"NEB\"],\n",
    "    description=\"Type:\",\n",
    "    disabled=False,\n",
    "    button_style=\"\",\n",
    "    tooltips=[\"Replica\", \"NEB\"],\n",
    "    icons=[\"check\", \"check\"],\n",
    "    style={\"description_width\": \"initial\"},\n",
    "    value=\"NEB\",\n",
    ")\n",
    "\n",
    "# Code.\n",
    "code_input_widget = awb.ComputationalResourcesWidget(\n",
    "    description=\"CP2K code:\", default_calc_job_plugin=\"cp2k\"\n",
    ")\n",
    "resources = computational_resources.ProcessResourcesWidget()\n",
    "\n",
    "# Global variables.\n",
    "label = \"CP2K_NEB\"\n",
    "the_workchain = Cp2kNebWorkChain\n",
    "\n",
    "# Protocol.\n",
    "protocol = ipw.Dropdown(\n",
    "    value=\"standard\",\n",
    "    options=[\n",
    "        (\"Standard\", \"standard\"),\n",
    "        (\"Low accuracy\", \"low_accuracy\"),\n",
    "        (\"Debug\", \"debug\"),\n",
    "    ],\n",
    "    description=\"Protocol:\",\n",
    "    style={\"description_width\": \"initial\"},\n",
    ")\n",
    "\n",
    "output = ipw.Output()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow_description = ipw.Text(\n",
    "    description=\"Workflow description:\",\n",
    "    placeholder=\"Provide the description here.\",\n",
    "    style={\"description_width\": \"initial\"},\n",
    "    layout={\"width\": \"70%\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ipw.dlink((code_input_widget, \"value\"), (input_details, \"selected_code\"))\n",
    "\n",
    "\n",
    "def prepare_inputs():\n",
    "    with output:\n",
    "        clear_output()\n",
    "    if not structure_selector.structure_node:\n",
    "        can_submit, msg = False, \"Select a structure furst.\"\n",
    "    elif not code_input_widget.value:\n",
    "        can_submit, msg = False, \"Select CP2K code.\"\n",
    "    else:\n",
    "        can_submit, msg, parameters = input_details.return_final_dictionary()\n",
    "\n",
    "    if not can_submit:\n",
    "        with output:\n",
    "            print(msg)\n",
    "            return\n",
    "\n",
    "    builder = the_workchain.get_builder()\n",
    "    builder.protocol = orm.Str(protocol.value)\n",
    "    builder.metadata.label = label\n",
    "    builder.metadata.description = workflow_description.value\n",
    "    builder.code = orm.load_code(code_input_widget.value)\n",
    "    builder.options = {\n",
    "        \"max_wallclock_seconds\": resources.walltime_seconds,\n",
    "        \"resources\": {\n",
    "            \"num_machines\": resources.nodes,\n",
    "            \"num_mpiprocs_per_machine\": resources.tasks_per_node,\n",
    "            \"num_cores_per_mpiproc\": resources.threads_per_task,\n",
    "        },\n",
    "    }\n",
    "\n",
    "    builder.structure = structure_selector.structure_node\n",
    "    builder.dft_params = orm.Dict(parameters[\"dft_params\"])\n",
    "    builder.sys_params = orm.Dict(parameters[\"sys_params\"])\n",
    "    if \"neb_params\" in parameters:\n",
    "        builder.neb_params = orm.Dict(parameters[\"neb_params\"])\n",
    "\n",
    "    if \"replica_uuids\" in parameters:\n",
    "        replicas = {}\n",
    "        for i in range(len(parameters[\"replica_uuids\"])):\n",
    "            name = \"replica_\" + str(i + 1).zfill(3)\n",
    "            replicas[name] = orm.load_node(parameters[\"replica_uuids\"][i])\n",
    "        builder.replicas = replicas\n",
    "\n",
    "    if \"restart_from\" in parameters:\n",
    "        builder.restart_from = orm.Str(parameters[\"restart_from\"])\n",
    "    elif replica_or_neb.value == \"Replica\":\n",
    "        # Check if a restart wfn is available for the replica chain calculations not restarting from a previous calculation.\n",
    "        wave_function = None\n",
    "        code = orm.load_code(code_input_widget.value)\n",
    "        if structure_selector.structure_node.is_stored:\n",
    "            wave_function = wfn.structure_available_wfn(\n",
    "                node=structure_selector.structure_node,\n",
    "                relative_replica_id=None,\n",
    "                current_hostname=code.computer.hostname,\n",
    "                return_path=False,\n",
    "                dft_params=parameters[\"dft_params\"],\n",
    "            )\n",
    "        if wave_function is not None:\n",
    "            print(f\"Restarting from wfn in folder: {wave_function.pk}\")\n",
    "            builder.parent_calc_folder = wave_function\n",
    "\n",
    "    return builder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "btn_output = ipw.Output()\n",
    "btn_submit_neb = awb.SubmitButtonWidget(\n",
    "    Cp2kNebWorkChain,\n",
    "    inputs_generator=prepare_inputs,\n",
    "    disable_after_submit=False,\n",
    "    append_output=True,\n",
    ")\n",
    "btn_submit_replica = awb.SubmitButtonWidget(\n",
    "    Cp2kReplicaWorkChain,\n",
    "    inputs_generator=prepare_inputs,\n",
    "    disable_after_submit=False,\n",
    "    append_output=True,\n",
    ")\n",
    "\n",
    "with btn_output:\n",
    "    display(btn_submit_neb)\n",
    "\n",
    "\n",
    "def on_replica_or_neb_change(change):\n",
    "    global label, the_workchain\n",
    "    with btn_output:\n",
    "        clear_output()\n",
    "        if change[\"new\"] == \"Replica\":\n",
    "            label = \"CP2K_Replica\"\n",
    "            the_workchain = Cp2kReplicaWorkChain\n",
    "            input_details.replica = True\n",
    "            input_details.neb = False\n",
    "            display(btn_submit_replica)\n",
    "        elif change[\"new\"] == \"NEB\":\n",
    "            label = \"CP2K_NEB\"\n",
    "            the_workchain = Cp2kNebWorkChain\n",
    "            input_details.replica = False\n",
    "            input_details.neb = True\n",
    "            display(btn_submit_neb)\n",
    "\n",
    "\n",
    "replica_or_neb.observe(on_replica_or_neb_change, names=\"value\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Resources estimation.\n",
    "resources_estimation_button = computational_resources.ResourcesEstimatorWidget()\n",
    "resources_estimation_button.link_to_resources_widget(resources)\n",
    "ipw.dlink((structure_selector, \"structure\"), (input_details, \"structure\"))\n",
    "ipw.dlink((input_details, \"details\"), (resources_estimation_button, \"details\"))\n",
    "ipw.dlink((input_details, \"uks\"), (resources_estimation_button, \"uks\"))\n",
    "ipw.dlink((input_details, \"neb\"), (resources, \"neb\"))\n",
    "ipw.dlink((input_details, \"n_replica_trait\"), (resources, \"n_replica_trait\"))\n",
    "ipw.dlink(\n",
    "    (input_details, \"n_replica_per_group_trait\"),\n",
    "    (resources, \"n_replica_per_group_trait\"),\n",
    ")\n",
    "ipw.dlink((resources, \"nproc_replica_trait\"), (input_details, \"nproc_replica_trait\"))\n",
    "_ = ipw.dlink(\n",
    "    (code_input_widget, \"value\"), (resources_estimation_button, \"selected_code\")\n",
    ")\n",
    "input_details.replica = False\n",
    "input_details.neb = True"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(protocol, replica_or_neb, input_details)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Code and resources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(code_input_widget, resources, resources_estimation_button)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Submit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(workflow_description, btn_output, output)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "d4d1e4263499bec80672ea0156c357c1ee493ec2b1c70f0acce89fc37c4a6abe"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}