Self-Consistent Field Methods

13. Self-Consistent Field Methods#

👩🏽‍💻 Assignment: Complete the Below Notebook#

To do this, complete the code blocks indicated by ### START YOUR CODE HERE and ### END YOUR CODE HERE .

The detailed steps you’ll follow are:

Read the theoretical background on the self-consistent field (SCF) method.
Implement the “core Hamiltonian” guess.
Implement the addition of an electric field to the Hamiltonian and compute the dipole moment and polarizability.

Motivation of Mean-Field (Average Potential) Methods#

In atomic units, the electronic Hamiltonian for an $N$-electron molecule, with $P$ atoms with atomic numbers $Z_A$ at positions $\mathbf{R}_A$, is given by $$\hat{H} = -\sum_{i=1}^N \frac{1}{2} \nabla_i^2 - \sum_{i=1}^N \sum_{A=1}^P \frac{Z_A}{|\mathbf{r}_i - \mathbf{R}_A|} + \sum_{i<j} \frac{1}{|\mathbf{r}_i - \mathbf{r}_j|}$$

This Hamiltonian cannot be solved by separation of variables because the electrons are coupled through the electron-electron repulsion term. Mean-field, or average potential, approximations decouple electrons by treating the effect of all other electrons on a given electron as an average potential (operator), $\hat{w}(\mathbf{r})$. (I.e., individual electrons feel only the effects of the “cloud” of other electrons.) With this revision, the electronic Hamiltonian becomes, $$\hat{H}_{\text{avg}} = -\sum_{i=1}^N \frac{1}{2} \nabla_i^2 + \sum_{i=1}^N v(\mathbf{r}_i) + \sum_{i=1}^N \hat{w}(\mathbf{r}_i)$$ where the external potential is defined as $$v(\mathbf{r}) = -\sum_{A=1}^P \frac{Z_A}{|\mathbf{r} - \mathbf{R}_A|}$$

The electronic Schrödinger equation for the mean-field Hamiltonian is $$\hat{H}_{\text{avg}} \Phi(\mathbf{r}_1, \mathbf{r}_2, \ldots, \mathbf{r}_N) = E \Phi(\mathbf{r}_1, \mathbf{r}_2, \ldots, \mathbf{r}_N)$$ which can be solved by separation of variables, yielding a set of one-electron equations, $$\left[-\frac{1}{2} \nabla^2 + v(\mathbf{r}) + \hat{w}(\mathbf{r})\right] \phi_p(\mathbf{r}) = \epsilon_p \phi_p(\mathbf{r})$$ where $\phi_p(\mathbf{r})$ is the $p$-th molecular orbital. However, the average potential $\hat{w}$ depends on the wavefunction $\Phi$, which is a Slater determinant of the molecular orbitals $\phi_p$. Making this dependence explicit, we have, $$\left[-\frac{1}{2} \nabla^2 + v(\mathbf{r}) + \hat{w}[\{\phi_q\}](\mathbf{r})\right] \phi_p(\mathbf{r}) = \epsilon_p \phi_p(\mathbf{r})$$ This is a *nonlinear* Schrödinger equation because the operator $\hat{w}$ depends on the orbitals ${\phi_q}$, which are the solutions to the equation. Since we cannot write the Hamiltonian without knowing the orbitals, and cannot know the orbitals without writing the Hamiltonian, we have a circular problem. The simplest self-consistent-field (SCF) method solves this problem by starting from an initial guess for the orbitals, constructing the Hamiltonian, solving for the orbitals, and then updating the Hamiltonian using those orbitals. This procedure is iterated until the orbitals used to construct the Hamiltonian are the same as the orbitals obtained from solving the Schrödinger equation, at which point we have a self-consistent solution.

Average Potential Models#

Hartree Model#

The Hartree model is the simplest average potential model. It approximates the average potential as the electrostatic potential generated by the total electron density, $\rho(\mathbf{r})$. Therefore, the Hartree average potential is given by $$w_{\text{H}}(\mathbf{r}) = \int \frac{\rho(\mathbf{r}')}{|\mathbf{r} - \mathbf{r}'|} d\mathbf{r}'$$ where $\rho(\mathbf{r})$ is the total electron density. The Hartree model does not account for the fact that an electron does not interact with itself, leading to an overestimation of the electron-electron repulsion energy. This is known as the self-interaction error.

Fermi-Amaldi Model#

The Fermi-Amaldi model is the simplest way to impose a self-interaction correction. It accounts for the fact that an $\alpha$-spin electron interacts with all of the $\beta$-spin electrons but only $N_\alpha - 1$ of the $\alpha$-spin electrons, and vice versa for $\beta$-spin electrons. Therefore, the Fermi-Amaldi average potential is given by $$ \begin{align*} w_{\text{FA};\alpha}(\mathbf{r}) &= \frac{N_\alpha-1}{N_\alpha} \int \frac{\rho_\alpha(\mathbf{r}')}{|\mathbf{r} - \mathbf{r}'|} d\mathbf{r}' + \int \frac{\rho_\beta(\mathbf{r}')}{|\mathbf{r} - \mathbf{r}'|} d\mathbf{r}'\\ w_{\text{FA};\beta}(\mathbf{r}) &= \frac{N_\beta-1}{N_\beta} \int \frac{\rho_\beta(\mathbf{r}')}{|\mathbf{r} - \mathbf{r}'|} d\mathbf{r}' + \int \frac{\rho_\alpha(\mathbf{r}')}{|\mathbf{r} - \mathbf{r}'|} d\mathbf{r}' \end{align*} $$ where $\rho_\alpha(\mathbf{r})$ and $\rho_\beta(\mathbf{r})$ are the electron densities for $\alpha$ and $\beta$ spins, respectively.

Kohn-Sham Density Functional Theory#

Kohn-Sham density functional theory (KS-DFT) is a more sophisticated average potential model that includes an exchange-correlation potential, $v_{\text{xc}}(\mathbf{r})$, to account for the effects of electron exchange and correlation. The KS-DFT average potential is given by $$\hat{w}_{\text{KS}}(\mathbf{r}) = w_{\text{H}}(\mathbf{r}) + v_{\text{xc}}(\mathbf{r})$$ where $v_{\text{xc}}(\mathbf{r})$ is the exchange-correlation potential, which is a functional of the electron (spin) densities $\rho_\alpha(\mathbf{r})$ and $\rho_\beta(\mathbf{r})$. The exact form of $v_{\text{xc}}$ is unknown, and various approximations have been developed to model it, such as the local density approximation (LDA) and the generalized gradient approximation (GGA).

The simplest—and one of the first—exchange-correlation functionals is the local spin density approximation (LSDA), which is a simple scaling of the exchange energy from the homogeneous electron gas. If only exchange effects are included, then one has $$v_{\text{xc};\alpha}(\mathbf{r}) = -\left(\frac{6}{\pi}\right)^{1/3} \rho_\alpha(\mathbf{r})^{1/3}$$v_{\text{xc};\beta}(\mathbf{r}) = -\left(\frac{6}{\pi}\right)^{1/3} \rho_\beta(\mathbf{r})^{1/3}$$

Hartree-Fock Theory#

Hartree-Fock (HF) theory is an average potential model that includes an exchange potential, $\hat{K}$, to account for the effects of electron exchange. The HF average potential is given by $$\hat{w}_{\text{HF}}(\mathbf{r}) = w_{\text{H}}(\mathbf{r}) + \hat{k}(\mathbf{r})$$ where $$\hat{k}(\mathbf{r}) \phi_p(\mathbf{r}) = -\sum_{q=1}^N \delta_{\sigma_q \sigma_p}\int \frac{\phi_q^*(\mathbf{r}') \phi_p(\mathbf{r}')}{|\mathbf{r} - \mathbf{r}'|}d\mathbf{r}' \phi_q(\mathbf{r}) $$ The exchange operator, accounting for the Pauli exclusion principle, only acts on orbitals with the same spin.

Initial Guess#

In the separated-atom limit, a molecule’s density is the sum of the isolated atomic densities. Even at (near) equilibrium geometries, this promolecular approximation is often a good starting point. When we do this, we say we are using the Harris functional as an initial guess. The Harris functional can be used with any pure density functional (including the Hartree, Fermi-Amaldi, and LSDA).

An easier alternative is the core Hamiltonian guess, which is the solution to the one-electron Schrödinger equation with only the kinetic energy and nuclear attraction terms. This guess isn’t as accurate, but it is easy to compute and for systems with weakly bound electrons, it sometimes performs better than the (raw) Harris functional guess. (An even better guess, however, is often the Harris functional guess but using the electron density of the molecular cation.)

External Packages Needed#

We will use several packages to construct the program.

AtomDB is a database of atomic properties. It can be used to generate the promolecular density for the initial guess.
IOData is a package for reading and writing data. We will use it to read the molecular geometry from an XYZ file.
GBasis is a package for evaluating Gaussian basis functions and integrals related to them.
QC-Grid is a package for evaluating numerical integrals (and derivatives).
QC-BFit is a package for fitting functions (like the electron density) to Gaussian basis sets. We use it to build fast (approximate) models for the Coulomb potential.
basis_set_exchange downloads Gaussian basis sets.

## Load Packages
!pip install qc-atomdb
!pip install qc-iodata
!pip install qc-gbasis
!pip install qc-grid
!pip install basis_set_exchange
!pip install qc-bfit

Collecting qc-atomdb

  Downloading qc_AtomDB-0.0.2.post5-py3-none-any.whl.metadata (7.0 kB)

Requirement already satisfied: numpy>=1.16 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from qc-atomdb) (2.4.4)
Requirement already satisfied: scipy>=1.4 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from qc-atomdb) (1.17.1)

Collecting msgpack>=1.0.0 (from qc-atomdb)
  Downloading msgpack-1.1.2-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl.metadata (8.1 kB)
Collecting msgpack-numpy>=0.4.8 (from qc-atomdb)
  Downloading msgpack_numpy-0.4.8-py2.py3-none-any.whl.metadata (5.0 kB)

Collecting h5py>=3.6.0 (from qc-atomdb)
  Downloading h5py-3.16.0-cp314-cp314-manylinux_2_28_x86_64.whl.metadata (3.0 kB)

