Helonium's Hartree-Fock program

Exordium

We will implement the Hartree-Fock¹ program from the classic Szabo-Ostlund text, a staple in quantum chemistry. If you have any experience in the field, chances are you know it well. If you don't, the BQN implementation will take you only a few cognitive units: the coarse mathematical description involves basis sets, the calculation of electronic integrals, and the self-consistent optimization of the Fock matrix. Using this program, we will compute the energy of the HeH\(^+\) molecule².

First, we import the required BQN system values and utility functions:

Sin‿Cos‿ATan‿Erf ← •math
Setplot‿Plot ← •Import "../bqn-utils/plots.bqn"

Some auxiliary functions are needed for the computation of the Boys function, the matrix and vector product, the spectrum of a matrix with shape 2‿2, and Löwdin's canonical orthogonalization. Additionally, we define a 1-modifier to compute the fixed point of a function:

E ← 1e¯6⊸<◶(1-÷⟜3)‿((π÷4)⊸÷×⟜Erf○√⊢)
P ← +˝∘×⎉1‿∞
D ← (⊢⋈·P´<∘⊢∾⋈)⟜([Cos‿Sin⋄Sin‿(-Cos)]{𝕎𝕩}¨·<2÷˜·ATan 2×0‿1⊸⊑÷·-´0‿0⊸⍉)
L ← (-⌾(1‿1⊸⊑)·÷∘√≢⥊2˙)⊸P(0¨⌾(0‿0⍉⌽˘)·÷∘√1+1‿¯1×⌽˘)
_fp ← {𝕊∘⊢⍟≢⟜𝔽𝕩}

Then, we define a function returning a namespace with the physical constants of the system³, as is customary in ab-initio methods. We have the flexibility of providing different molecular geometries:

System ← {𝕊𝕩:
  e1‿e2 ⇐ 2.0925‿1.24
  z1‿z2 ⇐ 2‿1
  r ⇐ 𝕩
}

Basis set

Basis sets are used to transform the PDEs into linear algebra problems. Physical intuition suggests that Slater type orbitals⁴ are a good choice for our Hamiltonian. However, the computation of the integrals is a lot easier if we approximate them using Gaussian functions⁵. Specifically, the STO-3G approach defines the approximating function as follows:

STO ← {
  e ← 0.109818‿0.405771‿2.22766 ××˜ 𝕩
  c ← 0.444635‿0.535328‿0.154329
  e⋈c×(2×e÷π)⋆3÷4
}

Electronic integrals

Constructing the integrals' tensor is complicated⁶ and is the main reason for the poor scaling of electronic structure methods. The \(1s\) orbitals are the simplest case, and here two types of integrals are analytical (S, T) while the rest already lacks a closed-form solution (V, ERI):

S ← {a‿b𝕊𝕩: (1.5⋆˜π÷a+b) × ⋆-𝕩×a(×÷+)b}
T ← {a‿b𝕊𝕩: f ← a(×÷+)b ⋄ ×´⟨1.5⋆˜π÷a+b, (3×f)-2×𝕩××˜f, ⋆-𝕩×f⟩}
V ← {a‿b‿z𝕊r‿s: ×´⟨-2×z×π÷a+b, E s×a+b, ⋆-r×a(×÷+)b⟩}
ERI ← {a‿b‿c‿d𝕊r1‿r2‿r3‿r4:
  r5 ← -´⟨a‿b ⋄ c‿d⟩ (+´∘×÷+´∘⊣)¨ ⟨r1‿r2 ⋄ r3‿r4⟩
  f1‿f2‿f3 ← a‿b ({(√∘+××)⋈(×÷+)}○(+´)∾<∘⋈○((×÷+)´)) c‿d
  ×´⟨f1÷˜2×π⋆5÷2, E f2××˜r5, ⋆-+´f3××˜-´¨⟨r1‿r2, r3‿r4⟩⟩
}

Derivation strategy

We need to compute the overlap (S), kinetic energy (T), nuclear attraction (V), and four-center (ERI) integrals. Crucially, the product of two Gaussians at different centers is proportional to a Gaussian at a scaled center. This property, combined with the Laplacian of a Gaussian, readily yields S and T. The remaining two sets are more complex: we combine the Gaussians as before, then transform to reciprocal space where the delta distribution arises and simplifies the problem to this integration by reduction:

\begin{equation*} I(x) = \int_0^{\infty}{{{e^ {- a\,k^2 }\,\sin \left(k\,x\right)}\over{k}}\;dk} \sim \text{Erf}(x) \end{equation*}

