contracted integral timing

timing_integral_reshape.ipy

# -*- mode: python -*-

from __future__ import print_function

import numpy as np


Nprim = 6
Ncntr = 2

D2 = np.random.random((Nprim, Ncntr))
Sprim = np.random.random((Nprim, Nprim))
Vprim = np.random.random((Nprim, Nprim, Nprim, Nprim))

D4 = np.tensordot(D2, D2, 0)

# 1-electron integral
S = np.einsum('ba,bc,cd', D2, Sprim, D2)

# 2-electron integral Option 1: form (Nprim * Ncntr * Nprim * Ncntr) tensor
def option1(D2, Vprim):
    D4 = np.tensordot(D2, D2, 0)
    V = np.einsum('badc,bedf,egfh->agch', D4, Vprim, D4)
    return V
V1 = option1(D2, Vprim)
%timeit option1(D2, Vprim)

# 2-electron integral Option 2: transform indices one-by-one
def option2(D2, Vprim):
    V = Vprim.copy()
    V = np.einsum('ba,bcde->acde', D2, V)    # i
    V = np.einsum('abcd,be->aecd', V, D2)    # j
    V = np.einsum('da,bcde->bcae', D2, V)    # k
    V = np.einsum('abcd,de->abce', V, D2)    # l
    return V
V2 = option2(D2, Vprim)
%timeit option2(D2, Vprim)

# 2-electron integral Option 3: reshape to 2-index matrices
def option3(D2, Vprim):
    Nprim, Ncntr = D2.shape
    Vprim = Vprim.reshape((Nprim**2, Nprim**2))
    D4 = np.tensordot(D2, D2, 0)
    # making sure to combine the correct axes
    D4 = np.swapaxes(D4, 1, 2)
    D4 = D4.reshape((Nprim**2, Ncntr**2))
    V = np.einsum('ba,bc,cd', D4, Vprim, D4)
    return V
V3 = option3(D2, Vprim)
%timeit option3(D2, Vprim)

print(V1.shape)
print(V2.shape)
print(V3.shape)

print(V1.reshape(Ncntr**2, Ncntr**2) / V3)
print(V2.reshape(Ncntr**2, Ncntr**2) / V3)
print(V3)

timing_integral_reshape.out

1000 loops, best of 3: 667 us per loop
10000 loops, best of 3: 33.8 us per loop
10000 loops, best of 3: 146 us per loop
(2, 2, 2, 2)
(2, 2, 2, 2)
(4, 4)
[[ 1.  1.  1.  1.]
 [ 1.  1.  1.  1.]
 [ 1.  1.  1.  1.]
 [ 1.  1.  1.  1.]]
[[ 1.  1.  1.  1.]
 [ 1.  1.  1.  1.]
 [ 1.  1.  1.  1.]
 [ 1.  1.  1.  1.]]
[[ 52.08672425  56.09980772  56.53323056  60.80065119]
 [ 55.70203934  60.3196274   60.48286959  65.42496997]
 [ 55.92963961  60.54072257  60.7119886   65.59192265]
 [ 59.64160155  64.85766581  64.76847957  70.34496875]]