{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3b1f2b27",
   "metadata": {},
   "source": [
    "## Initial Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e35f8232",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PLAMS working folder: /path/plams/examples/TuningRangeSeparation/plams_workdir\n"
     ]
    }
   ],
   "source": [
    "import multiprocessing\n",
    "import numpy as np\n",
    "from scm.plams import Settings, Results, MultiJob, JobRunner, config, Molecule, Atom, AMSJob, init\n",
    "\n",
    "# this line is not required in AMS2025+\n",
    "init()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f09740d7",
   "metadata": {},
   "source": [
    "## Helper Classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d0516495",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GammaResults(Results):\n",
    "\n",
    "    @staticmethod\n",
    "    def get_difference(job, jobplus):\n",
    "        \"\"\"Calculate the difference between HOMO and IP.\n",
    "        *jobplus* should be the counterpart of *job* with one less electron.\"\"\"\n",
    "        homo = job.results.readrkf(\"Properties\", \"HOMO\", file=\"engine\")\n",
    "        IP = jobplus.results.get_energy() - job.results.get_energy()\n",
    "        return IP + homo\n",
    "\n",
    "    def get_J(self):\n",
    "        N = GammaResults.get_difference(self.job.children[1], self.job.children[2])\n",
    "        A = GammaResults.get_difference(self.job.children[0], self.job.children[1])\n",
    "        return (N**2 + A**2) ** 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5b2877aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GammaJob(MultiJob):\n",
    "    _result_type = GammaResults\n",
    "\n",
    "    def __init__(self, molecule, gamma, charge, spins, **kwargs):\n",
    "        MultiJob.__init__(self, **kwargs)\n",
    "        self.molecule = molecule\n",
    "        self.charge = charge\n",
    "        self.spins = spins\n",
    "        self.gamma = gamma\n",
    "\n",
    "    def prerun(self):\n",
    "        charges = [self.charge - 1, self.charge, self.charge + 1]\n",
    "        for charge, spin in zip(charges, self.spins):\n",
    "            name = \"{}_charge_{}\".format(self.name, charge)\n",
    "            name = name.replace(\"-\", \"minus\")\n",
    "            newjob = AMSJob(name=name, molecule=self.molecule, settings=self.settings)\n",
    "            newjob.molecule.properties.charge = charge\n",
    "            newjob.settings.input.adf.xc.rangesep = \"gamma={:f}\".format(self.gamma)\n",
    "            if spin != 0:\n",
    "                newjob.settings.input.adf.unrestricted = True\n",
    "                newjob.settings.input.adf.SpinPolarization = spin\n",
    "\n",
    "            self.children.append(newjob)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d122a4ce",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def gamma_scan(gammas, settings, molecule, name=\"scan\", charge=0, spins=(1, 0, 1)):\n",
    "    \"\"\"Calculate values of J function for given range of gammas.\n",
    "\n",
    "    Arguments:\n",
    "    gammas   - list of gamma values to calculate the J function for\n",
    "    settings - Settings object for an ADF calculation\n",
    "    molecule - Molecule object with the system of interest\n",
    "    name     - base name of all the jobs\n",
    "    charge   - base charge of the system of interest. The J function is going to be\n",
    "               calculated based on two systems: with charge, and charge-1\n",
    "    spins    - values of spin polarization for jobs with, respectively, charge-1,\n",
    "               charge and charge +1\n",
    "\n",
    "    In other words, if charge=X and spins=(a,b,c) the three resulting jobs\n",
    "    are going to have the following values for charge and spin:\n",
    "\n",
    "    Charge=X-1  SpinPolarization=a\n",
    "    Charge=X    SpinPolarization=b\n",
    "    Charge=X+1  SpinPolarization=c\n",
    "\n",
    "    Returns a list of pairs (gamma, J) of the same length as the parameter *gammas*\n",
    "    \"\"\"\n",
    "    jobs = [\n",
    "        GammaJob(\n",
    "            molecule=molecule, settings=settings, gamma=g, charge=charge, spins=spins, name=name + \"_gamma_\" + str(g)\n",
    "        )\n",
    "        for g in gammas\n",
    "    ]\n",
    "    results = [j.run() for j in jobs]\n",
    "    js = [r.get_J() for r in results]\n",
    "    return list(zip(gammas, js))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6ec49cb",
   "metadata": {},
   "source": [
    "## Configure Parallel JobRunner\n",
    "Set up the default jobrunner to run in parallel, with as many jobs as there are cores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e2200ace",
   "metadata": {},
   "outputs": [],
   "source": [
    "config.default_jobrunner = JobRunner(parallel=True, maxjobs=multiprocessing.cpu_count())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53eb08e1",
   "metadata": {},
   "source": [
    "## Settings of the ADF calculations\n",
    "Configure settings object for the calculation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "927bac04",
   "metadata": {},
   "outputs": [],
   "source": [
    "s = Settings()\n",
    "s.input.ams.task = \"SinglePoint\"\n",
    "s.input.adf.basis.type = \"DZP\"\n",
    "s.input.adf.basis.core = \"None\"\n",
    "s.input.adf.xc.gga = \"PBE\"\n",
    "s.input.adf.xc.xcfun = True\n",
    "s.runscript.nproc = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a614ba3a",
   "metadata": {},
   "source": [
    "## Set Up Molecule\n",
    "Create a toy hydrogen dimer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a416b532",
   "metadata": {},
   "outputs": [],
   "source": [
    "mol = Molecule()\n",
    "mol.add_atom(Atom(symbol=\"H\", coords=(0, 0, -0.3540)))\n",
    "mol.add_atom(Atom(symbol=\"H\", coords=(0, 0, 0.3540)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "211f92f3",
   "metadata": {},
   "source": [
    "## Calculate Gamma Values\n",
    "Perform a scan of a few values for gamma. In practice, you want to scan a wider range and smaller step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d681e630",
   "metadata": {},
   "outputs": [],
   "source": [
    "gammas = np.around(np.arange(1.2, 1.9, 0.2), decimals=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d77cfda5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[11.02|10:02:41] JOB scan_gamma_1.2 STARTED\n",
      "[11.02|10:02:41] JOB scan_gamma_1.4 STARTED\n",
      "[11.02|10:02:41] JOB scan_gamma_1.6 STARTED\n",
      "[11.02|10:02:41] JOB scan_gamma_1.8 STARTED\n",
      "[11.02|10:02:41] JOB scan_gamma_1.2 RUNNING\n",
      "[11.02|10:02:41] Waiting for job scan_gamma_1.2 to finish\n",
      "[11.02|10:02:41] JOB scan_gamma_1.2/scan_gamma_1.2_charge_minus1 STARTED\n",
      "[11.02|10:02:41] JOB scan_gamma_1.2/scan_gamma_1.2_charge_0 STARTED\n",
      "[11.02|10:02:41] JOB scan_gamma_1.2/scan_gamma_1.2_charge_1 STARTED\n",
      "[11.02|10:02:41] JOB scan_gamma_1.4 RUNNING\n",
      "[11.02|10:02:41] JOB scan_gamma_1.4/scan_gamma_1.4_charge_minus1 STARTED\n",
      "[11.02|10:02:41] JOB scan_gamma_1.6 RUNNING\n",
      "[11.02|10:02:41] JOB scan_gamma_1.4/scan_gamma_1.4_charge_0 STARTED\n",
      "[11.02|10:02:41] JOB scan_gamma_1.4/scan_gamma_1.4_charge_1 STARTED\n",
      "[11.02|10:02:41] JOB scan_gamma_1.8 RUNNING\n",
      "[11.02|10:02:41] JOB scan_gamma_1.6/scan_gamma_1.6_charge_minus1 STARTED\n",
      "[11.02|10:02:41] JOB scan_gamma_1.6/scan_gamma_1.6_charge_0 STARTED\n",
      "[11.02|10:02:41] JOB scan_gamma_1.8/scan_gamma_1.8_charge_minus1 STARTED\n",
      "[11.02|10:02:41] JOB scan_gamma_1.6/scan_gamma_1.6_charge_1 STARTED\n",
      "[11.02|10:02:41] JOB scan_gamma_1.8/scan_gamma_1.8_charge_0 STARTED\n",
      "[11.02|10:02:41] JOB scan_gamma_1.8/scan_gamma_1.8_charge_1 STARTED\n",
      "[11.02|10:02:42] JOB scan_gamma_1.2/scan_gamma_1.2_charge_minus1 RUNNING\n",
      "[11.02|10:02:42] JOB scan_gamma_1.2/scan_gamma_1.2_charge_0 RUNNING\n",
      "[11.02|10:02:42] JOB scan_gamma_1.2/scan_gamma_1.2_charge_1 RUNNING\n",
      "[11.02|10:02:42] JOB scan_gamma_1.4/scan_gamma_1.4_charge_0 RUNNING\n",
      "[11.02|10:02:42] JOB scan_gamma_1.4/scan_gamma_1.4_charge_1 RUNNING\n",
      "[11.02|10:02:42] JOB scan_gamma_1.8/scan_gamma_1.8_charge_1 RUNNING\n",
      "[11.02|10:02:42] JOB scan_gamma_1.8/scan_gamma_1.8_charge_0 RUNNING\n",
      "[11.02|10:02:42] JOB scan_gamma_1.4/scan_gamma_1.4_charge_minus1 RUNNING\n",
      "[11.02|10:02:42] JOB scan_gamma_1.6/scan_gamma_1.6_charge_1 RUNNING\n",
      "[11.02|10:02:42] JOB scan_gamma_1.6/scan_gamma_1.6_charge_minus1 RUNNING\n",
      "[11.02|10:02:42] JOB scan_gamma_1.8/scan_gamma_1.8_charge_minus1 RUNNING\n",
      "[11.02|10:02:42] JOB scan_gamma_1.6/scan_gamma_1.6_charge_0 RUNNING\n",
      "[11.02|10:03:15] JOB scan_gamma_1.2/scan_gamma_1.2_charge_0 FINISHED\n",
      "[11.02|10:03:15] JOB scan_gamma_1.2/scan_gamma_1.2_charge_0 SUCCESSFUL\n",
      "[11.02|10:03:16] JOB scan_gamma_1.4/scan_gamma_1.4_charge_0 FINISHED\n",
      "[11.02|10:03:16] JOB scan_gamma_1.4/scan_gamma_1.4_charge_0 SUCCESSFUL\n",
      "[11.02|10:03:16] JOB scan_gamma_1.6/scan_gamma_1.6_charge_0 FINISHED\n",
      "[11.02|10:03:16] JOB scan_gamma_1.6/scan_gamma_1.6_charge_0 SUCCESSFUL\n",
      "[11.02|10:03:17] JOB scan_gamma_1.8/scan_gamma_1.8_charge_0 FINISHED\n",
      "[11.02|10:03:17] JOB scan_gamma_1.8/scan_gamma_1.8_charge_0 SUCCESSFUL\n",
      "[11.02|10:03:17] JOB scan_gamma_1.4/scan_gamma_1.4_charge_1 FINISHED\n",
      "[11.02|10:03:17] JOB scan_gamma_1.4/scan_gamma_1.4_charge_1 SUCCESSFUL\n",
      "[11.02|10:03:17] JOB scan_gamma_1.2/scan_gamma_1.2_charge_1 FINISHED\n",
      "[11.02|10:03:17] JOB scan_gamma_1.2/scan_gamma_1.2_charge_1 SUCCESSFUL\n",
      "[11.02|10:03:17] JOB scan_gamma_1.2/scan_gamma_1.2_charge_minus1 FINISHED\n",
      "[11.02|10:03:17] JOB scan_gamma_1.2/scan_gamma_1.2_charge_minus1 SUCCESSFUL\n",
      "[11.02|10:03:17] JOB scan_gamma_1.8/scan_gamma_1.8_charge_1 FINISHED\n",
      "[11.02|10:03:17] JOB scan_gamma_1.8/scan_gamma_1.8_charge_1 SUCCESSFUL\n",
      "[11.02|10:03:17] JOB scan_gamma_1.6/scan_gamma_1.6_charge_1 FINISHED\n",
      "[11.02|10:03:17] JOB scan_gamma_1.6/scan_gamma_1.6_charge_1 SUCCESSFUL\n",
      "[11.02|10:03:18] JOB scan_gamma_1.4/scan_gamma_1.4_charge_minus1 FINISHED\n",
      "[11.02|10:03:18] JOB scan_gamma_1.4/scan_gamma_1.4_charge_minus1 SUCCESSFUL\n",
      "[11.02|10:03:18] JOB scan_gamma_1.6/scan_gamma_1.6_charge_minus1 FINISHED\n",
      "[11.02|10:03:18] JOB scan_gamma_1.6/scan_gamma_1.6_charge_minus1 SUCCESSFUL\n",
      "[11.02|10:03:18] JOB scan_gamma_1.8/scan_gamma_1.8_charge_minus1 FINISHED\n",
      "[11.02|10:03:18] JOB scan_gamma_1.8/scan_gamma_1.8_charge_minus1 SUCCESSFUL\n",
      "[11.02|10:03:22] JOB scan_gamma_1.4 FINISHED\n",
      "[11.02|10:03:22] JOB scan_gamma_1.6 FINISHED\n",
      "[11.02|10:03:22] JOB scan_gamma_1.8 FINISHED\n",
      "[11.02|10:03:22] JOB scan_gamma_1.2 FINISHED\n",
      "[11.02|10:03:22] Pickling of scan_gamma_1.2 failed\n",
      "[11.02|10:03:22] Pickling of scan_gamma_1.4 failed\n",
      "[11.02|10:03:22] Pickling of scan_gamma_1.6 failed\n",
      "[11.02|10:03:22] JOB scan_gamma_1.2 SUCCESSFUL\n",
      "[11.02|10:03:22] JOB scan_gamma_1.4 SUCCESSFUL\n",
      "[11.02|10:03:22] Pickling of scan_gamma_1.8 failed\n",
      "[11.02|10:03:22] JOB scan_gamma_1.6 SUCCESSFUL\n",
      "[11.02|10:03:22] JOB scan_gamma_1.8 SUCCESSFUL\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/ormrodmorley/Documents/code/ams/amshome/bin.release/python3.8/lib/python3.8/site-packages/dill/_dill.py:1707: PicklingWarning: Cannot pickle <class '__main__.GammaResults'>: __main__.GammaResults has recursive self-references that trigger a RecursionError.\n",
      "  warnings.warn('Cannot pickle %r: %s.%s has recursive self-references that trigger a RecursionError.' % (obj, obj.__module__, obj_name), PicklingWarning)\n",
      "/Users/ormrodmorley/Documents/code/ams/amshome/bin.release/python3.8/lib/python3.8/site-packages/dill/_dill.py:1705: PicklingWarning: Cannot locate reference to <class '__main__.GammaJob'>.\n",
      "  warnings.warn('Cannot locate reference to %r.' % (obj,), PicklingWarning)\n",
      "/Users/ormrodmorley/Documents/code/ams/amshome/bin.release/python3.8/lib/python3.8/site-packages/dill/_dill.py:1707: PicklingWarning: Cannot pickle <class '__main__.GammaJob'>: __main__.GammaJob has recursive self-references that trigger a RecursionError.\n",
      "  warnings.warn('Cannot pickle %r: %s.%s has recursive self-references that trigger a RecursionError.' % (obj, obj.__module__, obj_name), PicklingWarning)\n",
      "/Users/ormrodmorley/Documents/code/ams/amshome/bin.release/python3.8/lib/python3.8/site-packages/dill/_dill.py:1705: PicklingWarning: Cannot locate reference to <class '__main__.GammaResults'>.\n",
      "  warnings.warn('Cannot locate reference to %r.' % (obj,), PicklingWarning)\n"
     ]
    }
   ],
   "source": [
    "results = gamma_scan(gammas, s, mol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5552817d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Results ==\n",
      "gamma \t J\n",
      "1.2000 \t 0.01139955\n",
      "1.4000 \t 0.00756624\n",
      "1.6000 \t 0.00858701\n",
      "1.8000 \t 0.01110892\n",
      "Optimal gamma value: 1.4000\n"
     ]
    }
   ],
   "source": [
    "print(\"== Results ==\")\n",
    "print(\"gamma \\t J\")\n",
    "for g, j in results:\n",
    "    print(\"{:.4f} \\t {:.8f}\".format(g, j))\n",
    "print(\"Optimal gamma value: {:.4f}\".format(min(results, key=lambda x: x[1])[0]))"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "executable": "/usr/bin/env plams",
   "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
}