Differentiable quantum chemistry with PySCF for molecules and materials at the mean-field level and beyond

PySCF uses contracted Gaussian basis functions, the radial parts of which can be expressed as

where $l$ is the angular momentum, and the three parameters $\mathbf{r_{0}}$, $\alpha_{i}$ and $C_{i}$ label the center, exponents and contraction coefficients, respectively. Differentiation of a basis function with respect to these variables leads to another Gaussian function. As such, derivatives of an ERI can be computed by a sequence of similar integral evaluations.

In the current implementation of PySCFAD, the program walks through the computation graph and determines the required intermediate ERIs. For example, computing the second order derivative of the one-electron overlap integral with respect to the basis function centers requires the evaluation of three integrals

where in the above, the permutation symmetries have been considered. In the implementation, these integrals are all evaluated analytically by the highly optimized integral library Libcint.[19] Thus, the ERIs themselves are treated as elementary functions during the AD procedure, instead of the lower-level arithmetic operations that compute them.

For the commonly used ERIs, up to fourth order derivatives with respect to the basis function centers are available through Libcint. If higher order derivatives are needed, one can extend Libcint by generating the code to compute the required intermediate integrals before runtime, with the accompanying automatic code generator. Only first order derivatives with respect to the exponents and contraction coefficients have been implemented so far. This should be sufficient for many purposes, such as for basis function optimization.

The procedure above, in a strict sense, is not fully differentiable, because the ERI evaluation is black-box. However, a general, efficient and fully differentiable implementation for the ERIs remains challenging. The advent of compiler level AD tools[20] may facilitate such developments in the future.

A semi-local exchange correlation (XC) density functional can be expressed as

where $\varepsilon_{\rm xc}$ is the XC energy per particle, and $\rho$ and $\tau$ label the electron density and the non-interacting kinetic energy density, respectively. The integration in Eq. 2 is usually carried out numerically over a quadrature grid due to the complexity of XC functionals:

where $w_{i}$ is the weight for the $i$-th grid point at position $\mathbf{r}_{i}$. The AD of $E_{\rm xc}$ is straightforward if one uses a fully differentiable implementation of $\varepsilon_{\rm xc}$. However, we do not pursue that strategy here, since density derivatives to finite order (usually up to fourth order) are sufficient in most practical scenarios and these have already been implemented in efficient density functional libraries such as Libxc.[21] Therefore, we take an approach similar to that introduced above for the AD of ERIs, where the derivatives of $\varepsilon_{\rm xc}$ with respect to the density variables are computed analytically, while other derivatives (e.g., the derivatives of the density variables) are generated by AD. In particular, Libxc provides analytic functional derivatives of $E_{\rm xc}$, from which the derivatives of $\varepsilon_{\rm xc}$ can be easily expressed. For example, the first order derivative of $\varepsilon_{\rm xc}$ with respect to $\rho$ within the local density approximation (LDA) is obtained as

where both $\varepsilon_{\rm xc}$ and $v_{\rm xc}$ (the XC potential) are computed analytically by Libxc. The higher order derivatives are obtained by recursion, e.g.,

where $f_{\rm xc}$ is the XC kernel.

Although Jax provides a differentiable implementation of the eigendecomposition, it does not handle degenerate eigenstates. Usually, in order to determine the derivatives of degenerate eigenstates, one needs to diagonalize the perturbations of the matrix being decomposed within the degenerate subspace until the degeneracy is removed. Such an approach, however, is not well defined for reverse-mode AD, because the basis in which to expand the degenerate eigenstates is undetermined until the back propagation is carried out. The development of a general differentiable eigensolver is beyond the scope of the current work. In the following, we only discuss a workaround for the AD of eigenvalue problems arising in mean-field calculations.

At the mean-field level, a gauge invariant quantity will only respond to perturbations that mix occupied and unoccupied states. Thus in many cases, the response of degenerate single-particle eigenstates, i.e. orbitals (which can only be either occupied or unoccupied for gapped systems) does not contribute to the derivatives, and can be ignored. Our implementation directly follows the standard analytic formalism derived from linear response theory.[22] For a generalized eigenvalue problem

where $\bm{\varepsilon}$ and $\mathbf{C}$ denote the eigenvalue and eigenvector matrices, respectively, their differentials can be expressed as

and

respectively. In the two equations above, $\circ$ represents element-wise multiplication,

