<commit_before>import numpy as np
import psi4

mol = psi4.geometry("""
O
H 1 1.1
H 1 1.1 2 105
""")

# Construct a Psi4 basis set
basis = psi4.core.BasisSet.build(mol, target="STO-3G")
nbf = basis.nbf()

# Compute integrals and convert the resulting matrices to numpy arrays
mints = psi4.core.MintsHelper(basis)
T = np.array(mints.ao_kinetic())
V = np.array(mints.ao_potential())
I = np.array(mints.ao_eri())

# One-electron hamiltonian is the sum of the kinetic and potential elements
H = T + V

# Build a symmetric random density matrix
D = np.random.rand(nbf, nbf)
D += D.T

# Build the fock matrix using loops
Floop = np.zeros((nbf, nbf))
for p in range(nbf):
    for q in range(nbf):
        Floop[p, q] = H[p, q]
        for r in range(nbf):
            for s in range(nbf):
                Floop[p, q] += 2 * I[p, q, r, s] * D[r, s]
                Floop[p, q] -=     I[p, r, q, s] * D[r, s]

# Build the fock matrix using einsum
J = np.einsum('pqrs,rs', I, D) 
K = np.einsum('prqs,rs', I, D)
F = H + 2 * J - K

# Make sure the correct answer is obtained
print('The loop and einsum fock builds match:    %s' % np.allclose(F, Floop))<commit_msg>Added timing feature to illustrate relative speed of Einsum vs. pure Python loops.<commit_after>import numpy as np
import psi4
import time

mol = psi4.geometry("""
O
H 1 1.1
H 1 1.1 2 105
""")

# Construct a Psi4 basis set
basis = psi4.core.BasisSet.build(mol, target="STO-3G")
nbf = basis.nbf()

# Compute integrals and convert the resulting matrices to numpy arrays
mints = psi4.core.MintsHelper(basis)
T = np.array(mints.ao_kinetic())
V = np.array(mints.ao_potential())
I = np.array(mints.ao_eri())

# One-electron hamiltonian is the sum of the kinetic and potential elements
H = T + V

# Build a symmetric random density matrix
D = np.random.rand(nbf, nbf)
D += D.T

# Build the fock matrix using loops, while keeping track of time
t_floop_0 = time.time()
Floop = np.zeros((nbf, nbf))
for p in range(nbf):
    for q in range(nbf):
        Floop[p, q] = H[p, q]
        for r in range(nbf):
            for s in range(nbf):
                Floop[p, q] += 2 * I[p, q, r, s] * D[r, s]
                Floop[p, q] -=     I[p, r, q, s] * D[r, s]

t_floop = time.time() - t_floop_0

# Build the fock matrix using einsum, while keeping track of time
t_einsum_0 = time.time()
J = np.einsum('pqrs,rs', I, D) 
K = np.einsum('prqs,rs', I, D)
F = H + 2 * J - K

t_einsum = time.time() - t_einsum_0

# Make sure the correct answer is obtained
print('The loop and einsum fock builds match:    %s\n' % np.allclose(F, Floop))
# Print out relative times for explicit loop vs einsum Fock builds
print('Time for loop Fock build:\t {:2.4f} seconds'.format(t_floop))
print('Time for einsum Fock build:\t {:2.4f} seconds'.format(t_einsum))
print('Fock builds with einsum are {:3.4f} times faster than Python loops!'.format(t_floop / t_einsum))
{'fraction_non_alphanumeric': 0.13525730180806675, 'fraction_numerical': 0.015994436717663423, 'mean_word_length': 3.2308823529411765, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5187225', 'n_tokens_mistral': 1167, 'n_tokens_neox': 1107, 'n_words': 374}