-
Notifications
You must be signed in to change notification settings - Fork 373
/
hartree_fock.py
425 lines (358 loc) · 16.5 KB
/
hartree_fock.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
"""
Module performs gradient based RHF, [WIP] UHF, [WIP] GHF
Module needs AO integrals
"""
# pylint: disable=C
from typing import Callable, Dict, List, Optional, Tuple, Union
from itertools import product
import numpy as np
import scipy as sp
from scipy.optimize.optimize import OptimizeResult
from openfermion.ops.representations import (InteractionOperator,
InteractionRDM,
general_basis_change)
import openfermion.linalg as linalg
def get_matrix_of_eigs(w: np.ndarray) -> np.ndarray:
"""
Transform the eigenvalues into a matrix corresponding
to summing the adjoint rep.
Args:
w: eigenvalues of C-matrix
Returns: new array of transformed eigenvalues
"""
transform_eigs = np.zeros((w.shape[0], w.shape[0]), dtype=np.complex128)
for i, j in product(range(w.shape[0]), repeat=2):
if np.isclose(abs(w[i] - w[j]), 0):
transform_eigs[i, j] = 1
else:
transform_eigs[i, j] = (np.exp(1j *
(w[i] - w[j])) - 1) / (1j *
(w[i] - w[j]))
return transform_eigs
class InputError(Exception):
pass
class HartreeFockFunctional():
"""
Implementation of the objective function code for Restricted Hartree-Fock
The object transforms a variety of input types into the appropriate output.
It does this by analyzing the type and size of the input based on its
knowledge of each type.
"""
def __init__(self,
*,
one_body_integrals: np.ndarray,
two_body_integrals: np.ndarray,
overlap: np.ndarray,
n_electrons: int,
model='rhf',
nuclear_repulsion: Optional[float] = 0.,
initial_orbitals: Optional[Union[None, Callable]] = None):
"""
Initialize functional
Args:
one_body_integrals: integrals in the atomic orbital basis for the
one-body potential.
two_body_integrals: integrals in the atomic obrital basis for the
two-body potential ordered according to
phi_{p}(r1)^{*}phi_{q}^{*}(r2) x
phi_{r}(r2)phi_{s}(r1)
overlap: overlap integrals in the atomic orbital basis
n_electrons: number of electrons total
model: Optional flag for performing restricted-, unrestricted-,
or generalized- hartree-fock.
nuclear_repulsion: Optional nuclear repulsion term. Energy is
shifted by this amount. default is 0.
initial_orbitals: Method for producing the initial orbitals from
the atomic orbitals. Default is defining the
core orbitals.
"""
if model not in ['rhf', 'uhf', 'ghf']:
raise InputError("{} is not rhf, uhf, or ghf".format(model))
self.model = model
self.obi = one_body_integrals
self.tbi = two_body_integrals
self.overlap = overlap
self.num_orbitals = one_body_integrals.shape[0]
self.num_electrons = n_electrons
self.constant_offset = nuclear_repulsion
self.hamiltonian = None
self.nocc = None
self.nvirt = None
self.occ = None
self.virt = None
if model == 'rhf':
self.nocc = self.num_electrons // 2
self.nvirt = self.num_orbitals - self.nocc
self.occ = list(range(self.nocc))
self.virt = list(range(self.nocc, self.nocc + self.nvirt))
elif model == 'uhf' or model == 'ghf':
self.nocc = self.num_electrons
self.nvirt = 2 * self.num_orbitals - self.nocc
self.occ = list(range(self.nocc))
self.virt = list(range(self.nocc, self.nocc + self.nvirt))
if initial_orbitals is None:
# use core orbitals
_, core_orbs = sp.linalg.eigh(one_body_integrals, b=overlap)
molecular_hamiltonian = generate_hamiltonian(
one_body_integrals=general_basis_change(self.obi, core_orbs,
(1, 0)),
two_body_integrals=general_basis_change(self.tbi, core_orbs,
(1, 1, 0, 0)),
constant=self.constant_offset)
self.hamiltonian = molecular_hamiltonian
else:
self.hamiltonian = initial_orbitals(self.obi, self.tbi,
self.num_electrons)
def rdms_from_rhf_opdm(self, opdm_aa: np.ndarray) -> InteractionRDM:
"""
Generate spin-orbital InteractionRDM object from the alpha-spin
opdm.
Args:
opdm_aa: single spin sector of the 1-particle denstiy matrix
Returns: InteractionRDM object for full spin-orbital 1-RDM and 2-RDM
"""
opdm = np.zeros((2 * self.num_orbitals, 2 * self.num_orbitals),
dtype=np.complex128)
opdm[::2, ::2] = opdm_aa
opdm[1::2, 1::2] = opdm_aa
tpdm = linalg.wedge(opdm, opdm, (1, 1), (1, 1))
rdms = InteractionRDM(opdm, 2 * tpdm)
return rdms
def energy_from_rhf_opdm(self, opdm_aa: np.ndarray) -> float:
"""
Compute the energy given a spin-up opdm
Args:
opdm_aa: spin-up opdm. Should be an n x n matrix where n is
the number of spatial orbitals
Returns: RHF energy
"""
rdms = self.rdms_from_rhf_opdm(opdm_aa)
return rdms.expectation(self.hamiltonian).real
def rhf_global_gradient(self, params: np.ndarray, alpha_opdm: np.ndarray):
"""
Compute rhf global gradient
Args:
params: rhf-parameters for rotation matrix.
alpha_opdm: 1-RDM corresponding to results of basis rotation
parameterized by `params'.
Returns: gradient vector the same size as the input `params'
"""
opdm = np.zeros((2 * self.num_orbitals, 2 * self.num_orbitals),
dtype=np.complex128)
opdm[::2, ::2] = alpha_opdm
opdm[1::2, 1::2] = alpha_opdm
tpdm = 2 * linalg.wedge(opdm, opdm, (1, 1), (1, 1))
# now go through and generate all the necessary Z, Y, Y_kl matrices
kappa_matrix = rhf_params_to_matrix(params,
len(self.occ) + len(self.virt),
self.occ, self.virt)
kappa_matrix_full = np.kron(kappa_matrix, np.eye(2))
w_full, v_full = np.linalg.eigh(
-1j * kappa_matrix_full) # so that kappa = i U lambda U^
eigs_scaled_full = get_matrix_of_eigs(w_full)
grad = np.zeros(self.nocc * self.nvirt, dtype=np.complex128)
# kdelta = np.eye(2 * self.num_orbitals)
# NOW GENERATE ALL TERMS ASSOCIATED WITH THE GRADIENT!!!!!!
for p in range(self.nocc * self.nvirt):
grad_params = np.zeros_like(params)
grad_params[p] = 1
Y = rhf_params_to_matrix(grad_params,
len(self.occ) + len(self.virt), self.occ,
self.virt)
Y_full = np.kron(Y, np.eye(2))
# Now rotate Y int othe basis that diagonalizes Z
Y_kl_full = v_full.conj().T.dot(Y_full).dot(v_full)
# now rotate Y_{kl} * (exp(i(l_{k} - l_{l})) - 1) / (i(l_{k} - l_{l}))
# into the original basis
pre_matrix_full = v_full.dot(eigs_scaled_full * Y_kl_full).dot(
v_full.conj().T)
grad_expectation = -1.0 * np.einsum(
'ab,pa,pb',
self.hamiltonian.one_body_tensor,
pre_matrix_full,
opdm,
optimize='optimal').real
grad_expectation += 1.0 * np.einsum(
'ab,bq,aq',
self.hamiltonian.one_body_tensor,
pre_matrix_full,
opdm,
optimize='optimal').real
grad_expectation += 1.0 * np.einsum(
'ijkl,pi,jpkl',
self.hamiltonian.two_body_tensor,
pre_matrix_full,
tpdm,
optimize='optimal').real
grad_expectation += -1.0 * np.einsum(
'ijkl,pj,ipkl',
self.hamiltonian.two_body_tensor,
pre_matrix_full,
tpdm,
optimize='optimal').real
grad_expectation += -1.0 * np.einsum(
'ijkl,kq,ijlq',
self.hamiltonian.two_body_tensor,
pre_matrix_full,
tpdm,
optimize='optimal').real
grad_expectation += 1.0 * np.einsum(
'ijkl,lq,ijkq',
self.hamiltonian.two_body_tensor,
pre_matrix_full,
tpdm,
optimize='optimal').real
grad[p] = grad_expectation
return grad
def generate_hamiltonian(one_body_integrals: np.ndarray,
two_body_integrals: np.ndarray,
constant: float,
EQ_TOLERANCE: Optional[float] = 1.0E-12
) -> InteractionOperator:
n_qubits = 2 * one_body_integrals.shape[0]
# Initialize Hamiltonian coefficients.
one_body_coefficients = np.zeros((n_qubits, n_qubits))
two_body_coefficients = np.zeros((n_qubits, n_qubits, n_qubits, n_qubits))
# Loop through integrals.
for p in range(n_qubits // 2):
for q in range(n_qubits // 2):
# Populate 1-body coefficients. Require p and q have same spin.
one_body_coefficients[2 * p, 2 * q] = one_body_integrals[p, q]
one_body_coefficients[2 * p + 1, 2 * q +
1] = one_body_integrals[p, q]
# Continue looping to prepare 2-body coefficients.
for r in range(n_qubits // 2):
for s in range(n_qubits // 2):
# Mixed spin
two_body_coefficients[2 * p, 2 * q + 1, 2 * r + 1, 2 *
s] = (two_body_integrals[p, q, r, s] /
2.)
two_body_coefficients[2 * p + 1, 2 * q, 2 * r, 2 * s +
1] = (two_body_integrals[p, q, r, s] /
2.)
# Same spin
two_body_coefficients[2 * p, 2 * q, 2 * r, 2 *
s] = (two_body_integrals[p, q, r, s] /
2.)
two_body_coefficients[2 * p + 1, 2 * q + 1, 2 * r +
1, 2 * s +
1] = (two_body_integrals[p, q, r, s] /
2.)
# Truncate.
one_body_coefficients[
np.absolute(one_body_coefficients) < EQ_TOLERANCE] = 0.
two_body_coefficients[
np.absolute(two_body_coefficients) < EQ_TOLERANCE] = 0.
# Cast to InteractionOperator class and return.
molecular_hamiltonian = InteractionOperator(constant, one_body_coefficients,
two_body_coefficients)
return molecular_hamiltonian
def rhf_params_to_matrix(parameters: np.ndarray,
num_orbitals: int,
occ: Optional[Union[None, List[int]]] = None,
virt: Optional[Union[None, List[int]]] = None
) -> np.ndarray:
"""
For restricted Hartree-Fock we have nocc * nvirt parameters. These are
provided as a list that is ordered by (virtuals) \times (occupied).
For example, for H4 we have 2 orbitals occupied and 2 virtuals
occupied = [0, 1] virtuals = [2, 3]
parameters = [(v_{0}, o_{0}), (v_{0}, o_{1}), (v_{1}, o_{0}), (v_{1}, o_{1})]
= [(2, 0), (2, 1), (3, 0), (3, 1)]
You can think of the tuples of elements of the upper right triangle of the
antihermitian matrix that specifies the c_{b, i} coefficients.
coefficient matrix
[[ c_{0, 0}, -c_{1, 0}, -c_{2, 0}, -c_{3, 0}],
[ c_{1, 0}, c_{1, 1}, -c_{2, 1}, -c_{3, 1}],
[ c_{2, 0}, c_{2, 1}, c_{2, 2}, -c_{3, 2}],
[ c_{3, 0}, c_{3, 1}, c_{3, 2}, c_{3, 3}]]
Since we are working with only non-redundant operators we know c_{i, i} = 0
and any c_{i, j} where i and j are both in occupied or both in virtual = 0.
Args:
parameters: array of parameters for kappa matrix
num_orbitals: total number of spatial orbitals
occ: (Optional) indices for doubly occupied sector
virt: (Optional) indices for virtual sector
Returns: np.ndarray kappa matrix
"""
if occ is None:
occ = range(num_orbitals // 2)
if virt is None:
virt = range(num_orbitals // 2, num_orbitals)
# check that parameters are a real array
if not np.allclose(parameters.imag, 0):
raise ValueError("parameters input must be real valued")
kappa = np.zeros((len(occ) + len(virt), len(occ) + len(virt)))
for idx, (v, o) in enumerate(product(virt, occ)):
kappa[v, o] = parameters[idx].real
kappa[o, v] = -parameters[idx].real
return kappa
def rhf_func_generator(rhf_func: HartreeFockFunctional,
init_occ_vec: Optional[Union[None, np.ndarray]] = None,
get_opdm_func: Optional[bool] = False
) -> Union[Tuple[Callable, Callable, Callable],
Tuple[Callable, Callable, Callable, Callable]]:
"""
Generate the energy, gradient, and unitary functions
Args:
rhf_func: objective function object.
init_occ_vec: (optional) vector for occupation numbers of
the alpha-opdm.
get_opdm_func: (optional) flag for returning Callable that returns
the final opdm.
Returns: functions for unitary, energy, gradient (in that order)
"""
if init_occ_vec is None:
initial_opdm = np.diag([1] * rhf_func.nocc + [0] * rhf_func.nvirt)
else:
initial_opdm = np.diag(init_occ_vec)
def energy(params):
u = unitary(params)
final_opdm_aa = u.dot(initial_opdm).dot(np.conjugate(u).T)
tenergy = rhf_func.energy_from_rhf_opdm(final_opdm_aa)
return tenergy
def gradient(params):
u = unitary(params)
final_opdm_aa = u.dot(initial_opdm).dot(np.conjugate(u).T)
return rhf_func.rhf_global_gradient(params, final_opdm_aa).real
def unitary(params):
kappa = rhf_params_to_matrix(params, rhf_func.nocc + rhf_func.nvirt,
rhf_func.occ, rhf_func.virt)
return sp.linalg.expm(kappa)
def get_opdm(params):
u = unitary(params)
return u.dot(initial_opdm).dot(np.conjugate(u).T)
if get_opdm_func:
return unitary, energy, gradient, get_opdm
return unitary, energy, gradient
def rhf_minimization(rhf_object: HartreeFockFunctional,
method: Optional[str] = 'CG',
initial_guess: Optional[Union[None, np.ndarray]] = None,
verbose: Optional[bool] = True,
sp_options: Optional[Union[None, Dict]] = None
) -> OptimizeResult:
"""
Perform Hartree-Fock energy minimization
Args:
rhf_object: An instantiation of the HartreeFockFunctional
method: (optional) scipy optimization method
initial_guess: (optional) initial rhf parameter vector. If None
zero vector is used.
verbose: (optional) turn on printing. This is passed to the
scipy 'disp' option.
sp_options:
Returns: scipy.optimize result object
"""
_, energy, gradient = rhf_func_generator(rhf_object)
if initial_guess is None:
init_guess = np.zeros(rhf_object.nocc * rhf_object.nvirt)
else:
init_guess = np.asarray(initial_guess).flatten()
sp_optimizer_options = {'disp': verbose}
if sp_options is not None:
sp_optimizer_options.update(sp_options)
return sp.optimize.minimize(energy,
init_guess,
jac=gradient,
method=method,
options=sp_optimizer_options)