Collecting importlib-resources>=3.0.0 (from qc-atomdb)
  Downloading importlib_resources-6.5.2-py3-none-any.whl.metadata (3.9 kB)
Collecting pooch>=1.8.1 (from qc-atomdb)
  Downloading pooch-1.9.0-py3-none-any.whl.metadata (10 kB)
Requirement already satisfied: platformdirs>=2.5.0 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from pooch>=1.8.1->qc-atomdb) (4.9.4)
Requirement already satisfied: packaging>=20.0 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from pooch>=1.8.1->qc-atomdb) (26.0)
Requirement already satisfied: requests>=2.19.0 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from pooch>=1.8.1->qc-atomdb) (2.33.0)
Requirement already satisfied: charset_normalizer<4,>=2 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from requests>=2.19.0->pooch>=1.8.1->qc-atomdb) (3.4.6)
Requirement already satisfied: idna<4,>=2.5 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from requests>=2.19.0->pooch>=1.8.1->qc-atomdb) (3.11)
Requirement already satisfied: urllib3<3,>=1.26 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from requests>=2.19.0->pooch>=1.8.1->qc-atomdb) (2.6.3)
Requirement already satisfied: certifi>=2023.5.7 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from requests>=2.19.0->pooch>=1.8.1->qc-atomdb) (2026.2.25)

Downloading qc_AtomDB-0.0.2.post5-py3-none-any.whl (62.0 MB)
?25l   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0.0/62.0 MB ? eta -:--:--

   ━━━━━╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 8.1/62.0 MB 41.5 MB/s eta 0:00:02

   ━━━━━━━━━━╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 15.7/62.0 MB 41.4 MB/s eta 0:00:02

   ━━━━━━━━━━━━━━━╸━━━━━━━━━━━━━━━━━━━━━━━━ 24.1/62.0 MB 41.1 MB/s eta 0:00:01

   ━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━━━━━━━━━━━ 31.7/62.0 MB 39.7 MB/s eta 0:00:01

   ━━━━━━━━━━━━━━━━━━━━━━━━━╸━━━━━━━━━━━━━━ 39.8/62.0 MB 40.5 MB/s eta 0:00:01

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━ 48.2/62.0 MB 40.6 MB/s eta 0:00:01

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸━━━ 56.6/62.0 MB 40.6 MB/s eta 0:00:01

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸ 61.9/62.0 MB 39.3 MB/s eta 0:00:01
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 62.0/62.0 MB 37.9 MB/s  0:00:01
?25hDownloading h5py-3.16.0-cp314-cp314-manylinux_2_28_x86_64.whl (5.4 MB)
?25l   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0.0/5.4 MB ? eta -:--:--

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 5.4/5.4 MB 326.6 MB/s  0:00:00
?25hDownloading importlib_resources-6.5.2-py3-none-any.whl (37 kB)
Downloading msgpack-1.1.2-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl (419 kB)
Downloading msgpack_numpy-0.4.8-py2.py3-none-any.whl (6.9 kB)
Downloading pooch-1.9.0-py3-none-any.whl (67 kB)

Installing collected packages: msgpack, importlib-resources, h5py, pooch, msgpack-numpy, qc-atomdb
?25l

   ━━━━━━━━━━━━━╺━━━━━━━━━━━━━━━━━━━━━━━━━━ 2/6 [h5py]

   ━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━━━━━━━━━━━ 3/6 [pooch]

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╺━━━━━━ 5/6 [qc-atomdb]

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╺━━━━━━ 5/6 [qc-atomdb]
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 6/6 [qc-atomdb]
?25h

Successfully installed h5py-3.16.0 importlib-resources-6.5.2 msgpack-1.1.2 msgpack-numpy-0.4.8 pooch-1.9.0 qc-atomdb-0.0.2.post5

Collecting qc-iodata

  Downloading qc_iodata-1.0.1-py3-none-any.whl.metadata (4.4 kB)
Requirement already satisfied: numpy>=1.26.4 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from qc-iodata) (2.4.4)
Requirement already satisfied: scipy>=1.13.1 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from qc-iodata) (1.17.1)
Requirement already satisfied: attrs>=21.3.0 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from qc-iodata) (26.1.0)
Downloading qc_iodata-1.0.1-py3-none-any.whl (3.4 MB)
?25l   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0.0/3.4 MB ? eta -:--:--

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 3.4/3.4 MB 39.5 MB/s  0:00:00
?25h

Installing collected packages: qc-iodata

Successfully installed qc-iodata-1.0.1

Collecting qc-gbasis

  Downloading qc_gbasis-0.1.0-py3-none-any.whl.metadata (7.4 kB)
Requirement already satisfied: scipy>=1.11.1 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from qc-gbasis) (1.17.1)
Requirement already satisfied: importlib-resources in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from qc-gbasis) (6.5.2)
Requirement already satisfied: sympy in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from qc-gbasis) (1.14.0)
Requirement already satisfied: numpy>=1.22 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from qc-gbasis) (2.4.4)
Requirement already satisfied: mpmath<1.4,>=1.1.0 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from sympy->qc-gbasis) (1.3.0)
Downloading qc_gbasis-0.1.0-py3-none-any.whl (97 kB)

Installing collected packages: qc-gbasis

Successfully installed qc-gbasis-0.1.0

Collecting qc-grid

  Downloading qc_grid-0.0.9.post1-py3-none-any.whl.metadata (4.5 kB)
Requirement already satisfied: numpy>=1.16 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from qc-grid) (2.4.4)

Collecting pytest>=8.0.0 (from qc-grid)
  Downloading pytest-9.0.2-py3-none-any.whl.metadata (7.6 kB)
Requirement already satisfied: scipy>=1.4 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from qc-grid) (1.17.1)
Requirement already satisfied: importlib_resources in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from qc-grid) (6.5.2)
Requirement already satisfied: sympy in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from qc-grid) (1.14.0)
Collecting iniconfig>=1.0.1 (from pytest>=8.0.0->qc-grid)
  Downloading iniconfig-2.3.0-py3-none-any.whl.metadata (2.5 kB)
Requirement already satisfied: packaging>=22 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from pytest>=8.0.0->qc-grid) (26.0)
Collecting pluggy<2,>=1.5 (from pytest>=8.0.0->qc-grid)
  Downloading pluggy-1.6.0-py3-none-any.whl.metadata (4.8 kB)
Requirement already satisfied: pygments>=2.7.2 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from pytest>=8.0.0->qc-grid) (2.20.0)
Requirement already satisfied: mpmath<1.4,>=1.1.0 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from sympy->qc-grid) (1.3.0)

Downloading qc_grid-0.0.9.post1-py3-none-any.whl (66.1 MB)
?25l   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0.0/66.1 MB ? eta -:--:--

   ━━━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━━━━━━━━━ 36.7/66.1 MB 199.5 MB/s eta 0:00:01

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 66.1/66.1 MB 172.7 MB/s  0:00:00
?25hDownloading pytest-9.0.2-py3-none-any.whl (374 kB)
Downloading pluggy-1.6.0-py3-none-any.whl (20 kB)
Downloading iniconfig-2.3.0-py3-none-any.whl (7.5 kB)

Installing collected packages: pluggy, iniconfig, pytest, qc-grid
?25l

   ━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━━━━━━━━━━━ 2/4 [pytest]

   ━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━━━━━━━━━━━ 2/4 [pytest]

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━ 3/4 [qc-grid]

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━ 3/4 [qc-grid]

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━ 3/4 [qc-grid]
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 4/4 [qc-grid]
?25h
Successfully installed iniconfig-2.3.0 pluggy-1.6.0 pytest-9.0.2 qc-grid-0.0.9.post1

Collecting basis_set_exchange

  Downloading basis_set_exchange-0.12-py3-none-any.whl.metadata (6.7 kB)
Requirement already satisfied: jsonschema in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from basis_set_exchange) (4.26.0)
Collecting argcomplete (from basis_set_exchange)
  Downloading argcomplete-3.6.3-py3-none-any.whl.metadata (16 kB)

Collecting regex (from basis_set_exchange)
  Downloading regex-2026.3.32-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl.metadata (40 kB)

Collecting unidecode (from basis_set_exchange)
  Downloading Unidecode-1.4.0-py3-none-any.whl.metadata (13 kB)
Requirement already satisfied: attrs>=22.2.0 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from jsonschema->basis_set_exchange) (26.1.0)
Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from jsonschema->basis_set_exchange) (2025.9.1)
Requirement already satisfied: referencing>=0.28.4 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from jsonschema->basis_set_exchange) (0.37.0)
Requirement already satisfied: rpds-py>=0.25.0 in /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages (from jsonschema->basis_set_exchange) (0.30.0)
Downloading basis_set_exchange-0.12-py3-none-any.whl (40.1 MB)
?25l   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0.0/40.1 MB ? eta -:--:--

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸ 40.1/40.1 MB 236.4 MB/s eta 0:00:01
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 40.1/40.1 MB 191.1 MB/s  0:00:00
?25hDownloading argcomplete-3.6.3-py3-none-any.whl (43 kB)
Downloading regex-2026.3.32-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl (800 kB)
?25l   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0.0/800.9 kB ? eta -:--:--

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 800.9/800.9 kB 166.9 MB/s  0:00:00
?25hDownloading Unidecode-1.4.0-py3-none-any.whl (235 kB)

Installing collected packages: unidecode, regex, argcomplete, basis_set_exchange
?25l

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0/4 [unidecode]

   ━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━━━━━━━━━━━ 2/4 [argcomplete]

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━ 3/4 [basis_set_exchange]

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━ 3/4 [basis_set_exchange]

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━ 3/4 [basis_set_exchange]

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━ 3/4 [basis_set_exchange]

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━ 3/4 [basis_set_exchange]

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━ 3/4 [basis_set_exchange]

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╺━━━━━━━━━ 3/4 [basis_set_exchange]

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 4/4 [basis_set_exchange]
?25h
Successfully installed argcomplete-3.6.3 basis_set_exchange-0.12 regex-2026.3.32 unidecode-1.4.0

Collecting qc-bfit

  Downloading qc_BFit-0.1.2-py3-none-any.whl.metadata (49 kB)
