submit_geometry_optimization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Submit geometry optimization"
   ]
  },
  {
   "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 urllib.parse as urlparse\n",
    "\n",
    "import ipywidgets as ipw\n",
    "from IPython.display import clear_output\n",
    "\n",
    "# AiiDA imports.\n",
    "%load_ext aiida\n",
    "%aiida\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": [
    "Cp2kGeoOptWorkChain = plugins.WorkflowFactory(\"nanotech_empa.cp2k.geo_opt\")"
   ]
  },
  {
   "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",
    "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=[\n",
    "        awb.BasicStructureEditor(title=\"Edit structure\"),\n",
    "        build_slab,\n",
    "        awb.BasicCellEditor(),\n",
    "        editors.InsertStructureWidget(title=\"Insert molecule\"),\n",
    "    ],\n",
    "    storable=True,\n",
    "    node_class=\"StructureData\",\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",
    "display(structure_selector)\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",
    "output = ipw.Output()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "url = urlparse.urlsplit(jupyter_notebook_url)\n",
    "parsed_url = urlparse.parse_qs(url.query)\n",
    "if \"structure_uuid\" in parsed_url:\n",
    "    structure_selector.input_structure = load_node(parsed_url[\"structure_uuid\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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",
    ")"
   ]
  },
  {
   "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_geometry_optimization():\n",
    "    with output:\n",
    "        clear_output()\n",
    "    if not structure_selector.structure_node:\n",
    "        can_submit, msg = False, \"Select a structure first.\"\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 = Cp2kGeoOptWorkChain.get_builder()\n",
    "    builder.protocol = orm.Str(protocol.value)\n",
    "    builder.metadata.label = (\n",
    "        \"CP2K_CellOpt\" if input_details.do_cell_opt else \"CP2K_GeoOpt\"\n",
    "    )\n",
    "    builder.metadata.description = workflow_description.value\n",
    "    code = orm.load_code(code_input_widget.value)\n",
    "    builder.code = code\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",
    "\n",
    "    # Check if a restart wfn is available.\n",
    "    wave_function = None\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",
    "    builder.dft_params = orm.Dict(parameters[\"dft_params\"])\n",
    "    builder.sys_params = orm.Dict(parameters[\"sys_params\"])\n",
    "\n",
    "    return builder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "btn_submit = awb.SubmitButtonWidget(\n",
    "    Cp2kGeoOptWorkChain,\n",
    "    inputs_generator=prepare_geometry_optimization,\n",
    "    disable_after_submit=False,\n",
    "    append_output=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Resources estimation.\n",
    "resources_estimation = computational_resources.ResourcesEstimatorWidget()\n",
    "resources_estimation.link_to_resources_widget(resources)\n",
    "ipw.dlink((input_details, \"details\"), (resources_estimation, \"details\"))\n",
    "ipw.dlink((input_details, \"uks\"), (resources_estimation, \"uks\"))\n",
    "_ = ipw.dlink((code_input_widget, \"value\"), (resources_estimation, \"selected_code\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(input_details, protocol)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Code and resources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(code_input_widget, resources, resources_estimation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Submit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(workflow_description, btn_submit, 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
}