and

Similar approaches have also been reported in other works.[9, 8]

Automatic differentiation of optimization problems or iterative solvers is usually done in one of two ways, namely, unrolling the iterations[1, 23] or implicit differentiation.[2, 24, 25]

In the first approach, the entire set of iterations is differentiated, leading to a memory complexity that scales linearly with the number of iterations for reverse-mode AD. In addition, the so-obtained derivatives are usually initial guess dependent. This can be understood with the example of computing the molecular orbital (MO) response (Eq. 8). Suppose in an extreme case that the self-consistent field (SCF) iteration takes the converged MOs as the initial guess, then the SCF will converge after one Fock matrix diagonalization. Due to the fact the SCF iteration is short-circuited (the Fock matrix is not rebuilt) the input MOs have no knowledge of their own derivatives, and the Fock matrix, which responds to the MO changes, will have the wrong derivative after this single SCF iteration, which in turn leads to the wrong MO response. In other words, self-consistency is not reached when solving Eq. 8, where $\partial\mathbf{F}$ depends on $\partial\mathbf{C}$; the quality of convergence of the self-consistent response equations is controlled only by the threshold of the SCF itself. Such errors can be corrected by increasing the number of SCF iterations, so that the Fock matrix response with respect to the MO changes is gradually restored, and self-consistency of Eq. 8 is achieved. To see this effect, we computed the nuclear gradient of N₂ molecule at the MP2 level, where the derivatives of the MO coefficients are evaluated by the scheme of unrolling the SCF iterations. Using the same set of converged MOs as the initial guess, we vary the number of SCF iterations and plot the gradient error, shown as the red curve in Fig. 1. It is clear that the computed gradient is inaccurate unless a sufficient number of SCF iterations are carried out to recover the correct MO response. Nevertheless, the method is problematic when performing AD for optimizations or for iterative solvers, as one has no direct control over the convergence of the computed derivatives.

In implicit differentiation, instead of differentiating the solver iterations, the optimality condition is implicitly differentiated. For example, the solution (denoted as $x^{\star}$) of the iterative solver should also be the root of some optimality condition

According to the implicit function theorem,[26, 2] $x^{\star}$ can be seen as an implicit function of $\theta$, and its derivative can be obtained by solving a set of linear equations

In the problem of computing the MO response, Eq. 12 simply corresponds to the coupled perturbed Hartree-Fock (CPHF) equations.[22] More interestingly, in the reverse-mode AD, one does not directly solve Eq. 12; instead, the vector Jacobian product (i.e., $v^{\top}J$ where $v$ is a vector and $J\equiv\frac{\partial x^{\star}}{\partial\theta}$) is computed. If we define $A\equiv\frac{\partial f}{\partial x^{\star}}$ and $B\equiv-\frac{\partial f}{\partial\theta}$, then the vector Jacobian product is obtained as

where

In this case, only one set of linear equations (Eq. 14) needs to be solved even if derivatives are evaluated with respect to multiple variables (i.e., when $B$ changes but not $A$ and $v$). Note that Eqs. 13 and 14 correspond exactly to the $Z$-vector approach[27] in conventional analytic derivative methods.

The advantages of the implicit differentiation approach are obvious. First, because only the optimality condition is differentiated, the computational complexity and memory footprint do not depend on the actual implementation of the solver or on the number of solver iterations. A notable benefit of this is that algorithms to accelerate convergence can be readily applied, such as the direct inversion of the iterative subspace[28] (DIIS) method, without any modification to the AD implementation. (Note, however, that Eq. 14 itself needs to be solved iteratively, which can be more expensive than the approach of unrolling the iterations, depending on the size and convergence of the problem). Second, the computed derivatives have errors that are governed only by the accuracy of the solution of Eq. 14, which can be controlled with a predefined convergence threshold. This can be seen from the blue curve in Fig. 1, where the implicit differentiation approach is applied for the same MP2 nuclear gradient calculation as discussed above.

Given a general implementation,[25] implicit differentiation can be applied to almost any iterative solver. Currently, we have adapted it for the AD of SCF iterations and the coupled cluster (CC) amplitude equations. This eliminates the need to explicitly implement and solve the CPHF equations and the CC $\Lambda$ equations.[29] Finally, it is interesting to mention that if the optimality condition being differentiated involves the fixed point problem of gauge variant quantities (e.g., wavefunctions), it is important to fix the gauge. In particular, the fixed point of the optimality condition should be identical to the variable used to evaluate the objective function.