Downloading qc_BFit-0.1.2-py3-none-any.whl (561 kB)
?25l   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0.0/561.1 kB ? eta -:--:--

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 561.1/561.1 kB 70.7 MB/s  0:00:00
?25h

Installing collected packages: qc-bfit

Successfully installed qc-bfit-0.1.2

# Load the Water molecule geometry from an XYZ file (using IOData)
from urllib.request import urlretrieve
from iodata import load_one

url = "https://raw.githubusercontent.com/theochem/iodata/main/iodata/test/data/water.xyz"
urlretrieve(url, "water.xyz")

h2o = load_one("water.xyz")
print(h2o.atnums)
print(h2o.atcoords)

[1 8 1]
[[ 1.48123726 -0.93019123 -0.        ]
 [-0.          0.11720081 -0.        ]
 [-1.48123726 -0.93019123 -0.        ]]

# Select a basis set (using basis_set_exchange)
import iodata.periodic
import iodata.utils
from gbasis.parsers import parse_gbs, make_contractions
import basis_set_exchange as bse
import tempfile, iodata


def get_ao_basis_from_bse(atnums, atcoords, basis_name, coord_types="spherical"):
    r"""Fetch the atomic orbital basis set from the online repository.

    The basis will be downloaded from the basis set exchange repository and returned in the gaussian94
    format. The basis set exchange repository is available at:
    https://www.basissetexchange.org/

    Parameters
    ----------
    atnums : list of int
        Atomic numbers of the atoms in the molecule.
    atcoords : numpy.ndarray[float, shape=(n_atoms, 3)]
        Coordinates of the atoms in the molecule.
    basis_name : str
        Name of the atomic orbital basis set to fetch.
    coord_types : {"cartesian"/"c", list/tuple of "cartesian"/"c" or "spherical"/"p", "spherical"/"p"}
            Types of the coordinate system for the contractions.
            If "cartesian" or "c", then all of the contractions are treated as Cartesian contractions.
            If "spherical" or "p", then all of the contractions are treated as spherical contractions.
            If list/tuple, then each entry must be a "cartesian" (or "c") or "spherical" (or "p") to
            specify the coordinate type of each `GeneralizedContractionShell` instance.
            Default value is "spherical".
    """
    local_atnums = [int(i) for i in atnums]
    atoms = [iodata.periodic.num2sym[i] for i in local_atnums]
    gbs_basis =bse.api.get_basis(name = basis_name, elements=local_atnums, fmt="gaussian94")
    # save basis to tmp file
    with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.gbs') as tmp_data:
        tmp_data.write(gbs_basis)
        fname = tmp_data.name

    basis_dict = parse_gbs(fname)
    ao_basis = make_contractions(basis_dict, atoms, atcoords, coord_types)
    return ao_basis

bse.get_all_basis_names()