The following namespace exports the corresponding integral arrays. Extending the code to an arbitrary number of atoms implies mapping over an array of coordinates, as opposed to fusing them in the implementation.

I ← {𝕊e1‿e2‿z1‿z2‿r:
  bs‿na‿nb ← (<∾·≢⊏∘>)⍉>STO¨ e1‿e2
  M ← {∾‿×({2: {nb(⋈˜/∘⋈˜)⊸⊔𝕎⌜˜𝕩}; 4: {𝕎⌜⍟3˜𝕩}}𝕩)¨○⊢<∘∾˘bs}

  sm‿hcore ⇐ {e𝕊c:
    mst ← ⌽⊸≍∾⟜0×˜r
    r1‿r2 ← <˘⍉⁼> (r⊸-⊸⋈˜×⟜r÷+)⌜´ ⊏bs
    mv ← ×˜∘{[0‿2,3‿1]⊏({0‿𝕨¨𝕩}⟜𝕩¨𝕨)∾⋈⟜⍉r⋈¨𝕩}´¨⟨0‿r, r1⟩‿⟨r‿0, r2⟩
    (⊑⋈·+´1⊸↓)+´∘⥊¨¨ c<⊸× ({e𝕏¨¨mst}¨S‿T) ∾ z1‿z2{e∾⟜𝕨⊸V¨¨𝕩}¨mv
  }´ M 2

  erim ⇐ {e𝕊c:
    meri ← (c⊸×⊣ERI¨⊢/˜·<¨≢∘⊣÷≢∘⊢)⟜{0‿r⊏˜⚇1↕na¨↕=𝕩} e
    =⊸{+˝∘⥊⎉𝕨 (2×↕𝕨)⍉⁼(na‿nb⥊˜na×𝕨)⥊𝕩} meri
  }´ M 4
}

Performance

The computation of the ERIs is expected to be the primary bottleneck, as there are N⋆4 of them—in our case, 16. The required tensors have a shape of 6¨↕4. As shown in the profile below, using an array-based strategy for the ERIs significantly improved their computational efficiency compared to the two-center integrals. For the latter, I increased the depth by grouping the tables (block matrices). The resulting code was significantly slower than replicating the elements to match each axis' length, like I do for the ERIs.

)profile {𝕊: I∘System 1.4632}¨↕1e4

Got 38006 samples
(self-hosted runtime1): 1067 samples
(REPL): 36939 samples:
    72│I ← {𝕊e1‿e2‿z1‿z2‿r:
    68│  bs‿na‿nb ← (<∾·≢⊏∘>)⍉>STO¨ e1‿e2
  2053│  M ← {∾‿×({2: {nb(⋈˜/∘⋈˜)⊸⊔𝕎⌜˜𝕩}; 4: {𝕎⌜⍟3˜𝕩}}𝕩)¨○⊢<∘∾˘bs}
      │
   245│  sm‿hcore ⇐ {e𝕊c:
    75│    mst ← ⌽⊸≍∾⟜0×˜r
  4181│    r1‿r2 ← <˘⍉⁼> (r⊸-⊸⋈˜×⟜r÷+)⌜´ ⊏bs
 16277│    mv ← ×˜∘{[0‿2,3‿1]⊏({0‿𝕨¨𝕩}⟜𝕩¨𝕨)∾⋈⟜⍉r⋈¨𝕩}´¨⟨0‿r, r1⟩‿⟨r‿0, r2⟩
  8830│    (⊑⋈·+´1⊸↓)+´∘⥊¨¨ c<⊸× ({e𝕏¨¨mst}¨S‿T) ∾ z1‿z2{e∾⟜𝕨⊸V¨¨𝕩}¨mv
  3708│  }´ M 2
      │
     8│  erim ⇐ {e𝕊c:
  1100│    meri ← (c⊸×⊣ERI¨⊢/˜·<¨≢∘⊣÷≢∘⊢)⟜{0‿r⊏˜⚇1↕na¨↕=𝕩} e
   318│    =⊸{+˝∘⥊⎉𝕨 (2×↕𝕨)⍉⁼(na‿nb⥊˜na×𝕨)⥊𝕩} meri
     4│  }´ M 4
      │}

Morals: Never underestimate the power of vectorization and reshaping operations are often computationally trivial.

Fock matrix

The following function constructs the Fock matrix, our approximation to the true Hamiltonian of the system:

