{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Suggest multiple active spaces\n",
    "\n",
    "## Example: the *para*-benzyne diradical"
   ]
  },
  {
   "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",
    "from pyscf.gto import Mole\n",
    "from pyscf.lib import logger\n",
    "from asf import ASFDMRG\n",
    "from asf.visualization import draw_pair_information\n",
    "from asf.preselection import MP2NatorbPreselection, MP2PairinfoPreselection\n",
    "from asf.scf import stable_scf\n",
    "from asf.utility import pictures_Jmol, show_mos_grid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Define the molecule as a PySCF object\n",
    "\n",
    "In the ASF package, molecular systems are represented as PySCF `Mole` objects. The first step is therefore to initialize the example molecule as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mol = Mole(\n",
    "    atom=\"pbenzyne.xyz\",\n",
    "    spin=0,\n",
    "    charge=0,\n",
    "    basis=\"def2-SVP\",\n",
    "    verbose=logger.NOTE,  # a less talkative option, to suppress printing of SCF iterations, etc.\n",
    ").build()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Perform a UHF calculation\n",
    "\n",
    "This is intentional even for singlet states, as a broken-symmetry solution helps with active space selection. The function `stable_scf` incorporates checks for convergence and stability of the solution. SCF calculations are restarted if an unconverged or unstable solution is encountered. The stability analysis often helps with nudging the orbitals towards a broken-symmetry solution, but it does not always work out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mf = stable_scf(mol)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Side note: the ASF package contains wrapper that perfom some of the steps outlined below (semi-)automatically in one function:\n",
    "\n",
    "- `asf.wrapper.sized_space_from_mol` - determines an optimal active space of given size starting from a `Mole` object,\n",
    "- `asf.wrapper.sized_space_from_scf` - determines an optimal active space of given size starting from an SCF solution,\n",
    "- `asf.wrapper.find_from_mol` - selects an active space based on a minimum entropy threshold starting from a `Mole` object, and\n",
    "- `asf.wrapper.find_from_scf` - selects an active space based on a minimum entropy threshold starting from an SCF solution.\n",
    "\n",
    "The `find_from_...` wrapper variants also support a limiting window for the active space size via `min_norb` and `max_norb` arguments."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Determine an initial active space\n",
    "\n",
    "This step concerns the definition of a subset of the MO space in which the final active space should be determined. While not mandatory, such a preselection step is recommended for larger systems. The ASF package implements two procedures to select an initial space:\n",
    "\n",
    "  - `asf.preselection.MP2PairinfoPreselection` - based on pair information from Møller-Plesset perturbation theory\n",
    "  - `asf.preselection.MP2NatorbPreselection` - based on MP2 natural occupation numbers\n",
    "\n",
    "Custom preselection approaches may be implemented by deriving from `asf.asfbase.Preselection`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If set to True, determine the initial active space using pair information derived from\n",
    "# Moller-Plesset theory (more complicated procedure).\n",
    "use_mp_pairinfo = True\n",
    "\n",
    "if use_mp_pairinfo:\n",
    "    pre = MP2PairinfoPreselection(mf)\n",
    "    active_space_pt = pre.select()\n",
    "# Otherwise, just use MP2 natural orbitals with a specific cutoff.\n",
    "else:\n",
    "    pre = MP2NatorbPreselection(mf, lower=0.02, upper=1.98)\n",
    "    active_space_pt = pre.select()\n",
    "\n",
    "print(active_space_pt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Run multi-reference single point calculation\n",
    "\n",
    "`ASFCI` / `ASFDMRG` are two important low-level interface classes for the functionality of the ASF to perform single-point calculations. Here, a DMRG calculation is set up using the initial active space previously determined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sf = ASFDMRG(\n",
    "    mol=mol,\n",
    "    nel=active_space_pt.nel,\n",
    "    mo_list=active_space_pt.mo_list,\n",
    "    mo_coeff=active_space_pt.mo_coeff,\n",
    "    maxM=250,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point, the DMRG interface is initialized including the default calculation settings. For a more extensive control of the FCI solver settings both `ASFDMRG` and `ASFCI` accept the `fcisolver_kwargs` argument (see available keys for [DMRG](https://block2.readthedocs.io/en/latest/user/dmrg-scf.html) and [CASCI](https://pyscf.org/pyscf_api_docs/pyscf.mcscf.html#pyscf.mcscf.casci.CASBase)). To start the single-point calculation, the `calculate` method must be called which is implemented by all of the ASF interfaces. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sf.calculate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Find active spaces\n",
    "\n",
    "With the finished single-point calculation finished, the determination of viable active spaces is available. The active space search relies on metrics such as single-orbital entropies and pair information from the diagonal of the two-body cumulant.\n",
    "The ASF interfaces provide several methods to determine active spaces:\n",
    "\n",
    "  - `find_many()` determining *multiple* reasonable active spaces sorted by number of active electrons and number of active orbitals.\n",
    "  - `find_one_sized()` which yields *one* active space matching a given number of active orbitals and/or a given number of active electrons.\n",
    "  - `find_one_entropy()` which yields *one* active space that has a minimal single-orbital entropy above a threshold that can be changed by the user (__recommended__).\n",
    "  - `find_one_generic()` which also yields *one* active space, but requires the user to define the selection criteria.\n",
    "\n",
    "This interactive example focuses on the usage of `find_many()`. See other tutorials for how `find_one_sized` or `find_one_entropy` can be employed.\n",
    "Internally, `find_many()` first determines a large number of candidate spaces (via `unfiltered_pairinfo_spaces()`) and then applies up to three selection criteria to select spaces most likely to converge to a reasonable solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "active_spaces = sf.find_many()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The resulting output is a dictionary whereby each value is a list of active space suggestions matching the number of electrons and number of orbitals in the key. Each item in the suggestion list is a `MOListInfo` dataclass containing the list of MO indices and additional metrics used for the selection procedure. For example:\n",
    "```python\n",
    "{\n",
    "    (2, 2): [\n",
    "        MOListInfo(\n",
    "            mo_list=[19, 20],\n",
    "            nel=2,\n",
    "            minimal=False,\n",
    "            pairinfo_sum=0.6640700019980257,\n",
    "            max_increment=0.0150994051180795,\n",
    "            min_decrement=0.9691123584300709,\n",
    "            min_entropy=0.6783514316563286,\n",
    "            min_edge_sum=0.34325572898426215,\n",
    "        )\n",
    "    ],\n",
    "    ...\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. Inspect suggested active spaces\n",
    "\n",
    "#### Molecular orbital indices\n",
    "To get an overview of the suggested active spaces, we print active space groups (in `(N, M)` notation) and the associated MO index lists. Note that the indices refer to the MO coefficient matrix `sf.mo_coeff` (and therefore also `active_space_pt.mo_coeff`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for nel, nmo in active_spaces.keys():\n",
    "    for mo_info in active_spaces[(nel, nmo)]:\n",
    "        print(f\"({nel}, {nmo}): {mo_info.mo_list}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Molecular orbital visualization\n",
    "\n",
    "The suggested active MOs can be easily visualized by tools such as Jmol. The ASF package provides a utility function `pictures_Jmol` that creates a temporary molden file and writes the MO plots to the current working directory. To save a molden file for interactive inspection of the molecular orbitals, use:\n",
    "```python\n",
    "from pyscf.tools import molden\n",
    "molden.from_mo(mol, 'mos.molden', mo_coeff=sf.mo_coeff)\n",
    "```\n",
    "Once `pictures_Jmol` has run, there should be several images named `mo_[index].png` in the current directory. The function `show_mos_grid` from the `utility` module arranges the orbital images in a grid. The MO index lists seen in the previous step can also be passed to only show the MOs of the active space of interest.\n",
    "**Note:** MO indices in Jmol are indexed from *one*, not from zero. Thus, the annotation from Jmol is disabled and `show_mos_grid` annotates each orbital image with the MO index that is consistent with the ASF (0-indexed)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If plots have been generated in a previous run of the Notebook, they will be reused here.\n",
    "# Set redo_plots = True, if new orbital images should be created regardless of existing files.\n",
    "redo_plots = False\n",
    "\n",
    "# Select for which active space molecular orbitals should be shown:\n",
    "space_of_interest = active_spaces[(12, 12)][0]\n",
    "\n",
    "def mo_images(name_glob: str = \"mo_*.png\") -> list[Path]:\n",
    "    \"\"\"Return the list of all MO images in this directory.\"\"\"\n",
    "    return sorted(Path.cwd().glob(name_glob))\n",
    "\n",
    "if redo_plots or not mo_images():\n",
    "    pictures_Jmol(mol=mol, mo_coeff=sf.mo_coeff, mo_list=sf.mo_list, mo_index_title=\" \")\n",
    "\n",
    "show_mos_grid(mo_images(), mo_list=space_of_interest.mo_list, columns=3, figsize=(10, 10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Pairinfo \n",
    "\n",
    "The ASF package also provides methods to visualize pair information which is a key metric in the active space search. Below, `draw_pair_information()` is employed to plot a graph of the pair information from the DMRG calculation. Note that the sign of the cumulant is inverted, so that the most energy-lowering interactions correspond to the largest pair information. In the resulting graph plot, the molecular orbitals are represented as nodes and the pair information as edges. Edges corresponding to energy-lowering interactions are colored red, while energy-raising interactions are shown in blue (cf. ['seismic'](https://matplotlib.org/stable/users/explain/colors/colormaps.html#diverging) color map). By default the color of the nodes indicates the summed pair information per orbital using a sequential color map (['RdPu'](https://matplotlib.org/stable/users/explain/colors/colormaps.html#diverging)), however any per-orbital-quantity, such as single-orbital entropies, can be passed instead via the `orbital_info` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pairinfo = -sf.diagonal_cumulant(full_space=False)\n",
    "draw_pair_information(pair_info=pairinfo, orbital_labels=sf.mo_list)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}