['2ZaP',
 '2ZaPa-NR',
 '2ZaPa-NR-CV',
 '3-21G',
 '3ZaP',
 '3ZaPa-NR',
 '3ZaPa-NR-CV',
 '4-31G',
 '4ZaP',
 '4ZaPa-NR',
 '4ZaPa-NR-CV',
 '5-21G',
 '5ZaP',
 '5ZaPa-NR',
 '5ZaPa-NR-CV',
 '6-21G',
 '6-31++G',
 '6-31++G*',
 '6-31++G**',
 '6-31++G**-J',
 '6-31+G',
 '6-31+G*',
 '6-31+G**',
 '6-31+G*-J',
 '6-311++G',
 '6-311++G(2d,2p)',
 '6-311++G(3df,3pd)',
 '6-311++G*',
 '6-311++G**',
 '6-311++G**-J',
 '6-311+G',
 '6-311+G(2d,p)',
 '6-311+G*',
 '6-311+G**',
 '6-311+G*-J',
 '6-311G',
 '6-311G(2df,2pd)',
 '6-311G(d,p)',
 '6-311G*',
 '6-311G**',
 '6-311G**-RIFIT',
 '6-311G-J',
 '6-311xxG(d,p)',
 '6-31G',
 '6-31G(2df,p)',
 '6-31G(3df,3pd)',
 '6-31G(d,p)',
 '6-31G*',
 '6-31G**',
 '6-31G**-RIFIT',
 '6-31G*-Blaudeau',
 '6-31G-Blaudeau',
 '6-31G-J',
 '6ZaP',
 '6ZaPa-NR',
 '7ZaPa-NR',
 'AHGBS-5',
 'AHGBS-7',
 'AHGBS-9',
 'AHGBSP1-5',
 'AHGBSP1-7',
 'AHGBSP1-9',
 'AHGBSP2-5',
 'AHGBSP2-7',
 'AHGBSP2-9',
 'AHGBSP3-5',
 'AHGBSP3-7',
 'AHGBSP3-9',
 'ANO-DK3',
 'ANO-R',
 'ANO-R0',
 'ANO-R1',
 'ANO-R2',
 'ANO-R3',
 'ANO-RCC',
 'ANO-RCC-MB',
 'ANO-RCC-VDZ',
 'ANO-RCC-VDZP',
 'ANO-RCC-VQZP',
 'ANO-RCC-VTZ',
 'ANO-RCC-VTZP',
 'ANO-VT-DZ',
 'ANO-VT-QZ',
 'ANO-VT-TZ',
 'ATZP-ZORA',
 'Ahlrichs TZV',
 'Ahlrichs VDZ',
 'Ahlrichs VTZ',
 'Ahlrichs pVDZ',
 'CADPAC-TZ2P',
 'CRENBL',
 'CRENBL ECP',
 'CRENBS',
 'CRENBS ECP',
 'Cologne DKH2',
 'DFO+-NRLMOL',
 'DFO-1',
 'DFO-1-BHS',
 'DFO-2',
 'DFO-NRLMOL',
 'DZ (Dunning-Hay)',
 'DZ + Double Rydberg (Dunning-Hay)',
 'DZ + Rydberg (Dunning-Hay)',
 'DZP (Dunning-Hay)',
 'DZP + Diffuse (Dunning-Hay)',
 'DZP + Rydberg (Dunning-Hay)',
 'FANO-5Z',
 'FANO-6Z',
 'FANO-DZ',
 'FANO-QZ',
 'FANO-TZ',
 'Grimme vDZP',
 'HGBS-5',
 'HGBS-7',
 'HGBS-9',
 'HGBSP1-5',
 'HGBSP1-7',
 'HGBSP1-9',
 'HGBSP2-5',
 'HGBSP2-7',
 'HGBSP2-9',
 'HGBSP3-5',
 'HGBSP3-7',
 'HGBSP3-9',
 'IGLO-II',
 'IGLO-III',
 'Koga unpolarized',
 'LANL08',
 'LANL08(d)',
 'LANL08(f)',
 'LANL08+',
 'LANL2DZ',
 'LANL2DZ ECP',
 'LANL2DZdp',
 'LANL2TZ',
 'LANL2TZ(f)',
 'LANL2TZ+',
 'MIDI',
 'MIDI!',
 'MIDIX',
 'MINI',
 'NASA Ames ANO',
 'NASA Ames ANO2',
 'NASA Ames cc-pCV5Z',
 'NASA Ames cc-pCVQZ',
 'NASA Ames cc-pCVTZ',
 'NASA Ames cc-pV5Z',
 'NASA Ames cc-pVQZ',
 'NASA Ames cc-pVTZ',
 'NLO-V',
 'NMR-DKH (TZ2P)',
 'ORP',
 'PAW-L05',
 'PAW-L1',
 'PAW-L1-contracted',
 'PAW-L2',
 'PAW-L2-contracted',
 'PB4-D',
 'PB4-F1',
 'PB4-F2',
 'PB5-D',
 'PB5-F',
 'PB5-G',
 'PB6-D',
 'PB6-F',
 'PB6-G',
 'PB6-H',
 'Partridge Uncontracted 1',
 'Partridge Uncontracted 2',
 'Partridge Uncontracted 3',
 'Partridge Uncontracted 4',
 'PsX-DZ',
 'PsX-QZ',
 'PsX-TZ',
 'Pt - mDZP',
 'Roos Augmented Double Zeta ANO',
 'Roos Augmented Triple Zeta ANO',
 'SARC-DKH2',
 'SARC-ZORA',
 'SARC2-QZV-DKH2',
 'SARC2-QZV-DKH2-JKFIT',
 'SARC2-QZV-ZORA',
 'SARC2-QZV-ZORA-JKFIT',
 'SARC2-QZVP-DKH2',
 'SARC2-QZVP-DKH2-JKFIT',
 'SARC2-QZVP-ZORA',
 'SARC2-QZVP-ZORA-JKFIT',
 'SBKJC Polarized (p,2d) - LFK',
 'SBKJC-ECP',
 'SBKJC-VDZ',
 'SBO4-DZ(d)-3G',
 'SBO4-DZ(d,p)-3G',
 'SBO4-SZ-3G',
 'STO-2G',
 'STO-3G',
 'STO-3G*',
 'STO-4G',
 'STO-5G',
 'STO-6G',
 'SV (Dunning-Hay)',
 'SV + Double Rydberg (Dunning-Hay)',
 'SV + Rydberg (Dunning-Hay)',
 'SVP (Dunning-Hay)',
 'SVP + Diffuse (Dunning-Hay)',
 'SVP + Diffuse + Rydberg (Dunning-Hay)',
 'SVP + Rydberg (Dunning-Hay)',
 'Sadlej pVTZ',
 'Sadlej+',
 'Sapporo-DKH3-DZP',
 'Sapporo-DKH3-DZP-2012',
 'Sapporo-DKH3-DZP-2012-diffuse',
 'Sapporo-DKH3-DZP-diffuse',
 'Sapporo-DKH3-QZP',
 'Sapporo-DKH3-QZP-2012',
 'Sapporo-DKH3-QZP-2012-diffuse',
 'Sapporo-DKH3-QZP-diffuse',
 'Sapporo-DKH3-TZP',
 'Sapporo-DKH3-TZP-2012',
 'Sapporo-DKH3-TZP-2012-diffuse',
 'Sapporo-DKH3-TZP-diffuse',
 'Sapporo-DZP',
 'Sapporo-DZP-2012',
 'Sapporo-DZP-2012-diffuse',
 'Sapporo-DZP-diffuse',
 'Sapporo-QZP',
 'Sapporo-QZP-2012',
 'Sapporo-QZP-2012-diffuse',
 'Sapporo-QZP-diffuse',
 'Sapporo-TZP',
 'Sapporo-TZP-2012',
 'Sapporo-TZP-2012-diffuse',
 'Sapporo-TZP-diffuse',
 'Scaled MINI',
 'Stuttgart RLC',
 'Stuttgart RLC ECP',
 'Stuttgart RSC 1997',
 'Stuttgart RSC 1997 ECP',
 'Stuttgart RSC ANO',
 'Stuttgart RSC Segmented + ECP',
 'TZ (Dunning-Hay)',
 'TZP-ZORA',
 'UGBS',
 'WTBS',
 'Wachters+f',
 'acv2z-J',
 'acv3z-J',
 'acv4z-J',
 'admm-1',
 'admm-2',
 'admm-3',
 'ano-pV5Z',
 'ano-pVDZ',
 'ano-pVQZ',
 'ano-pVTZ',
 'apr-cc-pV(Q+d)Z',
 'asigmaDZ',
 'asigmaQZ',
 'asigmaTZ',
 'aug-admm-1',
 'aug-admm-2',
 'aug-admm-3',
 'aug-ano-pV5Z',
 'aug-ano-pVDZ',
 'aug-ano-pVQZ',
 'aug-ano-pVTZ',
 'aug-cc-pCV5Z',
 'aug-cc-pCVDZ',
 'aug-cc-pCVDZ-DK',
 'aug-cc-pCVQZ',
 'aug-cc-pCVQZ-DK',
 'aug-cc-pCVTZ',
 'aug-cc-pCVTZ-DK',
 'aug-cc-pV(5+d)Z',
 'aug-cc-pV(D+d)Z',
 'aug-cc-pV(Q+d)Z',
 'aug-cc-pV(T+d)Z',
 'aug-cc-pV5Z',
 'aug-cc-pV5Z-DK',
 'aug-cc-pV5Z-OPTRI',
 'aug-cc-pV5Z-PP',
 'aug-cc-pV5Z-PP-OPTRI',
 'aug-cc-pV5Z-PP-RIFIT',
 'aug-cc-pV5Z-RIFIT',
 'aug-cc-pV6Z',
 'aug-cc-pV6Z-RIFIT',
 'aug-cc-pV7Z',
 'aug-cc-pVDZ',
 'aug-cc-pVDZ-DK',
 'aug-cc-pVDZ-DK3',
 'aug-cc-pVDZ-OPTRI',
 'aug-cc-pVDZ-PP',
 'aug-cc-pVDZ-PP-OPTRI',
 'aug-cc-pVDZ-PP-RIFIT',
 'aug-cc-pVDZ-RIFIT',
 'aug-cc-pVDZ-X2C',
 'aug-cc-pVQZ',
 'aug-cc-pVQZ-DK',
 'aug-cc-pVQZ-DK3',
 'aug-cc-pVQZ-OPTRI',
 'aug-cc-pVQZ-PP',
 'aug-cc-pVQZ-PP-OPTRI',
 'aug-cc-pVQZ-PP-RIFIT',
 'aug-cc-pVQZ-RIFIT',
 'aug-cc-pVQZ-X2C',
 'aug-cc-pVTZ',
 'aug-cc-pVTZ-DK',
 'aug-cc-pVTZ-DK3',
 'aug-cc-pVTZ-J',
 'aug-cc-pVTZ-OPTRI',
 'aug-cc-pVTZ-PP',
 'aug-cc-pVTZ-PP-OPTRI',
 'aug-cc-pVTZ-PP-RIFIT',
 'aug-cc-pVTZ-RIFIT',
 'aug-cc-pVTZ-X2C',
 'aug-cc-pwCV5Z',
 'aug-cc-pwCV5Z-DK',
 'aug-cc-pwCV5Z-PP',
 'aug-cc-pwCV5Z-PP-OPTRI',
 'aug-cc-pwCV5Z-PP-RIFIT',
 'aug-cc-pwCV5Z-RIFIT',
 'aug-cc-pwCVDZ',
 'aug-cc-pwCVDZ-DK3',
 'aug-cc-pwCVDZ-PP',
 'aug-cc-pwCVDZ-PP-OPTRI',
 'aug-cc-pwCVDZ-PP-RIFIT',
 'aug-cc-pwCVDZ-RIFIT',
 'aug-cc-pwCVDZ-X2C',
 'aug-cc-pwCVQZ',
 'aug-cc-pwCVQZ-DK',
 'aug-cc-pwCVQZ-DK3',
 'aug-cc-pwCVQZ-PP',
 'aug-cc-pwCVQZ-PP-OPTRI',
 'aug-cc-pwCVQZ-PP-RIFIT',
 'aug-cc-pwCVQZ-RIFIT',
 'aug-cc-pwCVQZ-X2C',
 'aug-cc-pwCVTZ',
 'aug-cc-pwCVTZ-DK',
 'aug-cc-pwCVTZ-DK3',
 'aug-cc-pwCVTZ-PP',
 'aug-cc-pwCVTZ-PP-OPTRI',
 'aug-cc-pwCVTZ-PP-RIFIT',
 'aug-cc-pwCVTZ-RIFIT',
 'aug-cc-pwCVTZ-X2C',
 'aug-ccX-5Z',
 'aug-ccX-DZ',
 'aug-ccX-QZ',
 'aug-ccX-TZ',
 'aug-mcc-pV5Z',
 'aug-mcc-pV6Z',
 'aug-mcc-pV7Z',
 'aug-mcc-pV8Z',
 'aug-mcc-pVQZ',
 'aug-mcc-pVTZ',
 'aug-pV7Z',
 'aug-pc-0',
 'aug-pc-1',
 'aug-pc-2',
 'aug-pc-3',
 'aug-pc-4',
 'aug-pcH-1',
 'aug-pcH-2',
 'aug-pcH-3',
 'aug-pcH-4',
 'aug-pcJ-0',
 'aug-pcJ-0_2006',
 'aug-pcJ-1',
 'aug-pcJ-1_2006',
 'aug-pcJ-2',
 'aug-pcJ-2_2006',
 'aug-pcJ-3',
 'aug-pcJ-3_2006',
 'aug-pcJ-4',
 'aug-pcJ-4_2006',
 'aug-pcS-0',
 'aug-pcS-1',
 'aug-pcS-2',
 'aug-pcS-3',
 'aug-pcS-4',
 'aug-pcSseg-0',
 'aug-pcSseg-1',
 'aug-pcSseg-2',
 'aug-pcSseg-3',
 'aug-pcSseg-4',
 'aug-pcX-1',
 'aug-pcX-2',
 'aug-pcX-3',
 'aug-pcX-4',
 'aug-pcseg-0',
 'aug-pcseg-1',
 'aug-pcseg-2',
 'aug-pcseg-3',
 'aug-pcseg-4',
 'aug-seg-cc-pV5Z-PP',
 'aug-seg-cc-pVDZ-PP',
 'aug-seg-cc-pVQZ-PP',
 'aug-seg-cc-pVTZ-PP',
 'aug-seg-cc-pwCV5Z-PP',
 'aug-seg-cc-pwCVDZ-PP',
 'aug-seg-cc-pwCVQZ-PP',
 'aug-seg-cc-pwCVTZ-PP',
 'binning 641',
 'binning 641(d)',
 'binning 641(df)',
 'binning 641+',
 'binning 641+(d)',
 'binning 641+(df)',
 'binning 962',
 'binning 962(d)',
 'binning 962(df)',
 'binning 962+',
 'binning 962+(d)',
 'binning 962+(df)',
 'cc-pCV5Z',
 'cc-pCVDZ',
 'cc-pCVDZ-DK',
 'cc-pCVDZ-F12',
 'cc-pCVDZ-F12-OPTRI',
 'cc-pCVDZ-F12-RIFIT',
 'cc-pCVQZ',
 'cc-pCVQZ-DK',
 'cc-pCVQZ-F12',
 'cc-pCVQZ-F12-OPTRI',
 'cc-pCVQZ-F12-RIFIT',
 'cc-pCVTZ',
 'cc-pCVTZ-DK',
 'cc-pCVTZ-F12',
 'cc-pCVTZ-F12-OPTRI',
 'cc-pCVTZ-F12-RIFIT',
 'cc-pV(5+d)Z',
 'cc-pV(D+d)Z',
 'cc-pV(Q+d)Z',
 'cc-pV(T+d)Z',
 'cc-pV5Z',
 'cc-pV5Z(fi/sf/fw)',
 'cc-pV5Z(fi/sf/lc)',
 'cc-pV5Z(fi/sf/sc)',
 'cc-pV5Z(pt/sf/fw)',
 'cc-pV5Z(pt/sf/lc)',
 'cc-pV5Z(pt/sf/sc)',
 'cc-pV5Z-DK',
 'cc-pV5Z-F12',
 'cc-pV5Z-F12(rev2)',
 'cc-pV5Z-JKFIT',
 'cc-pV5Z-PP',
 'cc-pV5Z-PP-RIFIT',
 'cc-pV5Z-RIFIT',
 'cc-pV6Z',
 'cc-pV6Z-RIFIT',
 'cc-pV8Z',
 'cc-pV9Z',
 'cc-pVDZ',
 'cc-pVDZ(fi/sf/fw)',
 'cc-pVDZ(fi/sf/lc)',
 'cc-pVDZ(fi/sf/sc)',
 'cc-pVDZ(pt/sf/fw)',
 'cc-pVDZ(pt/sf/lc)',
 'cc-pVDZ(pt/sf/sc)',
 'cc-pVDZ(seg-opt)',
 'cc-pVDZ-DK',
 'cc-pVDZ-DK3',
 'cc-pVDZ-F12',
 'cc-pVDZ-F12(rev2)',
 'cc-pVDZ-F12-OPTRI',
 'cc-pVDZ-F12-OPTRI+',
 'cc-pVDZ-PP',
 'cc-pVDZ-PP-RIFIT',
 'cc-pVDZ-RIFIT',
 'cc-pVDZ-X2C',
 'cc-pVQZ',
 'cc-pVQZ(fi/sf/fw)',
 'cc-pVQZ(fi/sf/lc)',
 'cc-pVQZ(fi/sf/sc)',
 'cc-pVQZ(pt/sf/fw)',
 'cc-pVQZ(pt/sf/lc)',
 'cc-pVQZ(pt/sf/sc)',
 'cc-pVQZ(seg-opt)',
 'cc-pVQZ-DK',
 'cc-pVQZ-DK3',
 'cc-pVQZ-F12',
 'cc-pVQZ-F12(rev2)',
 'cc-pVQZ-F12-OPTRI',
 'cc-pVQZ-F12-OPTRI+',
 'cc-pVQZ-JKFIT',
 'cc-pVQZ-PP',
 'cc-pVQZ-PP-RIFIT',
 'cc-pVQZ-RIFIT',
 'cc-pVQZ-X2C',
 'cc-pVTZ',
 'cc-pVTZ(fi/sf/fw)',
 'cc-pVTZ(fi/sf/lc)',
 'cc-pVTZ(fi/sf/sc)',
 'cc-pVTZ(pt/sf/fw)',
 'cc-pVTZ(pt/sf/lc)',
 'cc-pVTZ(pt/sf/sc)',
 'cc-pVTZ(seg-opt)',
 'cc-pVTZ-DK',
 'cc-pVTZ-DK3',
 'cc-pVTZ-F12',
 'cc-pVTZ-F12(rev2)',
 'cc-pVTZ-F12-OPTRI',
 'cc-pVTZ-F12-OPTRI+',
 'cc-pVTZ-JKFIT',
 'cc-pVTZ-PP',
 'cc-pVTZ-PP-RIFIT',
 'cc-pVTZ-RIFIT',
 'cc-pVTZ-X2C',
 'cc-pwCV5Z',
 'cc-pwCV5Z-DK',
 'cc-pwCV5Z-PP',
 'cc-pwCV5Z-PP-RIFIT',
 'cc-pwCV5Z-RIFIT',
 'cc-pwCVDZ',
 'cc-pwCVDZ-DK3',
 'cc-pwCVDZ-PP',
 'cc-pwCVDZ-PP-RIFIT',
 'cc-pwCVDZ-RIFIT',
 'cc-pwCVDZ-X2C',
 'cc-pwCVQZ',
 'cc-pwCVQZ-DK',
 'cc-pwCVQZ-DK3',
 'cc-pwCVQZ-PP',
 'cc-pwCVQZ-PP-RIFIT',
 'cc-pwCVQZ-RIFIT',
 'cc-pwCVQZ-X2C',
 'cc-pwCVTZ',
 'cc-pwCVTZ-DK',
 'cc-pwCVTZ-DK3',
 'cc-pwCVTZ-PP',
 'cc-pwCVTZ-PP-RIFIT',
 'cc-pwCVTZ-RIFIT',
 'cc-pwCVTZ-X2C',
 'ccJ-pV5Z',
 'ccJ-pVDZ',
 'ccJ-pVQZ',
 'ccJ-pVTZ',
 'ccX-5Z',
 'ccX-DZ',
 'ccX-QZ',
 'ccX-TZ',
 'ccemd-2',
 'ccemd-3',
 'coemd-2',
 'coemd-3',
 'coemd-4',
 'coemd-ref',
 'd-aug-cc-pV5Z',
 'd-aug-cc-pV6Z',
 'd-aug-cc-pVDZ',
 'd-aug-cc-pVQZ',
 'd-aug-cc-pVTZ',
 'deMon2k-DZVP-GGA',
 'def2-ECP',
 'def2-QZVP',
 'def2-QZVP-RIFIT',
 'def2-QZVPD',
 'def2-QZVPP',
 'def2-QZVPP-RIFIT',
 'def2-QZVPPD',
 'def2-QZVPPD-RIFIT',
 'def2-SV(P)',
 'def2-SV(P)-JKFIT',
 'def2-SV(P)-RIFIT',
 'def2-SVP',
 'def2-SVP-RIFIT',
 'def2-SVPD',
 'def2-SVPD-RIFIT',
 'def2-TZVP',
 'def2-TZVP-RIFIT',
 'def2-TZVPD',
 'def2-TZVPD-RIFIT',
 'def2-TZVPP',
 'def2-TZVPP-RIFIT',
 'def2-TZVPPD',
 'def2-TZVPPD-RIFIT',
 'def2-mTZVP',
 'def2-mTZVPP',
 'def2-mTZVPP-RIJ',
 'def2-universal-JFIT',
 'def2-universal-JKFIT',
 'dgauss-a1-dftjfit',
 'dgauss-a1-dftxfit',
 'dgauss-a2-dftjfit',
 'dgauss-a2-dftxfit',
 'dgauss-dzvp',
 'dgauss-dzvp2',
 'dgauss-tzvp',
 'dhf-ECP',
 'dhf-QZVP',
 'dhf-QZVPP',
 'dhf-SV(P)',
 'dhf-SVP',
 'dhf-TZVP',
 'dhf-TZVPP',
 'dyall-aae2z',
 'dyall-aae3z',
 'dyall-aae4z',
 'dyall-aae5z',
 'dyall-acv2z',
 'dyall-acv3z',
 'dyall-acv4z',
 'dyall-acv5z',
 'dyall-ae2z',
 'dyall-ae3z',
 'dyall-ae4z',
 'dyall-ae5z',
 'dyall-av2z',
 'dyall-av3z',
 'dyall-av4z',
 'dyall-av5z',
 'dyall-cv2z',
 'dyall-cv3z',
 'dyall-cv4z',
 'dyall-cv5z',
 'dyall-v2z',
 'dyall-v3z',
 'dyall-v4z',
 'dyall-v5z',
 'epc-10s10p10d10f',
 'epc-8s8p8d',
 'jgauss-dzp',
 'jgauss-qz2p',
 'jgauss-qzp',
 'jgauss-tzp1',
 'jgauss-tzp2',
 'jorge-5ZP',
 'jorge-5ZP-DKH',
 'jorge-5ZP-ZORA',
 'jorge-6ZP',
 'jorge-6ZP-DKH',
 'jorge-6ZP-ZORA',
 'jorge-A5ZP',
 'jorge-A6ZP',
 'jorge-ADZP',
 'jorge-AQZP',
 'jorge-ATZP',
 'jorge-ATZP-ZORA',
 'jorge-DZP',
 'jorge-DZP-DKH',
 'jorge-DZP-ZORA',
 'jorge-QZP',
 'jorge-QZP-DKH',
 'jorge-QZP-ZORA',
 'jorge-TZP',
 'jorge-TZP-DKH',
 'jorge-TZP-ZORA',
 'jul-cc-pV(D+d)Z',
 'jul-cc-pV(Q+d)Z',
 'jul-cc-pV(T+d)Z',
 'jun-cc-pV(D+d)Z',
 'jun-cc-pV(Q+d)Z',
 'jun-cc-pV(T+d)Z',
 'lcecp-0-QZVP',
 'lcecp-0-QZVPP',
 'lcecp-0-SV(P)',
 'lcecp-0-SVP',
 'lcecp-0-TZVP',
 'lcecp-0-TZVPP',
 'lcecp-1-QZVP',
 'lcecp-1-QZVPP',
 'lcecp-1-SV(P)',
 'lcecp-1-SVP',
 'lcecp-1-TZVP',
 'lcecp-1-TZVPP',
 'lcecp-2-QZVP',
 'lcecp-2-QZVPP',
 'lcecp-2-SV(P)',
 'lcecp-2-SVP',
 'lcecp-2-TZVP',
 'lcecp-2-TZVPP',
 'm6-31G',
 'm6-31G*',
 'maug-cc-pV(D+d)Z',
 'maug-cc-pV(Q+d)Z',
 'maug-cc-pV(T+d)Z',
 'may-cc-pV(Q+d)Z',
 'may-cc-pV(T+d)Z',
 'modified-LANL2DZ',
 'pSBKJC',
 'pV6Z',
 'pV7Z',
 'pc-0',
 'pc-1',
 'pc-2',
 'pc-3',
 'pc-4',
 'pcH-1',
 'pcH-2',
 'pcH-3',
 'pcH-4',
 'pcJ-0',
 'pcJ-0_2006',
 'pcJ-1',
 'pcJ-1_2006',
 'pcJ-2',
 'pcJ-2_2006',
 'pcJ-3',
 'pcJ-3_2006',
 'pcJ-4',
 'pcJ-4_2006',
 'pcS-0',
 'pcS-1',
 'pcS-2',
 'pcS-3',
 'pcS-4',
 'pcSseg-0',
 'pcSseg-1',
 'pcSseg-2',
 'pcSseg-3',
 'pcSseg-4',
 'pcX-1',
 'pcX-2',
 'pcX-3',
 'pcX-4',
 'pcemd-2',
 'pcemd-3',
 'pcemd-4',
 'pcseg-0',
 'pcseg-1',
 'pcseg-2',
 'pcseg-3',
 'pcseg-4',
 'pecJ-1',
 'pecJ-2',
 'pob-DZVP-rev2',
 'pob-TZVP',
 'pob-TZVP-rev2',
 's3-21G',
 's3-21G*',
 's6-31G',
 's6-31G*',
 'sap_grasp_large',
 'sap_grasp_small',
 'sap_helfem_large',
 'sap_helfem_small',
 'saug-ano-pV5Z',
 'saug-ano-pVDZ',
 'saug-ano-pVQZ',
 'saug-ano-pVTZ',
 'seg-cc-pV5Z-PP',
 'seg-cc-pVDZ-PP',
 'seg-cc-pVQZ-PP',
 'seg-cc-pVTZ-PP',
 'seg-cc-pwCV5Z-PP',
 'seg-cc-pwCVDZ-PP',
 'seg-cc-pwCVQZ-PP',
 'seg-cc-pwCVTZ-PP',
 'sigmaDZ',
 'sigmaDZHF',
 'sigmaQZ',
 'sigmaSZHF',
 'sigmaTZ',
 'sigmaTZHF',
 'un-ccemd-ref',
 'un-pcemd-ref',
 'x2c-JFIT',
 'x2c-JFIT-universal',
 'x2c-QZVPPall',
 'x2c-QZVPPall-2c',
 'x2c-QZVPPall-2c-s',
 'x2c-QZVPPall-s',
 'x2c-QZVPall',
 'x2c-QZVPall-2c',
 'x2c-QZVPall-2c-s',
 'x2c-QZVPall-s',
 'x2c-SV(P)all',
 'x2c-SV(P)all-2c',
 'x2c-SV(P)all-s',
 'x2c-SVPall',
 'x2c-SVPall-2c',
 'x2c-SVPall-s',
 'x2c-TZVPPall',
 'x2c-TZVPPall-2c',
 'x2c-TZVPPall-s',
 'x2c-TZVPall',
 'x2c-TZVPall-2c',
 'x2c-TZVPall-s']

