1import importlib
2import sys
3import time
4
5import os
6
7os.environ["JAX_PLATFORMS"] = "cpu"
8os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
9from jax import config
10
11config.update("jax_enable_x64", True)
12
13if importlib.util.find_spec("pyscf") is None:
14 print("pyscf is required to run this example, please install it")
15 sys.exit(0)
16
17import pyscf
18import numpy as np
19import autohf
20
21# helper functions
22def format_fixed_width_strings(strings):
23 """Format a list of strings to be fixed width."""
24 return ' '.join('{:>17}'.format(s) for s in strings)
25
26def format_fixed_width_floats(floats):
27 """Format a list of floats to be fixed width."""
28 return ' '.join('{: .10e}'.format(f) for f in floats)
29
30def modified_cholesky_direct(M, tol=1e-5, verbose=False, cmax=20):
31 """Modified cholesky decomposition of matrix.
32
33 See, e.g. [Motta17]_
34
35 Parameters
36 ----------
37 M : :class:`np.ndarray`
38 Positive semi-definite, symmetric matrix.
39 tol : float
40 Accuracy desired. Optional. Default : False.
41 verbose : bool
42 If true print out convergence progress. Optional. Default : False.
43 cmax : int
44 Number of cholesky vectors to store N_gamma = cmax M.
45
46 Returns
47 -------
48 chol_vecs : :class:`np.ndarray`
49 Matrix of cholesky vectors.
50 """
51 # matrix of residuals.
52 delta = np.copy(M.diagonal())
53 nchol_max = min(int(cmax*M.shape[0]**0.5),M.shape[1])
54 # index of largest diagonal element of residual matrix.
55 nu = np.argmax(np.abs(delta))
56 delta_max = delta[nu]
57 if verbose:
58 print(f"Performing Cholesky decomposition with tolerance {tol}")
59 print ("# max number of cholesky vectors = %d"%nchol_max)
60 header = ['iteration', 'max_residual', 'time']
61 print(format_fixed_width_strings(header))
62 init = [delta_max]
63 print('{:17d} '.format(0)+format_fixed_width_floats(init))
64 # Store for current approximation to input matrix.
65 Mapprox = np.zeros(M.shape[0], dtype=M.dtype)
66 chol_vecs = np.zeros((nchol_max, M.shape[0]), dtype=M.dtype)
67 nchol = 0
68 chol_vecs[0] = np.copy(M[:,nu])/delta_max**0.5
69 while abs(delta_max) > tol and nchol<nchol_max-1: # -1 because we already have one vector.
70 # Update cholesky vector
71 start = time.time()
72 Mapprox += chol_vecs[nchol]*chol_vecs[nchol].conj()
73 delta = M.diagonal() - Mapprox
74 nu = np.argmax(np.abs(delta))
75 delta_max = np.abs(delta[nu])
76 nchol += 1
77 Munu0 = np.dot(chol_vecs[:nchol,nu].conj(), chol_vecs[:nchol,:])
78 chol_vecs[nchol] = (M[:,nu] - Munu0) / (delta_max)**0.5
79 if verbose:
80 step_time = time.time() - start
81 out = [delta_max, step_time]
82 print('{:17d} '.format(nchol)+format_fixed_width_floats(out))
83
84 return np.array(chol_vecs[:nchol])
85
86def get_ortho_ao_mol(S, LINDEP_CUTOFF=0):
87 """Generate canonical orthogonalization transformation matrix.
88 Parameters
89 ----------
90 S : :class:`numpy.ndarray`
91 Overlap matrix.
92 LINDEP_CUTOFF : float
93 Linear dependency cutoff. Basis functions whose eigenvalues lie below
94 this value are removed from the basis set. Should be set in accordance
95 with value in :func:`pyscf.scf.addons.remove_linear_dep`.
96
97 Returns
98 -------
99 X : :class:`numpy.array`
100 Transformation matrix.
101 """
102 sdiag, Us = np.linalg.eigh(S)
103 X = Us[:,sdiag>LINDEP_CUTOFF] / np.sqrt(sdiag[sdiag>LINDEP_CUTOFF])
104 return X
105
106## End helper functions
107
108
109mol = pyscf.M(
110 atom="""
111 N 0 0 0.5488
112 N 0 0 -0.5488
113""",
114 basis="sto-3g",
115)
116
117onebody = mol.intor("int1e_kin") + mol.intor("int1e_nuc")
118overlap = mol.intor("int1e_ovlp")
119eri = mol.intor("int2e")
120nuclear_repulsion = mol.energy_nuc()
121
122X = get_ortho_ao_mol(overlap)
123
124norb = onebody.shape[0]
125print(f"{norb=}")
126
127chol_tol = 1e-8
128cholesky_operators = modified_cholesky_direct(
129 eri.reshape(norb**2, norb**2), tol=chol_tol, verbose=False
130).reshape(-1, norb, norb)
131
132onebody = X.T @ onebody @ X
133cholesky_operators = 1j * np.einsum("ij,njk,kl->nil", X.T, cholesky_operators, X)
134
135
136onebody += 0.5 * np.einsum("nij,njk", cholesky_operators, cholesky_operators).real
137
138# double along spin dimension
139cholesky_operators = np.stack([cholesky_operators, cholesky_operators], axis=1)
140onebody = np.stack([onebody, onebody], axis=0)
141
142
143mf = mol.UHF.run()
144
145
146T = np.zeros((norb, norb))
147
148H = autohf.AutoHFHamiltonian(T=(T, T), U=0)
149nelec = np.sum(mf.get_occ(), axis=1).astype(int)
150
151hf_settings = {
152 "verbose": False,
153 "steps": 2000,
154 "opt_method": "lbfgs",
155 "ansatz": "SD_ROT",
156 "batch_size": 4,
157 "gpu": False,
158 "nelec": nelec,
159 "noncollinear": False,
160 "state0_scale": 0.01,
161 "seed": 1234,
162 "measure_spin": True,
163}
164
165
166def molecular_energy(state, rdms):
167 return (
168 autohf.terms.cholesky.hf_onebody(rdms, onebody)
169 + autohf.terms.cholesky.hf_cholesky(rdms, cholesky_operators)
170 + nuclear_repulsion
171 )
172
173
174data, data_functions = autohf.solve_hf(
175 H, settings=hf_settings, custom_energy_term=molecular_energy
176)
177rounding_level = -np.log10(chol_tol).astype(int)
178# recall we have a cholesky cutoff, so we will be slightly off in evaluating the energy
179# practically this isn't a huge issue
180print("PySCF energy:", np.round(mf.energy_tot(), rounding_level))
181print("AutoHF :", np.round(data["E"], rounding_level))