F ← {𝕩.hcore + 𝕨 (+˝∘⥊⎉2⊣×⎉2⊢-2÷˜0‿3‿2⊸⍉⁼) 𝕩.erim}

Physical context

The Fock operator is an effective one-electron operator that arises after constrained minimization of the energy functional. The form of the functional is a consequence of the use of Slater determinants as wave functions.

\begin{equation*} \tilde{\mathcal{F}} \left[ \{\psi_i\} \right] = \sum_i h_i + \frac{1}{2} \sum_{i,j} (J_{ij} - K_{ij}) - \sum_{i,j} \lambda_{ij} \left( \langle \psi_i | \psi_j \rangle - \delta_{ij} \right) \end{equation*}

where \(h_i\) is the core Hamiltonian matrix, \(J_{ij}, K_{ij}\) are the Coulomb and exchange components of the ERI matrix, and \(\lambda_{ij}\) are Lagrange multipliers. To fully understand the derivation, consider the variational derivative of this functional with respect to the complex conjugate of the one-particle wave function \(\psi_i^*\):

\begin{align*} \lim_{\epsilon \to 0} \frac{\tilde{\mathcal{F}} \left[ \psi_k^* + \epsilon \delta \psi_k^* \right] - \tilde{\mathcal{F}} \left[ \psi_k^* \right]}{\epsilon} &= \langle \delta \psi_k | \hat{h} | \psi_k \rangle + \sum_j \left( \langle \delta \psi_k \psi_j | \frac{1}{r} | \psi_k \psi_j \rangle - \langle \delta \psi_k \psi_j | \frac{1}{r} | \psi_j \psi_k \rangle \right) - \sum_j \lambda_{kj} \langle \delta \psi_k | \psi_j \rangle \\ &= \int \left[ \hat{h} \psi_k(x) + \sum_j \left( \psi_k(x) \int \frac{|\psi_j(x')|^2}{|r - r'|} dx' - \psi_j(x) \int \frac{\psi_j^*(x') \psi_k(x')}{|r - r'|} dx' \right) \right. \left. - \sum_j \lambda_{kj} \psi_j(x) \right] \delta \psi_k^*(x) \, dx. \end{align*}

As discussed earlier, basis sets are used to discretize the Hartree-Fock problem. This process results in the Roothaan equations, which are implemented in the code below.

Self-consistent field

The final stage of the computation involves solving the pseudo-eigenvalue problem using a fixed-point iteration. This process is commonly known as the self-consistent field method, a term coined by D. R. Hartree.

SCF ← {
  ints ← I𝕩
  pm ← {2××⌜˜⊏⍉ xm(⊢P⟜⊑·D·P´⍉∘⊢<⊸∾⋈)˜𝕩F ints}_fp 0¨xm ← L ints.sm
  2÷˜+´∘⥊ pm (⊣⍉⊸×{𝕩.hcore + 𝕨F𝕩}) ints
}

If you are receptive and humble, mathematics will lead you by the hand⁷:

SCF∘System 1.4632

¯4.22752930421725

Compare the electronic energy with the one computed using the original Fortran program. In terms of performance, the CBQN implementation runs in 5 ms, which is of the same order of magnitude as the original program. Notably, the BQN version consists of just 45 lines of code, compared to 541 lines in the Fortran version.

Potential Energy Surface

Here we compute the system's PES. To do this, we need to add to the electronic energy above the nuclear repulsion energy. We also catch the error of non-converged calculations, instead of fiddling with convergence thresholds and different starting points:

PES ← 2⊸÷+SCF⎊∞∘System

Then we leverage my modified version of the •Plot namespace:

)r Setplot "line" ⋄ •Out¨ Plot´ (⊢/¨˜·<∞>⊢´)(⊢⋈PES¨) ↕∘⌈⌾((0.5+1e¯2×⊢)⁼)3

Footnotes:

Recasting of the TISE into a set of coupled integro-differential equations. Derived by optimizing the expectation value of the energy subject to normalization constraints, then discretizing it using a suitable basis set.

It may not look like much, but helonium was the very first molecule formed in the universe.

This program can compute the Hartree-Fock energy of any two-electron diatomic molecule.

⁴

STO: functions of the form \(r^le^{-\zeta r}Y_l^m(\theta, \phi)\). For \(1s\) orbitals the spherical harmonics integrate out to 1.

⁵

STO-nG: a non-linear least-squares fit of an STO as a weighted sum of n Gaussians.

⁶

See for example arXiv:2007.12057.

⁷

Paul A.M. Dirac, 27 November, 1975

⊑∘∞