# Load basis set
basis = get_ao_basis_from_bse(h2o.atnums, atcoords=h2o.atcoords, basis_name="STO-6G", coord_types = "spherical")

Secular Equations#

Given a set of basis functions $\{\chi_\lambda\}$, we can expand the molecular orbitals as linear combinations of the basis functions, $$\phi_p(\mathbf{r}) = \sum_\lambda c_{\lambda p} \chi_\lambda(\mathbf{r})$$ where $c_{\lambda p}$ is the coefficient of the $\lambda$-th basis function in the expansion of the $p$-th molecular orbital. Substituting this expansion into the one-electron equations and projecting onto a basis function $\chi_\kappa$, we obtain the secular equations, $$\sum_\lambda \left[ \int \chi_\kappa^*(\mathbf{r}) \left(-\frac{1}{2} \nabla^2 + v(\mathbf{r}) + \hat{w}[\{\phi_q\}](\mathbf{r})\right) \chi_\lambda(\mathbf{r}) d\mathbf{r} \right] c_{\lambda p} = \epsilon_p \sum_\lambda \left[ \int \chi_\kappa^*(\mathbf{r}) \chi_\lambda(\mathbf{r}) d\mathbf{r} \right] c_{\lambda p}$$ Defining the Hamiltonian matrix elements as $$F_{\kappa \lambda} = \int \chi_\kappa^*(\mathbf{r}) \left(-\frac{1}{2} \nabla^2 + v(\mathbf{r}) + \hat{w}[\{\phi_q\}](\mathbf{r})\right) \chi_\lambda(\mathbf{r}) d\mathbf{r}$$ and the overlap matrix elements as $$S_{\kappa \lambda} = \int \chi_\kappa^*(\mathbf{r}) \chi_\lambda(\mathbf{r}) d\mathbf{r}$$ the secular equations can be written in matrix form as $$\mathbf{F} \mathbf{c}_p = \epsilon_p \mathbf{S} \mathbf{c}_p$$ where $\mathbf{F}$ is the Fock matrix, $\mathbf{S}$ is the overlap matrix, and $\mathbf{c}_p$ is the vector of coefficients for the $p$-th molecular orbital. This is a generalized eigenvalue problem that can be solved to obtain the molecular orbitals and their corresponding energies. Note that one needs to (usually) solve for the $\alpha$- and $\beta$-spin orbitals separately, which leads to separate Fock matrices for each spin.