Orbital optimization in electronic structure methods provides the following advantages:

1. 1.

   Energies become variational with respect to orbital rotations, thus there is no need for orbital response when computing nuclear gradients.[30]

2. 2.

   Properties can be computed more easily because there are no orbital response contributions to the density matrices.

3. 3.

   The spurious poles in response functions for inexact methods such as coupled cluster theory can be removed.[31]

4. 4.

   Symmetry breaking problems may be described better.[32]

However, it is not always straightforward to implement analytic orbital gradients (and hessians for quadratically convergent optimization) if the underlying electronic structure method is complicated. In contrast, little effort is needed to obtain the orbital gradients and hessians using AD as long as the energy can be defined and implemented as differentiable with respect to orbital rotations.

As an example, in Fig. 2, we show the application of orbital optimization to the random phase approximation for the energy[33, 34, 35] (RPA) within the framework of PySCFAD. The base RPA method is implemented following Ren et al.,[36] where the correlation energy is evaluated by numerical integration over the imaginary frequency axis and density fitting is applied to reduce the computation cost. The total energy is then minimized with respect to the orbital rotation $e^{\mathbf{x}}$, where $\mathbf{x}$ is anti-hermitian and its upper triangular part corresponds to the variable “x” in Fig. 2. Note that only the energy function needs to be explicitly implemented (lines 24–31 in Fig. 2), whereas the orbital gradient and hessian (specifically, the hessian-vector product in the example) are obtained directly by AD (lines 33–36 in Fig. 2). The performance of the resulting orbital optimized (OO) RPA method is shown in Fig. 3. We plot the binding energy curves for the He₂ molecule computed at the RPA, OO-RPA and CCSD(T) levels. It is clear that the molecule is underbound at the RPA level, whereas applying orbital optimization corrects that and the binding energies are more consistent with the CCSD(T) results.

In general, ground-state dynamic (frequency-dependent) response properties are related to the derivatives of the quasienergy with respect to perturbations.[37] The quasienergy is defined as the time-averaged expectation value of $\hat{H}-i{\partial}/{\partial t}$ over the phase-isolated time-dependent wavefunction.[37] As such, although there is no difficulty applying AD to compute such derivatives, a straightforward application to the quasienergy in this form requires the time-dependent wavefunction itself. In the time-independent limit, however, the quasienergy reduces to the usual energy, thus the static response properties can be computed straightforwardly with AD from only time-independent quantities.

For example, the static Raman activity is related to the susceptibility[38, 39]

where $E$ is the ground-state energy, $\mathbf{R}$ denotes the nuclear coordinates, and $\bm{\varepsilon}$ represents the electric-field. In Fig. 4, we present an implementation of Raman activity at the CCSD level within the framework of PySCFAD. Again, only the energy function needs to be explicitly implemented (lines 14–24 in Fig. 4), while all the derivatives are obtained by AD (lines 27 and 32 in Fig. 4). In particular, the CPHF and CC $\Lambda$ equations are solved implicitly through the implicit differentiation procedure. Using this code, we compute the harmonic vibrational frequency, Raman activity and depolarization ratio for the BH molecule, and the results are displayed in Table 1. It can be seen that they agree very well with the reference analytic results, which validates our AD implementation.

Similarly, the IR intensity, which involves computing the nuclear derivative of the dipole moment, can be obtained in the same manner with AD. In Table 2, the IR intensities of the H₂O molecule at the CCSD level are displayed. Both AD and analytic evaluation give almost identical results.

Excitation energies can be identified from the poles of response functions, and can be obtained by solving the (generalized) eigenvalue problems which arise from the derivatives of quasienergy Lagrangians.[37] For example, in coupled cluster theory, the excitation energies can be computed as the eigenvalues of the CC Jacobian, which is defined as the second order derivative of the zeroth order Lagrangian with respect to the amplitudes and to the multipliers:

Note that $\frac{\partial L^{(0)}}{\partial\bm{\lambda}^{(0)}}$ gives the CC amplitude equations, which are already implemented in the ground state CC methods. Therefore, the CC Jacobian and subsequently the excitation energies can be obtained effortlessly if AD is applied to differentiate the amplitude equations. In Fig. 5, we show such an example of computing the excitation energies at the CCSD level, where only the CC amplitude equations are explicitly implemented (line 29 in Fig. 5), while the CC Jacobian is obtained via AD (line 33 Fig. 5) and is then diagonalized. The resulting excitation energies are identical to those from the equation-of-motion (EOM) method (see Table 3). Note that it is also possible to only compute the Jacobian vector product rather than the full Jacobian (similar to line 36 in Fig. 2), so that the diagonalization can be performed with Krylov subspace methods,[41] making larger calculations practical.

The strategy above is in principle applicable to any method, with the caveat that complications may arise depending on the particular ansatz used for representing the wavefunction. In general, the eigenvalue problem being solved has the form

where $\mathbf{E}^{[2]}$ is the second order derivative of the zeroth order energy or Lagrangian, and $\mathbf{S}^{[2]}$ may be related to the second order derivative of the wavefunction overlap, which may or may not be the identity.

The first order derivative coupling between two electronic states $\Psi_{I}$ and $\Psi_{J}$ is defined as[42, 43]

where $\mathbf{R}$ denotes the nuclear coordinates. Although derivative couplings are closely related to excited-state nuclear gradients, their analytic derivation and implementation may still be tedious and error-prone, depending on the complexity of the underlying electronic structure methods. However, such calculations can be greatly simplified by applying AD.

In Fig. 6, we give an example of computing the derivative coupling at the full configuration interaction[44] (FCI) level within PySCFAD. It should be noted that only a very simple function is explicitly implemented and then differentiated by AD, which is the wavefunction overlap $\langle\Psi_{I}(\mathbf{R}_{0})|\Psi_{J}(\mathbf{R})\rangle$. The $\mathbf{R}_{0}$ here emphasizes that $\Psi_{I}$ does not respond to the perturbation. In practice, all the variables corresponding to $\Psi_{I}$ are held constant when defining the objective function (e.g., “mol”, “fcivec” and “mf.mo_coeff” at lines 24–26 in Fig. 6).

The strategy above is applicable to any wavefunction method. In Table 4, we compare the derivative couplings computed at the configuration interaction singles[45] (CIS) and FCI levels. It is clear that the two methods give very different results. Nevertheless, our AD implementation produces exactly the same results as the reference implementation. In addition, second order derivative couplings can be readily obtained by performing AD on the same wavefunction overlap another time. Finally, with AD, it is even more straightforward to compute the Hellmann-Feynman part of the derivative coupling (which is translationally invariant,[46, 47] unlike the full derivative coupling)

where $\mathbf{H}$ is the Hamiltonian represented in a certain Hilbert space and $\mathbf{C}$ denotes the eigenvectors (i.e., CI vectors) of $\mathbf{H}$. Here, only the Hamiltonian needs to be differentiated (with the orbital response taken into account), but the CI vectors are treated as constants. In other words, it is not necessary to differentiate the CI eigensolver (e.g., the Davidson iterations[48]), which greatly simplifies the problem.

The PySCFAD framework also works seamlessly for solid calculations. Here, we give an example of computing the stress tensor to illustrate this.

The stress tensor $\bm{\sigma}$ is defined as the first order energy response to the infinitesimal strain deformation[50]

where $E$ and $V$ are the energy and volume of a unit cell, respectively, and the strain tensor $\bm{\varepsilon}$ defines the transformation of real space coordinates $\mathbf{R}$ according to

in which $\alpha$ and $\beta$ denote the Cartesian components. The strain deformation applies to both nuclear coordinates and lattice vectors, which also implicitly affects the derived quantities such as the reciprocal lattice vectors and the unit cell volume.

In Fig. 7, the stress tensor for a two-atom Si primitive cell is computed by AD at the Hartree-Fock level, where a mixed Gaussian and plane wave approach (plane-wave density fitting) is used.[51, 52] The results are plotted in Fig. 8 along with the corresponding finite difference values. Perfect agreement is obtained with an average discrepancy of $4\times 10^{-8}$ eV/ų between the AD and finite difference results. It should be noted that in this case, the plane wave basis (i.e., the uniform grid point) response to the strain deformation is non-negligible, which subsequently requires differentiation with respect to the grid points. However, all these complications are hidden from the user, and it suffices to modify the energy function (as shown in Fig. 7) for the gradient calculations.