{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simple functions for ground and excited states\n",
    "\n",
    "## Example: the electronic states of molecular oxygen\n",
    "\n",
    "This example notebook shows how to determine active spaces for the triplet ground state of oxygen, and its lowest excited singlet states. An emphasis is placed on simple wrapper functions to determine active spaces directly from a molecule definition.\n",
    "\n",
    "The following learning points are covered:\n",
    "\n",
    "- Selecting an active space using a simple wrapper function.\n",
    "- Selecting an active space using a wrapper function for multiple electronic states.\n",
    "- Using the selected active space in CASSCF calculations, including an example with state averaging.\n",
    "- Using suitable wrapper functions to choose active spaces of a specific size defined by the user."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Suppress PySCF warning...\n",
    "import pyscf\n",
    "pyscf.__config__.B3LYP_WITH_VWN5 = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "# The Mole class is used to define molecular information in PySCF.\n",
    "from pyscf.gto import Mole\n",
    "\n",
    "# logger contains definitions of verbosity levels for PySCF.\n",
    "from pyscf.lib import logger\n",
    "\n",
    "# Functionality for (state-averaged) CASSCF.\n",
    "from pyscf.mcscf import CASSCF, state_average_mix\n",
    "from pyscf.fci.direct_spin1 import FCISolver\n",
    "from pyscf.fci.addons import fix_spin\n",
    "\n",
    "# Wrapper functions to perform selection for variable and fixed active space sizes\n",
    "from asf.wrapper import find_from_mol, find_from_scf, sized_space_from_mol, sized_space_from_scf\n",
    "\n",
    "# Various utility functions...\n",
    "from asf.utility import compare_active_spaces, show_mos_grid, pictures_Jmol"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are two variants of the wrapper function to find general active spaces: `find_from_mol` and `find_from_scf`. The former accepts a `Mole` objects containing a molecular definition and performs the entire active space selection, including an SCF calculation. It is usually performed as a UHF calculation even if the molecule does not have any unpaired electrons. Stability analysis is used to help the UHF calculation converge to a broken-symmetry solution, which helps with the initial active space selection. It is also possible to start from a converged RHF or UHF solution directly using `find_from_scf`, which carries out the same steps as `find_from_mol` with the exception of the SCF calculation itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(find_from_mol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(find_from_scf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start with a molecule definition of oxygen in its (triplet) ground state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mol = Mole()\n",
    "mol.atom = \"O2.xyz\"\n",
    "mol.basis = \"def2-SVP\"\n",
    "mol.charge = 0\n",
    "mol.spin = 2\n",
    "# Set mol.verbose = logger.INFO to enable printing of SCF iterations and further output.\n",
    "mol.verbose = logger.NOTE\n",
    "mol.build()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All that is required now for an active space selection is one function call:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "active_space = find_from_mol(mol)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The returned object contains the number of active MOs and electrons, the matrix of MO coeffients and a list of active orbital indices referring to the aforementioned matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(active_space)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below, we plot the active orbitals that were found."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pictures_Jmol(\n",
    "    mol=mol, mo_coeff=active_space.mo_coeff, mo_list=active_space.mo_list, mo_index_title=\" \"\n",
    ")\n",
    "mo_images = sorted(dd for dd in Path.cwd().glob(\"mo_*.png\"))\n",
    "show_mos_grid(mo_images, mo_list=active_space.mo_list, columns=3, figsize=(10, 10))\n",
    "# clean up files\n",
    "for filename in mo_images:\n",
    "    Path.unlink(filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The selected active space can be used to set up a CASSCF calculation. Note that it is important to use the orbitals returned by the ASF in the object `active_space` - not the orbitals obtained from SCF or elsewhere!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mc = CASSCF(mol, ncas=active_space.norb, nelecas=active_space.nel)\n",
    "# Before carrying out the CASSCF calculation, the orbital coefficient matrix columns need to be\n",
    "# ordered as core - active - virtual from left to right.\n",
    "# Take care - by default, sort_mo assumes that the MO list starts indexing from one, not from zero.\n",
    "mo_guess = mc.sort_mo(active_space.mo_list, active_space.mo_coeff, base=0)\n",
    "_ = mc.kernel(mo_coeff=mo_guess)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A CASSCF calculation converging in a small number of iterations is often a good sign. In general, it is advisable to inspect converged CASSCF orbitals visually in order to verify that they indeed form the expected active space. In our case, that would mean comparing them to the guess orbitals determined by the active space finder.\n",
    "\n",
    "Here, however, we will employ a simple numerical criterion by performing a singular value decomposition of the overlap matrix between the guess and the converged active orbitals. The following function reports the smallest singular value - if it is close to one, that is a good sign. A value close to zero indicates that an orbital may have rotated out, while intermediate values indicate some degree of mixing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compare_active_spaces(mc, active_space.mo_coeff, active_space.mo_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The functions `find_from_mol` / `find_from_scf` can also be employed to find active spaces for state-averaged CASSCF calculations of multiple electronic states. This functionality is relatively simple, determining active spaces for each electronic state in isolation, and then forming the union of the MO sets. States can be specified in one of two ways:\n",
    "\n",
    "1. Using `states: int`. in that case, the active space will be determined for a number `states` of states with the multiplicity inferred from the Mole or SCF object.\n",
    "2. Using `states: list[tuple[int, int]]`. the first integer in each tuple specifies the number of electrons (`spin` in PySCF conventions), the second integer the number of electronic states of that spin. In this case, the union is formed over the active spaces for all states of all spin multiplicities requested.\n",
    "Below this will be demonstrated for one triplet and three singlet states of oxygen:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "active_space = find_from_mol(mol, states=[(0, 3), (2, 1)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The recommended active space is the same as for triplet oxygen; note that the table printed by the wrapper can be inspected for the orbital occupations in each state.\n",
    "\n",
    "Below, a state-averaged CASSCF calculation is performed with four states: the triplet, the two doubly degenerate lowest singlet states, and the third singlet state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mc = CASSCF(mol, ncas=active_space.norb, nelecas=active_space.nel)\n",
    "# The solvers for each spin need to be set up individually.\n",
    "# See PySCF examples for state averages and state average mixes.\n",
    "solver_t = FCISolver(mol)\n",
    "solver_t.spin = 2\n",
    "solver_s = FCISolver(mol)\n",
    "solver_s.spin = 0\n",
    "solver_s.nroots = 3\n",
    "solver_s = fix_spin(solver_s, ss=0)\n",
    "# All four states are weighted equally.\n",
    "mc_mix = state_average_mix(mc, [solver_t, solver_s], weights=[0.25] * 4)\n",
    "mo_guess = mc.sort_mo(active_space.mo_list, active_space.mo_coeff, base=0)\n",
    "_ = mc_mix.kernel(mo_coeff=mo_guess)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The main parameter controlling the size of the selected space is the entropy threshold, controlled via the parameter `entropy_threshold` of the `find_from_mol` and `find_from_scf` functions.\n",
    "\n",
    "If you know the size of the desired active space and only need to select the correct orbitals, it is possible to do so via the functions `sized_space_from_mol` and `sized_space_from_scf`. The only difference between the two functions is that the former performs an SCF calculation, while the latter expects an object with an SCF solution as its input. These wrapper functions currently support only one electronic state at a time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(sized_space_from_mol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(sized_space_from_scf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is possible to specify the size of an active space either via the number of orbitals by providing an integer, or by providing both the number of active electrons and orbitals as a tuple `(electrons, orbitals)`. If only the number of orbitals is provided, the number of electrons is determined by the active space finder.\n",
    "\n",
    "For example, the following function call constrains the number of active orbitals to four and chooses an active space with a suitable number of electrons, resulting in a (6, 4) space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sized_space_from_mol(mol, size=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Likewise, it is possible to specify six electrons and four orbitals explicitly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sized_space_from_mol(mol, size=(6, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the `sized_space` functions attempt to choose active spaces from a set of sensible suggestions determined internally. In consequence, if the user requests an active space size that is not considered sensible, the function will raise an `ActiveSpaceSelectionError` exception."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}