The Fock matrix is a sum of three parts, the kinetic energy matrix, the nuclear attraction matrix, and the average potential matrix, $$F_{\kappa \lambda} = T_{\kappa \lambda} + V_{\kappa \lambda} + W_{\kappa \lambda}$$ where $$T_{\kappa \lambda} = \int \chi_\kappa^*(\mathbf{r}) \left(-\frac{1}{2} \nabla^2\right) \chi_\lambda(\mathbf{r}) d\mathbf{r}$$V_{\kappa \lambda} = \int \chi_\kappa^*(\mathbf{r}) v(\mathbf{r}) \chi_\lambda(\mathbf{r}) d\mathbf{r}$$W_{\kappa \lambda} = \int \chi_\kappa^*(\mathbf{r}) \hat{w}[\{\phi_q\}](\mathbf{r}) \chi_\lambda(\mathbf{r}) d\mathbf{r}$$

import numpy as np
import scipy
from scipy.special import erf
import json
import atomdb
from iodata import load_one
from grid.onedgrid import GaussLegendre
from grid.rtransform import BeckeRTransform
from grid.becke import BeckeWeights
from grid.molgrid import MolGrid
from gbasis.integrals.overlap import overlap_integral
from gbasis.integrals.kinetic_energy import kinetic_energy_integral
from gbasis.integrals.nuclear_electron_attraction import nuclear_electron_attraction_integral
from gbasis.integrals.electron_repulsion import electron_repulsion_integral
from gbasis.evals.eval import evaluate_basis
from gbasis.evals.density import evaluate_density

# Make a decent(ish) grid.
def make_grid(atnums, atcoords, rad_points=100, preset="ultrafine"):
    '''Make a grid object from a molecule
    Args:
        atnums (list of int): Atomic numbers of the atoms in the molecule
        atcoords (numpy.ndarray): Coordinates of the atoms in the molecule
        rad_points (int): number of radial points
        preset (str): preset grid type
            "coarse", "medium", "fine", "veryfine", "ultrafine", "insane"
    Returns:
        mol_grid (Grid): grid object
    '''
    # create a radial grid with rad_points points
    oned_grid = GaussLegendre(rad_points)
    radial_grid = BeckeRTransform(0.0, R=1.5).transform_1d_grid(oned_grid)

    # create a pruned molecular grid with the preset angular convention
    mol_grid = MolGrid.from_preset(atnums, atcoords, preset, rgrid=radial_grid, store=True)

    return mol_grid

# Compute overlap integrals
def compute_overlap(ao_basis):
    r"""Compute the overlap matrix in the AO basis.

    Parameters
    ----------
    ao_basis : list of `GeneralizedContractionShell` instances
        Atomic orbital basis set for the molecule.

    Returns
    -------
    S : numpy.ndarray[float, shape=(n_ao, n_ao)]
        Overlap matrix in the AO basis.
    """

    return overlap_integral(ao_basis)

---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
Cell In[6], line 7
      3 from scipy.special import erf
      4 import json
      5 import atomdb
      6 from iodata import load_one
----> 7 from grid.onedgrid import GaussLegendre
      8 from grid.rtransform import BeckeRTransform
      9 from grid.becke import BeckeWeights
     10 from grid.molgrid import MolGrid

File /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages/grid/__init__.py:25
     21 """Grid Module."""
     22 # flake8: noqa
---> 25 from grid.angular import *
     26 from grid.atomgrid import *
     27 from grid.basegrid import *

File /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages/grid/angular.py:75
     72 import numpy as np
     73 from importlib_resources import files
---> 75 from grid.basegrid import Grid
     77 # Lebedev dictionary for grid's number of points (keys) and degrees (values)
     78 LEBEDEV_NPOINTS = {
     79     6: 3,
     80     18: 5,
   (...)    110     5810: 131,
    111 }

File /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages/grid/basegrid.py:25
     22 import numpy as np
     23 from scipy.spatial import cKDTree
---> 25 from grid.utils import (
     26     convert_cart_to_sph,
     27     generate_orders_horton_order,
     28     solid_harmonics,
     29 )
     32 class Grid:
     33     """Basic Grid class for grid information storage."""

File /opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages/grid/utils.py:23
     20 """Utility Module."""
     22 import numpy as np
---> 23 from scipy.special import sph_harm
     25 _bragg = np.array(
     26     [
     27         np.nan,  # index 0, place holder
   (...)    114     ]
    115 )
    117 _cambridge = np.array(
    118     [
    119         np.nan,  # index 0, place holder
   (...)    206     ]
    207 )

ImportError: cannot import name 'sph_harm' from 'scipy.special' (/opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/site-packages/scipy/special/__init__.py)

# Evaluate the Coulomb potential of a Gaussian and r^2 times a Gaussian (on a grid)

def coulomb_gaussian_s(r, a):
    """
    Coulomb potential of exp(-a r^2), for scalar or array r.
    Returns V(r) = (pi/a)^(3/2) * erf(sqrt(a) r) / r
    with the correct finite limit at r=0.
    """
    r = np.asarray(r, dtype=float)
    out = np.empty_like(r)

    mask0 = np.isclose(r, 0.0)
    mask = ~mask0

    out[mask] = (np.pi / a) ** 1.5 * erf(np.sqrt(a) * r[mask]) / r[mask]
    out[mask0] = 2.0 * np.pi / a
    return out


def coulomb_gaussian_r2(r, b):
    """
    Coulomb potential of r^2 exp(-b r^2), for scalar or array r.
    Returns
        V(r) = (3 pi^(3/2)/(2 b^(5/2))) * erf(sqrt(b) r)/r
             + (2 pi / b^2) * exp(-b r^2)
    with the correct finite limit at r=0.
    """
    r = np.asarray(r, dtype=float)
    out = np.empty_like(r)

    mask0 = np.isclose(r, 0.0)
    mask = ~mask0

    out[mask] = (
        (3.0 * np.pi**1.5 / (2.0 * b**2.5)) * erf(np.sqrt(b) * r[mask]) / r[mask]
        + (2.0 * np.pi / b**2) * np.exp(-b * r[mask] ** 2)
    )
    out[mask0] = 5.0 * np.pi / b**2
    return out

def coulomb_potential(r, c, a, d=None, b=None):
    """
    Coulomb potential of
        rho(r) = sum_i c_i exp(-a_i r^2) + sum_j d_j r^2 exp(-b_j r^2)

    Parameters
    ----------
    r : float or ndarray
        Radial coordinate(s).
    c, a : 1D arrays
        Coefficients and exponents for the pure Gaussians.
    d, b : 1D arrays, optional
        Coefficients and exponents for the r^2 Gaussians.

    Returns
    -------
    V : ndarray
        Coulomb potential evaluated at r.
    """
    r = np.asarray(r, dtype=float)
    V = np.zeros_like(r)

    c = np.asarray(c, dtype=float)
    a = np.asarray(a, dtype=float)
    for ci, ai in zip(c, a):
        V += ci * coulomb_gaussian_s(r, ai)

    if d is not None and b is not None:
        d = np.asarray(d, dtype=float)
        b = np.asarray(b, dtype=float)
        for dj, bj in zip(d, b):
            V += dj * coulomb_gaussian_r2(r, bj)

    return V

