Skip to content

Commit 50298a6

Browse files
edeustuaedeustua
committed
feat: Decouple einsum from np einsum
1 parent 62552ec commit 50298a6

File tree

125 files changed

+13843
-13713
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

125 files changed

+13843
-13713
lines changed

‎ccpy/cc/accd.py‎

Lines changed: 66 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
'''
44

55
import numpy as np
6+
from ccpy.utilities.linear_algebra import ccpy_einsum
67
# Modules for type checking
78
from typing import List, Tuple
89
from ccpy.models.operators import ClusterOperator
@@ -89,32 +90,32 @@ def update_t2a(T: ClusterOperator,
8990
d1, d2, d3, d4, d5 = acparray
9091

9192
# < ijab | (F T2)_C | 0 >
92-
dT.aa = -0.5 * np.einsum("mi,abmj->abij", H.a.oo, T.aa, optimize=True) # A(ij)
93-
dT.aa += 0.5 * np.einsum("ae,ebij->abij", H.a.vv, T.aa, optimize=True) # A(ab)
93+
dT.aa = -0.5 * ccpy_einsum("mi,abmj->abij", H.a.oo, T.aa) # A(ij)
94+
dT.aa += 0.5 * ccpy_einsum("ae,ebij->abij", H.a.vv, T.aa) # A(ab)
9495

9596
# < ijab | (V T2)_C | 0 >
96-
dT.aa += np.einsum("amie,ebmj->abij", H.aa.voov, T.aa, optimize=True) # A(ab)A(ij)
97-
dT.aa += np.einsum("amie,bejm->abij", H.ab.voov, T.ab, optimize=True) # A(ab)A(ij)
98-
dT.aa += 0.125 * np.einsum("mnij,abmn->abij", H.aa.oooo, T.aa, optimize=True) # 1
99-
dT.aa += 0.125 * np.einsum("abef,efij->abij", H.aa.vvvv, T.aa, optimize=True) # 1
97+
dT.aa += ccpy_einsum("amie,ebmj->abij", H.aa.voov, T.aa) # A(ab)A(ij)
98+
dT.aa += ccpy_einsum("amie,bejm->abij", H.ab.voov, T.ab) # A(ab)A(ij)
99+
dT.aa += 0.125 * ccpy_einsum("mnij,abmn->abij", H.aa.oooo, T.aa) # 1
100+
dT.aa += 0.125 * ccpy_einsum("abef,efij->abij", H.aa.vvvv, T.aa) # 1
100101

101102
# < ijab | (V T2**2)_C | 0 >
102103
# dT.aa += 0.5 * np.einsum(
103104
# "mnef,aeim,bfjn->abij", H.aa.oovv, T.aa, T.aa, optimize=True
104105
# ) # A(ij) [D1 + D2]
105-
dT.aa += d1 * 0.5 * np.einsum("mnef,aeim,bfjn->abij", H.ab.oovv, T.aa, T.aa, optimize=True) # A(ij) [D1]
106-
dT.aa -= d2 * 0.5 * np.einsum("mnfe,aeim,bfjn->abij", H.ab.oovv, T.aa, T.aa, optimize=True) # A(ij) [D2]
107-
dT.aa += d5 * 0.25 * 0.25 * np.einsum("mnef,efij,abmn->abij", H.aa.oovv, T.aa, T.aa, optimize=True) # 1 [D5]
108-
dT.aa -= d4 * 0.25 * np.einsum("mnef,abim,efjn->abij", H.aa.oovv, T.aa, T.aa, optimize=True) # A(ij) [D4]
109-
dT.aa -= d3 * 0.25 * np.einsum("mnef,aeij,bfmn->abij", H.aa.oovv, T.aa, T.aa, optimize=True) # A(ab) [D3]
110-
dT.aa += d1 * np.einsum("mnef,aeim,bfjn->abij", H.ab.oovv, T.aa, T.ab, optimize=True) # A(ij)A(ab) [D1]
111-
dT.aa -= d4 * 0.5 * np.einsum("mnef,abim,efjn->abij", H.ab.oovv, T.aa, T.ab, optimize=True) # A(ij) [D4]
112-
dT.aa -= d3 * 0.5 * np.einsum("mnef,aeij,bfmn->abij", H.ab.oovv, T.aa, T.ab, optimize=True) # A(ab) [D3]
106+
dT.aa += d1 * 0.5 * ccpy_einsum("mnef,aeim,bfjn->abij", H.ab.oovv, T.aa, T.aa) # A(ij) [D1]
107+
dT.aa -= d2 * 0.5 * ccpy_einsum("mnfe,aeim,bfjn->abij", H.ab.oovv, T.aa, T.aa) # A(ij) [D2]
108+
dT.aa += d5 * 0.25 * 0.25 * ccpy_einsum("mnef,efij,abmn->abij", H.aa.oovv, T.aa, T.aa) # 1 [D5]
109+
dT.aa -= d4 * 0.25 * ccpy_einsum("mnef,abim,efjn->abij", H.aa.oovv, T.aa, T.aa) # A(ij) [D4]
110+
dT.aa -= d3 * 0.25 * ccpy_einsum("mnef,aeij,bfmn->abij", H.aa.oovv, T.aa, T.aa) # A(ab) [D3]
111+
dT.aa += d1 * ccpy_einsum("mnef,aeim,bfjn->abij", H.ab.oovv, T.aa, T.ab) # A(ij)A(ab) [D1]
112+
dT.aa -= d4 * 0.5 * ccpy_einsum("mnef,abim,efjn->abij", H.ab.oovv, T.aa, T.ab) # A(ij) [D4]
113+
dT.aa -= d3 * 0.5 * ccpy_einsum("mnef,aeij,bfmn->abij", H.ab.oovv, T.aa, T.ab) # A(ab) [D3]
113114
# dT.aa += 0.5 * np.einsum(
114115
# "mnef,aeim,bfjn->abij", H.bb.oovv, T.ab, T.ab, optimize=True
115116
# ) # A(ij) [D1 + D2]
116-
dT.aa += d1 * 0.5 * np.einsum("mnef,aeim,bfjn->abij", H.ab.oovv, T.ab, T.ab, optimize=True) # A(ij) [D1]
117-
dT.aa -= d2 * 0.5 * np.einsum("mnfe,aeim,bfjn->abij", H.ab.oovv, T.ab, T.ab, optimize=True) # A(ij) [D2]
117+
dT.aa += d1 * 0.5 * ccpy_einsum("mnef,aeim,bfjn->abij", H.ab.oovv, T.ab, T.ab) # A(ij) [D1]
118+
dT.aa -= d2 * 0.5 * ccpy_einsum("mnfe,aeim,bfjn->abij", H.ab.oovv, T.ab, T.ab) # A(ij) [D2]
118119

119120
T.aa, dT.aa = cc_loops2.update_t2a(
120121
T.aa, dT.aa + 0.25 * H.aa.vvoo, H.a.oo, H.a.vv, shift
@@ -152,43 +153,43 @@ def update_t2b(T: ClusterOperator,
152153
d1, d2, d3, d4, d5 = acparray
153154

154155
# < ijab | (F T2)_C | 0 >
155-
dT.ab = -np.einsum("mi,abmj->abij", H.a.oo, T.ab, optimize=True)
156-
dT.ab += np.einsum("ae,ebij->abij", H.a.vv, T.ab, optimize=True)
157-
dT.ab -= np.einsum("mj,abim->abij", H.b.oo, T.ab, optimize=True)
158-
dT.ab += np.einsum("be,aeij->abij", H.b.vv, T.ab, optimize=True)
156+
dT.ab = -ccpy_einsum("mi,abmj->abij", H.a.oo, T.ab)
157+
dT.ab += ccpy_einsum("ae,ebij->abij", H.a.vv, T.ab)
158+
dT.ab -= ccpy_einsum("mj,abim->abij", H.b.oo, T.ab)
159+
dT.ab += ccpy_einsum("be,aeij->abij", H.b.vv, T.ab)
159160

160161
# < ijab | (V T2)_C | 0 >
161-
dT.ab += np.einsum("amie,ebmj->abij", H.aa.voov, T.ab, optimize=True)
162-
dT.ab += np.einsum("amie,ebmj->abij", H.ab.voov, T.bb, optimize=True)
163-
dT.ab += np.einsum("mbej,aeim->abij", H.ab.ovvo, T.aa, optimize=True)
164-
dT.ab += np.einsum("bmje,aeim->abij", H.bb.voov, T.ab, optimize=True)
165-
dT.ab -= np.einsum("mbie,aemj->abij", H.ab.ovov, T.ab, optimize=True)
166-
dT.ab -= np.einsum("amej,ebim->abij", H.ab.vovo, T.ab, optimize=True)
167-
dT.ab += np.einsum("mnij,abmn->abij", H.ab.oooo, T.ab, optimize=True)
168-
dT.ab += np.einsum("abef,efij->abij", H.ab.vvvv, T.ab, optimize=True)
162+
dT.ab += ccpy_einsum("amie,ebmj->abij", H.aa.voov, T.ab)
163+
dT.ab += ccpy_einsum("amie,ebmj->abij", H.ab.voov, T.bb)
164+
dT.ab += ccpy_einsum("mbej,aeim->abij", H.ab.ovvo, T.aa)
165+
dT.ab += ccpy_einsum("bmje,aeim->abij", H.bb.voov, T.ab)
166+
dT.ab -= ccpy_einsum("mbie,aemj->abij", H.ab.ovov, T.ab)
167+
dT.ab -= ccpy_einsum("amej,ebim->abij", H.ab.vovo, T.ab)
168+
dT.ab += ccpy_einsum("mnij,abmn->abij", H.ab.oooo, T.ab)
169+
dT.ab += ccpy_einsum("abef,efij->abij", H.ab.vvvv, T.ab)
169170

170171
# < ijab | (V T2**2)_C | 0 >
171-
# dT.ab += np.einsum("mnef,aeim,fbnj->abij", H.aa.oovv, T.aa, T.ab, optimize=True) # [D1 + D2]
172-
dT.ab += d1 * np.einsum("mnef,aeim,fbnj->abij", H.ab.oovv, T.aa, T.ab, optimize=True) # [D1]
173-
dT.ab -= d2 * np.einsum("mnfe,aeim,fbnj->abij", H.ab.oovv, T.aa, T.ab, optimize=True) # [D2]
174-
dT.ab -= d4 * 0.5 * np.einsum("mnef,efin,abmj->abij", H.aa.oovv, T.aa, T.ab, optimize=True) # [D4]
175-
dT.ab -= d3 * 0.5 * np.einsum("mnef,afmn,ebij->abij", H.aa.oovv, T.aa, T.ab, optimize=True) # [D3]
176-
177-
dT.ab += d1 * np.einsum("nmfe,aeim,fbnj->abij", H.ab.oovv, T.ab, T.ab, optimize=True) # [D1]
178-
dT.ab += d2 * np.einsum("mnef,ebin,afmj->abij", H.ab.oovv, T.ab, T.ab, optimize=True) # [D2]
179-
dT.ab += d4 * np.einsum("mnef,efij,abmn->abij", H.ab.oovv, T.ab, T.ab, optimize=True) # [D4]
180-
dT.ab -= d4 * np.einsum("mnef,efin,abmj->abij", H.ab.oovv, T.ab, T.ab, optimize=True) # [D4]
181-
dT.ab -= d5 * np.einsum("nmfe,fenj,abim->abij", H.ab.oovv, T.ab, T.ab, optimize=True) # [D5]
182-
dT.ab -= d3 * np.einsum("mnef,afmn,ebij->abij", H.ab.oovv, T.ab, T.ab, optimize=True) # [D3]
183-
dT.ab -= d3 * np.einsum("nmfe,fbnm,aeij->abij", H.ab.oovv, T.ab, T.ab, optimize=True) # [D3]
184-
185-
dT.ab += d1 * np.einsum("mnef,aeim,fbnj->abij", H.ab.oovv, T.aa, T.bb, optimize=True) # [D1]
186-
187-
# dT.ab += np.einsum("mnef,aeim,fbnj->abij", H.bb.oovv, T.ab, T.bb, optimize=True) # [D1 + D2]
188-
dT.ab += d1 * np.einsum("mnef,aeim,fbnj->abij", H.ab.oovv, T.ab, T.bb, optimize=True) # [D1]
189-
dT.ab -= d2 * np.einsum("mnfe,aeim,fbnj->abij", H.ab.oovv, T.ab, T.bb, optimize=True) # [D2]
190-
dT.ab -= d4 * 0.5 * np.einsum("mnef,efjn,abim->abij", H.bb.oovv, T.bb, T.ab, optimize=True) # [D4]
191-
dT.ab -= d3 * 0.5 * np.einsum("mnef,bfmn,aeij->abij", H.bb.oovv, T.bb, T.ab, optimize=True) # [D3]
172+
# dT.ab += ccpy_einsum("mnef,aeim,fbnj->abij", H.aa.oovv, T.aa, T.ab) # [D1 + D2]
173+
dT.ab += d1 * ccpy_einsum("mnef,aeim,fbnj->abij", H.ab.oovv, T.aa, T.ab) # [D1]
174+
dT.ab -= d2 * ccpy_einsum("mnfe,aeim,fbnj->abij", H.ab.oovv, T.aa, T.ab) # [D2]
175+
dT.ab -= d4 * 0.5 * ccpy_einsum("mnef,efin,abmj->abij", H.aa.oovv, T.aa, T.ab) # [D4]
176+
dT.ab -= d3 * 0.5 * ccpy_einsum("mnef,afmn,ebij->abij", H.aa.oovv, T.aa, T.ab) # [D3]
177+
178+
dT.ab += d1 * ccpy_einsum("nmfe,aeim,fbnj->abij", H.ab.oovv, T.ab, T.ab) # [D1]
179+
dT.ab += d2 * ccpy_einsum("mnef,ebin,afmj->abij", H.ab.oovv, T.ab, T.ab) # [D2]
180+
dT.ab += d4 * ccpy_einsum("mnef,efij,abmn->abij", H.ab.oovv, T.ab, T.ab) # [D4]
181+
dT.ab -= d4 * ccpy_einsum("mnef,efin,abmj->abij", H.ab.oovv, T.ab, T.ab) # [D4]
182+
dT.ab -= d5 * ccpy_einsum("nmfe,fenj,abim->abij", H.ab.oovv, T.ab, T.ab) # [D5]
183+
dT.ab -= d3 * ccpy_einsum("mnef,afmn,ebij->abij", H.ab.oovv, T.ab, T.ab) # [D3]
184+
dT.ab -= d3 * ccpy_einsum("nmfe,fbnm,aeij->abij", H.ab.oovv, T.ab, T.ab) # [D3]
185+
186+
dT.ab += d1 * ccpy_einsum("mnef,aeim,fbnj->abij", H.ab.oovv, T.aa, T.bb) # [D1]
187+
188+
# dT.ab += ccpy_einsum("mnef,aeim,fbnj->abij", H.bb.oovv, T.ab, T.bb) # [D1 + D2]
189+
dT.ab += d1 * ccpy_einsum("mnef,aeim,fbnj->abij", H.ab.oovv, T.ab, T.bb) # [D1]
190+
dT.ab -= d2 * ccpy_einsum("mnfe,aeim,fbnj->abij", H.ab.oovv, T.ab, T.bb) # [D2]
191+
dT.ab -= d4 * 0.5 * ccpy_einsum("mnef,efjn,abim->abij", H.bb.oovv, T.bb, T.ab) # [D4]
192+
dT.ab -= d3 * 0.5 * ccpy_einsum("mnef,bfmn,aeij->abij", H.bb.oovv, T.bb, T.ab) # [D3]
192193

193194
T.ab, dT.ab = cc_loops2.update_t2b(
194195
T.ab, dT.ab + H.ab.vvoo, H.a.oo, H.a.vv, H.b.oo, H.b.vv, shift
@@ -226,32 +227,32 @@ def update_t2c(T: ClusterOperator,
226227
d1, d2, d3, d4, d5 = acparray
227228

228229
# < ijab | (F T2)_C | 0 >
229-
dT.bb = -0.5 * np.einsum("mi,abmj->abij", H.b.oo, T.bb, optimize=True) # A(ij)
230-
dT.bb += 0.5 * np.einsum("ae,ebij->abij", H.b.vv, T.bb, optimize=True) # A(ab)
230+
dT.bb = -0.5 * ccpy_einsum("mi,abmj->abij", H.b.oo, T.bb) # A(ij)
231+
dT.bb += 0.5 * ccpy_einsum("ae,ebij->abij", H.b.vv, T.bb) # A(ab)
231232

232233
# < ijab | (V T2)_C | 0 >
233-
dT.bb += np.einsum("amie,ebmj->abij", H.bb.voov, T.bb, optimize=True) # A(ab)A(ij)
234-
dT.bb += np.einsum("maei,ebmj->abij", H.ab.ovvo, T.ab, optimize=True) # A(ab)A(ij)
235-
dT.bb += 0.125 * np.einsum("mnij,abmn->abij", H.bb.oooo, T.bb, optimize=True) # 1
236-
dT.bb += 0.125 * np.einsum("abef,efij->abij", H.bb.vvvv, T.bb, optimize=True) # 1
234+
dT.bb += ccpy_einsum("amie,ebmj->abij", H.bb.voov, T.bb) # A(ab)A(ij)
235+
dT.bb += ccpy_einsum("maei,ebmj->abij", H.ab.ovvo, T.ab) # A(ab)A(ij)
236+
dT.bb += 0.125 * ccpy_einsum("mnij,abmn->abij", H.bb.oooo, T.bb) # 1
237+
dT.bb += 0.125 * ccpy_einsum("abef,efij->abij", H.bb.vvvv, T.bb) # 1
237238

238239
# < ijab | (V T2**2)_C | 0 >
239240
# dT.bb += 0.5 * np.einsum(
240241
# "mnef,aeim,bfjn->abij", H.bb.oovv, T.bb, T.bb, optimize=True
241242
# ) # A(ij) [D1 + D2]
242-
dT.bb += d1 * 0.5 * np.einsum("nmfe,aeim,bfjn->abij", H.ab.oovv, T.bb, T.bb, optimize=True) # A(ij) [D1]
243-
dT.bb -= d2 * 0.5 * np.einsum("nmef,aeim,bfjn->abij", H.ab.oovv, T.bb, T.bb, optimize=True) # A(ij) [D2]
244-
dT.bb += d5 * 0.25 * 0.25 * np.einsum("mnef,efij,abmn->abij", H.bb.oovv, T.bb, T.bb, optimize=True) # 1 [D5]
245-
dT.bb -= d4 * 0.25 * np.einsum("mnef,abim,efjn->abij", H.bb.oovv, T.bb, T.bb, optimize=True) # A(ij) [D4]
246-
dT.bb -= d3 * 0.25 * np.einsum("mnef,aeij,bfmn->abij", H.bb.oovv, T.bb, T.bb, optimize=True) # A(ab) [D3]
247-
dT.bb += d1 * np.einsum("nmfe,aeim,fbnj->abij", H.ab.oovv, T.bb, T.ab, optimize=True) # A(ij)A(ab) [D1]
248-
dT.bb -= d4 * 0.5 * np.einsum("nmfe,abim,fenj->abij", H.ab.oovv, T.bb, T.ab, optimize=True) # A(ij) [D4]
249-
dT.bb -= d3 * 0.5 * np.einsum("nmfe,aeij,fbnm->abij", H.ab.oovv, T.bb, T.ab, optimize=True) # A(ab) [D3]
243+
dT.bb += d1 * 0.5 * ccpy_einsum("nmfe,aeim,bfjn->abij", H.ab.oovv, T.bb, T.bb) # A(ij) [D1]
244+
dT.bb -= d2 * 0.5 * ccpy_einsum("nmef,aeim,bfjn->abij", H.ab.oovv, T.bb, T.bb) # A(ij) [D2]
245+
dT.bb += d5 * 0.25 * 0.25 * ccpy_einsum("mnef,efij,abmn->abij", H.bb.oovv, T.bb, T.bb) # 1 [D5]
246+
dT.bb -= d4 * 0.25 * ccpy_einsum("mnef,abim,efjn->abij", H.bb.oovv, T.bb, T.bb) # A(ij) [D4]
247+
dT.bb -= d3 * 0.25 * ccpy_einsum("mnef,aeij,bfmn->abij", H.bb.oovv, T.bb, T.bb) # A(ab) [D3]
248+
dT.bb += d1 * ccpy_einsum("nmfe,aeim,fbnj->abij", H.ab.oovv, T.bb, T.ab) # A(ij)A(ab) [D1]
249+
dT.bb -= d4 * 0.5 * ccpy_einsum("nmfe,abim,fenj->abij", H.ab.oovv, T.bb, T.ab) # A(ij) [D4]
250+
dT.bb -= d3 * 0.5 * ccpy_einsum("nmfe,aeij,fbnm->abij", H.ab.oovv, T.bb, T.ab) # A(ab) [D3]
250251
# dT.bb += 0.5 * np.einsum(
251252
# "mnef,aeim,bfjn->abij", H.bb.oovv, T.ab, T.ab, optimize=True
252253
# ) # A(ij) [D1 + D2]
253-
dT.bb += d1 * 0.5 * np.einsum("nmfe,eami,fbnj->abij", H.ab.oovv, T.ab, T.ab, optimize=True) # A(ij) [D1]
254-
dT.bb -= d2 * 0.5 * np.einsum("nmef,eami,fbnj->abij", H.ab.oovv, T.ab, T.ab, optimize=True) # A(ij) [D2]
254+
dT.bb += d1 * 0.5 * ccpy_einsum("nmfe,eami,fbnj->abij", H.ab.oovv, T.ab, T.ab) # A(ij) [D1]
255+
dT.bb -= d2 * 0.5 * ccpy_einsum("nmef,eami,fbnj->abij", H.ab.oovv, T.ab, T.ab) # A(ij) [D2]
255256

256257
T.bb, dT.bb = cc_loops2.update_t2c(
257258
T.bb, dT.bb + 0.25 * H.bb.vvoo, H.b.oo, H.b.vv, shift

0 commit comments

Comments
 (0)