submit_gw_ic.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Submit GW calculation for molecules followed by the IC correction"
   ]
  },
  {
   "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",
    "# AiiDAlab imports.\n",
    "import aiidalab_widgets_base as awb\n",
    "from aiida import orm, plugins\n",
    "\n",
    "# Custom imports.\n",
    "from surfaces_tools.widgets import (\n",
    "    analyze_structure,\n",
    "    cdxml,\n",
    "    computational_resources,\n",
    "    inputs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Cp2kAdsorbedGwIcWorkChain = plugins.WorkflowFactory(\"nanotech_empa.cp2k.ads_gw_ic\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Structure selector.\n",
    "input_details = inputs.InputDetails()\n",
    "\n",
    "structure_selector = awb.StructureManagerWidget(\n",
    "    importers=[\n",
    "        awb.StructureBrowserWidget(title=\"AiiDA database\"),\n",
    "        awb.StructureUploadWidget(title=\"Import from computer\"),\n",
    "        awb.SmilesWidget(title=\"From SMILES\"),\n",
    "        cdxml.CdxmlUpload2GnrWidget(title=\"CDXML\"),\n",
    "    ],\n",
    "    storable=False,\n",
    "    node_class=\"StructureData\",\n",
    ")\n",
    "display(structure_selector)\n",
    "ipw.dlink((structure_selector, \"structure\"), (input_details, \"structure\"))\n",
    "\n",
    "# Code.\n",
    "resources = awb.ComputationalResourcesWidget(\n",
    "    description=\"CP2K code:\", default_calc_job_plugin=\"cp2k\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# gw_type\n",
    "gw_type = ipw.Dropdown(\n",
    "    options=[\"gpw_std\", \"gapw_std\", \"gapw_hq\"],\n",
    "    value=\"gpw_std\",\n",
    "    description=\"Protocol:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "\n",
    "# Multiplicity\n",
    "multiplicity = ipw.IntText(value=1, description=\"Multiplicity\", disabled=False)\n",
    "# UKS\n",
    "uks = ipw.Checkbox(value=False, description=\"UKS\", disabled=False, indent=False)\n",
    "\n",
    "# spin guess\n",
    "spins_up = ipw.Text(\n",
    "    value=\"\", placeholder=\"1 2 10..13\", description=\"Spins U:\", disabled=False\n",
    ")\n",
    "\n",
    "spins_down = ipw.Text(\n",
    "    value=\"\", placeholder=\"3 4 14..17\", description=\"Spins D:\", disabled=False\n",
    ")\n",
    "\n",
    "# geometry mode\n",
    "geo_mode = ipw.Dropdown(\n",
    "    options=[\"ads_geo\", \"gas_opt\"],\n",
    "    value=\"ads_geo\",\n",
    "    description=\"Geo mode:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "# geometry mode\n",
    "ads_height = ipw.FloatText(\n",
    "    value=0.0,\n",
    "    description=\"adsorption height in Å:\",\n",
    "    disabled=False,\n",
    "    style={\"description_width\": \"initial\"},\n",
    ")\n",
    "# description\n",
    "description = ipw.Text(\n",
    "    value=\"\",\n",
    "    placeholder=\"Calculation description\",\n",
    "    description=\"Description:\",\n",
    "    disabled=False,\n",
    "    style={\"description_width\": \"initial\"},\n",
    ")\n",
    "\n",
    "\n",
    "# description adsorption height\n",
    "html_ads_height = ipw.HTML(\n",
    "    value=\"\"\"\n",
    "                <p style=\"font-weight:400;\">If you specify a value >0\n",
    "                    <font style=\"font-style:italic;font-weight:400;\">(mandatory in case there is no slab)</font> it will override the height extracted from the geometry.\n",
    "                </p>\n",
    "                <p>The substrate surface is at <font style=\"font-style:italic;font-weight:600;\">geometric center of the molecule - adsorption height</font></p>\n",
    "               \"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_builder_gw():\n",
    "    builder = Cp2kAdsorbedGwIcWorkChain.get_builder()\n",
    "    builder.metadata.description = description.value\n",
    "    builder.code = orm.load_code(resources.value)\n",
    "    builder.geometry_mode = orm.Str(geo_mode.value)\n",
    "\n",
    "    # Override automatic adsorption height.\n",
    "    if ads_height.value > 0.0:\n",
    "        builder.ads_height = orm.Float(ads_height.value)\n",
    "\n",
    "    ase_geom = structure_selector.structure\n",
    "\n",
    "    # Spin guess.\n",
    "    mag_list = [0 for t in ase_geom]\n",
    "    if uks.value:\n",
    "        for i in awb.utils.string_range_to_list(spins_up.value)[0]:\n",
    "            mag_list[i] = 1\n",
    "        for i in string_range_to_list(spins_down.value)[0]:\n",
    "            mag_list[i] = -1\n",
    "\n",
    "        builder.multiplicity = orm.Int(multiplicity.value)\n",
    "\n",
    "    builder.structure = structure_selector.structure_node\n",
    "    builder.magnetization_per_site = orm.List(list=mag_list)\n",
    "\n",
    "    builder.protocol = orm.Str(gw_type.value)\n",
    "\n",
    "    builder.geometry_mode = orm.Str(geo_mode.value)\n",
    "\n",
    "    builder.options.scf = {\n",
    "        \"max_wallclock_seconds\": resources_scf.walltime_seconds,\n",
    "        \"resources\": {\n",
    "            \"num_machines\": resources_scf.nodes,\n",
    "            \"num_mpiprocs_per_machine\": resources_scf.tasks_per_node,\n",
    "            \"num_cores_per_mpiproc\": resources_scf.threads_per_task,\n",
    "        },\n",
    "    }\n",
    "    builder.options.gw = {\n",
    "        \"max_wallclock_seconds\": resources_gw.walltime_seconds,\n",
    "        \"resources\": {\n",
    "            \"num_machines\": resources_gw.nodes,\n",
    "            \"num_mpiprocs_per_machine\": resources_gw.tasks_per_node,\n",
    "            \"num_cores_per_mpiproc\": resources_gw.threads_per_task,\n",
    "        },\n",
    "    }\n",
    "    builder.options.ic = {\n",
    "        \"max_wallclock_seconds\": resources_ic.walltime_seconds,\n",
    "        \"resources\": {\n",
    "            \"num_machines\": resources_ic.nodes,\n",
    "            \"num_mpiprocs_per_machine\": resources_ic.tasks_per_node,\n",
    "            \"num_cores_per_mpiproc\": resources_ic.threads_per_task,\n",
    "        },\n",
    "    }\n",
    "\n",
    "    return builder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def after_submission(_=None):\n",
    "    structure_selector.value = None\n",
    "\n",
    "\n",
    "btn_submit_gw = awb.SubmitButtonWidget(\n",
    "    Cp2kAdsorbedGwIcWorkChain, inputs_generator=get_builder_gw\n",
    ")\n",
    "btn_submit_gw.btn_submit.disabled = True\n",
    "\n",
    "btn_submit_gw.on_submitted(after_submission)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = ipw.Output()\n",
    "\n",
    "\n",
    "def update_all(_=None):\n",
    "    btn_submit_gw.btn_submit.disabled = False\n",
    "    # check system\n",
    "    only_one_molecule = input_details.details[\"system_type\"] == \"SlabXY\"\n",
    "    only_one_molecule = (\n",
    "        only_one_molecule or input_details.details[\"system_type\"] == \"Molecule\"\n",
    "    )\n",
    "    only_one_molecule = (\n",
    "        only_one_molecule and len(input_details.details[\"all_molecules\"]) == 1\n",
    "    )\n",
    "    msg = \"GW for this system not implemented\"\n",
    "    if only_one_molecule:\n",
    "        btn_submit_gw.btn_submit.disabled = False\n",
    "        msg = \"\"\n",
    "    spins_up.value = awb.utils.list_to_string_range(input_details.details[\"spins_up\"])\n",
    "    spins_down.value = awb.utils.list_to_string_range(\n",
    "        input_details.details[\"spins_down\"]\n",
    "    )\n",
    "    with output:\n",
    "        clear_output()\n",
    "        print(msg)\n",
    "\n",
    "        if uks.value:\n",
    "            to_display = [\n",
    "                gw_type,\n",
    "                geo_mode,\n",
    "                ipw.HBox([ads_height, html_ads_height]),\n",
    "                uks,\n",
    "                spins_up,\n",
    "                spins_down,\n",
    "                multiplicity,\n",
    "                description,\n",
    "            ]\n",
    "        else:\n",
    "            to_display = [\n",
    "                gw_type,\n",
    "                geo_mode,\n",
    "                ipw.HBox([ads_height, html_ads_height]),\n",
    "                uks,\n",
    "            ]\n",
    "        display(ipw.VBox(to_display))\n",
    "\n",
    "\n",
    "structure_selector.observe(update_all, names=\"structure\")\n",
    "uks.observe(update_all, names=\"value\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Resources.\n",
    "resources_scf = computational_resources.ProcessResourcesWidget()\n",
    "resources_gw = computational_resources.ProcessResourcesWidget()\n",
    "resources_ic = computational_resources.ProcessResourcesWidget()\n",
    "\n",
    "# Resources estimation.\n",
    "resources_estimation_scf = computational_resources.ResourcesEstimatorWidget(\n",
    "    calculation_type=\"dft\"\n",
    ")\n",
    "resources_estimation_gw = computational_resources.ResourcesEstimatorWidget(\n",
    "    calculation_type=\"gw\"\n",
    ")\n",
    "resources_estimation_ic = computational_resources.ResourcesEstimatorWidget(\n",
    "    calculation_type=\"gw_ic\"\n",
    ")\n",
    "\n",
    "# Link resources widgets.\n",
    "resources_estimation_scf.link_to_resources_widget(resources_scf)\n",
    "resources_estimation_gw.link_to_resources_widget(resources_gw)\n",
    "resources_estimation_ic.link_to_resources_widget(resources_ic)\n",
    "\n",
    "\n",
    "def extract_details_on_adsorbed_molecule(all_atoms):\n",
    "    structure_analyzer = analyze_structure.StructureAnalyzer()\n",
    "    structure_analyzer.structure = all_atoms\n",
    "    if \"all_molecules\" in structure_analyzer.details:\n",
    "        molecules_indices = [\n",
    "            item\n",
    "            for sublist in structure_analyzer.details[\"all_molecules\"]\n",
    "            for item in sublist\n",
    "        ]\n",
    "        structure_analyzer.structure = all_atoms[molecules_indices]\n",
    "        return structure_analyzer.details\n",
    "    else:\n",
    "        return {}\n",
    "\n",
    "\n",
    "# Link viewer to resources estimation widgets.\n",
    "ipw.dlink(\n",
    "    (structure_selector, \"structure\"),\n",
    "    (resources_estimation_scf, \"details\"),\n",
    "    transform=extract_details_on_adsorbed_molecule,\n",
    ")\n",
    "ipw.dlink(\n",
    "    (structure_selector, \"structure\"),\n",
    "    (resources_estimation_gw, \"details\"),\n",
    "    transform=extract_details_on_adsorbed_molecule,\n",
    ")\n",
    "ipw.dlink(\n",
    "    (structure_selector, \"structure\"),\n",
    "    (resources_estimation_ic, \"details\"),\n",
    "    transform=extract_details_on_adsorbed_molecule,\n",
    ")\n",
    "\n",
    "# Link code selector to resources estimation widgets.\n",
    "_ = ipw.dlink((resources, \"value\"), (resources_estimation_scf, \"selected_code\"))\n",
    "_ = ipw.dlink((resources, \"value\"), (resources_estimation_gw, \"selected_code\"))\n",
    "_ = ipw.dlink((resources, \"value\"), (resources_estimation_ic, \"selected_code\"))\n",
    "\n",
    "# Link UKS\n",
    "_ = ipw.dlink((uks, \"value\"), (resources_estimation_scf, \"uks\"))\n",
    "_ = ipw.dlink((uks, \"value\"), (resources_estimation_gw, \"uks\"))\n",
    "_ = ipw.dlink((uks, \"value\"), (resources_estimation_ic, \"uks\"))\n",
    "\n",
    "# Estimate all resources\n",
    "estimate_nodes_button = ipw.Button(\n",
    "    description=\"Estimate resources\", button_style=\"warning\"\n",
    ")\n",
    "estimate_nodes_button.on_click(resources_estimation_scf.estimate_resources)\n",
    "estimate_nodes_button.on_click(resources_estimation_gw.estimate_resources)\n",
    "estimate_nodes_button.on_click(resources_estimation_ic.estimate_resources)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(\n",
    "    output,\n",
    "    description,\n",
    "    ipw.HBox(\n",
    "        [\n",
    "            ipw.VBox(\n",
    "                [\n",
    "                    ipw.HTML(\"DFT resources\"),\n",
    "                    resources_scf,\n",
    "                ]\n",
    "            ),\n",
    "            ipw.VBox(\n",
    "                [\n",
    "                    ipw.HTML(\"GW resources\"),\n",
    "                    resources_gw,\n",
    "                ]\n",
    "            ),\n",
    "            ipw.VBox(\n",
    "                [\n",
    "                    ipw.HTML(\"IC resources\"),\n",
    "                    resources_ic,\n",
    "                ]\n",
    "            ),\n",
    "        ]\n",
    "    ),\n",
    "    estimate_nodes_button,\n",
    "    resources,\n",
    "    btn_submit_gw,\n",
    ")"
   ]
  }
 ],
 "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
}