from jax import vjp


def compute_harris_lsda(atnums, atcoords, ao_basis, n_electrons):
    r"""Compute the LSDA Harris-functional guess for the exchange-correlation potential matrix.

    The LSDA Harris-functional guess is computed by evaluating the Fock matrix using the promolecular density
    as the input density. The exchange-correlation potential is computed using the local spin-density approximation (LSDA).

    Parameters
    ----------
    atnums : list of int
        Atomic numbers of the atoms in the molecule.
    atcoords : numpy.ndarray[float, shape=(n_atoms, 3)]
        Coordinates of the atoms in the molecule.
    ao_basis : list of `GeneralizedContractionShell` instances
        Atomic orbital basis set for the molecule.
    n_electrons : numpy.ndarray[int, shape=(2,)]
        Number of electrons in the molecule (of each spin channel).
    Returns
    -------
    harris_guess : numpy.ndarray[float, shape=(n_ao, n_ao, 2)]
        LSDA exchange-correlation potential matrix elements. The first component of the last dimension
        corresponds to the alpha spin, and the second component corresponds to the beta spin.
    """
    # We use `AtomDB` to compute the promolecular density of the molecule. We assume
    # that the spin-densities are proportional to the total density.

    # Make a grid. No need to be super accurate here.
    mol_grid = make_grid(atnums, atcoords, rad_points=50, preset="medium")

    # Compute the promolecular density on the grid.
    # evaluate the promolecular density on the grid
    promol = atomdb.make_promolecule(atnums, atcoords, dataset="gaussian")
    promol_dens = promol.density(mol_grid.points, log=True)

    # Get spin-densities.
    n_a, n_b = n_electrons
    #The promolecular density is always neutral, so the number of electrons *expected* is the
    #sum of the atomic numbers.
    n_total = np.sum(atnums)


    # Components of the density are n_alpha/n_total and n_beta/n_total times the total density.
    promol_dens_a = promol_dens * n_a / (n_total)
    promol_dens_b = promol_dens * n_b / (n_total)

    # Compute the LSDA exchange-correlation potential on the grid.
    vxc_a = (6/np.pi)**(1/3) * promol_dens_a ** (1/3)    # LSDA exchange potential for alpha spin
    vxc_b = (6/np.pi)**(1/3) * promol_dens_b ** (1/3)    # LSDA exchange potential for beta spin

    # Evaluate the exchange-correlation potential matrix elements in the AO basis.
    phi = evaluate_basis(basis, mol_grid.points)
    # AO matrices for alpha and beta
    Fxc = np.empty((phi.shape[0], phi.shape[0], 2))
    Fxc[:, :, 0] = np.einsum("ip,p,jp->ij", phi, mol_grid.weights * vxc_a, phi)
    Fxc[:, :, 1] = np.einsum("ip,p,jp->ij", phi, mol_grid.weights * vxc_b, phi)

    return Fxc

def compute_harris_hartree(atnums, atcoords, ao_basis, n_electrons):
    r"""Compute the Hartree Harris-functional guess for the Fock matrix.

    The Hartree Harris-functional guess is computed by evaluating the Fock matrix using the promolecular density
    as the input density.

    Parameters
    ----------
    atnums : list of int
        Atomic numbers of the atoms in the molecule.
    atcoords : numpy.ndarray[float, shape=(n_atoms, 3)]
        Coordinates of the atoms in the molecule.
    ao_basis : list of `GeneralizedContractionShell` instances
        Atomic orbital basis set for the molecule.
    n_electrons : numpy.ndarray[int, shape=(2,)]
        Number of electrons in the molecule (of each spin channel).
    Returns
    -------
    hartree_pot : numpy.ndarray[float, shape=(n_ao, n_ao)]
        Hartree potential matrix elements.
    """
    n_a, n_b = n_electrons
    n_total = sum(atnums)

    # We use `BFit` to obtain an approximate density with the same number of electrons as the promolecular density,
    # expanded in s-type Gaussians. The Coulomb potential is then easy to evaluate on a grid.
    # Make a grid. No need to be super accurate here.
    mol_grid = make_grid(atnums, atcoords, rad_points=50, preset="medium")

    # Load atomic fit data from BFit.
    url = "https://raw.githubusercontent.com/theochem/BFit/master/bfit/data/kl_slsqp_results.json"
    urlretrieve(url, "kl_slsqp_results.json")
    with open("kl_slsqp_results.json") as file:
        data = json.load(file)

    # For each atom, we compute the Coulomb potential on the grid and sum them up to get the total Hartree potential.
    # Loop over the atoms
    hartree_pot = np.zeros((mol_grid.points.shape[0],))
    for atnum, atcoord in zip(atnums, atcoords):
        # Get the BFit model for the atom
        # Convert the atomic number to the element symbol
        element = iodata.periodic.num2sym[atnum]
        # BFit uses lowercase element symbols, so we convert to lowercase.
        element = element.lower()
        num_s = data[element]["num_s"]
        num_p = data[element]["num_p"]
        coefficients = np.array(data[element]["coeffs"])
        exponents = np.array(data[element]["exps"])

        # For every grid point, compute its distance to the atom in question.
        dist = np.linalg.norm(mol_grid.points - atcoord, axis=1)

        # Evaluate the Coulomb potential from the atom at each grid point using the BFit model.
        # Break the coefficients and exponents into s-type and p-type contributions.
        # Our code uses unnormalized Gaussians but BFit produces normalized Gaussians, so we need to
        # multiply the coefficients by the normalization factors (cf. eqs. (22) and (23)).
        c_s = coefficients[:num_s] * (exponents[:num_s] / np.pi )**(3/2)
        a_s = exponents[:num_s]
        c_p = coefficients[num_s:num_s+num_p] * 2/3 *exponents[num_s:num_s+num_p]**(5/2) / np.pi**(3/2)
        a_p = exponents[num_s:num_s+num_p]
        hartree_pot += coulomb_potential(dist, c_s, a_s, d=c_p, b=a_p)

    # Evaluate the Hartree potential matrix elements in the AO basis.
    phi = evaluate_basis(ao_basis, mol_grid.points)
    vj = hartree_pot * (n_a +n_b) / n_total  # Scale the Hartree potential if the molecule is not neutral.
    # AO matrices for alpha and beta
    fj = np.empty((phi.shape[0], phi.shape[0]))
    fj[:, :] = np.einsum("ip,p,jp->ij", phi, mol_grid.weights * vj, phi)

    return fj

def compute_harris_guess(atnums, atcoords, ao_basis, n_electrons, guess_type='fermi-amaldi'):
    r"""Compute the Harris-functional guess for the Fock matrix.

    The Harris-functional guess is computed by evaluating the Fock matrix using the promolecular density
    as the input density.

    Parameters
    ----------
    atnums : list of int
        Atomic numbers of the atoms in the molecule.
    atcoords : numpy.ndarray[float, shape=(n_atoms, 3)]
        Coordinates of the atoms in the molecule.
    ao_basis : list of `GeneralizedContractionShell` instances
        Atomic orbital basis set for the molecule.
    n_electrons : numpy.ndarray[int, shape=(2,)]
        Number of electrons in the molecule (of each spin channel).
    guess_type : {"fermi-amaldi", "lsda", "hartree"}
        Type of Harris-functional guess to compute. Choices are LSDA (exchange-only), Fermi-Amaldi, and Hartree.

    Returns
    -------
    harris_guess : numpy.ndarray[float, shape=(n_ao, n_ao, 2)]
        Harris-functional guess for the Fock matrix. The first component of the last dimension corresponds to the
        alpha spin, and the second component corresponds to the beta spin.
    """

    n_a, n_b = n_electrons
    kin_energy = kinetic_energy_integral(ao_basis)
    nuc_attraction = nuclear_electron_attraction_integral(ao_basis, atcoords, atnums)
    # Assign the hartree potentials for a and b based on the last axis of compute_harris_hartree.
    vj = compute_harris_hartree(atnums, atcoords, ao_basis, n_electrons)

    fock = np.empty((kin_energy.shape[0], kin_energy.shape[1], 2))
    fock[:, :, 0] = kin_energy + nuc_attraction + vj
    fock[:, :, 1] = fock[:, :, 0]

    if guess_type == "lsda":
        fock += compute_harris_lsda(atnums, atcoords, ao_basis, n_electrons)
    elif guess_type == "fermi-amaldi":
        fock[:, :, 0] -= (1/n_a) * vj
        fock[:, :, 1] -= (1/n_b) * vj
    #make sure there is no weird case where the user choose a forbidden option
    elif guess_type != "hartree":
        raise ValueError(f"Invalid guess type: {guess_type}. Must be one of 'fermi-amaldi', 'lsda', or 'hartree'.")

    return fock

# Solve for the eigenvalues and eigenvectors of the Harris guess Fock matrix to get the initial guess density matrix.
import scipy
from sympy.abc import phi

s = compute_overlap(basis)
n_electrons = np.array([5, 5])  # Number of electrons in the molecule (of each spin channel).
fock = compute_harris_guess(h2o.atnums, h2o.atcoords, basis, n_electrons, guess_type='fermi-amaldi')
eps_a, phi_a = scipy.linalg.eigh(fock[:, :, 0], s)  # Solve the generalized eigenvalue problem for the alpha spin channel
eps_b, phi_b = scipy.linalg.eigh(fock[:, :, 1], s)  # Solve the generalized eigenvalue problem for the beta spin channel

# Put the eigenvalues and MOs into a 2-d array for easier indexing.
eps = np.empty((eps_a.shape[0], 2))
eps[:, 0] = eps_a
eps[:, 1] = eps_b
phi = np.empty((phi_a.shape[0], phi_a.shape[1], 2))
phi[:, :, 0] = phi_a
phi[:, :, 1] = phi_b
print(phi.shape)

print(eps)

5 5 number of electrons in each spin channel
0.2 0.2 number of electrons in each spin channel
(7, 7, 2)
[[-14.3825598  -14.3825598 ]
 [ -1.12906637  -1.12906637]
 [ -0.8205876   -0.8205876 ]
 [ -0.44593379  -0.44593379]
 [ -0.34309427  -0.34309427]
 [  0.11801745   0.11801745]
 [  0.19562679   0.19562679]]

