{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "08edae93",
   "metadata": {},
   "source": [
    "## Initial Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "396e9b99",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PLAMS working folder: /path/plams/examples/ExcitationsWorkflow/plams_workdir\n"
     ]
    }
   ],
   "source": [
    "from scm.plams import AMSResults, Units, add_to_class, Settings, read_molecules, AMSJob, init\n",
    "\n",
    "# this line is not required in AMS2025+\n",
    "init()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a63d1ac",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## Helper Functions\n",
    "Set up a couple of useful functions for extracting results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c94154ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "@add_to_class(AMSResults)\n",
    "def get_excitations(results):\n",
    "    \"\"\"Returns excitation energies (in eV) and oscillator strengths (in Debye).\"\"\"\n",
    "    if results.job.ok():\n",
    "        exci_energies_au = results.readrkf(\"Excitations SS A\", \"excenergies\", file=\"engine\")\n",
    "        oscillator_str_au = results.readrkf(\"Excitations SS A\", \"oscillator strengths\", file=\"engine\")\n",
    "        # The results are stored in atomic units. Convert them to more convenient units:\n",
    "        exci_energies = Units.convert(exci_energies_au, \"au\", \"eV\")\n",
    "        oscillator_str = Units.convert(oscillator_str_au, \"au\", \"Debye\")\n",
    "        return exci_energies, oscillator_str\n",
    "    else:\n",
    "        return [], []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "812f3721",
   "metadata": {},
   "outputs": [],
   "source": [
    "@add_to_class(AMSResults)\n",
    "def has_good_excitations(results, min_energy, max_energy, oscillator_str_threshold=1e-4):\n",
    "    \"\"\"Returns True if there is at least one excitation with non-vanishing oscillator strenght\n",
    "    in the energy range [min_energy, max_energy]. Unit for min_energy and max energy: eV.\"\"\"\n",
    "    exci_energies, oscillator_str = results.get_excitations()\n",
    "    for e, o in zip(exci_energies, oscillator_str):\n",
    "        if min_energy < e < max_energy and o > oscillator_str_threshold:\n",
    "            return True\n",
    "    return False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63520b81",
   "metadata": {},
   "source": [
    "## Calculation settings\n",
    "\n",
    "Configure the settings for the various jobs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "911a62ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Settings for geometry optimization with the AMS driver:\n",
    "go_sett = Settings()\n",
    "go_sett.input.ams.Task = \"GeometryOptimization\"\n",
    "go_sett.input.ams.GeometryOptimization.Convergence.Gradients = 1.0e-4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6f625c6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Settings for single point calculation with the AMS driver\n",
    "sp_sett = Settings()\n",
    "sp_sett.input.ams.Task = \"SinglePoint\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0dcca668",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Settings for the DFTB engine (including excitations)\n",
    "dftb_sett = Settings()\n",
    "dftb_sett.input.dftb.Model = \"SCC-DFTB\"\n",
    "dftb_sett.input.dftb.ResourcesDir = \"QUASINANO2015\"\n",
    "dftb_sett.input.dftb.Properties.Excitations.TDDFTB.calc = \"singlet\"\n",
    "dftb_sett.input.dftb.Properties.Excitations.TDDFTB.lowest = 10\n",
    "dftb_sett.input.dftb.Occupation.Temperature = 5.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "33932ecc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Settings for the geometry optimization with the ADF engine\n",
    "adf_sett = Settings()\n",
    "adf_sett.input.adf.Basis.Type = \"DZP\"\n",
    "adf_sett.input.adf.NumericalQuality = \"Basic\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3eb2b834",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# Settings for the excitation calculation using the ADF engine\n",
    "adf_exci_sett = Settings()\n",
    "adf_exci_sett.input.adf.Basis.Type = \"TZP\"\n",
    "adf_exci_sett.input.adf.XC.GGA = \"PBE\"\n",
    "adf_exci_sett.input.adf.NumericalQuality = \"Basic\"\n",
    "adf_exci_sett.input.adf.Symmetry = \"NoSym\"\n",
    "adf_exci_sett.input.adf.Excitations.lowest = 10\n",
    "adf_exci_sett.input.adf.Excitations.OnlySing = \"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dc547d5",
   "metadata": {},
   "source": [
    "## Load Molecules\n",
    "Import all xyz files in the folder 'molecules'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "abc9ab9d",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "molecules = read_molecules(\"molecules\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d2a078d",
   "metadata": {},
   "source": [
    "## DFTB Prescreen\n",
    "Perform an initial prescreen of all molecules with DFTB."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "07ad9788",
   "metadata": {},
   "outputs": [],
   "source": [
    "promising_molecules = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "efe03566",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[10.02|15:15:10] JOB DFTB_H2O STARTED\n",
      "[10.02|15:15:10] JOB DFTB_H2O RUNNING\n",
      "[10.02|15:15:12] JOB DFTB_H2O FINISHED\n",
      "[10.02|15:15:12] JOB DFTB_H2O SUCCESSFUL\n",
      "[10.02|15:15:12] JOB DFTB_NH3 STARTED\n",
      "[10.02|15:15:12] JOB DFTB_NH3 RUNNING\n",
      "[10.02|15:15:21] WARNING: Job DFTB_NH3 finished with nonzero return code\n",
      "[10.02|15:15:21] JOB DFTB_NH3 CRASHED\n",
      "[10.02|15:15:21] Error message for job DFTB_NH3 was:\n",
      "\tAtoms 3 and 4 are extremely close (<0.001 Bohr). Input error? If this was intended, set the System%AllowCloseAtoms option to True.\n",
      "[10.02|15:15:21] JOB DFTB_S2Cl2 STARTED\n",
      "[10.02|15:15:21] JOB DFTB_S2Cl2 RUNNING\n",
      "[10.02|15:15:22] JOB DFTB_S2Cl2 FINISHED\n",
      "[10.02|15:15:22] JOB DFTB_S2Cl2 SUCCESSFUL\n",
      "[10.02|15:15:22] JOB DFTB_AlF3 STARTED\n",
      "[10.02|15:15:22] JOB DFTB_AlF3 RUNNING\n",
      "[10.02|15:15:25] JOB DFTB_AlF3 FINISHED\n",
      "[10.02|15:15:25] JOB DFTB_AlF3 SUCCESSFUL\n",
      "[10.02|15:15:25] JOB DFTB_CSCl2 STARTED\n",
      "[10.02|15:15:25] JOB DFTB_CSCl2 RUNNING\n",
      "[10.02|15:15:26] JOB DFTB_CSCl2 FINISHED\n",
      "[10.02|15:15:26] JOB DFTB_CSCl2 SUCCESSFUL\n"
     ]
    }
   ],
   "source": [
    "for name, mol in molecules.items():\n",
    "    dftb_job = AMSJob(name=\"DFTB_\" + name, molecule=mol, settings=go_sett + dftb_sett)\n",
    "    dftb_job.run()\n",
    "\n",
    "    if dftb_job.results.has_good_excitations(1, 6):\n",
    "        promising_molecules[name] = dftb_job.results.get_main_molecule()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9b7877ce",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 2 promising molecules with DFTB\n"
     ]
    }
   ],
   "source": [
    "print(f\"Found {len(promising_molecules)} promising molecules with DFTB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59c0fe21",
   "metadata": {},
   "source": [
    "## Optimization and excitations calculation with ADF\n",
    "For each of the molecules identified in the prescreen, run a further calculation with ADF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0e9a4be",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[10.02|15:15:26] JOB ADF_GO_S2Cl2 STARTED\n",
      "[10.02|15:15:26] JOB ADF_GO_S2Cl2 RUNNING\n",
      "[10.02|15:15:37] JOB ADF_GO_S2Cl2 FINISHED\n",
      "[10.02|15:15:37] JOB ADF_GO_S2Cl2 SUCCESSFUL\n",
      "[10.02|15:15:37] JOB ADF_exci_S2Cl2 STARTED\n",
      "[10.02|15:15:37] JOB ADF_exci_S2Cl2 RUNNING\n",
      "[10.02|15:15:47] JOB ADF_exci_S2Cl2 FINISHED\n",
      "[10.02|15:15:47] JOB ADF_exci_S2Cl2 SUCCESSFUL\n",
      "Molecule S2Cl2 has excitation(s) satysfying our criteria!\n",
      "  Atoms: \n",
      "    1         S      -0.658306      -0.316643       0.909151\n",
      "    2         S      -0.658306       0.316643      -0.909151\n",
      "    3        Cl       0.758306       0.752857       2.053019\n",
      "    4        Cl       0.758306      -0.752857      -2.053019\n",
      "\n",
      "Excitation energy [eV], oscillator strength:\n",
      "  3.4107,   0.0114\n",
      "  3.5386,   0.0160\n",
      "  3.5400,   0.0011\n",
      "  3.9864,   0.1105\n",
      "  4.3225,   0.0049\n",
      "  4.3513,   0.2551\n",
      "  4.7544,   0.0011\n",
      "  4.9414,   0.0105\n",
      "  5.3188,   0.0036\n",
      "  5.3272,   0.0721\n",
      "[10.02|15:15:47] JOB ADF_GO_CSCl2 STARTED\n",
      "[10.02|15:15:47] JOB ADF_GO_CSCl2 RUNNING\n",
      "[10.02|15:15:58] JOB ADF_GO_CSCl2 FINISHED\n",
      "[10.02|15:15:58] JOB ADF_GO_CSCl2 SUCCESSFUL\n",
      "[10.02|15:15:58] JOB ADF_exci_CSCl2 STARTED\n",
      "[10.02|15:15:58] JOB ADF_exci_CSCl2 RUNNING\n"
     ]
    }
   ],
   "source": [
    "for name, mol in promising_molecules.items():\n",
    "    adf_go_job = AMSJob(name=\"ADF_GO_\" + name, molecule=mol, settings=go_sett + adf_sett)\n",
    "    adf_go_job.run()\n",
    "\n",
    "    optimized_mol = adf_go_job.results.get_main_molecule()\n",
    "\n",
    "    adf_exci_job = AMSJob(name=\"ADF_exci_\" + name, molecule=optimized_mol, settings=sp_sett + adf_exci_sett)\n",
    "    adf_exci_job.run()\n",
    "\n",
    "    if adf_exci_job.results.has_good_excitations(2, 4):\n",
    "        print(f\"Molecule {name} has excitation(s) satysfying our criteria!\")\n",
    "        print(optimized_mol)\n",
    "        exci_energies, oscillator_str = adf_exci_job.results.get_excitations()\n",
    "        print(\"Excitation energy [eV], oscillator strength:\")\n",
    "        for e, o in zip(exci_energies, oscillator_str):\n",
    "            print(f\"{e:8.4f}, {o:8.4f}\")"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "executable": "/usr/bin/env amspython",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  },
  "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}