from sympy import N


def compute_density_matrix(eps, phi, n_electrons):
    r"""Compute the density matrix from the eigenvalues and eigenvectors of the Fock matrix.

    Parameters
    ----------
    eps : numpy.ndarray[float, shape=(n_ao, 2)]
        Eigenvalues of the Fock matrix.
    phi : numpy.ndarray[float, shape=(n_ao, n_ao, 2)]
        Eigenvectors of the Fock matrix.
    n_electrons : numpy.ndarray[int, shape=(2,)]
        Number of electrons in the molecule (of a given spin channel).

    Returns
    -------
    density_matrix : numpy.ndarray[float, shape=(n_ao, n_ao, 2)]
        Density matrix in the AO basis.
    """
    n_a, n_b = n_electrons
    # Compute the density matrix as D = sum_i |phi_i><phi_i| where i runs over the occupied orbitals.
    density_matrix = np.zeros((phi.shape[0], phi.shape[0], 2))
    density_matrix[:, :, 0] = np.einsum("pi,qi->pq", phi[:, :n_a, 0], phi[:, :n_a, 0])
    density_matrix[:, :, 1] = np.einsum("pi,qi->pq", phi[:, :n_b, 1], phi[:, :n_b, 1])

    return density_matrix

def check_normalization(density_matrix, n_electrons, overlap):
    r"""Check the normalization of the density matrix.

    Parameters
    ----------
    density_matrix : numpy.ndarray[float, shape=(n_ao, n_ao, 2)]
        Density matrix in the AO basis.
    n_electrons : numpy.ndarray[int, shape=(2,)]
        Number of electrons in the molecule (of each spin channel).
    overlap : numpy.ndarray[float, shape=(n_ao, n_ao)]
        Overlap matrix in the AO basis.

    Returns
    -------
    is_normalized : bool
        True if the density matrix is normalized, False otherwise.
    """
    n_a, n_b = n_electrons
    # The number of electrons can be computed as N = Tr(D S) where D is the density matrix and S is the overlap matrix.
    n_a_computed = np.trace(density_matrix[:, :, 0] @ overlap)
    n_b_computed = np.trace(density_matrix[:, :, 1] @ overlap)
    print(n_a_computed, n_b_computed, "number of electrons in each spin channel computed from the density matrix")

    return np.isclose(n_a_computed, n_a) and np.isclose(n_b_computed, n_b)

print(check_normalization(compute_density_matrix(eps, phi, n_electrons), n_electrons, s ))

4.999999999999998 4.999999999999998 number of electrons in each spin channel computed from the density matrix
True

def get_eri_ao(basis):
    """
    Return AO two-electron integrals in chemist notation:
        eri[mu, nu, lam, sig] = (mu nu | lam sig)
    """
    return electron_repulsion_integral(basis, notation="chemist")

vee = get_eri_ao(basis)

def build_j_ao(D, eri):
    """
    Build the AO Coulomb matrix

        J_mn = sum_ls D_ls (mn|ls)

    Parameters
    ----------
    D : (nao, nao) ndarray
        AO density matrix.
    eri : (nao, nao, nao, nao) ndarray
        AO ERI tensor in chemist notation.

    Returns
    -------
    J : (nao, nao) ndarray
        AO Coulomb matrix.
    """
    D = np.asarray(D)
    return np.einsum("ls,mnls->mn", D, eri, optimize=True)


def build_k_ao(D, eri):
    """
    Build the AO exchange matrix

        K_mn = sum_ls D_ls (ml|ns)

    Parameters
    ----------
    D : (nao, nao) ndarray
        AO density matrix.
    eri : (nao, nao, nao, nao) ndarray
        AO ERI tensor in chemist notation.

    Returns
    -------
    K : (nao, nao) ndarray
        AO exchange matrix.
    """
    D = np.asarray(D)
    return np.einsum("ls,mlns->mn", D, eri, optimize=True)

def build_fock_ao(h_core, density_matrix, eri, overlap, n_electrons, calc_type = 'fermi-amaldi'):
    """
    Build the AO Fock matrix for a given density matrix and ERI tensor.

    Parameters
    ----------
    h_core : (nao, nao) ndarray
        AO core Hamiltonian matrix.
    density_matrix : (nao, nao, 2) ndarray
        AO density matrix for both spin channels.
    eri : (nao, nao, nao, nao) ndarray
        AO ERI tensor in chemist notation.
    overlap : (nao, nao) ndarray
        AO overlap matrix.
    n_electrons : array-like of shape (2,)
        Number of electrons in the molecule (of each spin channel).
    calc_type : str
        Type of calculation to perform. Choices are 'fermi-amaldi', 'lsda', and 'hartree'.

    Returns
    -------
    F_a : (nao, nao) ndarray
        AO Fock matrix for the alpha spin channel.
    F_b : (nao, nao) ndarray
        AO Fock matrix for the beta spin channel.
    """
    D_a = density_matrix[:, :, 0]
    D_b = density_matrix[:, :, 1]
    J_a = build_j_ao(D_a, eri)
    J_b = build_j_ao(D_b, eri)
    K_a = build_k_ao(D_a, eri)
    K_b = build_k_ao(D_b, eri)

    # For consistency check that the number of electrons targetted is consistent with the density matrix given.
    check_normalization(density_matrix, n_electrons, overlap)

    if calc_type == 'fermi-amaldi':
        F_a = h_core + J_a + J_b - (1/n_electrons[0]) * J_a
        F_b = h_core + J_a + J_b - (1/n_electrons[1]) * J_b
    elif calc_type == 'hartree':
        F_a = h_core + J_a + J_b
        F_b = h_core + J_a + J_b
    elif calc_type == 'lsda':
        raise NotImplementedError("LSDA calculation is not implemented yet.")
    else:
        raise ValueError(f"Invalid calculation type: {calc_type}. Must be one of 'fermi-amaldi', 'lsda', or 'hartree'.")

    return F_a, F_b

def scf_cycle(h_core, density_matrix, eri, s, n_electrons, calc_type='hartree', max_iter=100, tol=1e-4):
    """
    Perform the SCF cycle to self-consistency.

    Parameters
    ----------
    h_core : (nao, nao) ndarray
        AO core Hamiltonian matrix.
    density_matrix : (nao, nao, 2) ndarray
        Initial AO density matrix for both spin channels.
    eri : (nao, nao, nao, nao) ndarray
        AO ERI tensor in chemist notation.
    s : (nao, nao) ndarray
        AO overlap matrix.
    n_electrons : array-like of shape (2,)
        Number of electrons in the molecule (of each spin channel).
    calc_type : str
        Type of calculation to perform. Choices are 'fermi-amaldi', 'lsda', and 'hartree'.
    max_iter : int
        Maximum number of SCF iterations to perform.
    tol : float
        Convergence tolerance for the density matrix.

    Returns
    -------
    eps : (n_ao, 2) ndarray
        Eigenvalues of the final Fock matrix.
    phi : (n_ao, n_ao, 2) ndarray
        Eigenvectors of the final Fock matrix.
    """
    nao = h_core.shape[0]
    for iteration in range(max_iter):
        F_a, F_b = build_fock_ao(h_core, density_matrix, eri, s, n_electrons, calc_type)
        eps_a, phi_a = scipy.linalg.eigh(F_a, s)
        eps_b, phi_b = scipy.linalg.eigh(F_b, s)

        # Put the eigenvalues and MOs into a 2-d array for easier indexing.
        eps = np.empty((eps_a.shape[0], 2))
        eps[:, 0] = eps_a
        eps[:, 1] = eps_b
        phi = np.empty((phi_a.shape[0], phi_a.shape[1], 2))
        phi[:, :, 0] = phi_a
        phi[:, :, 1] = phi_b

        # Compute the new density matrix.
        density_matrix_new = compute_density_matrix(eps, phi, n_electrons)

        # Check for convergence
        delta_dm = np.linalg.norm(density_matrix - density_matrix_new)
        print(f"Iteration {iteration+1}: delta_dm = {delta_dm:.6e}")
        if delta_dm < tol:
            print(f"SCF converged in {iteration+1} iterations.")
            return eps, phi

        density_matrix = density_matrix_new *0.1 + density_matrix *0.9  # Simple damping to help convergence.

    raise RuntimeError("SCF did not converge within the maximum number of iterations.")

h_core = kinetic_energy_integral(basis) + nuclear_electron_attraction_integral(basis, h2o.atcoords, h2o.atnums)
eri = get_eri_ao(basis)
eps, phi = scf_cycle(h_core, compute_density_matrix(eps, phi, n_electrons), eri, s, n_electrons, calc_type='fermi-amaldi')
print("Final eigenvalues (orbital energies):", eps)

5.0 5.0 number of electrons in each spin channel computed from the density matrix
Iteration 1: delta_dm = 3.746842e-06
SCF converged in 1 iterations.
Final eigenvalues (orbital energies): [[-17.35918998 -17.35918998]
 [ -0.83428658  -0.83428658]
 [ -0.42319285  -0.42319285]
 [ -0.14136572  -0.14136572]
 [ -0.07400932  -0.07400932]
 [  0.27785362   0.27785362]
 [  0.40140387   0.40140387]]

# Write a function to solve the self-consistent field equations using the Harris-functional guess.

def solve_scf(atnums, atcoords, ao_basis, n_electrons, guess_type='fermi-amaldi', method='fermi-amaldi', max_iter=100, tol=1e-6):
    r"""Solve the self-consistent field equations using the Harris-functional guess.

    Parameters
    ----------
    atnums : list of int
        Atomic numbers of the atoms in the molecule.
    atcoords : numpy.ndarray[float, shape=(n_atoms, 3)]
        Coordinates of the atoms in the molecule.
    ao_basis : list of `GeneralizedContractionShell` instances
        Atomic orbital basis set for the molecule.
    n_electrons : numpy.ndarray[int, shape=(2,)]
        Number of electrons in the molecule (of each spin channel).
    guess_type : {"fermi-amaldi", "lsda", "hartree"}
        Type of Harris-functional guess to compute. Choices are LSDA (exchange-only), Fermi-Amaldi, and Hartree.

# This