diff --git a/ebcc/codegen/GCCSDwtwp.py b/ebcc/codegen/GCCSDwtwp.py index ffe055ca..8e0ef679 100644 --- a/ebcc/codegen/GCCSDwtwp.py +++ b/ebcc/codegen/GCCSDwtwp.py @@ -2,14 +2,14 @@ Code generated by `albert`: https://github.com/obackhouse/albert - * date: 2024-09-27T21:44:58.220797 - * python version: 3.10.12 (main, Sep 11 2024, 15:47:36) [GCC 11.4.0] + * date: 2024-09-29T17:18:28.820515 + * python version: 3.10.15 (main, Sep 9 2024, 03:02:45) [GCC 11.4.0] * albert version: 0.0.0 - * caller: /home/ollie/git/albert/albert/codegen/einsum.py - * node: ollie-desktop + * caller: /opt/hostedtoolcache/Python/3.10.15/x64/lib/python3.10/site-packages/albert/codegen/einsum.py + * node: fv-az1788-690 * system: Linux * processor: x86_64 - * release: 6.8.0-40-generic + * release: 6.8.0-1014-azure """ from ebcc import numpy as np @@ -19,7 +19,7 @@ def energy(f=None, t1=None, t2=None, v=None, **kwargs): """ - Code generated by `albert` 0.0.0 on 2024-09-27T21:44:58.384057. + Code generated by `albert` 0.0.0 on 2024-09-29T17:18:29.071491. Parameters ---------- @@ -39,7 +39,7 @@ def energy(f=None, t1=None, t2=None, v=None, **kwargs): """ tmp0 = np.copy(t2) - tmp0 += einsum(t1, (0, 1), t1, (2, 3), (0, 2, 1, 3)) * 2 + tmp0 += einsum(t1, (0, 1), t1, (2, 3), (2, 0, 3, 1)) * 2 e_cc = einsum(f.ov, (0, 1), t1, (0, 1), ()) e_cc += einsum(tmp0, (0, 1, 2, 3), v.oovv, (0, 1, 2, 3), ()) * 0.25 del tmp0 @@ -48,7 +48,7 @@ def energy(f=None, t1=None, t2=None, v=None, **kwargs): def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): """ - Code generated by `albert` 0.0.0 on 2024-09-27T21:45:33.101020. + Code generated by `albert` 0.0.0 on 2024-09-29T17:19:16.670741. Parameters ---------- @@ -81,235 +81,235 @@ def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): sv = np.ones((t1.shape[1],), dtype=bool) sO = space.active[space.correlated][space.occupied[space.correlated]] sV = space.active[space.correlated][space.virtual[space.correlated]] - t2_ooVv = np.copy(t2[np.ix_(so, so, sV, sv)]) - t1_oV = np.copy(t1[np.ix_(so, sV)]) - t2_ooVV = np.copy(t2[np.ix_(so, so, sV, sV)]) - t2_OOVv = np.copy(t2[np.ix_(sO, sO, sV, sv)]) - t1_Ov = np.copy(t1[np.ix_(sO, sv)]) + t2_OOvv = np.copy(t2[np.ix_(sO, sO, sv, sv)]) + t2_OoVv = np.copy(t2[np.ix_(sO, so, sV, sv)]) t1_ov = np.copy(t1[np.ix_(so, sv)]) + t1_Ov = np.copy(t1[np.ix_(sO, sv)]) t2_oovv = np.copy(t2[np.ix_(so, so, sv, sv)]) - t2_Oovv = np.copy(t2[np.ix_(sO, so, sv, sv)]) t2_OoVV = np.copy(t2[np.ix_(sO, so, sV, sV)]) - t2_OOvv = np.copy(t2[np.ix_(sO, sO, sv, sv)]) - t2_OoVv = np.copy(t2[np.ix_(sO, so, sV, sv)]) + t2_ooVv = np.copy(t2[np.ix_(so, so, sV, sv)]) + t2_Oovv = np.copy(t2[np.ix_(sO, so, sv, sv)]) + t1_oV = np.copy(t1[np.ix_(so, sV)]) + t2_OOVv = np.copy(t2[np.ix_(sO, sO, sV, sv)]) + t2_ooVV = np.copy(t2[np.ix_(so, so, sV, sV)]) - tmp110 = einsum(v.oovv, (0, 1, 2, 3), t1_Ov, (4, 3), (0, 1, 2, 4)) - tmp146 = einsum(t1_Ov, (0, 1), tmp110, (2, 3, 1, 4), (3, 2, 4, 0)) - tmp143 = einsum(t1_oV, (0, 1), v.ooVV, (2, 0, 3, 4), (2, 1, 3, 4)) - tmp137 = einsum(v.OOvv, (0, 1, 2, 3), t1_Ov, (4, 3), (2, 4, 0, 1)) - tmp134 = einsum(v.ooVv, (0, 1, 2, 3), t1, (1, 3), (0, 2)) * -1 - tmp126 = einsum(v.oVvv, (0, 1, 2, 3), t1_Ov, (4, 3), (0, 2, 4, 1)) - tmp149 = einsum(t1_oV, (0, 1), tmp110, (2, 0, 3, 4), (2, 3, 4, 1)) + tmp110 = einsum(t1_Ov, (0, 1), v.oovv, (2, 3, 4, 1), (2, 3, 4, 0)) + tmp140 = einsum(v.OoVv, (0, 1, 2, 3), t1_Ov, (4, 3), (1, 4, 0, 2)) tmp115 = einsum(t1_oV, (0, 1), v.oovv, (2, 0, 3, 4), (2, 3, 4, 1)) + tmp137 = einsum(t1_Ov, (0, 1), v.OOvv, (2, 3, 4, 1), (4, 0, 2, 3)) + tmp123 = einsum(t1_oV, (0, 1), v.Ovoo, (2, 3, 4, 0), (4, 3, 2, 1)) * -1 + tmp126 = einsum(v.Vovv, (0, 1, 2, 3), t1_Ov, (4, 3), (1, 2, 4, 0)) * -1 + tmp146 = einsum(t1_Ov, (0, 1), tmp110, (2, 3, 1, 4), (3, 2, 4, 0)) + tmp1 = einsum(t1, (0, 1), v.oovv, (2, 0, 3, 1), (2, 3)) + tmp143 = einsum(v.VVoo, (0, 1, 2, 3), t1_oV, (3, 4), (2, 4, 0, 1)) + tmp120 = einsum(v.Ovoo, (0, 1, 2, 3), t1_Ov, (4, 1), (2, 3, 4, 0)) * -1 tmp131 = einsum(v.Oovv, (0, 1, 2, 3), t1, (1, 3), (2, 0)) * -1 - tmp140 = einsum(t1_Ov, (0, 1), v.OoVv, (2, 3, 4, 1), (3, 0, 2, 4)) - tmp1 = einsum(v.oovv, (0, 1, 2, 3), t1, (1, 3), (0, 2)) - tmp120 = einsum(t1_Ov, (0, 1), v.Ovoo, (2, 1, 3, 4), (3, 4, 0, 2)) * -1 - tmp123 = einsum(v.Ovoo, (0, 1, 2, 3), t1_oV, (3, 4), (2, 1, 0, 4)) * -1 + tmp149 = einsum(t1_oV, (0, 1), tmp110, (2, 0, 3, 4), (2, 3, 4, 1)) + tmp134 = einsum(t1, (0, 1), v.Vvoo, (2, 1, 3, 0), (3, 2)) * -1 tmp2 = einsum(v.oovv, (0, 1, 2, 3), t1, (4, 3), (4, 0, 1, 2)) - tmp23 = einsum(t1, (0, 1), v.ovvv, (2, 3, 4, 1), (0, 2, 3, 4)) - tmp18 = einsum(t1, (0, 1), v.ooov, (2, 3, 4, 1), (0, 2, 3, 4)) + tmp18 = einsum(v.ooov, (0, 1, 2, 3), t1, (4, 3), (4, 0, 1, 2)) + tmp23 = einsum(v.ovvv, (0, 1, 2, 3), t1, (4, 3), (4, 0, 1, 2)) + tmp104 = einsum(v.Vovv, (0, 1, 2, 3), t2_OOvv, (4, 5, 2, 3), (1, 4, 5, 0)) * -1 + tmp87 = einsum(v.Oovv, (0, 1, 2, 3), t2_Oovv, (4, 1, 2, 3), (4, 0)) + tmp100 = einsum(t2_OoVv, (0, 1, 2, 3), v.Ovoo, (4, 3, 5, 1), (5, 0, 4, 2)) * -1 + tmp106 = einsum(v.Vovv, (0, 1, 2, 3), t2_OoVv, (4, 1, 5, 3), (2, 4, 5, 0)) * -1 + tmp141 = einsum(tmp140, (0, 1, 2, 3), t1_oV, (0, 4), (1, 2, 4, 3)) + del tmp140 + tmp118 = einsum(tmp115, (0, 1, 2, 3), t2_OoVv, (4, 0, 5, 1), (2, 4, 3, 5)) + tmp138 = einsum(t1_Ov, (0, 1), tmp137, (1, 2, 3, 4), (2, 0, 4, 3)) + del tmp137 + tmp102 = einsum(t2_ooVV, (0, 1, 2, 3), v.Ovoo, (4, 5, 0, 1), (5, 4, 2, 3)) * -1 + tmp56 = einsum(f.Ov, (0, 1), t1_Ov, (2, 1), (0, 2)) + tmp97 = einsum(t2_ooVV, (0, 1, 2, 3), v.VVoo, (4, 5, 0, 1), (2, 3, 4, 5)) + tmp124 = einsum(tmp123, (0, 1, 2, 3), t1_oV, (0, 4), (1, 2, 3, 4)) * -1 + del tmp123 + tmp129 = einsum(tmp126, (0, 1, 2, 3), t1_oV, (0, 4), (1, 2, 4, 3)) tmp147 = einsum(t1_oV, (0, 1), tmp146, (0, 2, 3, 4), (2, 3, 4, 1)) * -1 del tmp146 + tmp108 = einsum(t2_OOVv, (0, 1, 2, 3), tmp1, (4, 3), (4, 0, 1, 2)) + tmp116 = einsum(tmp115, (0, 1, 2, 3), t2_OOvv, (4, 5, 2, 1), (0, 4, 5, 3)) * -1 + del tmp115 + tmp85 = einsum(v.VoVV, (0, 1, 2, 3), t1_oV, (1, 4), (4, 0, 2, 3)) * -1 tmp111 = einsum(t2_OoVv, (0, 1, 2, 3), tmp110, (1, 4, 3, 5), (4, 5, 0, 2)) - tmp93 = einsum(t2_OoVv, (0, 1, 2, 3), v.OoVv, (4, 1, 5, 3), (0, 4, 2, 5)) - tmp87 = einsum(t2_Oovv, (0, 1, 2, 3), v.Oovv, (4, 1, 2, 3), (0, 4)) - tmp144 = einsum(tmp143, (0, 1, 2, 3), t1_oV, (0, 4), (1, 4, 3, 2)) + tmp91 = einsum(t2_ooVv, (0, 1, 2, 3), v.Vvoo, (4, 3, 0, 1), (2, 4)) + tmp144 = einsum(t1_oV, (0, 1), tmp143, (0, 2, 3, 4), (2, 1, 4, 3)) del tmp143 - tmp138 = einsum(t1_Ov, (0, 1), tmp137, (1, 2, 3, 4), (2, 0, 4, 3)) - del tmp137 - tmp77 = einsum(t1_Ov, (0, 1), v.OvOO, (2, 1, 3, 4), (0, 2, 3, 4)) * -1 - tmp69 = einsum(t1_Ov, (0, 1), v.OvoV, (2, 1, 3, 4), (3, 0, 2, 4)) * -1 - tmp79 = einsum(t1_oV, (0, 1), v.OoOV, (2, 0, 3, 4), (2, 3, 1, 4)) * -1 - tmp95 = einsum(t2_OOVv, (0, 1, 2, 3), v.OOVv, (4, 5, 6, 3), (0, 1, 4, 5, 2, 6)) - tmp75 = einsum(v.OoOv, (0, 1, 2, 3), t1, (1, 3), (0, 2)) - tmp135 = einsum(tmp134, (0, 1), t1_oV, (0, 2), (2, 1)) * -1 - del tmp134 - tmp127 = einsum(t1_Ov, (0, 1), tmp126, (2, 1, 3, 4), (2, 3, 0, 4)) * -1 - tmp100 = einsum(v.Ovoo, (0, 1, 2, 3), t2_OoVv, (4, 3, 5, 1), (2, 4, 0, 5)) * -1 - tmp150 = einsum(t1_oV, (0, 1), tmp149, (0, 2, 3, 4), (2, 3, 4, 1)) * -1 - del tmp149 - tmp116 = einsum(t2_OOvv, (0, 1, 2, 3), tmp115, (4, 3, 2, 5), (4, 0, 1, 5)) * -1 - tmp71 = einsum(t1_oV, (0, 1), v.OvoV, (2, 3, 0, 4), (3, 2, 1, 4)) * -1 - tmp60 = einsum(f.ov, (0, 1), t2_OOVv, (2, 3, 4, 1), (0, 2, 3, 4)) * -1 - tmp118 = einsum(t2_OoVv, (0, 1, 2, 3), tmp115, (1, 3, 4, 5), (4, 0, 5, 2)) - del tmp115 - tmp104 = einsum(v.oVvv, (0, 1, 2, 3), t2_OOvv, (4, 5, 2, 3), (0, 4, 5, 1)) - tmp129 = einsum(t1_oV, (0, 1), tmp126, (0, 2, 3, 4), (2, 3, 1, 4)) - del tmp126 - tmp113 = einsum(tmp110, (0, 1, 2, 3), t2_ooVV, (1, 0, 4, 5), (2, 3, 4, 5)) * -1 + tmp113 = einsum(t2_ooVV, (0, 1, 2, 3), tmp110, (1, 0, 4, 5), (4, 5, 2, 3)) * -1 del tmp110 - tmp132 = einsum(t1_Ov, (0, 1), tmp131, (1, 2), (0, 2)) * -1 - del tmp131 - tmp89 = einsum(t2_OOvv, (0, 1, 2, 3), v.OOvv, (4, 5, 2, 3), (0, 1, 4, 5)) - tmp141 = einsum(t1_oV, (0, 1), tmp140, (0, 2, 3, 4), (2, 3, 1, 4)) - del tmp140 - tmp102 = einsum(v.Ovoo, (0, 1, 2, 3), t2_ooVV, (2, 3, 4, 5), (1, 0, 4, 5)) * -1 - tmp56 = einsum(t1_Ov, (0, 1), f.Ov, (2, 1), (2, 0)) - tmp106 = einsum(v.oVvv, (0, 1, 2, 3), t2_OoVv, (4, 0, 5, 3), (2, 4, 5, 1)) - tmp30 = einsum(v.OoVV, (0, 1, 2, 3), t3, (4, 5, 0, 6, 2, 3), (1, 4, 5, 6)) * -1 - tmp73 = einsum(t1_Ov, (0, 1), v.VVvv, (2, 3, 4, 1), (4, 0, 2, 3)) - tmp58 = einsum(t1_oV, (0, 1), f.oV, (0, 2), (2, 1)) - tmp83 = einsum(t1_Ov, (0, 1), v.OVVv, (2, 3, 4, 1), (0, 2, 4, 3)) * -1 - tmp81 = einsum(t1, (0, 1), v.oVVv, (0, 2, 3, 1), (2, 3)) * -1 - tmp67 = einsum(t1_oV, (0, 1), v.OOoo, (2, 3, 4, 0), (4, 2, 3, 1)) - tmp108 = einsum(t2_OOVv, (0, 1, 2, 3), tmp1, (4, 3), (4, 0, 1, 2)) - tmp97 = einsum(v.ooVV, (0, 1, 2, 3), t2_ooVV, (0, 1, 4, 5), (4, 5, 2, 3)) - tmp85 = einsum(t1_oV, (0, 1), v.oVVV, (0, 2, 3, 4), (1, 2, 3, 4)) - tmp91 = einsum(v.ooVv, (0, 1, 2, 3), t2_ooVv, (0, 1, 4, 3), (4, 2)) + tmp93 = einsum(v.OoVv, (0, 1, 2, 3), t2_OoVv, (4, 1, 5, 3), (4, 0, 5, 2)) tmp121 = einsum(t1_oV, (0, 1), tmp120, (2, 0, 3, 4), (2, 3, 4, 1)) del tmp120 - tmp124 = einsum(t1_oV, (0, 1), tmp123, (0, 2, 3, 4), (2, 3, 4, 1)) * -1 - del tmp123 + tmp69 = einsum(v.OvVo, (0, 1, 2, 3), t1_Ov, (4, 1), (3, 4, 0, 2)) + tmp132 = einsum(tmp131, (0, 1), t1_Ov, (2, 0), (2, 1)) * -1 + del tmp131 + tmp150 = einsum(tmp149, (0, 1, 2, 3), t1_oV, (0, 4), (1, 2, 3, 4)) * -1 + del tmp149 + tmp71 = einsum(v.OvVo, (0, 1, 2, 3), t1_oV, (3, 4), (1, 0, 4, 2)) + tmp127 = einsum(t1_Ov, (0, 1), tmp126, (2, 1, 3, 4), (2, 3, 0, 4)) * -1 + del tmp126 + tmp75 = einsum(v.OoOv, (0, 1, 2, 3), t1, (1, 3), (0, 2)) + tmp60 = einsum(f.ov, (0, 1), t2_OOVv, (2, 3, 4, 1), (0, 2, 3, 4)) * -1 + tmp30 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OoVV, (2, 6, 4, 5), (6, 0, 1, 3)) * -1 + tmp79 = einsum(v.OoOV, (0, 1, 2, 3), t1_oV, (1, 4), (0, 2, 4, 3)) * -1 + tmp73 = einsum(v.VVvv, (0, 1, 2, 3), t1_Ov, (4, 3), (2, 4, 0, 1)) + tmp77 = einsum(v.OvOO, (0, 1, 2, 3), t1_Ov, (4, 1), (4, 0, 2, 3)) * -1 + tmp81 = einsum(v.VoVv, (0, 1, 2, 3), t1, (1, 3), (0, 2)) + tmp89 = einsum(v.OOvv, (0, 1, 2, 3), t2_OOvv, (4, 5, 2, 3), (4, 5, 0, 1)) + tmp83 = einsum(v.OVVv, (0, 1, 2, 3), t1_Ov, (4, 3), (4, 0, 2, 1)) * -1 + tmp67 = einsum(t1_oV, (0, 1), v.OOoo, (2, 3, 4, 0), (4, 2, 3, 1)) + tmp95 = einsum(t2_OOVv, (0, 1, 2, 3), v.OOVv, (4, 5, 6, 3), (0, 1, 4, 5, 2, 6)) + tmp58 = einsum(t1_oV, (0, 1), f.Vo, (2, 0), (2, 1)) + tmp135 = einsum(tmp134, (0, 1), t1_oV, (0, 2), (2, 1)) * -1 + del tmp134 tmp28 = einsum(v.OOVv, (0, 1, 2, 3), t1, (4, 3), (4, 0, 1, 2)) * -1 + tmp43 = einsum(t2, (0, 1, 2, 3), tmp1, (4, 3), (0, 1, 4, 2)) * -1 + tmp0 = einsum(f.ov, (0, 1), t1, (2, 1), (0, 2)) tmp4 = einsum(v.ooov, (0, 1, 2, 3), t1, (1, 3), (0, 2)) + tmp37 = einsum(t2, (0, 1, 2, 3), v.oovv, (4, 5, 2, 3), (0, 1, 4, 5)) + tmp3 = einsum(t2, (0, 1, 2, 3), v.oovv, (4, 1, 2, 3), (0, 4)) + tmp40 = einsum(t2, (0, 1, 2, 3), v.oovv, (4, 1, 5, 3), (0, 4, 2, 5)) + tmp46 = einsum(tmp2, (0, 1, 2, 3), t2, (4, 1, 5, 3), (0, 4, 2, 5)) + tmp20 = einsum(t2, (0, 1, 2, 3), v.ooov, (4, 1, 5, 3), (0, 4, 5, 2)) tmp25 = einsum(t2, (0, 1, 2, 3), v.ovvv, (4, 5, 2, 3), (0, 1, 4, 5)) - tmp10 = einsum(t2, (0, 1, 2, 3), f.ov, (4, 3), (4, 0, 1, 2)) - tmp46 = einsum(t2, (0, 1, 2, 3), tmp2, (4, 1, 5, 3), (4, 0, 5, 2)) - tmp51 = einsum(t1, (0, 1), tmp23, (2, 3, 4, 1), (2, 0, 3, 4)) * -1 - tmp37 = einsum(v.oovv, (0, 1, 2, 3), t2, (4, 5, 2, 3), (4, 5, 0, 1)) tmp33 = einsum(v.ovov, (0, 1, 2, 3), t1, (4, 3), (4, 0, 2, 1)) - tmp45 = einsum(tmp2, (0, 1, 2, 3), t1, (4, 3), (0, 4, 2, 1)) - tmp38 = einsum(t2, (0, 1, 2, 3), v.oovv, (0, 1, 4, 3), (2, 4)) - tmp20 = einsum(t2, (0, 1, 2, 3), v.ooov, (4, 1, 5, 3), (0, 4, 5, 2)) - tmp3 = einsum(v.oovv, (0, 1, 2, 3), t2, (4, 1, 2, 3), (4, 0)) tmp5 = einsum(v.ovvv, (0, 1, 2, 3), t1, (0, 3), (1, 2)) - tmp43 = einsum(t2, (0, 1, 2, 3), tmp1, (4, 3), (0, 1, 4, 2)) * -1 + tmp10 = einsum(t2, (0, 1, 2, 3), f.ov, (4, 3), (4, 0, 1, 2)) + tmp6 = einsum(tmp1, (0, 1), t1, (2, 1), (2, 0)) + tmp45 = einsum(tmp2, (0, 1, 2, 3), t1, (4, 3), (0, 4, 2, 1)) tmp49 = einsum(tmp18, (0, 1, 2, 3), t1, (1, 4), (0, 2, 3, 4)) * -1 - tmp6 = einsum(t1, (0, 1), tmp1, (2, 1), (0, 2)) - tmp0 = einsum(f.ov, (0, 1), t1, (2, 1), (0, 2)) - tmp40 = einsum(t2, (0, 1, 2, 3), v.oovv, (4, 1, 5, 3), (0, 4, 2, 5)) - tmp148 = einsum(t2_OoVV, (0, 1, 2, 3), tmp147, (1, 4, 5, 6), (5, 4, 0, 6, 2, 3)) * -1 - del tmp147 - tmp112 = einsum(t2_OoVV, (0, 1, 2, 3), tmp111, (1, 4, 5, 6), (4, 0, 5, 2, 3, 6)) * -1 - del tmp111 - tmp94 = einsum(tmp93, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 3), (0, 4, 5, 2, 6, 7)) - del tmp93 + tmp51 = einsum(t1, (0, 1), tmp23, (2, 3, 4, 1), (2, 0, 3, 4)) * -1 + tmp38 = einsum(t2, (0, 1, 2, 3), v.oovv, (0, 1, 4, 3), (2, 4)) + tmp105 = einsum(t2_OoVV, (0, 1, 2, 3), tmp104, (1, 4, 5, 6), (4, 5, 0, 2, 3, 6)) * -1 + del tmp104 tmp88 = einsum(tmp87, (0, 1), t3, (2, 3, 1, 4, 5, 6), (0, 2, 3, 4, 5, 6)) del tmp87 - tmp145 = einsum(tmp144, (0, 1, 2, 3), t3, (4, 5, 6, 7, 2, 3), (4, 5, 6, 0, 1, 7)) - del tmp144 - tmp139 = einsum(tmp138, (0, 1, 2, 3), t3, (4, 2, 3, 5, 6, 7), (1, 0, 4, 5, 6, 7)) * -1 - del tmp138 - tmp63 = einsum(t2_OOVv, (0, 1, 2, 3), v.OvVV, (4, 3, 5, 6), (0, 1, 4, 2, 5, 6)) - tmp78 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp77, (6, 7, 2, 1), (6, 0, 7, 3, 4, 5)) * -1 - del tmp77 - tmp70 = einsum(t2_OoVV, (0, 1, 2, 3), tmp69, (1, 4, 5, 6), (4, 0, 5, 2, 3, 6)) * -1 - del tmp69 - tmp80 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp79, (2, 6, 7, 5), (0, 1, 6, 7, 3, 4)) - del tmp79 - tmp96 = einsum(tmp95, (0, 1, 2, 3, 4, 5), t3, (6, 3, 2, 7, 8, 5), (0, 1, 6, 4, 7, 8)) * -1 - del tmp95 - tmp76 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp75, (2, 6), (0, 1, 6, 3, 4, 5)) - del tmp75 - tmp136 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp135, (6, 5), (0, 1, 2, 6, 3, 4)) - del tmp135 - tmp66 = einsum(v.VVVV, (0, 1, 2, 3), t3, (4, 5, 6, 7, 2, 3), (4, 5, 6, 7, 0, 1)) - tmp128 = einsum(tmp127, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (1, 2, 4, 5, 6, 3)) * -1 - del tmp127 - tmp101 = einsum(t2_OoVV, (0, 1, 2, 3), tmp100, (1, 4, 5, 6), (0, 4, 5, 2, 3, 6)) * -1 + tmp101 = einsum(tmp100, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (4, 1, 2, 5, 6, 3)) * -1 del tmp100 - tmp151 = einsum(t2_OOVv, (0, 1, 2, 3), tmp150, (3, 4, 5, 6), (4, 0, 1, 5, 6, 2)) - del tmp150 - tmp117 = einsum(t2_OoVV, (0, 1, 2, 3), tmp116, (1, 4, 5, 6), (4, 5, 0, 6, 2, 3)) - del tmp116 - tmp72 = einsum(tmp71, (0, 1, 2, 3), t2_OOVv, (4, 5, 6, 0), (4, 5, 1, 2, 6, 3)) * -1 - del tmp71 - tmp61 = einsum(t2_OoVV, (0, 1, 2, 3), tmp60, (1, 4, 5, 6), (5, 4, 0, 6, 2, 3)) - del tmp60 - tmp119 = einsum(tmp118, (0, 1, 2, 3), t2_OOVv, (4, 5, 6, 0), (4, 5, 1, 2, 6, 3)) * -1 - del tmp118 - tmp105 = einsum(t2_OoVV, (0, 1, 2, 3), tmp104, (1, 4, 5, 6), (4, 5, 0, 2, 3, 6)) * -1 - del tmp104 - tmp130 = einsum(t2_OOVv, (0, 1, 2, 3), tmp129, (3, 4, 5, 6), (4, 0, 1, 5, 2, 6)) - del tmp129 - tmp114 = einsum(tmp113, (0, 1, 2, 3), t2_OOVv, (4, 5, 6, 0), (1, 4, 5, 2, 3, 6)) - del tmp113 - tmp62 = einsum(v.OOoV, (0, 1, 2, 3), t2_OoVV, (4, 2, 5, 6), (4, 0, 1, 5, 6, 3)) * -1 - tmp133 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp132, (6, 2), (6, 0, 1, 3, 4, 5)) - del tmp132 - tmp90 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp89, (6, 7, 2, 1), (7, 6, 0, 3, 4, 5)) - del tmp89 + tmp107 = einsum(t2_OOVv, (0, 1, 2, 3), tmp106, (3, 4, 5, 6), (0, 1, 4, 2, 5, 6)) * -1 + del tmp106 tmp142 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp141, (6, 2, 7, 5), (6, 0, 1, 7, 3, 4)) del tmp141 - tmp65 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OVOV, (6, 5, 2, 7), (0, 1, 6, 3, 4, 7)) + tmp119 = einsum(t2_OOVv, (0, 1, 2, 3), tmp118, (3, 4, 5, 6), (0, 1, 4, 5, 2, 6)) * -1 + del tmp118 + tmp139 = einsum(tmp138, (0, 1, 2, 3), t3, (4, 2, 3, 5, 6, 7), (1, 0, 4, 5, 6, 7)) * -1 + del tmp138 tmp103 = einsum(tmp102, (0, 1, 2, 3), t2_OOVv, (4, 5, 6, 0), (4, 5, 1, 2, 3, 6)) * -1 del tmp102 - tmp57 = einsum(tmp56, (0, 1), t3, (2, 3, 0, 4, 5, 6), (1, 2, 3, 4, 5, 6)) + tmp57 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp56, (2, 6), (6, 0, 1, 3, 4, 5)) del tmp56 - tmp107 = einsum(t2_OOVv, (0, 1, 2, 3), tmp106, (3, 4, 5, 6), (0, 1, 4, 2, 5, 6)) * -1 - del tmp106 - tmp99 = einsum(t2_OoVV, (0, 1, 2, 3), tmp30, (1, 4, 5, 6), (0, 5, 4, 2, 3, 6)) - tmp74 = einsum(t2_OOVv, (0, 1, 2, 3), tmp73, (3, 4, 5, 6), (4, 0, 1, 2, 5, 6)) - del tmp73 - tmp59 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp58, (5, 6), (0, 1, 2, 6, 3, 4)) - del tmp58 - tmp55 = einsum(t3, (0, 1, 2, 3, 4, 5), f.VV, (6, 5), (0, 1, 2, 6, 3, 4)) - tmp54 = einsum(f.OO, (0, 1), t3, (2, 3, 1, 4, 5, 6), (0, 2, 3, 4, 5, 6)) - tmp84 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp83, (6, 2, 5, 7), (6, 0, 1, 3, 4, 7)) - del tmp83 - tmp82 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp81, (6, 5), (0, 1, 2, 3, 4, 6)) - del tmp81 - tmp68 = einsum(t2_OoVV, (0, 1, 2, 3), tmp67, (1, 4, 5, 6), (0, 4, 5, 6, 2, 3)) - del tmp67 - tmp109 = einsum(t2_OoVV, (0, 1, 2, 3), tmp108, (1, 4, 5, 6), (4, 5, 0, 6, 2, 3)) - del tmp108 - tmp64 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OOOO, (6, 7, 1, 2), (0, 6, 7, 3, 4, 5)) tmp98 = einsum(tmp97, (0, 1, 2, 3), t3, (4, 5, 6, 7, 3, 2), (4, 5, 6, 1, 0, 7)) del tmp97 + tmp125 = einsum(tmp124, (0, 1, 2, 3), t2_OOVv, (4, 5, 6, 0), (4, 5, 1, 2, 3, 6)) * -1 + del tmp124 + tmp130 = einsum(t2_OOVv, (0, 1, 2, 3), tmp129, (3, 4, 5, 6), (4, 0, 1, 5, 2, 6)) + del tmp129 + tmp148 = einsum(t2_OoVV, (0, 1, 2, 3), tmp147, (1, 4, 5, 6), (5, 4, 0, 6, 2, 3)) * -1 + del tmp147 + tmp109 = einsum(t2_OoVV, (0, 1, 2, 3), tmp108, (1, 4, 5, 6), (4, 5, 0, 6, 2, 3)) + del tmp108 + tmp66 = einsum(v.VVVV, (0, 1, 2, 3), t3, (4, 5, 6, 7, 2, 3), (4, 5, 6, 7, 0, 1)) + tmp55 = einsum(f.VV, (0, 1), t3, (2, 3, 4, 5, 6, 1), (2, 3, 4, 0, 5, 6)) + tmp117 = einsum(t2_OoVV, (0, 1, 2, 3), tmp116, (1, 4, 5, 6), (4, 5, 0, 6, 2, 3)) + del tmp116 + tmp54 = einsum(f.OO, (0, 1), t3, (2, 3, 1, 4, 5, 6), (0, 2, 3, 4, 5, 6)) tmp86 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp85, (6, 7, 5, 4), (0, 1, 2, 6, 3, 7)) * -1 del tmp85 - tmp92 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp91, (6, 5), (0, 1, 2, 6, 3, 4)) + tmp65 = einsum(v.OVOV, (0, 1, 2, 3), t3, (4, 5, 2, 6, 7, 1), (4, 5, 0, 6, 7, 3)) + tmp112 = einsum(tmp111, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (1, 4, 2, 5, 6, 3)) * -1 + del tmp111 + tmp92 = einsum(tmp91, (0, 1), t3, (2, 3, 4, 5, 6, 1), (2, 3, 4, 0, 5, 6)) del tmp91 + tmp145 = einsum(tmp144, (0, 1, 2, 3), t3, (4, 5, 6, 7, 2, 3), (4, 5, 6, 0, 1, 7)) + del tmp144 + tmp114 = einsum(t2_OOVv, (0, 1, 2, 3), tmp113, (3, 4, 5, 6), (4, 0, 1, 5, 6, 2)) + del tmp113 + tmp94 = einsum(tmp93, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 3), (0, 4, 5, 2, 6, 7)) + del tmp93 tmp122 = einsum(tmp121, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (1, 4, 2, 3, 5, 6)) del tmp121 - tmp125 = einsum(tmp124, (0, 1, 2, 3), t2_OOVv, (4, 5, 6, 0), (4, 5, 1, 2, 3, 6)) * -1 - del tmp124 - tmp16 = einsum(v.OvVV, (0, 1, 2, 3), t3, (4, 5, 0, 6, 2, 3), (1, 4, 5, 6)) * -1 + tmp70 = einsum(tmp69, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (1, 4, 2, 5, 6, 3)) * -1 + del tmp69 + tmp133 = einsum(tmp132, (0, 1), t3, (2, 3, 1, 4, 5, 6), (0, 2, 3, 4, 5, 6)) + del tmp132 + tmp151 = einsum(tmp150, (0, 1, 2, 3), t2_OOVv, (4, 5, 6, 0), (1, 4, 5, 2, 3, 6)) + del tmp150 + tmp63 = einsum(t2_OOVv, (0, 1, 2, 3), v.OvVV, (4, 3, 5, 6), (0, 1, 4, 2, 5, 6)) + tmp72 = einsum(t2_OOVv, (0, 1, 2, 3), tmp71, (3, 4, 5, 6), (0, 1, 4, 5, 2, 6)) * -1 + del tmp71 + tmp128 = einsum(t2_OoVV, (0, 1, 2, 3), tmp127, (1, 4, 5, 6), (4, 5, 0, 2, 3, 6)) * -1 + del tmp127 + tmp76 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp75, (2, 6), (0, 1, 6, 3, 4, 5)) + del tmp75 + tmp61 = einsum(tmp60, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (1, 2, 4, 3, 5, 6)) * -1 + del tmp60 + tmp64 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OOOO, (6, 7, 1, 2), (0, 6, 7, 3, 4, 5)) + tmp99 = einsum(tmp30, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (4, 2, 1, 5, 6, 3)) + tmp80 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp79, (2, 6, 7, 5), (0, 1, 6, 7, 3, 4)) + del tmp79 + tmp74 = einsum(t2_OOVv, (0, 1, 2, 3), tmp73, (3, 4, 5, 6), (4, 0, 1, 2, 5, 6)) + del tmp73 + tmp78 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp77, (6, 7, 2, 1), (6, 0, 7, 3, 4, 5)) * -1 + del tmp77 + tmp62 = einsum(v.OOVo, (0, 1, 2, 3), t2_OoVV, (4, 3, 5, 6), (4, 0, 1, 5, 6, 2)) + tmp82 = einsum(tmp81, (0, 1), t3, (2, 3, 4, 5, 6, 1), (2, 3, 4, 5, 6, 0)) + del tmp81 + tmp90 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp89, (6, 7, 2, 1), (7, 6, 0, 3, 4, 5)) + del tmp89 + tmp84 = einsum(tmp83, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 2), (0, 4, 5, 6, 7, 3)) + del tmp83 + tmp68 = einsum(tmp67, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (4, 1, 2, 3, 5, 6)) + del tmp67 + tmp96 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp95, (6, 7, 1, 2, 8, 5), (6, 7, 0, 8, 3, 4)) + del tmp95 + tmp59 = einsum(tmp58, (0, 1), t3, (2, 3, 4, 5, 6, 0), (2, 3, 4, 1, 5, 6)) + del tmp58 + tmp136 = einsum(tmp135, (0, 1), t3, (2, 3, 4, 5, 6, 1), (2, 3, 4, 0, 5, 6)) + del tmp135 + tmp16 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OvVV, (2, 6, 4, 5), (6, 0, 1, 3)) * -1 tmp31 = einsum(tmp30, (0, 1, 2, 3), t1, (0, 4), (4, 1, 2, 3)) del tmp30 - tmp15 = einsum(v.OOoV, (0, 1, 2, 3), t3, (4, 0, 1, 5, 6, 3), (2, 4, 5, 6)) - tmp29 = einsum(tmp28, (0, 1, 2, 3), t3, (4, 2, 1, 5, 6, 3), (0, 4, 5, 6)) * -1 + tmp29 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp28, (6, 2, 1, 5), (6, 0, 3, 4)) * -1 del tmp28 + tmp15 = einsum(v.OOVo, (0, 1, 2, 3), t3, (4, 0, 1, 5, 6, 2), (3, 4, 5, 6)) * -1 tmp27 = einsum(v.OoVv, (0, 1, 2, 3), t1, (1, 3), (0, 2)) + tmp44 = einsum(tmp43, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) + del tmp43 + tmp9 = einsum(tmp0, (0, 1), t2, (2, 0, 3, 4), (1, 2, 3, 4)) + tmp8 = einsum(f.vv, (0, 1), t2, (2, 3, 4, 1), (2, 3, 0, 4)) + tmp14 = einsum(t2, (0, 1, 2, 3), v.ovov, (4, 3, 1, 5), (0, 4, 2, 5)) tmp17 = einsum(t2, (0, 1, 2, 3), tmp4, (1, 4), (0, 4, 2, 3)) * -1 - tmp19 = einsum(t2, (0, 1, 2, 3), tmp18, (4, 1, 0, 5), (4, 5, 2, 3)) * -1 - del tmp18 - tmp26 = einsum(t1, (0, 1), tmp25, (2, 3, 0, 4), (2, 3, 1, 4)) - del tmp25 - tmp12 = einsum(t1, (0, 1), v.ooov, (2, 3, 0, 4), (2, 3, 1, 4)) - tmp11 = einsum(t1, (0, 1), tmp10, (0, 2, 3, 4), (2, 3, 1, 4)) - del tmp10 - tmp47 = einsum(t1, (0, 1), tmp46, (2, 3, 0, 4), (2, 3, 1, 4)) + tmp48 = einsum(tmp37, (0, 1, 2, 3), t1, (2, 4), (0, 1, 3, 4)) * -1 + tmp36 = einsum(tmp3, (0, 1), t2, (2, 1, 3, 4), (2, 0, 3, 4)) + tmp12 = einsum(v.ooov, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) + tmp41 = einsum(tmp40, (0, 1, 2, 3), t2, (4, 1, 5, 3), (0, 4, 2, 5)) + del tmp40 + tmp47 = einsum(tmp46, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) del tmp46 - tmp14 = einsum(t2, (0, 1, 2, 3), v.ovov, (4, 3, 1, 5), (0, 4, 2, 5)) - tmp52 = einsum(t1, (0, 1), tmp51, (2, 3, 0, 4), (3, 2, 1, 4)) * -1 - del tmp51 - tmp48 = einsum(t1, (0, 1), tmp37, (2, 3, 0, 4), (2, 3, 4, 1)) * -1 - tmp34 = einsum(t1, (0, 1), tmp33, (2, 0, 3, 4), (2, 3, 1, 4)) - del tmp33 - tmp53 = einsum(tmp45, (0, 1, 2, 3), t1, (3, 4), (0, 1, 2, 4)) - tmp39 = einsum(t2, (0, 1, 2, 3), tmp38, (4, 3), (0, 1, 2, 4)) - del tmp38 tmp21 = einsum(t1, (0, 1), tmp20, (2, 0, 3, 4), (2, 3, 1, 4)) del tmp20 - tmp36 = einsum(t2, (0, 1, 2, 3), tmp3, (4, 1), (0, 4, 2, 3)) - tmp8 = einsum(t2, (0, 1, 2, 3), f.vv, (4, 3), (0, 1, 4, 2)) + tmp26 = einsum(tmp25, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) + del tmp25 + tmp19 = einsum(tmp18, (0, 1, 2, 3), t2, (2, 1, 4, 5), (0, 3, 4, 5)) * -1 + del tmp18 + tmp34 = einsum(t1, (0, 1), tmp33, (2, 0, 3, 4), (2, 3, 1, 4)) + del tmp33 tmp22 = einsum(tmp5, (0, 1), t2, (2, 3, 4, 1), (2, 3, 4, 0)) * -1 - tmp7 = einsum(t2, (0, 1, 2, 3), f.oo, (4, 1), (4, 0, 2, 3)) tmp24 = einsum(t2, (0, 1, 2, 3), tmp23, (4, 1, 5, 3), (4, 0, 2, 5)) * -1 del tmp23 - tmp44 = einsum(tmp43, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) - del tmp43 + tmp11 = einsum(tmp10, (0, 1, 2, 3), t1, (0, 4), (1, 2, 4, 3)) + del tmp10 + tmp7 = einsum(f.oo, (0, 1), t2, (2, 1, 3, 4), (0, 2, 3, 4)) + tmp13 = einsum(v.ovvv, (0, 1, 2, 3), t1, (4, 1), (4, 0, 2, 3)) + tmp42 = einsum(tmp6, (0, 1), t2, (2, 1, 3, 4), (0, 2, 3, 4)) * -1 + tmp53 = einsum(tmp45, (0, 1, 2, 3), t1, (3, 4), (0, 1, 2, 4)) tmp50 = einsum(t1, (0, 1), tmp49, (2, 0, 3, 4), (2, 3, 4, 1)) * -1 del tmp49 - tmp42 = einsum(t2, (0, 1, 2, 3), tmp6, (4, 1), (4, 0, 2, 3)) * -1 - tmp9 = einsum(tmp0, (0, 1), t2, (2, 0, 3, 4), (1, 2, 3, 4)) + tmp35 = einsum(t1, (0, 1), v.vvvv, (2, 3, 4, 1), (0, 2, 3, 4)) + tmp52 = einsum(t1, (0, 1), tmp51, (2, 3, 0, 4), (3, 2, 1, 4)) * -1 + del tmp51 tmp32 = einsum(v.oooo, (0, 1, 2, 3), t1, (3, 4), (0, 1, 2, 4)) - tmp35 = einsum(v.vvvv, (0, 1, 2, 3), t1, (4, 3), (4, 0, 1, 2)) - tmp13 = einsum(v.ovvv, (0, 1, 2, 3), t1, (4, 1), (4, 0, 2, 3)) - tmp41 = einsum(t2, (0, 1, 2, 3), tmp40, (4, 1, 5, 3), (0, 4, 2, 5)) - del tmp40 + tmp39 = einsum(tmp38, (0, 1), t2, (2, 3, 4, 1), (2, 3, 4, 0)) + del tmp38 t3new = np.copy(np.transpose(tmp54, (1, 2, 0, 4, 5, 3))) * -1 t3new += np.transpose(tmp54, (1, 0, 2, 4, 5, 3)) t3new += np.transpose(tmp54, (0, 2, 1, 4, 5, 3)) @@ -813,7 +813,7 @@ def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): t2new += np.transpose(tmp26, (0, 1, 3, 2)) * 0.5 t2new += tmp26 * -0.5 del tmp26 - t2new += einsum(t1, (0, 1), tmp32, (2, 3, 0, 4), (2, 3, 1, 4)) + t2new += einsum(tmp32, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) del tmp32 t2new += np.transpose(tmp34, (0, 1, 3, 2)) * -1 t2new += np.transpose(tmp34, (1, 0, 3, 2)) @@ -846,7 +846,7 @@ def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): t2new += np.transpose(tmp47, (1, 0, 3, 2)) * -1 t2new += np.transpose(tmp47, (0, 1, 3, 2)) del tmp47 - t2new += einsum(t1, (0, 1), tmp48, (2, 3, 0, 4), (2, 3, 1, 4)) * 0.5 + t2new += einsum(tmp48, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) * 0.5 del tmp48 t2new += np.transpose(tmp50, (0, 1, 3, 2)) t2new += np.transpose(tmp50, (1, 0, 3, 2)) * -1 @@ -854,37 +854,37 @@ def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): t2new += np.transpose(tmp52, (0, 1, 3, 2)) t2new += tmp52 * -1 del tmp52 - t2new += einsum(t1, (0, 1), tmp53, (2, 3, 0, 4), (3, 2, 1, 4)) * -1 + t2new += einsum(tmp53, (0, 1, 2, 3), t1, (2, 4), (1, 0, 4, 3)) * -1 del tmp53 - t1new_OV = einsum(v.OOVV, (0, 1, 2, 3), t3, (4, 0, 1, 5, 2, 3), (4, 5)) * 0.25 + t1new_OV = einsum(t3, (0, 1, 2, 3, 4, 5), v.OOVV, (1, 2, 4, 5), (0, 3)) * 0.25 t1new = np.copy(f.ov) t1new += einsum(f.oo, (0, 1), t1, (1, 2), (0, 2)) * -1 t1new += einsum(f.vv, (0, 1), t1, (2, 1), (2, 0)) - t1new += einsum(f.ov, (0, 1), t2, (2, 0, 3, 1), (2, 3)) + t1new += einsum(t2, (0, 1, 2, 3), f.ov, (1, 3), (0, 2)) t1new += einsum(tmp0, (0, 1), t1, (0, 2), (1, 2)) * -1 del tmp0 - t1new += einsum(v.ovov, (0, 1, 2, 3), t1, (2, 1), (0, 3)) * -1 + t1new += einsum(t1, (0, 1), v.ovov, (2, 1, 0, 3), (2, 3)) * -1 t1new += einsum(t2, (0, 1, 2, 3), v.ovoo, (4, 3, 0, 1), (4, 2)) * -0.5 t1new += einsum(v.ovvv, (0, 1, 2, 3), t2, (4, 0, 2, 3), (4, 1)) * -0.5 t1new += einsum(tmp1, (0, 1), t2, (2, 0, 3, 1), (2, 3)) del tmp1 - t1new += einsum(tmp2, (0, 1, 2, 3), t2, (2, 1, 4, 3), (0, 4)) * -0.5 + t1new += einsum(t2, (0, 1, 2, 3), tmp2, (4, 1, 0, 3), (4, 2)) * -0.5 del tmp2 t1new += einsum(tmp3, (0, 1), t1, (1, 2), (0, 2)) * -0.5 del tmp3 - t1new += einsum(tmp4, (0, 1), t1, (0, 2), (1, 2)) * -1 + t1new += einsum(t1, (0, 1), tmp4, (0, 2), (2, 1)) * -1 del tmp4 t1new += einsum(tmp5, (0, 1), t1, (2, 1), (2, 0)) * -1 del tmp5 t1new += einsum(tmp6, (0, 1), t1, (1, 2), (0, 2)) * -1 del tmp6 - t2new += _inflate(t2new.shape, np.ix_(so, sO, sV, sV), t2new_oOVV) + t2new += _inflate(t2new.shape, np.ix_(sO, sO, sV, sV), t2new_OOVV) + t2new += _inflate(t2new.shape, np.ix_(sO, sO, sv, sV), t2new_OOvV) t2new += _inflate(t2new.shape, np.ix_(sO, so, sV, sV), t2new_OoVV) + t2new += _inflate(t2new.shape, np.ix_(so, sO, sV, sV), t2new_oOVV) t1new += _inflate(t1new.shape, np.ix_(sO, sV), t1new_OV) t2new += _inflate(t2new.shape, np.ix_(sO, sO, sV, sv), t2new_OOVv) - t2new += _inflate(t2new.shape, np.ix_(sO, sO, sv, sV), t2new_OOvV) - t2new += _inflate(t2new.shape, np.ix_(sO, sO, sV, sV), t2new_OOVV) return {f"t1new": t1new, f"t2new": t2new, f"t3new": t3new} diff --git a/ebcc/codegen/RCCSDwtwp.py b/ebcc/codegen/RCCSDwtwp.py index a237bb7b..ba33f50c 100644 --- a/ebcc/codegen/RCCSDwtwp.py +++ b/ebcc/codegen/RCCSDwtwp.py @@ -2,14 +2,14 @@ Code generated by `albert`: https://github.com/obackhouse/albert - * date: 2024-09-27T21:44:54.649193 - * python version: 3.10.12 (main, Sep 11 2024, 15:47:36) [GCC 11.4.0] + * date: 2024-09-29T17:16:22.166169 + * python version: 3.10.15 (main, Sep 9 2024, 03:02:45) [GCC 11.4.0] * albert version: 0.0.0 - * caller: /home/ollie/git/albert/albert/codegen/einsum.py - * node: ollie-desktop + * caller: /opt/hostedtoolcache/Python/3.10.15/x64/lib/python3.10/site-packages/albert/codegen/einsum.py + * node: fv-az1788-690 * system: Linux * processor: x86_64 - * release: 6.8.0-40-generic + * release: 6.8.0-1014-azure """ from ebcc import numpy as np @@ -19,7 +19,7 @@ def energy(f=None, t1=None, t2=None, v=None, **kwargs): """ - Code generated by `albert` 0.0.0 on 2024-09-27T21:44:55.123874. + Code generated by `albert` 0.0.0 on 2024-09-29T17:16:22.869247. Parameters ---------- @@ -41,7 +41,7 @@ def energy(f=None, t1=None, t2=None, v=None, **kwargs): tmp0 = np.copy(np.transpose(v.ovov, (0, 2, 3, 1))) tmp0 += np.transpose(v.ovov, (0, 2, 1, 3)) * -0.5 tmp1 = np.copy(f.ov) - tmp1 += einsum(tmp0, (0, 1, 2, 3), t1, (0, 3), (1, 2)) + tmp1 += einsum(t1, (0, 1), tmp0, (0, 2, 3, 1), (2, 3)) del tmp0 e_cc = einsum(t2, (0, 1, 2, 3), v.ovov, (0, 2, 1, 3), ()) * 2 e_cc += einsum(t2, (0, 1, 2, 3), v.ovov, (0, 3, 1, 2), ()) * -1 @@ -52,7 +52,7 @@ def energy(f=None, t1=None, t2=None, v=None, **kwargs): def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): """ - Code generated by `albert` 0.0.0 on 2024-09-27T21:46:20.711211. + Code generated by `albert` 0.0.0 on 2024-09-29T17:18:26.069150. Parameters ---------- @@ -85,435 +85,435 @@ def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): sv = np.ones((t1.shape[1],), dtype=bool) sO = space.active[space.correlated][space.occupied[space.correlated]] sV = space.active[space.correlated][space.virtual[space.correlated]] - t2_oovV = np.copy(t2[np.ix_(so, so, sv, sV)]) - t2_oovv = np.copy(t2[np.ix_(so, so, sv, sv)]) - t2_OovV = np.copy(t2[np.ix_(sO, so, sv, sV)]) + t2_OoVv = np.copy(t2[np.ix_(sO, so, sV, sv)]) + t1_oV = np.copy(t1[np.ix_(so, sV)]) t1_Ov = np.copy(t1[np.ix_(sO, sv)]) - t2_ooVV = np.copy(t2[np.ix_(so, so, sV, sV)]) - t2_OOvV = np.copy(t2[np.ix_(sO, sO, sv, sV)]) - t2_OoVV = np.copy(t2[np.ix_(sO, so, sV, sV)]) t2_Oovv = np.copy(t2[np.ix_(sO, so, sv, sv)]) t2_OOvv = np.copy(t2[np.ix_(sO, sO, sv, sv)]) + t2_OovV = np.copy(t2[np.ix_(sO, so, sv, sV)]) + t2_OoVV = np.copy(t2[np.ix_(sO, so, sV, sV)]) + t2_oovv = np.copy(t2[np.ix_(so, so, sv, sv)]) + t2_ooVV = np.copy(t2[np.ix_(so, so, sV, sV)]) + t2_ooVv = np.copy(t2[np.ix_(so, so, sV, sv)]) + t2_OOvV = np.copy(t2[np.ix_(sO, sO, sv, sV)]) t1_ov = np.copy(t1[np.ix_(so, sv)]) - t1_oV = np.copy(t1[np.ix_(so, sV)]) - t2_OoVv = np.copy(t2[np.ix_(sO, so, sV, sv)]) - tmp173 = einsum(t1_Ov, (0, 1), v.ovov, (2, 3, 4, 1), (2, 4, 3, 0)) + tmp173 = einsum(v.ovov, (0, 1, 2, 3), t1_Ov, (4, 3), (0, 2, 1, 4)) + tmp294 = einsum(v.Vvov, (0, 1, 2, 3), t1_Ov, (4, 1), (2, 3, 4, 0)) tmp218 = einsum(v.ovov, (0, 1, 2, 3), t1_oV, (2, 4), (0, 1, 3, 4)) - tmp294 = einsum(v.ovVv, (0, 1, 2, 3), t1_Ov, (4, 3), (0, 1, 4, 2)) - tmp181 = einsum(v.Ovov, (0, 1, 2, 3), t1, (2, 1), (3, 0)) - tmp307 = einsum(t1_Ov, (0, 1), tmp173, (2, 3, 1, 4), (2, 3, 0, 4)) - tmp154 = einsum(t1_oV, (0, 1), v.oVoV, (2, 3, 0, 4), (2, 1, 3, 4)) - tmp187 = einsum(t1_Ov, (0, 1), v.OvoV, (2, 1, 3, 4), (3, 0, 2, 4)) - tmp304 = einsum(v.OvOv, (0, 1, 2, 3), t1_Ov, (4, 3), (1, 4, 0, 2)) - tmp190 = einsum(t1_Ov, (0, 1), v.OVov, (2, 3, 4, 1), (4, 0, 2, 3)) - tmp176 = einsum(t1_Ov, (0, 1), v.Ooov, (2, 3, 4, 1), (4, 3, 0, 2)) - tmp148 = einsum(v.ovoV, (0, 1, 2, 3), t1, (2, 1), (0, 3)) - tmp297 = einsum(t1_Ov, (0, 1), v.ovVv, (2, 1, 3, 4), (2, 4, 0, 3)) - tmp14 = einsum(v.ovov, (0, 1, 2, 3), t1, (2, 1), (0, 3)) - tmp184 = einsum(v.Ovov, (0, 1, 2, 3), t1, (2, 3), (1, 0)) tmp1 = einsum(t1, (0, 1), v.ovov, (2, 3, 0, 1), (2, 3)) + tmp187 = einsum(v.OvVo, (0, 1, 2, 3), t1_Ov, (4, 1), (3, 4, 0, 2)) + tmp148 = einsum(t1, (0, 1), v.Voov, (2, 0, 3, 1), (3, 2)) + tmp14 = einsum(v.ovov, (0, 1, 2, 3), t1, (2, 1), (0, 3)) + tmp154 = einsum(t1_oV, (0, 1), v.VoVo, (2, 3, 4, 0), (3, 1, 2, 4)) + tmp307 = einsum(tmp173, (0, 1, 2, 3), t1_Ov, (4, 2), (0, 1, 4, 3)) + tmp176 = einsum(t1_Ov, (0, 1), v.Ooov, (2, 3, 4, 1), (4, 3, 0, 2)) tmp134 = einsum(t1_oV, (0, 1), v.Ooov, (2, 0, 3, 4), (3, 4, 2, 1)) - tmp151 = einsum(v.oVov, (0, 1, 2, 3), t1, (2, 3), (0, 1)) - tmp193 = einsum(t1_oV, (0, 1), tmp173, (0, 2, 3, 4), (2, 3, 4, 1)) - tmp75 = einsum(v.ovvv, (0, 1, 2, 3), t1, (4, 3), (4, 0, 1, 2)) - tmp4 = einsum(t1, (0, 1), v.ovov, (2, 3, 4, 1), (0, 2, 4, 3)) - tmp116 = einsum(v.ovov, (0, 1, 2, 3), t2, (4, 5, 1, 3), (4, 5, 0, 2)) - tmp62 = einsum(t1, (0, 1), v.ooov, (2, 3, 4, 1), (0, 2, 3, 4)) - tmp212 = einsum(t2_OovV, (0, 1, 2, 3), v.Ooov, (4, 5, 1, 2), (5, 0, 4, 3)) + tmp184 = einsum(v.Ovov, (0, 1, 2, 3), t1, (2, 3), (1, 0)) + tmp297 = einsum(v.Vvov, (0, 1, 2, 3), t1_Ov, (4, 3), (2, 1, 4, 0)) + tmp181 = einsum(v.Ovov, (0, 1, 2, 3), t1, (2, 1), (3, 0)) + tmp193 = einsum(tmp173, (0, 1, 2, 3), t1_oV, (0, 4), (1, 2, 3, 4)) + tmp151 = einsum(t1, (0, 1), v.Voov, (2, 3, 0, 1), (3, 2)) + tmp190 = einsum(v.OVov, (0, 1, 2, 3), t1_Ov, (4, 3), (2, 4, 0, 1)) + tmp304 = einsum(v.OvOv, (0, 1, 2, 3), t1_Ov, (4, 3), (1, 4, 0, 2)) + tmp4 = einsum(v.ovov, (0, 1, 2, 3), t1, (4, 3), (4, 0, 2, 1)) + tmp116 = einsum(t2, (0, 1, 2, 3), v.ovov, (4, 2, 5, 3), (0, 1, 4, 5)) + tmp62 = einsum(v.ooov, (0, 1, 2, 3), t1, (4, 3), (4, 0, 1, 2)) + tmp75 = einsum(t1, (0, 1), v.ovvv, (2, 3, 4, 1), (0, 2, 3, 4)) + tmp122 = einsum(v.OoOo, (0, 1, 2, 3), t1_oV, (3, 4), (1, 0, 2, 4)) + tmp313 = einsum(tmp294, (0, 1, 2, 3), t1_Ov, (4, 1), (0, 4, 2, 3)) + tmp231 = einsum(tmp218, (0, 1, 2, 3), t2_OovV, (4, 0, 1, 5), (2, 4, 3, 5)) + tmp275 = einsum(t2_OoVv, (0, 1, 2, 3), tmp173, (1, 4, 3, 5), (4, 5, 0, 2)) + tmp216 = einsum(tmp1, (0, 1), t2_OOvV, (2, 3, 1, 4), (0, 2, 3, 4)) + tmp188 = einsum(tmp187, (0, 1, 2, 3), t1_oV, (0, 4), (1, 2, 4, 3)) + del tmp187 + tmp260 = einsum(v.Vvov, (0, 1, 2, 3), t2_OoVv, (4, 2, 5, 1), (3, 4, 5, 0)) + tmp292 = einsum(t1_Ov, (0, 1), v.VvVv, (2, 3, 4, 1), (3, 0, 2, 4)) tmp300 = einsum(v.OvOv, (0, 1, 2, 3), t2_OOvv, (4, 5, 1, 3), (4, 5, 0, 2)) - tmp227 = einsum(tmp218, (0, 1, 2, 3), t2_OoVv, (4, 0, 5, 2), (1, 4, 3, 5)) - tmp167 = einsum(t1_Ov, (0, 1), v.OvOO, (2, 1, 3, 4), (0, 2, 3, 4)) + tmp167 = einsum(v.OvOO, (0, 1, 2, 3), t1_Ov, (4, 1), (4, 0, 2, 3)) tmp219 = einsum(tmp218, (0, 1, 2, 3), t2_OovV, (4, 0, 2, 5), (1, 4, 3, 5)) + tmp149 = einsum(t1_oV, (0, 1), tmp148, (0, 2), (1, 2)) + del tmp148 + tmp285 = einsum(t1_Ov, (0, 1), v.OvVV, (2, 1, 3, 4), (0, 2, 3, 4)) + tmp53 = einsum(v.OVVo, (0, 1, 2, 3), t3, (4, 5, 0, 2, 6, 1), (3, 4, 5, 6)) + tmp203 = einsum(v.OvOV, (0, 1, 2, 3), t2_OOvV, (4, 5, 1, 6), (4, 5, 0, 2, 6, 3)) + tmp227 = einsum(t2_OoVv, (0, 1, 2, 3), tmp218, (1, 4, 3, 5), (4, 0, 5, 2)) + tmp161 = einsum(v.VoVv, (0, 1, 2, 3), t1, (1, 3), (0, 2)) + tmp214 = einsum(tmp14, (0, 1), t2_OOvV, (2, 3, 1, 4), (0, 2, 3, 4)) + tmp281 = einsum(v.OoVv, (0, 1, 2, 3), t1_Ov, (4, 3), (1, 4, 0, 2)) + tmp223 = einsum(v.OVVo, (0, 1, 2, 3), t3, (4, 5, 0, 6, 1, 2), (3, 4, 5, 6)) + tmp201 = einsum(t2_OOvV, (0, 1, 2, 3), f.ov, (4, 2), (4, 0, 1, 3)) + tmp265 = einsum(t1_oV, (0, 1), v.VoVV, (2, 0, 3, 4), (1, 2, 3, 4)) + tmp239 = einsum(v.OvVo, (0, 1, 2, 3), t2_OovV, (4, 3, 1, 5), (4, 0, 5, 2)) + tmp142 = einsum(v.Voov, (0, 1, 2, 3), t2_ooVv, (2, 1, 4, 3), (4, 0)) + tmp155 = einsum(tmp154, (0, 1, 2, 3), t1_oV, (0, 4), (1, 4, 3, 2)) + del tmp154 + tmp165 = einsum(f.Ov, (0, 1), t1_Ov, (2, 1), (0, 2)) + tmp308 = einsum(t1_oV, (0, 1), tmp307, (0, 2, 3, 4), (2, 4, 3, 1)) + del tmp307 + tmp171 = einsum(t2_Oovv, (0, 1, 2, 3), v.Ovov, (4, 2, 1, 3), (0, 4)) + tmp252 = einsum(v.OVov, (0, 1, 2, 3), t2_OovV, (4, 2, 3, 5), (4, 0, 5, 1)) + tmp130 = einsum(t1_oV, (0, 1), v.OOVo, (2, 3, 4, 0), (2, 3, 1, 4)) + tmp128 = einsum(t1_oV, (0, 1), v.OoOV, (2, 0, 3, 4), (2, 3, 1, 4)) + tmp179 = einsum(t1_oV, (0, 1), tmp176, (0, 2, 3, 4), (2, 3, 4, 1)) tmp279 = einsum(t2_OovV, (0, 1, 2, 3), tmp173, (1, 4, 2, 5), (4, 5, 0, 3)) - tmp265 = einsum(v.oVVV, (0, 1, 2, 3), t1_oV, (0, 4), (4, 1, 2, 3)) - tmp210 = einsum(t2_OoVv, (0, 1, 2, 3), v.Ooov, (4, 5, 1, 3), (5, 0, 4, 2)) - tmp242 = einsum(v.OVov, (0, 1, 2, 3), t1_oV, (2, 4), (3, 0, 4, 1)) - tmp258 = einsum(t2_OoVv, (0, 1, 2, 3), v.ovVv, (1, 3, 4, 5), (5, 0, 2, 4)) - tmp221 = einsum(t2_OovV, (0, 1, 2, 3), v.Ooov, (4, 1, 5, 2), (5, 0, 4, 3)) - tmp126 = einsum(t1, (0, 1), v.OOov, (2, 3, 0, 1), (2, 3)) - tmp124 = einsum(t1, (0, 1), v.OoOv, (2, 0, 3, 1), (2, 3)) - tmp275 = einsum(t2_OoVv, (0, 1, 2, 3), tmp173, (1, 4, 3, 5), (4, 5, 0, 2)) + tmp126 = einsum(v.OOov, (0, 1, 2, 3), t1, (2, 3), (0, 1)) + tmp246 = einsum(t1_oV, (0, 1), v.OoVv, (2, 0, 3, 4), (4, 2, 1, 3)) + tmp254 = einsum(t2_OoVv, (0, 1, 2, 3), v.OVov, (4, 5, 1, 3), (0, 4, 2, 5)) + tmp283 = einsum(t2_OovV, (0, 1, 2, 3), tmp173, (4, 1, 2, 5), (4, 5, 0, 3)) + tmp229 = einsum(tmp218, (0, 1, 2, 3), t2_OoVv, (4, 0, 5, 1), (2, 4, 3, 5)) + tmp159 = einsum(v.VVov, (0, 1, 2, 3), t1, (2, 3), (0, 1)) + tmp55 = einsum(v.OVVo, (0, 1, 2, 3), t3, (4, 5, 0, 6, 2, 1), (3, 4, 5, 6)) + tmp144 = einsum(t2_ooVv, (0, 1, 2, 3), v.Voov, (4, 0, 1, 3), (2, 4)) + tmp146 = einsum(t2_ooVV, (0, 1, 2, 3), v.VoVo, (4, 1, 5, 0), (3, 2, 4, 5)) + tmp140 = einsum(f.Vo, (0, 1), t1_oV, (1, 2), (0, 2)) + tmp132 = einsum(t2_ooVV, (0, 1, 2, 3), v.Ooov, (4, 0, 1, 5), (5, 4, 3, 2)) + tmp244 = einsum(v.Vvov, (0, 1, 2, 3), t2_OovV, (4, 2, 1, 5), (3, 4, 5, 0)) tmp295 = einsum(t1_oV, (0, 1), tmp294, (0, 2, 3, 4), (2, 3, 1, 4)) + del tmp294 + tmp135 = einsum(t1_oV, (0, 1), tmp134, (0, 2, 3, 4), (2, 3, 1, 4)) + del tmp134 tmp302 = einsum(tmp218, (0, 1, 2, 3), t2_OOvv, (4, 5, 1, 2), (0, 4, 5, 3)) - tmp182 = einsum(t1_Ov, (0, 1), tmp181, (1, 2), (0, 2)) + del tmp218 + tmp242 = einsum(t1_oV, (0, 1), v.OVov, (2, 3, 0, 4), (4, 2, 1, 3)) + tmp185 = einsum(tmp184, (0, 1), t1_Ov, (2, 0), (2, 1)) + del tmp184 + tmp225 = einsum(v.OVVo, (0, 1, 2, 3), t3, (4, 0, 5, 6, 1, 2), (3, 4, 5, 6)) + tmp258 = einsum(v.Vvov, (0, 1, 2, 3), t2_OoVv, (4, 2, 5, 3), (1, 4, 5, 0)) + tmp298 = einsum(t1_oV, (0, 1), tmp297, (0, 2, 3, 4), (2, 3, 1, 4)) + del tmp297 + tmp182 = einsum(tmp181, (0, 1), t1_Ov, (2, 0), (2, 1)) del tmp181 - tmp254 = einsum(v.OVov, (0, 1, 2, 3), t2_OoVv, (4, 2, 5, 3), (4, 0, 5, 1)) - tmp308 = einsum(t1_oV, (0, 1), tmp307, (0, 2, 3, 4), (2, 4, 3, 1)) - del tmp307 - tmp130 = einsum(v.OOoV, (0, 1, 2, 3), t1_oV, (2, 4), (0, 1, 4, 3)) - tmp155 = einsum(t1_oV, (0, 1), tmp154, (0, 2, 3, 4), (2, 1, 4, 3)) - del tmp154 - tmp188 = einsum(tmp187, (0, 1, 2, 3), t1_oV, (0, 4), (1, 2, 4, 3)) - del tmp187 - tmp283 = einsum(t2_OovV, (0, 1, 2, 3), tmp173, (4, 1, 2, 5), (4, 5, 0, 3)) - tmp305 = einsum(tmp304, (0, 1, 2, 3), t1_Ov, (4, 0), (1, 4, 3, 2)) - del tmp304 - tmp252 = einsum(v.OVov, (0, 1, 2, 3), t2_OovV, (4, 2, 3, 5), (4, 0, 5, 1)) - tmp244 = einsum(v.ovVv, (0, 1, 2, 3), t2_OovV, (4, 0, 3, 5), (1, 4, 5, 2)) - tmp203 = einsum(t2_OOvV, (0, 1, 2, 3), v.OvOV, (4, 2, 5, 6), (0, 1, 4, 5, 3, 6)) - tmp191 = einsum(tmp190, (0, 1, 2, 3), t1_oV, (0, 4), (1, 2, 4, 3)) - tmp285 = einsum(t1_Ov, (0, 1), v.OvVV, (2, 1, 3, 4), (0, 2, 3, 4)) - tmp146 = einsum(v.oVoV, (0, 1, 2, 3), t2_ooVV, (0, 2, 4, 5), (5, 4, 3, 1)) - tmp142 = einsum(t2_oovV, (0, 1, 2, 3), v.oVov, (0, 4, 1, 2), (3, 4)) + tmp124 = einsum(v.OoOv, (0, 1, 2, 3), t1, (1, 3), (0, 2)) + tmp169 = einsum(t2_Oovv, (0, 1, 2, 3), v.Ovov, (4, 3, 1, 2), (0, 4)) + tmp289 = einsum(v.OVVv, (0, 1, 2, 3), t1_Ov, (4, 3), (4, 0, 2, 1)) tmp177 = einsum(tmp176, (0, 1, 2, 3), t1_oV, (1, 4), (0, 2, 3, 4)) - tmp281 = einsum(t1_Ov, (0, 1), v.OoVv, (2, 3, 4, 1), (3, 0, 2, 4)) - tmp201 = einsum(t2_OOvV, (0, 1, 2, 3), f.ov, (4, 2), (4, 0, 1, 3)) - tmp231 = einsum(tmp218, (0, 1, 2, 3), t2_OovV, (4, 0, 1, 5), (2, 4, 3, 5)) - tmp122 = einsum(t1_oV, (0, 1), v.OoOo, (2, 3, 4, 0), (3, 2, 4, 1)) - tmp179 = einsum(t1_oV, (0, 1), tmp176, (0, 2, 3, 4), (2, 3, 4, 1)) del tmp176 - tmp174 = einsum(tmp173, (0, 1, 2, 3), t2_ooVV, (1, 0, 4, 5), (2, 3, 5, 4)) - tmp208 = einsum(t2_OoVv, (0, 1, 2, 3), v.Ooov, (4, 1, 5, 3), (5, 0, 4, 2)) - tmp161 = einsum(v.oVVv, (0, 1, 2, 3), t1, (0, 3), (1, 2)) - tmp149 = einsum(t1_oV, (0, 1), tmp148, (0, 2), (1, 2)) - del tmp148 - tmp289 = einsum(t1_Ov, (0, 1), v.OVVv, (2, 3, 4, 1), (0, 2, 4, 3)) - tmp171 = einsum(v.Ovov, (0, 1, 2, 3), t2_Oovv, (4, 2, 1, 3), (4, 0)) - tmp55 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OVoV, (2, 5, 6, 4), (6, 0, 1, 3)) - tmp298 = einsum(tmp297, (0, 1, 2, 3), t1_oV, (0, 4), (1, 2, 4, 3)) - del tmp297 - tmp311 = einsum(v.ovVv, (0, 1, 2, 3), t2_OOvv, (4, 5, 1, 3), (0, 4, 5, 2)) - tmp239 = einsum(v.OvoV, (0, 1, 2, 3), t2_OovV, (4, 2, 1, 5), (4, 0, 5, 3)) - tmp214 = einsum(t2_OOvV, (0, 1, 2, 3), tmp14, (4, 2), (4, 0, 1, 3)) + tmp174 = einsum(tmp173, (0, 1, 2, 3), t2_ooVV, (0, 1, 4, 5), (2, 3, 4, 5)) + tmp262 = einsum(v.Vvov, (0, 1, 2, 3), t2_OovV, (4, 2, 3, 5), (1, 4, 5, 0)) + tmp250 = einsum(v.OvVo, (0, 1, 2, 3), t2_OoVv, (4, 3, 5, 1), (4, 0, 5, 2)) tmp277 = einsum(t2_OoVv, (0, 1, 2, 3), tmp173, (4, 1, 3, 5), (4, 5, 0, 2)) del tmp173 - tmp185 = einsum(tmp184, (0, 1), t1_Ov, (2, 0), (2, 1)) - del tmp184 - tmp165 = einsum(t1_Ov, (0, 1), f.Ov, (2, 1), (2, 0)) - tmp53 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OVoV, (2, 5, 6, 3), (6, 0, 1, 4)) - tmp250 = einsum(v.OvoV, (0, 1, 2, 3), t2_OoVv, (4, 2, 5, 1), (4, 0, 5, 3)) - tmp313 = einsum(t1_Ov, (0, 1), tmp294, (2, 1, 3, 4), (2, 0, 3, 4)) - del tmp294 - tmp169 = einsum(v.Ovov, (0, 1, 2, 3), t2_Oovv, (4, 2, 3, 1), (4, 0)) - tmp229 = einsum(t2_OoVv, (0, 1, 2, 3), tmp218, (1, 3, 4, 5), (4, 0, 5, 2)) - del tmp218 - tmp262 = einsum(v.ovVv, (0, 1, 2, 3), t2_OovV, (4, 0, 1, 5), (3, 4, 5, 2)) - tmp246 = einsum(v.OoVv, (0, 1, 2, 3), t1_oV, (1, 4), (3, 0, 4, 2)) - tmp216 = einsum(t2_OOvV, (0, 1, 2, 3), tmp1, (4, 2), (4, 0, 1, 3)) - tmp225 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OVoV, (1, 4, 6, 5), (6, 0, 2, 3)) - tmp128 = einsum(v.OoOV, (0, 1, 2, 3), t1_oV, (1, 4), (0, 2, 4, 3)) - tmp135 = einsum(tmp134, (0, 1, 2, 3), t1_oV, (0, 4), (1, 2, 4, 3)) - del tmp134 - tmp132 = einsum(t2_ooVV, (0, 1, 2, 3), v.Ooov, (4, 1, 0, 5), (5, 4, 2, 3)) - tmp159 = einsum(t1, (0, 1), v.ovVV, (0, 1, 2, 3), (2, 3)) - tmp144 = einsum(t2_oovV, (0, 1, 2, 3), v.ovoV, (0, 2, 1, 4), (3, 4)) - tmp152 = einsum(tmp151, (0, 1), t1_oV, (0, 2), (2, 1)) - del tmp151 tmp194 = einsum(t1_oV, (0, 1), tmp193, (0, 2, 3, 4), (2, 3, 4, 1)) del tmp193 - tmp260 = einsum(t2_OoVv, (0, 1, 2, 3), v.ovVv, (1, 4, 5, 3), (4, 0, 2, 5)) - tmp223 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OVoV, (2, 4, 6, 5), (6, 0, 1, 3)) - tmp140 = einsum(t1_oV, (0, 1), f.oV, (0, 2), (2, 1)) - tmp292 = einsum(t1_Ov, (0, 1), v.VvVv, (2, 3, 4, 1), (3, 0, 2, 4)) - tmp64 = einsum(t1, (0, 1), v.OvOV, (2, 1, 3, 4), (0, 2, 3, 4)) + tmp152 = einsum(t1_oV, (0, 1), tmp151, (0, 2), (1, 2)) + del tmp151 + tmp311 = einsum(v.Vvov, (0, 1, 2, 3), t2_OOvv, (4, 5, 3, 1), (2, 4, 5, 0)) + tmp208 = einsum(t2_OoVv, (0, 1, 2, 3), v.Ooov, (4, 1, 5, 3), (5, 0, 4, 2)) + tmp191 = einsum(t1_oV, (0, 1), tmp190, (0, 2, 3, 4), (2, 3, 1, 4)) + tmp305 = einsum(tmp304, (0, 1, 2, 3), t1_Ov, (4, 0), (1, 4, 3, 2)) + del tmp304 + tmp210 = einsum(t2_OoVv, (0, 1, 2, 3), v.Ooov, (4, 5, 1, 3), (5, 0, 4, 2)) + tmp221 = einsum(t2_OovV, (0, 1, 2, 3), v.Ooov, (4, 1, 5, 2), (5, 0, 4, 3)) + tmp212 = einsum(v.Ooov, (0, 1, 2, 3), t2_OovV, (4, 2, 3, 5), (1, 4, 0, 5)) + tmp64 = einsum(v.OvOV, (0, 1, 2, 3), t1, (4, 1), (4, 0, 2, 3)) tmp38 = einsum(v.OVov, (0, 1, 2, 3), t1, (2, 3), (0, 1)) - tmp36 = einsum(v.OvoV, (0, 1, 2, 3), t1, (2, 1), (0, 3)) - tmp76 = einsum(v.ovvv, (0, 1, 2, 3), t1, (4, 1), (4, 0, 2, 3)) - tmp45 = einsum(t2, (0, 1, 2, 3), f.ov, (4, 3), (4, 0, 1, 2)) - tmp83 = einsum(v.ooov, (0, 1, 2, 3), t2, (4, 1, 3, 5), (4, 0, 2, 5)) - tmp101 = einsum(t1, (0, 1), tmp75, (2, 3, 1, 4), (0, 2, 3, 4)) - tmp99 = einsum(tmp4, (0, 1, 2, 3), t1, (4, 3), (0, 4, 2, 1)) - tmp7 = einsum(t2, (0, 1, 2, 3), v.ovov, (4, 2, 1, 3), (0, 4)) - tmp48 = einsum(t2, (0, 1, 2, 3), v.ovov, (4, 5, 1, 2), (0, 4, 3, 5)) - tmp26 = einsum(v.ooov, (0, 1, 2, 3), t1, (1, 3), (0, 2)) - tmp114 = einsum(v.ovvv, (0, 1, 2, 3), t2, (4, 5, 1, 3), (4, 5, 0, 2)) + tmp36 = einsum(v.OvVo, (0, 1, 2, 3), t1, (3, 1), (0, 2)) + tmp79 = einsum(t2, (0, 1, 2, 3), tmp4, (4, 5, 1, 2), (4, 0, 5, 3)) + tmp94 = einsum(v.ovov, (0, 1, 2, 3), t2, (4, 2, 5, 1), (4, 0, 5, 3)) + tmp87 = einsum(t2, (0, 1, 2, 3), v.ooov, (4, 5, 1, 3), (0, 4, 5, 2)) + tmp27 = einsum(v.ooov, (0, 1, 2, 3), t1, (2, 3), (0, 1)) + tmp83 = einsum(t2, (0, 1, 2, 3), v.ooov, (4, 1, 5, 2), (0, 4, 5, 3)) + tmp45 = einsum(f.ov, (0, 1), t2, (2, 3, 4, 1), (0, 2, 3, 4)) + tmp91 = einsum(v.ooov, (0, 1, 2, 3), t2, (4, 2, 3, 5), (4, 0, 1, 5)) + tmp118 = einsum(tmp116, (0, 1, 2, 3), t1, (3, 4), (0, 1, 2, 4)) tmp89 = einsum(t2, (0, 1, 2, 3), v.ooov, (4, 1, 5, 3), (0, 4, 5, 2)) - tmp11 = einsum(v.ovvv, (0, 1, 2, 3), t1, (0, 1), (2, 3)) - tmp118 = einsum(tmp116, (0, 1, 2, 3), t1, (2, 4), (1, 0, 3, 4)) + tmp11 = einsum(t1, (0, 1), v.ovvv, (0, 1, 2, 3), (2, 3)) + tmp40 = einsum(v.ovov, (0, 1, 2, 3), t2, (0, 2, 4, 3), (4, 1)) + tmp109 = einsum(tmp4, (0, 1, 2, 3), t2, (4, 1, 5, 3), (0, 4, 2, 5)) + tmp13 = einsum(t1, (0, 1), v.ovvv, (0, 2, 3, 1), (2, 3)) + tmp111 = einsum(t2, (0, 1, 2, 3), tmp4, (4, 1, 5, 2), (4, 0, 5, 3)) + tmp9 = einsum(t2, (0, 1, 2, 3), v.ovov, (4, 3, 1, 2), (0, 4)) + tmp15 = einsum(tmp14, (0, 1), t1, (2, 1), (2, 0)) + tmp0 = einsum(t1, (0, 1), f.ov, (2, 1), (2, 0)) + tmp76 = einsum(t1, (0, 1), v.ovvv, (2, 1, 3, 4), (0, 2, 3, 4)) + tmp47 = einsum(t2, (0, 1, 2, 3), v.ovov, (4, 5, 1, 3), (0, 4, 2, 5)) + tmp48 = einsum(t2, (0, 1, 2, 3), v.ovov, (4, 5, 1, 2), (0, 4, 3, 5)) + tmp107 = einsum(tmp4, (0, 1, 2, 3), t2, (4, 2, 5, 3), (0, 4, 1, 5)) tmp57 = einsum(t2, (0, 1, 2, 3), tmp14, (4, 3), (0, 1, 4, 2)) - tmp107 = einsum(t2, (0, 1, 2, 3), tmp4, (4, 5, 1, 3), (4, 0, 5, 2)) + tmp7 = einsum(t2, (0, 1, 2, 3), v.ovov, (4, 2, 1, 3), (0, 4)) + tmp16 = einsum(tmp1, (0, 1), t1, (2, 1), (2, 0)) + tmp26 = einsum(v.ooov, (0, 1, 2, 3), t1, (1, 3), (0, 2)) + tmp114 = einsum(t2, (0, 1, 2, 3), v.ovvv, (4, 2, 5, 3), (0, 1, 4, 5)) + tmp85 = einsum(t2, (0, 1, 2, 3), v.ovov, (4, 2, 1, 5), (0, 4, 3, 5)) + tmp99 = einsum(t1, (0, 1), tmp4, (2, 3, 4, 1), (2, 0, 4, 3)) tmp73 = einsum(v.oovv, (0, 1, 2, 3), t1, (4, 3), (4, 0, 1, 2)) - tmp94 = einsum(t2, (0, 1, 2, 3), v.ovov, (4, 3, 1, 5), (0, 4, 2, 5)) - tmp15 = einsum(t1, (0, 1), tmp14, (2, 1), (0, 2)) - tmp9 = einsum(t2, (0, 1, 2, 3), v.ovov, (4, 3, 1, 2), (0, 4)) - tmp91 = einsum(v.ooov, (0, 1, 2, 3), t2, (4, 2, 3, 5), (4, 0, 1, 5)) - tmp71 = einsum(tmp62, (0, 1, 2, 3), t1, (3, 4), (0, 2, 1, 4)) + tmp71 = einsum(t1, (0, 1), tmp62, (2, 3, 4, 0), (2, 4, 3, 1)) tmp59 = einsum(t2, (0, 1, 2, 3), tmp1, (4, 3), (0, 1, 4, 2)) - tmp42 = einsum(v.ovov, (0, 1, 2, 3), t2, (0, 2, 3, 4), (4, 1)) - tmp87 = einsum(t2, (0, 1, 2, 3), v.ooov, (4, 5, 1, 3), (0, 4, 5, 2)) - tmp79 = einsum(t2, (0, 1, 2, 3), tmp4, (4, 5, 1, 2), (4, 0, 5, 3)) - tmp85 = einsum(t2, (0, 1, 2, 3), v.ovov, (4, 2, 1, 5), (0, 4, 3, 5)) - tmp40 = einsum(t2, (0, 1, 2, 3), v.ovov, (0, 2, 1, 4), (3, 4)) - tmp13 = einsum(t1, (0, 1), v.ovvv, (0, 2, 3, 1), (2, 3)) - tmp109 = einsum(t2, (0, 1, 2, 3), tmp4, (4, 1, 5, 3), (4, 0, 5, 2)) - tmp111 = einsum(tmp4, (0, 1, 2, 3), t2, (4, 1, 3, 5), (0, 4, 2, 5)) - tmp47 = einsum(t2, (0, 1, 2, 3), v.ovov, (4, 5, 1, 3), (0, 4, 2, 5)) - tmp27 = einsum(t1, (0, 1), v.ooov, (2, 3, 0, 1), (2, 3)) - tmp16 = einsum(t1, (0, 1), tmp1, (2, 1), (0, 2)) - tmp0 = einsum(t1, (0, 1), f.ov, (2, 1), (2, 0)) - tmp213 = einsum(t2_OoVV, (0, 1, 2, 3), tmp212, (1, 4, 5, 6), (4, 0, 5, 6, 3, 2)) - del tmp212 - tmp301 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp300, (6, 7, 1, 2), (7, 6, 0, 3, 5, 4)) - tmp228 = einsum(tmp227, (0, 1, 2, 3), t2_OOvV, (4, 5, 0, 6), (1, 4, 5, 2, 3, 6)) + tmp101 = einsum(tmp75, (0, 1, 2, 3), t1, (4, 2), (4, 0, 1, 3)) + tmp42 = einsum(v.ovov, (0, 1, 2, 3), t2, (0, 2, 4, 1), (4, 3)) + tmp237 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OOVV, (6, 2, 7, 4), (0, 1, 6, 3, 5, 7)) + tmp123 = einsum(t2_OoVV, (0, 1, 2, 3), tmp122, (1, 4, 5, 6), (0, 5, 4, 6, 3, 2)) + del tmp122 + tmp314 = einsum(t2_OoVV, (0, 1, 2, 3), tmp313, (1, 4, 5, 6), (4, 5, 0, 3, 2, 6)) + del tmp313 + tmp232 = einsum(t2_OOvV, (0, 1, 2, 3), tmp231, (2, 4, 5, 6), (4, 0, 1, 5, 6, 3)) + del tmp231 + tmp276 = einsum(t2_OoVV, (0, 1, 2, 3), tmp275, (1, 4, 5, 6), (4, 5, 0, 6, 3, 2)) + del tmp275 + tmp217 = einsum(t2_OoVV, (0, 1, 2, 3), tmp216, (1, 4, 5, 6), (0, 4, 5, 3, 2, 6)) + del tmp216 + tmp199 = einsum(tmp188, (0, 1, 2, 3), t3, (4, 1, 5, 6, 7, 3), (0, 4, 5, 2, 6, 7)) + tmp261 = einsum(tmp260, (0, 1, 2, 3), t2_OOvV, (4, 5, 0, 6), (1, 4, 5, 2, 6, 3)) + del tmp260 + tmp293 = einsum(t2_OOvV, (0, 1, 2, 3), tmp292, (2, 4, 5, 6), (4, 0, 1, 3, 6, 5)) + del tmp292 + tmp310 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp300, (6, 7, 0, 2), (7, 6, 1, 3, 5, 4)) * -1 + tmp168 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp167, (6, 2, 7, 0), (6, 1, 7, 3, 5, 4)) * -1 + tmp220 = einsum(t2_OOvV, (0, 1, 2, 3), tmp219, (2, 4, 5, 6), (4, 0, 1, 5, 6, 3)) + del tmp219 + tmp120 = einsum(t3, (0, 1, 2, 3, 4, 5), f.OO, (6, 2), (6, 0, 1, 3, 5, 4)) + tmp150 = einsum(tmp149, (0, 1), t3, (2, 3, 4, 5, 6, 1), (2, 4, 3, 0, 5, 6)) + tmp286 = einsum(tmp285, (0, 1, 2, 3), t3, (4, 5, 1, 6, 3, 7), (0, 4, 5, 6, 7, 2)) + tmp206 = einsum(tmp53, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (4, 1, 2, 6, 5, 3)) + tmp205 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp203, (6, 7, 1, 2, 8, 5), (6, 7, 0, 8, 3, 4)) + tmp228 = einsum(t2_OOvV, (0, 1, 2, 3), tmp227, (2, 4, 5, 6), (4, 0, 1, 5, 6, 3)) del tmp227 - tmp310 = einsum(tmp300, (0, 1, 2, 3), t3, (2, 4, 3, 5, 6, 7), (1, 0, 4, 5, 7, 6)) * -1 - del tmp300 + tmp137 = einsum(v.OOOO, (0, 1, 2, 3), t3, (4, 3, 1, 5, 6, 7), (4, 0, 2, 5, 7, 6)) + tmp233 = einsum(v.OoOV, (0, 1, 2, 3), t2_OoVV, (4, 1, 5, 6), (4, 0, 2, 6, 5, 3)) + tmp162 = einsum(tmp161, (0, 1), t3, (2, 3, 4, 5, 6, 0), (2, 4, 3, 5, 6, 1)) + tmp215 = einsum(tmp214, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (4, 1, 2, 6, 5, 3)) + del tmp214 + tmp157 = einsum(t3, (0, 1, 2, 3, 4, 5), f.VV, (6, 5), (0, 2, 1, 6, 3, 4)) + tmp282 = einsum(tmp281, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (1, 4, 2, 6, 5, 3)) + del tmp281 + tmp158 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OVOV, (6, 7, 2, 5), (0, 1, 6, 3, 4, 7)) + tmp224 = einsum(tmp223, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (4, 1, 2, 6, 5, 3)) + del tmp223 + tmp202 = einsum(tmp201, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (4, 1, 2, 6, 5, 3)) + del tmp201 + tmp266 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp265, (6, 5, 7, 3), (0, 2, 1, 6, 4, 7)) * -1 + tmp240 = einsum(tmp239, (0, 1, 2, 3), t3, (4, 5, 1, 6, 3, 7), (0, 4, 5, 2, 6, 7)) + tmp163 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp142, (6, 4), (0, 2, 1, 6, 3, 5)) + tmp268 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp155, (6, 7, 3, 5), (0, 2, 1, 6, 7, 4)) + tmp166 = einsum(tmp165, (0, 1), t3, (2, 3, 0, 4, 5, 6), (1, 2, 3, 4, 6, 5)) + tmp309 = einsum(t2_OoVV, (0, 1, 2, 3), tmp308, (1, 4, 5, 6), (5, 4, 0, 6, 3, 2)) + del tmp308 + tmp172 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp171, (6, 2), (6, 0, 1, 3, 5, 4)) + tmp253 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp252, (6, 2, 7, 5), (6, 0, 1, 7, 3, 4)) tmp196 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp167, (6, 1, 7, 2), (6, 0, 7, 3, 5, 4)) - tmp220 = einsum(tmp219, (0, 1, 2, 3), t2_OOvV, (4, 5, 0, 6), (1, 4, 5, 2, 3, 6)) - del tmp219 + tmp257 = einsum(tmp239, (0, 1, 2, 3), t3, (4, 1, 5, 6, 7, 3), (0, 4, 5, 2, 6, 7)) + tmp131 = einsum(tmp130, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 3), (4, 5, 0, 2, 6, 7)) + tmp129 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp128, (6, 2, 7, 5), (0, 1, 6, 7, 3, 4)) + tmp180 = einsum(tmp179, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (1, 4, 2, 3, 6, 5)) + del tmp179 + tmp236 = einsum(tmp203, (0, 1, 2, 3, 4, 5), t3, (6, 2, 3, 7, 5, 8), (0, 1, 6, 4, 7, 8)) tmp280 = einsum(t2_OoVV, (0, 1, 2, 3), tmp279, (1, 4, 5, 6), (4, 5, 0, 6, 3, 2)) del tmp279 - tmp266 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp265, (6, 3, 7, 5), (0, 2, 1, 6, 4, 7)) - tmp211 = einsum(t2_OoVV, (0, 1, 2, 3), tmp210, (1, 4, 5, 6), (4, 0, 5, 6, 3, 2)) - del tmp210 - tmp243 = einsum(tmp242, (0, 1, 2, 3), t2_OOvV, (4, 5, 0, 6), (4, 5, 1, 2, 6, 3)) - del tmp242 - tmp259 = einsum(t2_OOvV, (0, 1, 2, 3), tmp258, (2, 4, 5, 6), (4, 0, 1, 5, 3, 6)) - del tmp258 - tmp222 = einsum(t2_OoVV, (0, 1, 2, 3), tmp221, (1, 4, 5, 6), (4, 0, 5, 6, 3, 2)) - del tmp221 + tmp270 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp265, (6, 5, 7, 4), (0, 2, 1, 6, 3, 7)) tmp138 = einsum(tmp126, (0, 1), t3, (2, 1, 3, 4, 5, 6), (2, 3, 0, 4, 6, 5)) - tmp125 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp124, (6, 2), (0, 1, 6, 3, 5, 4)) - tmp269 = einsum(t3, (0, 1, 2, 3, 4, 5), v.VVVV, (6, 5, 7, 4), (0, 2, 1, 3, 6, 7)) - tmp276 = einsum(t2_OoVV, (0, 1, 2, 3), tmp275, (1, 4, 5, 6), (4, 5, 0, 6, 3, 2)) - del tmp275 - tmp296 = einsum(tmp295, (0, 1, 2, 3), t2_OOvV, (4, 5, 0, 6), (1, 4, 5, 2, 6, 3)) + tmp247 = einsum(tmp246, (0, 1, 2, 3), t2_OOvV, (4, 5, 0, 6), (4, 5, 1, 2, 6, 3)) + del tmp246 + tmp271 = einsum(tmp254, (0, 1, 2, 3), t3, (4, 5, 1, 6, 3, 7), (0, 4, 5, 2, 6, 7)) + tmp127 = einsum(tmp126, (0, 1), t3, (2, 3, 1, 4, 5, 6), (2, 3, 0, 4, 6, 5)) + del tmp126 + tmp284 = einsum(t2_OoVV, (0, 1, 2, 3), tmp283, (1, 4, 5, 6), (4, 5, 0, 6, 3, 2)) + del tmp283 + tmp139 = einsum(tmp130, (0, 1, 2, 3), t3, (4, 1, 5, 6, 7, 3), (4, 5, 0, 2, 6, 7)) + tmp230 = einsum(t2_OOvV, (0, 1, 2, 3), tmp229, (2, 4, 5, 6), (4, 0, 1, 5, 6, 3)) + del tmp229 + tmp143 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp142, (6, 5), (0, 2, 1, 6, 3, 4)) + del tmp142 + tmp160 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp159, (6, 5), (0, 2, 1, 3, 4, 6)) + tmp198 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp171, (6, 1), (6, 0, 2, 3, 5, 4)) + del tmp171 + tmp255 = einsum(tmp254, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 3), (0, 4, 5, 2, 6, 7)) + tmp207 = einsum(t2_OoVV, (0, 1, 2, 3), tmp55, (1, 4, 5, 6), (0, 4, 5, 3, 2, 6)) + tmp291 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp285, (6, 1, 7, 5), (6, 0, 2, 3, 4, 7)) + tmp200 = einsum(tmp167, (0, 1, 2, 3), t3, (4, 3, 1, 5, 6, 7), (0, 4, 2, 5, 7, 6)) + del tmp167 + tmp145 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp144, (6, 5), (0, 2, 1, 6, 3, 4)) + tmp267 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp146, (6, 7, 5, 3), (0, 2, 1, 7, 6, 4)) + tmp141 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp140, (5, 6), (0, 2, 1, 6, 3, 4)) + tmp133 = einsum(t2_OOvV, (0, 1, 2, 3), tmp132, (2, 4, 5, 6), (0, 1, 4, 5, 6, 3)) + del tmp132 + tmp164 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp144, (6, 4), (0, 2, 1, 6, 3, 5)) + del tmp144 + tmp301 = einsum(tmp300, (0, 1, 2, 3), t3, (4, 2, 3, 5, 6, 7), (1, 0, 4, 5, 7, 6)) + del tmp300 + tmp245 = einsum(t2_OOvV, (0, 1, 2, 3), tmp244, (2, 4, 5, 6), (4, 0, 1, 5, 3, 6)) + del tmp244 + tmp296 = einsum(t2_OOvV, (0, 1, 2, 3), tmp295, (2, 4, 5, 6), (4, 0, 1, 5, 3, 6)) del tmp295 + tmp204 = einsum(tmp203, (0, 1, 2, 3, 4, 5), t3, (3, 6, 2, 7, 8, 5), (0, 1, 6, 4, 7, 8)) * -1 + tmp136 = einsum(t2_OOvV, (0, 1, 2, 3), tmp135, (2, 4, 5, 6), (0, 1, 4, 5, 6, 3)) + del tmp135 + tmp287 = einsum(tmp188, (0, 1, 2, 3), t3, (4, 5, 1, 6, 3, 7), (0, 4, 5, 2, 6, 7)) + tmp147 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp146, (6, 7, 5, 4), (0, 2, 1, 6, 7, 3)) + del tmp146 tmp303 = einsum(t2_OoVV, (0, 1, 2, 3), tmp302, (1, 4, 5, 6), (5, 4, 0, 6, 3, 2)) del tmp302 + tmp243 = einsum(tmp242, (0, 1, 2, 3), t2_OOvV, (4, 5, 0, 6), (4, 5, 1, 2, 6, 3)) + del tmp242 + tmp186 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp185, (6, 2), (6, 0, 1, 3, 5, 4)) + tmp226 = einsum(t2_OoVV, (0, 1, 2, 3), tmp225, (1, 4, 5, 6), (0, 5, 4, 3, 2, 6)) * -1 + del tmp225 + tmp259 = einsum(t2_OOvV, (0, 1, 2, 3), tmp258, (2, 4, 5, 6), (4, 0, 1, 5, 3, 6)) + del tmp258 + tmp299 = einsum(t2_OOvV, (0, 1, 2, 3), tmp298, (2, 4, 5, 6), (4, 0, 1, 5, 3, 6)) + del tmp298 tmp183 = einsum(tmp182, (0, 1), t3, (2, 3, 1, 4, 5, 6), (0, 2, 3, 4, 6, 5)) - tmp271 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp254, (6, 2, 7, 4), (6, 0, 1, 7, 3, 5)) - tmp309 = einsum(tmp308, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (2, 1, 4, 3, 6, 5)) - del tmp308 - tmp131 = einsum(tmp130, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 3), (4, 5, 0, 2, 6, 7)) - tmp268 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp155, (6, 7, 5, 3), (0, 2, 1, 6, 7, 4)) * -1 - tmp287 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp188, (6, 2, 7, 4), (6, 0, 1, 7, 3, 5)) - tmp284 = einsum(t2_OoVV, (0, 1, 2, 3), tmp283, (1, 4, 5, 6), (4, 5, 0, 6, 3, 2)) - del tmp283 - tmp315 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp305, (6, 7, 0, 2), (7, 6, 1, 3, 5, 4)) * -1 - tmp253 = einsum(tmp252, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 3), (0, 4, 5, 2, 6, 7)) - tmp245 = einsum(tmp244, (0, 1, 2, 3), t2_OOvV, (4, 5, 0, 6), (1, 4, 5, 2, 6, 3)) - del tmp244 - tmp205 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp203, (6, 7, 1, 2, 8, 5), (6, 7, 0, 8, 3, 4)) - tmp192 = einsum(tmp191, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 3), (0, 4, 5, 2, 6, 7)) - tmp286 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp285, (6, 2, 7, 4), (6, 0, 1, 3, 5, 7)) - tmp147 = einsum(tmp146, (0, 1, 2, 3), t3, (4, 5, 6, 7, 2, 3), (4, 6, 5, 1, 0, 7)) + tmp121 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OOOO, (6, 0, 7, 2), (1, 6, 7, 3, 5, 4)) + tmp241 = einsum(v.OVVv, (0, 1, 2, 3), t2_OOvV, (4, 5, 3, 6), (4, 5, 0, 6, 2, 1)) + tmp125 = einsum(tmp124, (0, 1), t3, (2, 3, 1, 4, 5, 6), (2, 3, 0, 4, 6, 5)) + tmp156 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp155, (6, 7, 4, 5), (0, 2, 1, 7, 6, 3)) + del tmp155 + tmp235 = einsum(tmp203, (0, 1, 2, 3, 4, 5), t3, (6, 3, 2, 7, 5, 8), (0, 1, 6, 4, 7, 8)) + tmp170 = einsum(tmp169, (0, 1), t3, (2, 3, 1, 4, 5, 6), (0, 2, 3, 4, 6, 5)) tmp273 = einsum(tmp265, (0, 1, 2, 3), t3, (4, 5, 6, 7, 1, 3), (4, 6, 5, 0, 7, 2)) - tmp163 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp142, (6, 4), (0, 2, 1, 6, 3, 5)) - tmp121 = einsum(v.OOOO, (0, 1, 2, 3), t3, (3, 4, 1, 5, 6, 7), (4, 0, 2, 5, 7, 6)) * -1 - tmp178 = einsum(tmp177, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (1, 4, 2, 3, 6, 5)) + del tmp265 + tmp238 = einsum(tmp130, (0, 1, 2, 3), t3, (4, 5, 1, 6, 3, 7), (4, 5, 0, 2, 6, 7)) + tmp256 = einsum(v.OOVV, (0, 1, 2, 3), t3, (4, 1, 5, 6, 7, 3), (4, 5, 0, 6, 7, 2)) + tmp269 = einsum(t3, (0, 1, 2, 3, 4, 5), v.VVVV, (6, 5, 7, 4), (0, 2, 1, 3, 6, 7)) + tmp274 = einsum(t2_OoVV, (0, 1, 2, 3), tmp190, (1, 4, 5, 6), (4, 0, 5, 3, 2, 6)) + del tmp190 + tmp264 = einsum(t3, (0, 1, 2, 3, 4, 5), v.VVVV, (6, 5, 7, 3), (0, 2, 1, 4, 6, 7)) * -1 + tmp290 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp289, (6, 2, 7, 5), (6, 0, 1, 3, 4, 7)) + tmp178 = einsum(t2_OoVV, (0, 1, 2, 3), tmp177, (1, 4, 5, 6), (4, 0, 5, 6, 3, 2)) del tmp177 - tmp238 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp130, (6, 2, 7, 4), (0, 1, 6, 7, 3, 5)) - tmp306 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp305, (6, 7, 1, 2), (7, 6, 0, 3, 5, 4)) - del tmp305 - tmp282 = einsum(t2_OoVV, (0, 1, 2, 3), tmp281, (1, 4, 5, 6), (4, 0, 5, 3, 2, 6)) - del tmp281 - tmp202 = einsum(t2_OoVV, (0, 1, 2, 3), tmp201, (1, 4, 5, 6), (0, 4, 5, 3, 2, 6)) - del tmp201 - tmp127 = einsum(tmp126, (0, 1), t3, (2, 3, 1, 4, 5, 6), (2, 3, 0, 4, 6, 5)) - del tmp126 - tmp232 = einsum(tmp231, (0, 1, 2, 3), t2_OOvV, (4, 5, 0, 6), (1, 4, 5, 2, 3, 6)) - del tmp231 - tmp120 = einsum(f.OO, (0, 1), t3, (2, 3, 1, 4, 5, 6), (0, 2, 3, 4, 6, 5)) - tmp123 = einsum(t2_OoVV, (0, 1, 2, 3), tmp122, (1, 4, 5, 6), (0, 5, 4, 6, 3, 2)) - del tmp122 - tmp180 = einsum(tmp179, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (1, 4, 2, 3, 6, 5)) - del tmp179 - tmp175 = einsum(t2_OOvV, (0, 1, 2, 3), tmp174, (2, 4, 5, 6), (4, 0, 1, 6, 5, 3)) + tmp175 = einsum(tmp174, (0, 1, 2, 3), t2_OOvV, (4, 5, 0, 6), (1, 4, 5, 3, 2, 6)) del tmp174 - tmp288 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp285, (6, 2, 7, 5), (6, 0, 1, 3, 4, 7)) - tmp235 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp203, (6, 7, 2, 1, 8, 4), (6, 7, 0, 8, 3, 5)) - tmp209 = einsum(tmp208, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (1, 4, 2, 3, 6, 5)) - del tmp208 - tmp162 = einsum(tmp161, (0, 1), t3, (2, 3, 4, 5, 6, 0), (2, 4, 3, 5, 6, 1)) - tmp150 = einsum(tmp149, (0, 1), t3, (2, 3, 4, 5, 6, 1), (2, 4, 3, 0, 5, 6)) - tmp290 = einsum(tmp289, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 3), (0, 4, 5, 6, 7, 2)) - tmp198 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp171, (6, 1), (6, 0, 2, 3, 5, 4)) - tmp272 = einsum(tmp254, (0, 1, 2, 3), t3, (4, 1, 5, 6, 3, 7), (0, 4, 5, 2, 6, 7)) - tmp270 = einsum(tmp265, (0, 1, 2, 3), t3, (4, 5, 6, 7, 3, 1), (4, 6, 5, 0, 7, 2)) - del tmp265 - tmp267 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp146, (6, 7, 5, 3), (0, 2, 1, 7, 6, 4)) - del tmp146 - tmp207 = einsum(t2_OoVV, (0, 1, 2, 3), tmp55, (1, 4, 5, 6), (0, 4, 5, 3, 2, 6)) - tmp299 = einsum(tmp298, (0, 1, 2, 3), t2_OOvV, (4, 5, 0, 6), (1, 4, 5, 2, 6, 3)) - del tmp298 - tmp312 = einsum(t2_OoVV, (0, 1, 2, 3), tmp311, (1, 4, 5, 6), (4, 5, 0, 3, 2, 6)) - del tmp311 + tmp197 = einsum(tmp169, (0, 1), t3, (2, 1, 3, 4, 5, 6), (0, 2, 3, 4, 6, 5)) + del tmp169 + tmp263 = einsum(tmp262, (0, 1, 2, 3), t2_OOvV, (4, 5, 0, 6), (1, 4, 5, 2, 6, 3)) + del tmp262 + tmp251 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp250, (6, 2, 7, 5), (6, 0, 1, 7, 3, 4)) tmp249 = einsum(tmp239, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 3), (0, 4, 5, 2, 6, 7)) - tmp256 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OOVV, (6, 1, 7, 5), (0, 2, 6, 3, 4, 7)) - tmp237 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OOVV, (6, 2, 7, 4), (0, 1, 6, 3, 5, 7)) - tmp257 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp239, (6, 1, 7, 5), (6, 0, 2, 7, 3, 4)) - tmp255 = einsum(tmp254, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 3), (0, 4, 5, 2, 6, 7)) + tmp272 = einsum(tmp254, (0, 1, 2, 3), t3, (4, 1, 5, 6, 3, 7), (0, 4, 5, 2, 6, 7)) del tmp254 - tmp274 = einsum(t2_OoVV, (0, 1, 2, 3), tmp190, (1, 4, 5, 6), (4, 0, 5, 3, 2, 6)) - del tmp190 - tmp215 = einsum(tmp214, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (4, 1, 2, 6, 5, 3)) - del tmp214 tmp278 = einsum(t2_OoVV, (0, 1, 2, 3), tmp277, (1, 4, 5, 6), (4, 5, 0, 6, 3, 2)) del tmp277 - tmp204 = einsum(tmp203, (0, 1, 2, 3, 4, 5), t3, (3, 6, 2, 7, 8, 5), (0, 1, 6, 4, 7, 8)) * -1 - tmp186 = einsum(tmp185, (0, 1), t3, (2, 3, 1, 4, 5, 6), (0, 2, 3, 4, 6, 5)) - tmp166 = einsum(tmp165, (0, 1), t3, (2, 3, 0, 4, 5, 6), (1, 2, 3, 4, 6, 5)) - tmp172 = einsum(tmp171, (0, 1), t3, (2, 3, 1, 4, 5, 6), (0, 2, 3, 4, 6, 5)) - del tmp171 - tmp264 = einsum(t3, (0, 1, 2, 3, 4, 5), v.VVVV, (6, 3, 7, 5), (0, 2, 1, 4, 6, 7)) - tmp248 = einsum(v.OOVV, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 3), (4, 5, 0, 6, 7, 2)) - tmp206 = einsum(t2_OoVV, (0, 1, 2, 3), tmp53, (1, 4, 5, 6), (0, 4, 5, 3, 2, 6)) - tmp251 = einsum(tmp250, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 3), (0, 4, 5, 2, 6, 7)) - tmp314 = einsum(tmp313, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (1, 2, 4, 6, 5, 3)) - del tmp313 - tmp199 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp188, (6, 1, 7, 5), (6, 0, 2, 7, 3, 4)) - tmp170 = einsum(tmp169, (0, 1), t3, (2, 3, 1, 4, 5, 6), (0, 2, 3, 4, 6, 5)) - tmp137 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OOOO, (6, 2, 7, 1), (0, 6, 7, 3, 5, 4)) - tmp197 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp169, (6, 1), (6, 0, 2, 3, 5, 4)) - del tmp169 - tmp241 = einsum(v.OVVv, (0, 1, 2, 3), t2_OOvV, (4, 5, 3, 6), (4, 5, 0, 6, 2, 1)) + tmp195 = einsum(tmp194, (0, 1, 2, 3), t2_OOvV, (4, 5, 0, 6), (1, 4, 5, 3, 2, 6)) + del tmp194 + tmp153 = einsum(tmp152, (0, 1), t3, (2, 3, 4, 5, 6, 1), (2, 4, 3, 0, 5, 6)) + tmp288 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp285, (6, 2, 7, 5), (6, 0, 1, 3, 4, 7)) + tmp312 = einsum(tmp311, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (1, 2, 4, 6, 5, 3)) + del tmp311 + tmp209 = einsum(t2_OoVV, (0, 1, 2, 3), tmp208, (1, 4, 5, 6), (4, 0, 5, 6, 3, 2)) + del tmp208 + tmp192 = einsum(tmp191, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 3), (0, 4, 5, 2, 6, 7)) + tmp306 = einsum(tmp305, (0, 1, 2, 3), t3, (4, 2, 3, 5, 6, 7), (1, 0, 4, 5, 7, 6)) + tmp211 = einsum(tmp210, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (1, 4, 2, 3, 6, 5)) + del tmp210 tmp234 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp203, (6, 7, 2, 1, 8, 5), (6, 7, 0, 8, 3, 4)) - tmp236 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp203, (6, 7, 1, 2, 8, 4), (6, 7, 0, 8, 3, 5)) del tmp203 - tmp230 = einsum(t2_OOvV, (0, 1, 2, 3), tmp229, (2, 4, 5, 6), (4, 0, 1, 5, 6, 3)) - del tmp229 - tmp158 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OVOV, (6, 7, 2, 5), (0, 1, 6, 3, 4, 7)) - tmp263 = einsum(tmp262, (0, 1, 2, 3), t2_OOvV, (4, 5, 0, 6), (1, 4, 5, 2, 6, 3)) - del tmp262 - tmp247 = einsum(t2_OOvV, (0, 1, 2, 3), tmp246, (2, 4, 5, 6), (0, 1, 4, 5, 3, 6)) - del tmp246 - tmp217 = einsum(t2_OoVV, (0, 1, 2, 3), tmp216, (1, 4, 5, 6), (0, 4, 5, 3, 2, 6)) - del tmp216 - tmp291 = einsum(tmp285, (0, 1, 2, 3), t3, (4, 1, 5, 6, 7, 3), (0, 4, 5, 6, 7, 2)) - tmp226 = einsum(tmp225, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (4, 1, 2, 6, 5, 3)) - del tmp225 - tmp200 = einsum(tmp167, (0, 1, 2, 3), t3, (4, 3, 1, 5, 6, 7), (0, 4, 2, 5, 7, 6)) - tmp129 = einsum(tmp128, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 3), (4, 5, 0, 2, 6, 7)) - tmp136 = einsum(t2_OOvV, (0, 1, 2, 3), tmp135, (2, 4, 5, 6), (0, 1, 4, 5, 6, 3)) - del tmp135 - tmp133 = einsum(t2_OOvV, (0, 1, 2, 3), tmp132, (2, 4, 5, 6), (0, 1, 4, 5, 6, 3)) - del tmp132 - tmp160 = einsum(tmp159, (0, 1), t3, (2, 3, 4, 5, 6, 1), (2, 4, 3, 5, 6, 0)) - tmp157 = einsum(t3, (0, 1, 2, 3, 4, 5), f.VV, (6, 5), (0, 2, 1, 6, 3, 4)) - tmp240 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp239, (6, 2, 7, 4), (6, 0, 1, 7, 3, 5)) - tmp233 = einsum(v.OoOV, (0, 1, 2, 3), t2_OoVV, (4, 1, 5, 6), (4, 0, 2, 6, 5, 3)) - tmp164 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp144, (6, 4), (0, 2, 1, 6, 3, 5)) - tmp143 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp142, (6, 5), (0, 2, 1, 6, 3, 4)) - del tmp142 + tmp248 = einsum(v.OOVV, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 3), (4, 5, 0, 6, 7, 2)) + tmp222 = einsum(tmp221, (0, 1, 2, 3), t2_OoVV, (4, 0, 5, 6), (1, 4, 2, 3, 6, 5)) + del tmp221 + tmp213 = einsum(t2_OoVV, (0, 1, 2, 3), tmp212, (1, 4, 5, 6), (4, 0, 5, 6, 3, 2)) + del tmp212 tmp189 = einsum(tmp188, (0, 1, 2, 3), t3, (4, 5, 1, 6, 7, 3), (0, 4, 5, 2, 6, 7)) - tmp156 = einsum(tmp155, (0, 1, 2, 3), t3, (4, 5, 6, 7, 2, 3), (4, 6, 5, 1, 0, 7)) - del tmp155 - tmp153 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp152, (6, 5), (0, 2, 1, 6, 3, 4)) - tmp195 = einsum(t2_OOvV, (0, 1, 2, 3), tmp194, (2, 4, 5, 6), (4, 0, 1, 6, 5, 3)) - del tmp194 - tmp261 = einsum(t2_OOvV, (0, 1, 2, 3), tmp260, (2, 4, 5, 6), (4, 0, 1, 5, 3, 6)) - del tmp260 - tmp139 = einsum(tmp130, (0, 1, 2, 3), t3, (4, 1, 5, 6, 7, 3), (4, 5, 0, 2, 6, 7)) - tmp168 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp167, (6, 2, 7, 0), (6, 1, 7, 3, 5, 4)) * -1 - del tmp167 - tmp224 = einsum(t2_OoVV, (0, 1, 2, 3), tmp223, (1, 4, 5, 6), (0, 4, 5, 3, 2, 6)) - del tmp223 - tmp145 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp144, (6, 5), (0, 2, 1, 6, 3, 4)) - del tmp144 - tmp141 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp140, (5, 6), (0, 2, 1, 6, 3, 4)) - tmp293 = einsum(tmp292, (0, 1, 2, 3), t2_OOvV, (4, 5, 0, 6), (1, 4, 5, 6, 3, 2)) - del tmp292 - tmp54 = einsum(t1, (0, 1), tmp53, (0, 2, 3, 4), (1, 2, 3, 4)) + tmp315 = einsum(tmp305, (0, 1, 2, 3), t3, (2, 4, 3, 5, 6, 7), (1, 0, 4, 5, 7, 6)) * -1 + del tmp305 + tmp50 = einsum(v.OVVv, (0, 1, 2, 3), t3, (4, 5, 0, 6, 2, 1), (3, 4, 5, 6)) + tmp56 = einsum(tmp55, (0, 1, 2, 3), t1, (0, 4), (4, 1, 2, 3)) + del tmp55 + tmp54 = einsum(tmp53, (0, 1, 2, 3), t1, (0, 4), (4, 1, 2, 3)) del tmp53 tmp49 = einsum(v.OVVv, (0, 1, 2, 3), t3, (4, 5, 0, 2, 6, 1), (3, 4, 5, 6)) - tmp56 = einsum(t1, (0, 1), tmp55, (0, 2, 3, 4), (1, 2, 3, 4)) - del tmp55 - tmp50 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OVVv, (2, 5, 4, 6), (6, 0, 1, 3)) - tmp30 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OoOV, (0, 6, 2, 5), (6, 1, 3, 4)) - tmp65 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp64, (6, 0, 2, 5), (6, 1, 3, 4)) + tmp31 = einsum(v.OoOV, (0, 1, 2, 3), t3, (4, 0, 2, 5, 6, 3), (1, 4, 5, 6)) tmp66 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp64, (6, 1, 2, 5), (6, 0, 3, 4)) + tmp30 = einsum(v.OVOo, (0, 1, 2, 3), t3, (0, 4, 2, 5, 6, 1), (3, 4, 5, 6)) * -1 + tmp65 = einsum(tmp64, (0, 1, 2, 3), t3, (2, 4, 1, 5, 6, 3), (0, 4, 5, 6)) * -1 del tmp64 - tmp31 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OVOo, (2, 5, 1, 6), (6, 0, 3, 4)) + tmp35 = einsum(f.OV, (0, 1), t3, (2, 3, 0, 4, 5, 1), (2, 3, 4, 5)) tmp39 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp38, (2, 5), (0, 1, 3, 4)) del tmp38 - tmp37 = einsum(t3, (0, 1, 2, 3, 4, 5), tmp36, (2, 5), (0, 1, 3, 4)) + tmp37 = einsum(tmp36, (0, 1), t3, (2, 3, 0, 4, 5, 1), (2, 3, 4, 5)) del tmp36 - tmp35 = einsum(t3, (0, 1, 2, 3, 4, 5), f.OV, (2, 5), (0, 1, 3, 4)) - tmp28 = einsum(f.oo, (0, 1), t2, (2, 1, 3, 4), (0, 2, 3, 4)) - tmp77 = einsum(tmp76, (0, 1, 2, 3), t2, (4, 1, 3, 5), (0, 4, 5, 2)) - tmp46 = einsum(t1, (0, 1), tmp45, (0, 2, 3, 4), (2, 3, 1, 4)) - del tmp45 + tmp63 = einsum(t2, (0, 1, 2, 3), tmp62, (4, 5, 1, 0), (4, 5, 3, 2)) + del tmp62 + tmp80 = einsum(tmp79, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) + del tmp79 tmp81 = einsum(v.ooov, (0, 1, 2, 3), t1, (1, 4), (0, 2, 4, 3)) - tmp82 = einsum(v.oovv, (0, 1, 2, 3), t2, (4, 1, 3, 5), (4, 0, 5, 2)) + tmp95 = einsum(tmp94, (0, 1, 2, 3), t2, (4, 1, 5, 3), (0, 4, 2, 5)) + del tmp94 + tmp88 = einsum(tmp87, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) + del tmp87 + tmp33 = einsum(t2, (0, 1, 2, 3), tmp27, (4, 1), (0, 4, 2, 3)) + tmp28 = einsum(t2, (0, 1, 2, 3), f.oo, (4, 1), (4, 0, 2, 3)) + tmp29 = einsum(v.oooo, (0, 1, 2, 3), t2, (3, 1, 4, 5), (0, 2, 5, 4)) tmp84 = einsum(tmp83, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) del tmp83 - tmp102 = einsum(t1, (0, 1), tmp101, (2, 3, 0, 4), (2, 3, 1, 4)) - del tmp101 - tmp100 = einsum(t2, (0, 1, 2, 3), tmp99, (4, 5, 1, 0), (5, 4, 2, 3)) - tmp68 = einsum(tmp7, (0, 1), t2, (2, 1, 3, 4), (2, 0, 3, 4)) - tmp29 = einsum(t2, (0, 1, 2, 3), v.oooo, (4, 1, 5, 0), (4, 5, 3, 2)) + tmp46 = einsum(tmp45, (0, 1, 2, 3), t1, (0, 4), (1, 2, 4, 3)) + del tmp45 + tmp92 = einsum(tmp91, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) + del tmp91 + tmp119 = einsum(tmp118, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) + del tmp118 + tmp90 = einsum(tmp89, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) + del tmp89 + tmp51 = einsum(t2, (0, 1, 2, 3), tmp11, (4, 3), (0, 1, 2, 4)) + tmp41 = einsum(t2, (0, 1, 2, 3), tmp40, (4, 3), (0, 1, 2, 4)) + del tmp40 + tmp110 = einsum(tmp109, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) + del tmp109 + tmp52 = einsum(t2, (0, 1, 2, 3), tmp13, (3, 4), (0, 1, 2, 4)) + tmp112 = einsum(tmp111, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) + del tmp111 + tmp105 = einsum(tmp75, (0, 1, 2, 3), t2, (4, 1, 5, 2), (0, 4, 5, 3)) + tmp67 = einsum(t2, (0, 1, 2, 3), tmp9, (4, 1), (0, 4, 2, 3)) + tmp69 = einsum(tmp15, (0, 1), t2, (2, 1, 3, 4), (0, 2, 3, 4)) + tmp61 = einsum(tmp0, (0, 1), t2, (2, 0, 3, 4), (1, 2, 3, 4)) + tmp77 = einsum(t2, (0, 1, 2, 3), tmp76, (4, 1, 5, 2), (4, 0, 3, 5)) + tmp98 = einsum(v.vvvv, (0, 1, 2, 3), t1, (4, 3), (4, 0, 1, 2)) + tmp97 = einsum(tmp47, (0, 1, 2, 3), t2, (4, 1, 5, 3), (0, 4, 2, 5)) + tmp78 = einsum(tmp4, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) tmp117 = einsum(tmp116, (0, 1, 2, 3), t2, (3, 2, 4, 5), (1, 0, 4, 5)) del tmp116 tmp96 = einsum(t2, (0, 1, 2, 3), tmp48, (4, 1, 5, 3), (0, 4, 2, 5)) - tmp106 = einsum(t2, (0, 1, 2, 3), tmp75, (4, 1, 2, 5), (4, 0, 3, 5)) - tmp103 = einsum(tmp99, (0, 1, 2, 3), t1, (2, 4), (1, 0, 3, 4)) - del tmp99 - tmp104 = einsum(tmp76, (0, 1, 2, 3), t2, (4, 1, 5, 3), (0, 4, 5, 2)) - del tmp76 - tmp32 = einsum(tmp26, (0, 1), t2, (2, 1, 3, 4), (2, 0, 3, 4)) - tmp115 = einsum(t1, (0, 1), tmp114, (2, 3, 0, 4), (2, 3, 1, 4)) - del tmp114 - tmp90 = einsum(tmp89, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) - del tmp89 - tmp63 = einsum(tmp62, (0, 1, 2, 3), t2, (2, 3, 4, 5), (0, 1, 4, 5)) - del tmp62 - tmp51 = einsum(tmp11, (0, 1), t2, (2, 3, 4, 1), (2, 3, 4, 0)) - tmp119 = einsum(tmp118, (0, 1, 2, 3), t1, (2, 4), (1, 0, 3, 4)) - del tmp118 + tmp34 = einsum(v.oooo, (0, 1, 2, 3), t1, (3, 4), (0, 1, 2, 4)) + tmp113 = einsum(v.vvvv, (0, 1, 2, 3), t2, (4, 5, 1, 3), (4, 5, 0, 2)) + tmp108 = einsum(tmp107, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) + del tmp107 tmp58 = einsum(tmp57, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) del tmp57 - tmp108 = einsum(t1, (0, 1), tmp107, (2, 3, 0, 4), (2, 3, 1, 4)) - del tmp107 + tmp68 = einsum(t2, (0, 1, 2, 3), tmp7, (4, 1), (0, 4, 2, 3)) + tmp70 = einsum(tmp16, (0, 1), t2, (2, 1, 3, 4), (0, 2, 3, 4)) + tmp32 = einsum(tmp26, (0, 1), t2, (2, 1, 3, 4), (2, 0, 3, 4)) + tmp115 = einsum(tmp114, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) + del tmp114 + tmp93 = einsum(t2, (0, 1, 2, 3), tmp85, (4, 1, 5, 3), (0, 4, 2, 5)) + tmp100 = einsum(t2, (0, 1, 2, 3), tmp99, (4, 5, 0, 1), (5, 4, 3, 2)) tmp74 = einsum(tmp73, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) del tmp73 - tmp95 = einsum(tmp94, (0, 1, 2, 3), t2, (4, 1, 5, 3), (0, 4, 2, 5)) - del tmp94 - tmp69 = einsum(t2, (0, 1, 2, 3), tmp15, (4, 1), (4, 0, 2, 3)) - tmp67 = einsum(tmp9, (0, 1), t2, (2, 1, 3, 4), (2, 0, 3, 4)) - tmp92 = einsum(t1, (0, 1), tmp91, (2, 3, 0, 4), (2, 3, 1, 4)) - del tmp91 - tmp72 = einsum(t1, (0, 1), tmp71, (2, 0, 3, 4), (2, 3, 1, 4)) + tmp44 = einsum(f.vv, (0, 1), t2, (2, 3, 4, 1), (2, 3, 0, 4)) + tmp103 = einsum(t1, (0, 1), tmp99, (2, 3, 0, 4), (3, 2, 4, 1)) + del tmp99 + tmp72 = einsum(tmp71, (0, 1, 2, 3), t1, (1, 4), (0, 2, 4, 3)) del tmp71 tmp60 = einsum(tmp59, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) del tmp59 + tmp102 = einsum(tmp101, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) + del tmp101 + tmp106 = einsum(t2, (0, 1, 2, 3), tmp75, (4, 1, 2, 5), (4, 0, 3, 5)) + tmp82 = einsum(t2, (0, 1, 2, 3), v.oovv, (4, 1, 5, 2), (0, 4, 3, 5)) + tmp104 = einsum(t2, (0, 1, 2, 3), tmp76, (4, 1, 5, 3), (4, 0, 2, 5)) + del tmp76 tmp43 = einsum(tmp42, (0, 1), t2, (2, 3, 4, 1), (2, 3, 4, 0)) del tmp42 - tmp98 = einsum(t1, (0, 1), v.vvvv, (2, 3, 4, 1), (0, 2, 3, 4)) - tmp88 = einsum(tmp87, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) - del tmp87 - tmp105 = einsum(tmp75, (0, 1, 2, 3), t2, (4, 1, 5, 2), (0, 4, 5, 3)) - tmp113 = einsum(v.vvvv, (0, 1, 2, 3), t2, (4, 5, 1, 3), (4, 5, 0, 2)) - tmp80 = einsum(t1, (0, 1), tmp79, (2, 3, 0, 4), (2, 3, 1, 4)) - del tmp79 - tmp93 = einsum(t2, (0, 1, 2, 3), tmp85, (4, 1, 5, 3), (0, 4, 2, 5)) - tmp41 = einsum(t2, (0, 1, 2, 3), tmp40, (4, 3), (0, 1, 2, 4)) - del tmp40 - tmp52 = einsum(tmp13, (0, 1), t2, (2, 3, 4, 0), (2, 3, 4, 1)) - tmp110 = einsum(tmp109, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) - del tmp109 - tmp112 = einsum(t1, (0, 1), tmp111, (2, 3, 0, 4), (2, 3, 1, 4)) - del tmp111 - tmp97 = einsum(tmp47, (0, 1, 2, 3), t2, (4, 1, 5, 3), (0, 4, 2, 5)) - tmp33 = einsum(tmp27, (0, 1), t2, (2, 1, 3, 4), (2, 0, 3, 4)) tmp86 = einsum(t2, (0, 1, 2, 3), v.oovv, (4, 1, 5, 3), (0, 4, 2, 5)) - tmp78 = einsum(tmp4, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) - tmp34 = einsum(v.oooo, (0, 1, 2, 3), t1, (3, 4), (0, 1, 2, 4)) - tmp70 = einsum(t2, (0, 1, 2, 3), tmp16, (4, 1), (4, 0, 2, 3)) - tmp44 = einsum(f.vv, (0, 1), t2, (2, 3, 4, 1), (2, 3, 0, 4)) - tmp61 = einsum(t2, (0, 1, 2, 3), tmp0, (1, 4), (4, 0, 2, 3)) - tmp20 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OVOV, (2, 3, 1, 5), (0, 4)) - tmp23 = einsum(v.OVOV, (0, 1, 2, 3), t3, (0, 4, 2, 1, 5, 3), (4, 5)) - tmp21 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OVOV, (2, 4, 1, 5), (0, 3)) - tmp22 = einsum(t3, (0, 1, 2, 3, 4, 5), v.OVOV, (2, 5, 1, 4), (0, 3)) + tmp21 = einsum(v.OVOV, (0, 1, 2, 3), t3, (4, 0, 2, 5, 3, 1), (4, 5)) + tmp20 = einsum(v.OVOV, (0, 1, 2, 3), t3, (4, 0, 2, 3, 5, 1), (4, 5)) + tmp22 = einsum(v.OVOV, (0, 1, 2, 3), t3, (4, 0, 2, 5, 1, 3), (4, 5)) + tmp23 = einsum(v.OVOV, (0, 1, 2, 3), t3, (0, 4, 2, 3, 5, 1), (4, 5)) * -1 + tmp2 = einsum(t2, (0, 1, 2, 3), v.ovvv, (1, 3, 4, 2), (0, 4)) + tmp24 = einsum(t2, (0, 1, 2, 3), tmp1, (1, 2), (0, 3)) + tmp12 = einsum(tmp11, (0, 1), t1, (2, 1), (2, 0)) + del tmp11 + tmp10 = einsum(t1, (0, 1), tmp9, (2, 0), (2, 1)) + del tmp9 tmp17 = einsum(t1, (0, 1), tmp16, (2, 0), (2, 1)) del tmp16 - tmp10 = einsum(tmp9, (0, 1), t1, (1, 2), (0, 2)) - del tmp9 - tmp24 = einsum(tmp1, (0, 1), t2, (2, 0, 1, 3), (2, 3)) - tmp19 = einsum(v.ooov, (0, 1, 2, 3), t2, (2, 1, 3, 4), (0, 4)) - tmp6 = einsum(tmp4, (0, 1, 2, 3), t2, (2, 1, 3, 4), (0, 4)) tmp8 = einsum(tmp7, (0, 1), t1, (1, 2), (0, 2)) del tmp7 - tmp5 = einsum(tmp4, (0, 1, 2, 3), t2, (1, 2, 3, 4), (0, 4)) - del tmp4 + tmp18 = einsum(t2, (0, 1, 2, 3), v.ooov, (4, 1, 0, 3), (4, 2)) + tmp6 = einsum(t2, (0, 1, 2, 3), tmp4, (4, 0, 1, 3), (4, 2)) + tmp19 = einsum(t2, (0, 1, 2, 3), v.ooov, (4, 0, 1, 3), (4, 2)) + tmp25 = einsum(t2, (0, 1, 2, 3), tmp1, (1, 3), (0, 2)) tmp3 = einsum(t2, (0, 1, 2, 3), v.ovvv, (1, 2, 4, 3), (0, 4)) - tmp2 = einsum(t2, (0, 1, 2, 3), v.ovvv, (1, 3, 4, 2), (0, 4)) - tmp25 = einsum(tmp1, (0, 1), t2, (2, 0, 3, 1), (2, 3)) - tmp12 = einsum(tmp11, (0, 1), t1, (2, 1), (2, 0)) - del tmp11 - tmp18 = einsum(v.ooov, (0, 1, 2, 3), t2, (1, 2, 3, 4), (0, 4)) + tmp5 = einsum(t2, (0, 1, 2, 3), tmp4, (4, 1, 0, 3), (4, 2)) + del tmp4 t3new = np.copy(np.transpose(tmp120, (1, 2, 0, 4, 5, 3))) t3new += np.transpose(tmp121, (1, 0, 2, 4, 5, 3)) * -1 del tmp121 @@ -525,7 +525,7 @@ def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): t3new += np.transpose(tmp129, (1, 0, 2, 5, 4, 3)) * -1 t3new += np.transpose(tmp133, (1, 0, 2, 5, 3, 4)) t3new += np.transpose(tmp136, (1, 0, 2, 5, 3, 4)) - t3new += einsum(f.OO, (0, 1), t3, (2, 1, 3, 4, 5, 6), (2, 0, 3, 4, 5, 6)) * -1 + t3new += einsum(t3, (0, 1, 2, 3, 4, 5), f.OO, (6, 1), (0, 6, 2, 3, 4, 5)) * -1 t3new += np.transpose(tmp137, (1, 2, 0, 4, 5, 3)) t3new += np.transpose(tmp123, (1, 2, 0, 5, 4, 3)) * -1 t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp124, (6, 1), (0, 6, 2, 3, 4, 5)) @@ -572,7 +572,7 @@ def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): t3new += np.transpose(tmp153, (1, 2, 0, 3, 5, 4)) * -2 del tmp153 t3new += einsum(t3, (0, 1, 2, 3, 4, 5), f.VV, (6, 4), (0, 1, 2, 3, 6, 5)) - t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp140, (4, 6), (0, 1, 2, 3, 6, 5)) * -1 + t3new += einsum(tmp140, (0, 1), t3, (2, 3, 4, 5, 0, 6), (2, 3, 4, 5, 1, 6)) * -1 del tmp140 t3new += einsum(tmp159, (0, 1), t3, (2, 3, 4, 5, 1, 6), (2, 3, 4, 5, 0, 6)) * 2 del tmp159 @@ -582,9 +582,9 @@ def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): del tmp163 t3new += np.transpose(tmp164, (1, 2, 0, 4, 3, 5)) * 2 del tmp164 - t3new += einsum(tmp149, (0, 1), t3, (2, 3, 4, 5, 1, 6), (2, 3, 4, 5, 0, 6)) + t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp149, (6, 4), (0, 1, 2, 3, 6, 5)) del tmp149 - t3new += einsum(tmp152, (0, 1), t3, (2, 3, 4, 5, 1, 6), (2, 3, 4, 5, 0, 6)) * -2 + t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp152, (6, 4), (0, 1, 2, 3, 6, 5)) * -2 del tmp152 t3new += np.transpose(tmp157, (0, 2, 1, 4, 5, 3)) del tmp157 @@ -607,7 +607,7 @@ def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): t3new += np.transpose(tmp192, (1, 2, 0, 4, 5, 3)) * -1 t3new += np.transpose(tmp192, (2, 1, 0, 5, 4, 3)) * -1 t3new += np.transpose(tmp195, (2, 1, 0, 5, 4, 3)) - t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp165, (1, 6), (0, 6, 2, 3, 4, 5)) * -1 + t3new += einsum(tmp165, (0, 1), t3, (2, 0, 3, 4, 5, 6), (2, 1, 3, 4, 5, 6)) * -1 del tmp165 t3new += np.transpose(tmp196, (2, 0, 1, 4, 5, 3)) t3new += np.transpose(tmp197, (2, 0, 1, 3, 5, 4)) * -1 @@ -619,7 +619,7 @@ def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): t3new += np.transpose(tmp178, (2, 0, 1, 5, 4, 3)) * -1 t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp182, (6, 1), (0, 6, 2, 3, 4, 5)) del tmp182 - t3new += einsum(tmp185, (0, 1), t3, (2, 1, 3, 4, 5, 6), (2, 0, 3, 4, 5, 6)) * -2 + t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp185, (6, 1), (0, 6, 2, 3, 4, 5)) * -2 del tmp185 t3new += np.transpose(tmp199, (2, 0, 1, 4, 5, 3)) * -1 t3new += np.transpose(tmp195, (1, 0, 2, 5, 3, 4)) * -1 @@ -940,25 +940,25 @@ def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): t3new += np.transpose(tmp158, (0, 2, 1, 4, 5, 3)) * -1 t3new += np.transpose(tmp158, (0, 2, 1, 3, 5, 4)) t3new += einsum(t3, (0, 1, 2, 3, 4, 5), v.OVOV, (6, 7, 2, 4), (0, 6, 1, 3, 7, 5)) * -1 - t3new += einsum(t3, (0, 1, 2, 3, 4, 5), v.OOVV, (6, 1, 7, 4), (0, 6, 2, 3, 7, 5)) * -1 - t3new += einsum(tmp128, (0, 1, 2, 3), t3, (4, 1, 5, 6, 3, 7), (4, 0, 5, 6, 2, 7)) * -1 + t3new += einsum(v.OOVV, (0, 1, 2, 3), t3, (4, 1, 5, 6, 3, 7), (4, 0, 5, 6, 2, 7)) * -1 + t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp128, (6, 1, 7, 4), (0, 6, 2, 3, 7, 5)) * -1 t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp130, (6, 1, 7, 4), (0, 6, 2, 3, 7, 5)) del tmp130 t3new += np.transpose(tmp129, (0, 2, 1, 5, 3, 4)) t3new += np.transpose(tmp129, (0, 2, 1, 4, 3, 5)) * -1 del tmp129 - t3new += einsum(tmp128, (0, 1, 2, 3), t3, (4, 5, 1, 6, 3, 7), (4, 0, 5, 6, 2, 7)) + t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp128, (6, 2, 7, 4), (0, 6, 1, 3, 7, 5)) del tmp128 t3new += np.transpose(tmp251, (1, 0, 2, 5, 3, 4)) t3new += np.transpose(tmp251, (1, 0, 2, 4, 3, 5)) * -1 - t3new += einsum(tmp250, (0, 1, 2, 3), t3, (4, 5, 1, 6, 3, 7), (4, 0, 5, 6, 2, 7)) + t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp250, (6, 2, 7, 4), (0, 6, 1, 3, 7, 5)) t3new += np.transpose(tmp255, (1, 0, 2, 5, 3, 4)) * -2 t3new += np.transpose(tmp255, (1, 0, 2, 4, 3, 5)) * 2 t3new += np.transpose(tmp271, (1, 0, 2, 4, 3, 5)) * -2 del tmp271 t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp239, (6, 1, 7, 4), (0, 6, 2, 3, 7, 5)) del tmp239 - t3new += einsum(tmp250, (0, 1, 2, 3), t3, (4, 1, 5, 6, 3, 7), (4, 0, 5, 6, 2, 7)) * -1 + t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp250, (6, 1, 7, 4), (0, 6, 2, 3, 7, 5)) * -1 del tmp250 t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp252, (6, 1, 7, 4), (0, 6, 2, 3, 7, 5)) * -1 t3new += np.transpose(tmp272, (2, 0, 1, 4, 3, 5)) * -2 @@ -1142,10 +1142,10 @@ def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): del tmp289 t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp188, (6, 1, 7, 4), (0, 6, 2, 3, 7, 5)) del tmp188 - t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp191, (6, 1, 7, 4), (0, 6, 2, 3, 7, 5)) * -1 + t3new += einsum(tmp191, (0, 1, 2, 3), t3, (4, 1, 5, 6, 3, 7), (4, 0, 5, 6, 2, 7)) * -1 t3new += np.transpose(tmp192, (1, 0, 2, 5, 3, 4)) t3new += np.transpose(tmp192, (1, 0, 2, 4, 3, 5)) * -1 - t3new += einsum(t3, (0, 1, 2, 3, 4, 5), tmp191, (6, 2, 7, 4), (0, 6, 1, 3, 7, 5)) + t3new += einsum(tmp191, (0, 1, 2, 3), t3, (4, 5, 1, 6, 3, 7), (4, 0, 5, 6, 2, 7)) del tmp191 t3new += np.transpose(tmp288, (0, 2, 1, 5, 4, 3)) * -1 t3new += np.transpose(tmp290, (0, 2, 1, 5, 4, 3)) @@ -1408,7 +1408,7 @@ def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): del tmp29 t2new += tmp32 t2new += tmp33 * -2 - t2new += einsum(tmp34, (0, 1, 2, 3), t1, (1, 4), (2, 0, 3, 4)) + t2new += einsum(tmp34, (0, 1, 2, 3), t1, (1, 4), (0, 2, 4, 3)) del tmp34 t2new += np.transpose(tmp28, (0, 1, 3, 2)) * -1 del tmp28 @@ -1497,12 +1497,12 @@ def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): del tmp84 t2new += np.transpose(tmp75, (0, 1, 3, 2)) del tmp75 - t2new += einsum(tmp98, (0, 1, 2, 3), t1, (4, 2), (4, 0, 1, 3)) + t2new += einsum(t1, (0, 1), tmp98, (2, 3, 1, 4), (0, 2, 3, 4)) del tmp98 t2new += tmp100 del tmp100 t2new += np.transpose(tmp102, (1, 0, 3, 2)) * -1 - t2new += einsum(t1, (0, 1), tmp103, (2, 3, 0, 4), (2, 3, 1, 4)) + t2new += einsum(tmp103, (0, 1, 2, 3), t1, (2, 4), (0, 1, 4, 3)) del tmp103 t2new += np.transpose(tmp104, (0, 1, 3, 2)) * -1 t2new += np.transpose(tmp105, (0, 1, 3, 2)) * 2 @@ -1588,9 +1588,9 @@ def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): t1new += tmp17 * -2 del tmp17 t1new += einsum(t1, (0, 1), f.oo, (2, 0), (2, 1)) * -1 - t1new += einsum(t2, (0, 1, 2, 3), f.ov, (1, 2), (0, 3)) * -1 + t1new += einsum(f.ov, (0, 1), t2, (2, 0, 1, 3), (2, 3)) * -1 t1new += einsum(t2, (0, 1, 2, 3), f.ov, (1, 3), (0, 2)) * 2 - t1new += einsum(v.oovv, (0, 1, 2, 3), t1, (1, 3), (0, 2)) * -1 + t1new += einsum(t1, (0, 1), v.oovv, (2, 0, 3, 1), (2, 3)) * -1 t1new += tmp18 del tmp18 t1new += tmp19 * -2 @@ -1604,15 +1604,15 @@ def update_amps(f=None, t1=None, t2=None, t3=None, v=None, **kwargs): del tmp25 t1new += einsum(t1, (0, 1), tmp26, (2, 0), (2, 1)) del tmp26 - t1new += einsum(t1, (0, 1), tmp27, (2, 0), (2, 1)) * -2 + t1new += einsum(tmp27, (0, 1), t1, (1, 2), (0, 2)) * -2 del tmp27 - t2new += _inflate(t2new.shape, np.ix_(sO, sO, sv, sV), t2new_OOvV) - t2new += _inflate(t2new.shape, np.ix_(sO, so, sV, sV), t2new_OoVV) + t2new += _inflate(t2new.shape, np.ix_(sO, sO, sV, sV), t2new_OOVV) t1new += _inflate(t1new.shape, np.ix_(sO, sV), t1new_OV) - t2new += _inflate(t2new.shape, np.ix_(so, sO, sV, sV), t2new_oOVV) t2new += _inflate(t2new.shape, np.ix_(sO, sO, sV, sv), t2new_OOVv) - t2new += _inflate(t2new.shape, np.ix_(sO, sO, sV, sV), t2new_OOVV) + t2new += _inflate(t2new.shape, np.ix_(sO, so, sV, sV), t2new_OoVV) + t2new += _inflate(t2new.shape, np.ix_(sO, sO, sv, sV), t2new_OOvV) + t2new += _inflate(t2new.shape, np.ix_(so, sO, sV, sV), t2new_oOVV) return {f"t1new": t1new, f"t2new": t2new, f"t3new": t3new} diff --git a/ebcc/codegen/bootstrap_CC2.py b/ebcc/codegen/bootstrap_CC2.py index 72df1ba4..339414d6 100644 --- a/ebcc/codegen/bootstrap_CC2.py +++ b/ebcc/codegen/bootstrap_CC2.py @@ -5,21 +5,24 @@ import sys import pdaggerq -from albert.qc._pdaggerq import import_from_pdaggerq -from albert.qc.index import Index +from albert.qc._pdaggerq import remove_reference_energy, remove_reference_energy_eom +from albert.qc.spin import ghf_to_uhf, ghf_to_rhf +from albert.qc import ghf, uhf, rhf from albert.tensor import Tensor +from albert.index import Index +from albert.code._ebcc import EBCCCodeGenerator +from albert.misc import Stopwatch +from albert.opt.tools import _tensor_info -from ebcc.codegen.bootstrap_common import * +from ebcc.codegen.bootstrap_common import get_energy, get_amplitudes, get_rdm1, get_rdm2, get_eom # Get the spin case spin = sys.argv[1] # Set up the code generators code_generators = { - "einsum": EinsumCodeGen( + "einsum": EBCCCodeGenerator( stdout=open(f"{spin[0].upper()}CC2.py", "w"), - name_generator=name_generators[spin], - spin=spin, ), } @@ -39,22 +42,17 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms = pq.fully_contracted_strings() - terms = remove_hf_energy(terms) + terms = remove_reference_energy(terms) # Get the energy in albert format - expr = import_from_pdaggerq(terms) - expr = spin_integrate(expr, spin) - output = tuple(Tensor(name="e_cc") for _ in range(len(expr))) - output, expr = optimise(output, expr, spin, strategy="exhaust") - returns = (Tensor(name="e_cc"),) + output_expr, returns = get_energy(terms, spin) # Generate the energy code for codegen in code_generators.values(): codegen( "energy", returns, - output, - expr, + output_expr, ) with Stopwatch("T amplitudes"): @@ -75,37 +73,15 @@ terms_t2 = pq.fully_contracted_strings() # Get the T amplitudes in albert format - terms = [terms_t1, terms_t2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"t{n+1}new") - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"t{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns = get_amplitudes([terms_t1, terms_t2], spin) # Generate the T amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "t1new = Namespace()\nt2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_amps", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) with Stopwatch("L amplitudes"): @@ -138,37 +114,15 @@ terms_l2 = pq.fully_contracted_strings() # Get the L amplitudes in albert format - terms = [terms_l1, terms_l2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_l_amplitude_outputs(expr_n, f"l{n+1}new") - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"l{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="opt") + output_expr, returns = get_amplitudes([terms_l1, terms_l2], spin, strategy="opt" if spin == "uhf" else "exhaust", which="l") # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "l1new = Namespace()\nl2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_lams", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) with Stopwatch("1RDM"): @@ -182,37 +136,40 @@ terms[sectors, indices] = pq.fully_contracted_strings() # Get the 1RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [("oo", "ij"), ("ov", "ia"), ("vo", "ai"), ("vv", "ab")]: - for index_spins in get_density_spins(1, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - if spin == "rhf": - expr_n = tuple(e * 2 for e in expr_n) - output_n = get_density_outputs(expr_n, f"d", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"d"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns, deltas, deltas_sources = get_rdm1(terms, spin) # Generate the 1RDM code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": get_density_einsum_preamble(1, spin), - "postamble": get_density_einsum_postamble(1, spin), - } - else: - kwargs = {} + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing + if spin != "uhf": + codegen.write("rdm1 = np.block([[rdm1.oo, rdm1.ov], [rdm1.vo, rdm1.vv]])") + else: + codegen.write("rdm1.aa = np.block([[rdm1.aa.oo, rdm1.aa.ov], [rdm1.aa.vo, rdm1.aa.vv]])") + codegen.write("rdm1.bb = np.block([[rdm1.bb.oo, rdm1.bb.ov], [rdm1.bb.vo, rdm1.bb.vv]])") + codegen( "make_rdm1_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("2RDM"): @@ -245,50 +202,42 @@ terms[sectors, indices] = pq.fully_contracted_strings() # Get the 2RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [ - ("oooo", "ijkl"), - ("ooov", "ijka"), - ("oovo", "ijak"), - ("ovoo", "iajk"), - ("vooo", "aijk"), - ("oovv", "ijab"), - ("ovov", "iajb"), - ("ovvo", "iabj"), - ("voov", "aijb"), - ("vovo", "aibj"), - ("vvoo", "abij"), - ("ovvv", "iabc"), - ("vovv", "aibc"), - ("vvov", "abic"), - ("vvvo", "abci"), - ("vvvv", "abcd"), - ]: - for index_spins in get_density_spins(2, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_density_outputs(expr_n, f"Γ", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"Γ"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="trav") + output_expr, returns, deltas, deltas_sources = get_rdm2(terms, spin, strategy="trav" if spin == "uhf" else "exhaust") # Generate the 2RDM code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": get_density_einsum_preamble(2, spin), - "postamble": get_density_einsum_postamble(2, spin), - } + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing, handle transpose + if spin != "uhf": + codegen.write("rdm2 = pack_2e(rdm2.oooo, rdm2.ooov, rdm2.oovo, rdm2.ovoo, rdm2.vooo, rdm2.oovv, rdm2.ovov, rdm2.ovvo, rdm2.voov, rdm2.vovo, rdm2.vvoo, rdm2.ovvv, rdm2.vovv, rdm2.vvov, rdm2.vvvo, rdm2.vvvv).transpose((0, 2, 1, 3))") + else: + codegen.write("rdm2.aaaa = pack_2e(rdm2.aaaa.oooo, rdm2.aaaa.ooov, rdm2.aaaa.oovo, rdm2.aaaa.ovoo, rdm2.aaaa.vooo, rdm2.aaaa.oovv, rdm2.aaaa.ovov, rdm2.aaaa.ovvo, rdm2.aaaa.voov, rdm2.aaaa.vovo, rdm2.aaaa.vvoo, rdm2.aaaa.ovvv, rdm2.aaaa.vovv, rdm2.aaaa.vvov, rdm2.aaaa.vvvo, rdm2.aaaa.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.aabb = pack_2e(rdm2.abab.oooo, rdm2.abab.ooov, rdm2.abab.oovo, rdm2.abab.ovoo, rdm2.abab.vooo, rdm2.abab.oovv, rdm2.abab.ovov, rdm2.abab.ovvo, rdm2.abab.voov, rdm2.abab.vovo, rdm2.abab.vvoo, rdm2.abab.ovvv, rdm2.abab.vovv, rdm2.abab.vvov, rdm2.abab.vvvo, rdm2.abab.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.bbbb = pack_2e(rdm2.bbbb.oooo, rdm2.bbbb.ooov, rdm2.bbbb.oovo, rdm2.bbbb.ovoo, rdm2.bbbb.vooo, rdm2.bbbb.oovv, rdm2.bbbb.ovov, rdm2.bbbb.ovvo, rdm2.bbbb.voov, rdm2.bbbb.vovo, rdm2.bbbb.vvoo, rdm2.bbbb.ovvv, rdm2.bbbb.vovv, rdm2.bbbb.vvov, rdm2.bbbb.vvvo, rdm2.bbbb.vvvv).transpose((0, 2, 1, 3))") + codegen.write("del rdm2.abab") + codegen( "make_rdm2_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("IP-EOM"): @@ -301,7 +250,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -312,60 +261,27 @@ pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("EA-EOM"): @@ -378,7 +294,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -389,60 +305,27 @@ pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -456,7 +339,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -467,60 +350,46 @@ pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_matvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) with Stopwatch("L-IP-EOM"): @@ -533,7 +402,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -544,60 +413,27 @@ pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip") # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("L-EA-EOM"): @@ -610,7 +446,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -621,60 +457,27 @@ pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea") # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -688,7 +491,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -699,60 +502,46 @@ pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee") # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_lmatvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) for codegen in code_generators.values(): diff --git a/ebcc/codegen/bootstrap_CC3.py b/ebcc/codegen/bootstrap_CC3.py index d1e05c7d..d21effa5 100644 --- a/ebcc/codegen/bootstrap_CC3.py +++ b/ebcc/codegen/bootstrap_CC3.py @@ -3,23 +3,26 @@ """ import sys -sys.setrecursionlimit(100000) import pdaggerq -from albert.qc._pdaggerq import import_from_pdaggerq +from albert.qc._pdaggerq import remove_reference_energy, remove_reference_energy_eom +from albert.qc.spin import ghf_to_uhf, ghf_to_rhf +from albert.qc import ghf, uhf, rhf from albert.tensor import Tensor +from albert.index import Index +from albert.code._ebcc import EBCCCodeGenerator +from albert.misc import Stopwatch +from albert.opt.tools import _tensor_info -from ebcc.codegen.bootstrap_common import * +from ebcc.codegen.bootstrap_common import get_energy, get_amplitudes, get_rdm1, get_rdm2, get_eom # Get the spin case spin = sys.argv[1] # Set up the code generators code_generators = { - "einsum": EinsumCodeGen( + "einsum": EBCCCodeGenerator( stdout=open(f"{spin[0].upper()}CC3.py", "w"), - name_generator=name_generators[spin], - spin=spin, ), } @@ -39,22 +42,17 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms = pq.fully_contracted_strings() - terms = remove_hf_energy(terms) + terms = remove_reference_energy(terms) # Get the energy in albert format - expr = import_from_pdaggerq(terms) - expr = spin_integrate(expr, spin) - output = tuple(Tensor(name="e_cc") for _ in range(len(expr))) - output, expr = optimise(output, expr, spin, strategy="exhaust") - returns = (Tensor(name="e_cc"),) + output_expr, returns = get_energy(terms, spin) # Generate the energy code for codegen in code_generators.values(): codegen( "energy", returns, - output, - expr, + output_expr, ) with Stopwatch("T amplitudes"): @@ -88,37 +86,15 @@ terms_t3 = pq.fully_contracted_strings() # Get the T amplitudes in albert format - terms = [terms_t1, terms_t2, terms_t3] - expr = [] - output = [] - returns = [] - for n in range(3): - for index_spins in get_amplitude_spins(n + 1, spin): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"t{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"t{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns = get_amplitudes([terms_t1, terms_t2, terms_t3], spin) # Generate the T amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "t1new = Namespace()\nt2new = Namespace()\nt3new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_amps", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) for codegen in code_generators.values(): diff --git a/ebcc/codegen/bootstrap_CCD.py b/ebcc/codegen/bootstrap_CCD.py index 7de0d540..8dc678c2 100644 --- a/ebcc/codegen/bootstrap_CCD.py +++ b/ebcc/codegen/bootstrap_CCD.py @@ -5,21 +5,24 @@ import sys import pdaggerq -from albert.qc._pdaggerq import import_from_pdaggerq +from albert.qc._pdaggerq import remove_reference_energy, remove_reference_energy_eom +from albert.qc.spin import ghf_to_uhf, ghf_to_rhf +from albert.qc import ghf, uhf, rhf from albert.tensor import Tensor -from albert.algebra import Mul +from albert.index import Index +from albert.code._ebcc import EBCCCodeGenerator +from albert.misc import Stopwatch +from albert.opt.tools import _tensor_info -from ebcc.codegen.bootstrap_common import * +from ebcc.codegen.bootstrap_common import get_energy, get_amplitudes, get_rdm1, get_rdm2, get_eom # Get the spin case spin = sys.argv[1] # Set up the code generators code_generators = { - "einsum": EinsumCodeGen( + "einsum": EBCCCodeGenerator( stdout=open(f"{spin[0].upper()}CCD.py", "w"), - name_generator=name_generators[spin], - spin=spin, ), } @@ -39,22 +42,17 @@ pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms = pq.fully_contracted_strings() - terms = remove_hf_energy(terms) + terms = remove_reference_energy(terms) # Get the energy in albert format - expr = import_from_pdaggerq(terms) - expr = spin_integrate(expr, spin) - output = tuple(Tensor(name="e_cc") for _ in range(len(expr))) - output, expr = optimise(output, expr, spin, strategy="exhaust") - returns = (Tensor(name="e_cc"),) + output_expr, returns = get_energy(terms, spin) # Generate the energy code for codegen in code_generators.values(): codegen( "energy", returns, - output, - expr, + output_expr, ) with Stopwatch("T amplitudes"): @@ -67,35 +65,15 @@ terms = pq.fully_contracted_strings() # Get the T amplitudes in albert format - expr = [] - output = [] - returns = [] - for index_spins in get_amplitude_spins(2, spin): - indices = default_indices["o"][:2] + default_indices["v"][:2] - expr_n = import_from_pdaggerq(terms, index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"t2new") - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"t2new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns = get_amplitudes([terms], spin, orders=[2]) # Generate the T amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "t2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_amps", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) with Stopwatch("L amplitudes"): @@ -114,35 +92,15 @@ terms = pq.fully_contracted_strings() # Get the L amplitudes in albert format - expr = [] - output = [] - returns = [] - for index_spins in get_amplitude_spins(2, spin): - indices = default_indices["v"][:2] + default_indices["o"][:2] - expr_n = import_from_pdaggerq(terms, index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_l_amplitude_outputs(expr_n, f"l2new") - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"l2new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="opt") + get_amplitudes([terms], spin, strategy="opt" if spin == "uhf" else "exhaust", which="l", orders=[2]) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "l2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_lams", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) with Stopwatch("1RDM"): @@ -156,51 +114,40 @@ terms[sectors, indices] = pq.fully_contracted_strings() # Get the 1RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [("oo", "ij"), ("ov", "ia"), ("vo", "ai"), ("vv", "ab")]: - for index_spins in get_density_spins(1, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - if not (isinstance(expr_n, int) and expr_n == 0): - expr_n = spin_integrate(expr_n, spin) - if spin == "rhf": - expr_n = tuple(e * 2 for e in expr_n) - output_n = get_density_outputs(expr_n, f"d", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"d"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns, deltas, deltas_sources = get_rdm1(terms, spin) # Generate the 1RDM code for name, codegen in code_generators.items(): - if name == "einsum": - def get_postamble(n, spin, name="rdm{n}"): - nm = name.format(n=n) - postamble = "" - if spin != "uhf": - for occ in ("ov", "vo"): - shape = ", ".join(f"t2.shape[{'0' if o == 'o' else '-1'}]" for o in occ) - postamble += f"{nm}.{occ} = np.zeros(({shape}))\n" - else: - for s in "ab": - for occ in ("ov", "vo"): - shape = ", ".join(f"t2.{s}{s}{s}{s}.shape[{'0' if o == 'o' else '-1'}]" for o in occ) - postamble += f"{nm}.{s}{s}.{occ} = np.zeros(({shape}))\n" - return postamble + get_density_einsum_postamble(n, spin) - kwargs = { - "preamble": get_density_einsum_preamble(1, spin), - "postamble": get_postamble(1, spin), - } - else: - kwargs = {} + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing + if spin != "uhf": + codegen.write("rdm1 = np.block([[rdm1.oo, rdm1.ov], [rdm1.vo, rdm1.vv]])") + else: + codegen.write("rdm1.aa = np.block([[rdm1.aa.oo, rdm1.aa.ov], [rdm1.aa.vo, rdm1.aa.vv]])") + codegen.write("rdm1.bb = np.block([[rdm1.bb.oo, rdm1.bb.ov], [rdm1.bb.vo, rdm1.bb.vv]])") + codegen( "make_rdm1_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("2RDM"): @@ -233,64 +180,42 @@ def get_postamble(n, spin, name="rdm{n}"): terms[sectors, indices] = pq.fully_contracted_strings() # Get the 2RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [ - ("oooo", "ijkl"), - ("ooov", "ijka"), - ("oovo", "ijak"), - ("ovoo", "iajk"), - ("vooo", "aijk"), - ("oovv", "ijab"), - ("ovov", "iajb"), - ("ovvo", "iabj"), - ("voov", "aijb"), - ("vovo", "aibj"), - ("vvoo", "abij"), - ("ovvv", "iabc"), - ("vovv", "aibc"), - ("vvov", "abic"), - ("vvvo", "abci"), - ("vvvv", "abcd"), - ]: - for index_spins in get_density_spins(2, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - if not (isinstance(expr_n, int) and expr_n == 0): - expr_n = spin_integrate(expr_n, spin) - output_n = get_density_outputs(expr_n, f"Γ", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"Γ"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="trav") + output_expr, returns, deltas, deltas_sources = get_rdm2(terms, spin, strategy="trav" if spin == "uhf" else "exhaust") # Generate the 2RDM code for name, codegen in code_generators.items(): - if name == "einsum": - def get_postamble(n, spin, name="rdm{n}"): - nm = name.format(n=n) - postamble = "" - if spin != "uhf": - for occ in ("ooov", "oovo", "ovoo", "vooo", "ovvv", "vovv", "vvov", "vvvo"): - shape = ", ".join(f"t2.shape[{'0' if o == 'o' else '-1'}]" for o in occ) - postamble += f"{nm}.{occ} = np.zeros(({shape}))\n" - else: - for s1, s2 in [("a", "a"), ("a", "b"), ("b", "b")]: - for occ in ("ooov", "oovo", "ovoo", "vooo", "ovvv", "vovv", "vvov", "vvvo"): - shape = ", ".join(f"t2.{s}{s}{s}{s}.shape[{'0' if o == 'o' else '-1'}]" for o, s in zip(occ, s1+s2+s1+s2)) - postamble += f"{nm}.{s1}{s2}{s1}{s2}.{occ} = np.zeros(({shape}))\n" - return postamble + get_density_einsum_postamble(n, spin) - kwargs = { - "preamble": get_density_einsum_preamble(2, spin), - "postamble": get_postamble(2, spin), - } + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing, handle transpose + if spin != "uhf": + codegen.write("rdm2 = pack_2e(rdm2.oooo, rdm2.ooov, rdm2.oovo, rdm2.ovoo, rdm2.vooo, rdm2.oovv, rdm2.ovov, rdm2.ovvo, rdm2.voov, rdm2.vovo, rdm2.vvoo, rdm2.ovvv, rdm2.vovv, rdm2.vvov, rdm2.vvvo, rdm2.vvvv).transpose((0, 2, 1, 3))") + else: + codegen.write("rdm2.aaaa = pack_2e(rdm2.aaaa.oooo, rdm2.aaaa.ooov, rdm2.aaaa.oovo, rdm2.aaaa.ovoo, rdm2.aaaa.vooo, rdm2.aaaa.oovv, rdm2.aaaa.ovov, rdm2.aaaa.ovvo, rdm2.aaaa.voov, rdm2.aaaa.vovo, rdm2.aaaa.vvoo, rdm2.aaaa.ovvv, rdm2.aaaa.vovv, rdm2.aaaa.vvov, rdm2.aaaa.vvvo, rdm2.aaaa.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.aabb = pack_2e(rdm2.abab.oooo, rdm2.abab.ooov, rdm2.abab.oovo, rdm2.abab.ovoo, rdm2.abab.vooo, rdm2.abab.oovv, rdm2.abab.ovov, rdm2.abab.ovvo, rdm2.abab.voov, rdm2.abab.vovo, rdm2.abab.vvoo, rdm2.abab.ovvv, rdm2.abab.vovv, rdm2.abab.vvov, rdm2.abab.vvvo, rdm2.abab.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.bbbb = pack_2e(rdm2.bbbb.oooo, rdm2.bbbb.ooov, rdm2.bbbb.oovo, rdm2.bbbb.ovoo, rdm2.bbbb.vooo, rdm2.bbbb.oovv, rdm2.bbbb.ovov, rdm2.bbbb.ovvo, rdm2.bbbb.voov, rdm2.bbbb.vovo, rdm2.bbbb.vvoo, rdm2.bbbb.ovvv, rdm2.bbbb.vovv, rdm2.bbbb.vvov, rdm2.bbbb.vvvo, rdm2.bbbb.vvvv).transpose((0, 2, 1, 3))") + codegen.write("del rdm2.abab") + codegen( "make_rdm2_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("IP-EOM"): @@ -303,7 +228,7 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -314,60 +239,27 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("EA-EOM"): @@ -380,7 +272,7 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -391,60 +283,27 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -458,7 +317,7 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -469,60 +328,46 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_matvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) with Stopwatch("L-IP-EOM"): @@ -535,7 +380,7 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -546,60 +391,27 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip") # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("L-EA-EOM"): @@ -612,7 +424,7 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -623,60 +435,27 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea") # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -690,7 +469,7 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -701,60 +480,46 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee") # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_lmatvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) for codegen in code_generators.values(): diff --git a/ebcc/codegen/bootstrap_CCSD.py b/ebcc/codegen/bootstrap_CCSD.py index 79a7974a..364ad5de 100644 --- a/ebcc/codegen/bootstrap_CCSD.py +++ b/ebcc/codegen/bootstrap_CCSD.py @@ -5,32 +5,24 @@ import sys import pdaggerq -from albert.qc._pdaggerq import import_from_pdaggerq, remove_reference_energy -from albert.qc.spin import ghf_to_uhf, ghf_to_rhf, get_amplitude_spins, get_density_spins +from albert.qc._pdaggerq import remove_reference_energy, remove_reference_energy_eom +from albert.qc.spin import ghf_to_uhf, ghf_to_rhf from albert.qc import ghf, uhf, rhf from albert.tensor import Tensor from albert.index import Index from albert.code._ebcc import EBCCCodeGenerator from albert.misc import Stopwatch -from albert.opt import optimise +from albert.opt.tools import _tensor_info + +from ebcc.codegen.bootstrap_common import get_energy, get_amplitudes, get_rdm1, get_rdm2, get_eom # Get the spin case spin = sys.argv[1] -spin_integrate = { - "ghf": lambda *args, **kwargs: args, - "uhf": ghf_to_uhf, - "rhf": lambda *args, **kwargs: (ghf_to_rhf(*args, **kwargs),), -}[spin] -spin_module = { - "ghf": ghf, - "uhf": uhf, - "rhf": rhf, -}[spin] # Set up the code generators code_generators = { "einsum": EBCCCodeGenerator( - #stdout=open(f"{spin[0].upper()}CCSD.py", "w"), + stdout=open(f"{spin[0].upper()}CCSD.py", "w"), ), } @@ -53,11 +45,7 @@ terms = remove_reference_energy(terms) # Get the energy in albert format - expr = import_from_pdaggerq(terms) - expr = spin_integrate(expr) - output = tuple(Tensor(name="e_cc") for _ in range(len(expr))) - output_expr = optimise(output, expr, strategy="exhaust") - returns = (Tensor(name="e_cc"),) + output_expr, returns = get_energy(terms, spin) # Generate the energy code for codegen in code_generators.values(): @@ -67,103 +55,75 @@ output_expr, ) -#with Stopwatch("T amplitudes"): -# # Get the T1 contractions in pdaggerq format -# pq.clear() -# pq.set_left_operators([["e1(i,a)"]]) -# pq.add_st_operator(1.0, ["f"], ["t1", "t2"]) -# pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) -# pq.simplify() -# terms_t1 = pq.fully_contracted_strings() -# -# # Get the T2 contractions in pdaggerq format -# pq.clear() -# pq.set_left_operators([["e2(i,j,b,a)"]]) -# pq.add_st_operator(1.0, ["f"], ["t1", "t2"]) -# pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) -# pq.simplify() -# terms_t2 = pq.fully_contracted_strings() -# -# # Get the T amplitudes in albert format -# terms = [terms_t1, terms_t2] -# expr = [] -# output = [] -# returns = [] -# for n in range(2): -# indices = "ijkl"[: n + 1] + "abcd"[: n + 1] -# for index_spins in get_amplitude_spins(indices[: n + 1], indices[n + 1 :], spin): -# expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) -# expr_n = spin_integrate(expr_n) -# output_n = [Tensor(*tuple(sorted(e.external_indices, key=lambda i: indices.index(i.name))), name=f"t{n+1}new") for e in expr_n] -# returns_n = (output_n[0],) -# expr.extend(expr_n) -# output.extend(output_n) -# returns.extend(returns_n) -# output_expr = optimise(output, expr, strategy="exhaust") -# -# # Generate the T amplitude code -# for name, codegen in code_generators.items(): -# codegen( -# "update_amps", -# returns, -# output_expr, -# as_dict=True, -# ) -# -#with Stopwatch("L amplitudes"): -# # Get the L1 contractions in pdaggerq format -# pq.clear() -# pq.set_left_operators([["1"]]) -# pq.set_right_operators([["1"]]) -# pq.add_st_operator(1.0, ["f", "e1(a,i)"], ["t1", "t2"]) -# pq.add_st_operator(1.0, ["v", "e1(a,i)"], ["t1", "t2"]) -# pq.set_left_operators([["l1"], ["l2"]]) -# pq.add_st_operator(1.0, ["f", "e1(a,i)"], ["t1", "t2"]) -# pq.add_st_operator(1.0, ["v", "e1(a,i)"], ["t1", "t2"]) -# pq.add_st_operator(-1.0, ["e1(a,i)", "f"], ["t1", "t2"]) -# pq.add_st_operator(-1.0, ["e1(a,i)", "v"], ["t1", "t2"]) -# pq.simplify() -# terms_l1 = pq.fully_contracted_strings() -# -# # Get the L2 contractions in pdaggerq format -# pq.clear() -# pq.set_left_operators([["1"]]) -# pq.set_right_operators([["1"]]) -# pq.add_st_operator(1.0, ["f", "e2(a,b,j,i)"], ["t1", "t2"]) -# pq.add_st_operator(1.0, ["v", "e2(a,b,j,i)"], ["t1", "t2"]) -# pq.set_left_operators([["l1"], ["l2"]]) -# pq.add_st_operator(1.0, ["f", "e2(a,b,j,i)"], ["t1", "t2"]) -# pq.add_st_operator(1.0, ["v", "e2(a,b,j,i)"], ["t1", "t2"]) -# pq.add_st_operator(-1.0, ["e2(a,b,j,i)", "f"], ["t1", "t2"]) -# pq.add_st_operator(-1.0, ["e2(a,b,j,i)", "v"], ["t1", "t2"]) -# pq.simplify() -# terms_l2 = pq.fully_contracted_strings() -# -# # Get the L amplitudes in albert format -# terms = [terms_l1, terms_l2] -# expr = [] -# output = [] -# returns = [] -# for n in range(2): -# indices = "abcd"[: n + 1] + "ijkl"[: n + 1] -# for index_spins in get_amplitude_spins(indices[: n + 1], indices[n + 1 :], spin): -# expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) -# expr_n = spin_integrate(expr_n) -# output_n = [Tensor(*tuple(sorted(e.external_indices, key=lambda i: indices.index(i.name))), name=f"l{n+1}new") for e in expr_n] -# returns_n = (output_n[0],) -# expr.extend(expr_n) -# output.extend(output_n) -# returns.extend(returns_n) -# output_expr = optimise(output, expr, strategy="opt") -# -# # Generate the L amplitude code -# for name, codegen in code_generators.items(): -# codegen( -# "update_lams", -# returns, -# output_expr, -# as_dict=True, -# ) +with Stopwatch("T amplitudes"): + # Get the T1 contractions in pdaggerq format + pq.clear() + pq.set_left_operators([["e1(i,a)"]]) + pq.add_st_operator(1.0, ["f"], ["t1", "t2"]) + pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) + pq.simplify() + terms_t1 = pq.fully_contracted_strings() + + # Get the T2 contractions in pdaggerq format + pq.clear() + pq.set_left_operators([["e2(i,j,b,a)"]]) + pq.add_st_operator(1.0, ["f"], ["t1", "t2"]) + pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) + pq.simplify() + terms_t2 = pq.fully_contracted_strings() + + # Get the T amplitudes in albert format + output_expr, returns = get_amplitudes([terms_t1, terms_t2], spin) + + # Generate the T amplitude code + for name, codegen in code_generators.items(): + codegen( + "update_amps", + returns, + output_expr, + as_dict=True, + ) + +with Stopwatch("L amplitudes"): + # Get the L1 contractions in pdaggerq format + pq.clear() + pq.set_left_operators([["1"]]) + pq.set_right_operators([["1"]]) + pq.add_st_operator(1.0, ["f", "e1(a,i)"], ["t1", "t2"]) + pq.add_st_operator(1.0, ["v", "e1(a,i)"], ["t1", "t2"]) + pq.set_left_operators([["l1"], ["l2"]]) + pq.add_st_operator(1.0, ["f", "e1(a,i)"], ["t1", "t2"]) + pq.add_st_operator(1.0, ["v", "e1(a,i)"], ["t1", "t2"]) + pq.add_st_operator(-1.0, ["e1(a,i)", "f"], ["t1", "t2"]) + pq.add_st_operator(-1.0, ["e1(a,i)", "v"], ["t1", "t2"]) + pq.simplify() + terms_l1 = pq.fully_contracted_strings() + + # Get the L2 contractions in pdaggerq format + pq.clear() + pq.set_left_operators([["1"]]) + pq.set_right_operators([["1"]]) + pq.add_st_operator(1.0, ["f", "e2(a,b,j,i)"], ["t1", "t2"]) + pq.add_st_operator(1.0, ["v", "e2(a,b,j,i)"], ["t1", "t2"]) + pq.set_left_operators([["l1"], ["l2"]]) + pq.add_st_operator(1.0, ["f", "e2(a,b,j,i)"], ["t1", "t2"]) + pq.add_st_operator(1.0, ["v", "e2(a,b,j,i)"], ["t1", "t2"]) + pq.add_st_operator(-1.0, ["e2(a,b,j,i)", "f"], ["t1", "t2"]) + pq.add_st_operator(-1.0, ["e2(a,b,j,i)", "v"], ["t1", "t2"]) + pq.simplify() + terms_l2 = pq.fully_contracted_strings() + + # Get the L amplitudes in albert format + output_expr, returns = get_amplitudes([terms_l1, terms_l2], spin, strategy="opt" if spin == "uhf" else "exhaust", which="l") + + # Generate the L amplitude code + for name, codegen in code_generators.items(): + codegen( + "update_lams", + returns, + output_expr, + as_dict=True, + ) with Stopwatch("1RDM"): # Get the 1RDM contractions in pdaggerq format @@ -176,38 +136,40 @@ terms[sectors, indices] = pq.fully_contracted_strings() # Get the 1RDM in albert format - expr = [] - output = [] - returns = [] - deltas = [] - deltas_sources = [] - for sectors, indices in [("oo", "ij"), ("ov", "ia"), ("vo", "ai"), ("vv", "ab")]: - for index_spins in get_density_spins(indices, spin): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - expr_n = spin_integrate(expr_n) - if spin == "rhf": - expr_n = tuple(e * 2 for e in expr_n) - output_n = [Tensor(*tuple(sorted(e.external_indices, key=lambda i: indices.index(i.name))), name=f"d") for e in expr_n] - returns_n = (output_n[0],) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - if len(set(sectors)) == 1: - delta = spin_module.Delta(*tuple(sorted(expr_n[0].external_indices, key=lambda i: indices.index(i.name)))) - deltas.append(delta) - deltas_sources.append(next(expr_n[0].search_leaves(spin_module.T1))) - output_expr = optimise(output, expr, strategy="exhaust") + output_expr, returns, deltas, deltas_sources = get_rdm1(terms, spin) # Generate the 1RDM code for name, codegen in code_generators.items(): - for delta, delta_source in zip(deltas, deltas_sources): - codegen.tensor_declaration( - delta, is_identity=True, shape_source=delta_source, shape_source_index=0 - ) + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing + if spin != "uhf": + codegen.write("rdm1 = np.block([[rdm1.oo, rdm1.ov], [rdm1.vo, rdm1.vv]])") + else: + codegen.write("rdm1.aa = np.block([[rdm1.aa.oo, rdm1.aa.ov], [rdm1.aa.vo, rdm1.aa.vv]])") + codegen.write("rdm1.bb = np.block([[rdm1.bb.oo, rdm1.bb.ov], [rdm1.bb.vo, rdm1.bb.vv]])") + codegen( "make_rdm1_f", returns, output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("2RDM"): @@ -240,53 +202,42 @@ terms[sectors, indices] = pq.fully_contracted_strings() # Get the 2RDM in albert format - expr = [] - output = [] - returns = [] - deltas = [] - deltas_sources = [] - for sectors, indices in [ - ("oooo", "ijkl"), - ("ooov", "ijka"), - ("oovo", "ijak"), - ("ovoo", "iajk"), - ("vooo", "aijk"), - ("oovv", "ijab"), - ("ovov", "iajb"), - ("ovvo", "iabj"), - ("voov", "aijb"), - ("vovo", "aibj"), - ("vvoo", "abij"), - ("ovvv", "iabc"), - ("vovv", "aibc"), - ("vvov", "abic"), - ("vvvo", "abci"), - ("vvvv", "abcd"), - ]: - for index_spins in get_density_spins(indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - expr_n = spin_integrate(expr_n) - output_n = [Tensor(*tuple(sorted(e.external_indices, key=lambda i: indices.index(i.name))), name=f"Γ") for e in expr_n] - returns_n = (output_n[0],) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - if len(set(sectors)) == 1: - delta = spin_module.Delta(*tuple(sorted(expr_n[0].external_indices, key=lambda i: indices.index(i.name)))) - deltas.append(delta) - deltas_sources.append(next(expr_n[0].search_leaves(spin_module.T1))) - output_expr = optimise(output, expr, strategy="trav") + output_expr, returns, deltas, deltas_sources = get_rdm2(terms, spin, strategy="trav" if spin == "uhf" else "exhaust") # Generate the 2RDM code for name, codegen in code_generators.items(): - for delta, delta_source in zip(deltas, deltas_sources): - codegen.tensor_declaration( - delta, is_identity=True, shape_source=delta_source, shape_source_index=0 - ) + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing, handle transpose + if spin != "uhf": + codegen.write("rdm2 = pack_2e(rdm2.oooo, rdm2.ooov, rdm2.oovo, rdm2.ovoo, rdm2.vooo, rdm2.oovv, rdm2.ovov, rdm2.ovvo, rdm2.voov, rdm2.vovo, rdm2.vvoo, rdm2.ovvv, rdm2.vovv, rdm2.vvov, rdm2.vvvo, rdm2.vvvv).transpose((0, 2, 1, 3))") + else: + codegen.write("rdm2.aaaa = pack_2e(rdm2.aaaa.oooo, rdm2.aaaa.ooov, rdm2.aaaa.oovo, rdm2.aaaa.ovoo, rdm2.aaaa.vooo, rdm2.aaaa.oovv, rdm2.aaaa.ovov, rdm2.aaaa.ovvo, rdm2.aaaa.voov, rdm2.aaaa.vovo, rdm2.aaaa.vvoo, rdm2.aaaa.ovvv, rdm2.aaaa.vovv, rdm2.aaaa.vvov, rdm2.aaaa.vvvo, rdm2.aaaa.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.aabb = pack_2e(rdm2.abab.oooo, rdm2.abab.ooov, rdm2.abab.oovo, rdm2.abab.ovoo, rdm2.abab.vooo, rdm2.abab.oovv, rdm2.abab.ovov, rdm2.abab.ovvo, rdm2.abab.voov, rdm2.abab.vovo, rdm2.abab.vvoo, rdm2.abab.ovvv, rdm2.abab.vovv, rdm2.abab.vvov, rdm2.abab.vvvo, rdm2.abab.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.bbbb = pack_2e(rdm2.bbbb.oooo, rdm2.bbbb.ooov, rdm2.bbbb.oovo, rdm2.bbbb.ovoo, rdm2.bbbb.vooo, rdm2.bbbb.oovv, rdm2.bbbb.ovov, rdm2.bbbb.ovvo, rdm2.bbbb.voov, rdm2.bbbb.vovo, rdm2.bbbb.vvoo, rdm2.bbbb.ovvv, rdm2.bbbb.vovv, rdm2.bbbb.vvov, rdm2.bbbb.vvvo, rdm2.bbbb.vvvv).transpose((0, 2, 1, 3))") + codegen.write("del rdm2.abab") + codegen( "make_rdm2_f", returns, output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("IP-EOM"): @@ -299,7 +250,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -310,60 +261,27 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("EA-EOM"): @@ -376,7 +294,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -387,60 +305,27 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -454,7 +339,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -465,61 +350,46 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "postamble": "r2new.baba = np.transpose(r2new.abab, (1, 0, 3, 2))" if spin == "uhf" else None, # FIXME - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_matvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) with Stopwatch("L-IP-EOM"): @@ -532,7 +402,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -543,60 +413,27 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip") # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("L-EA-EOM"): @@ -609,7 +446,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -620,60 +457,27 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea") # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -687,7 +491,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -698,128 +502,48 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee") # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "postamble": "r2new.baba = np.transpose(r2new.abab, (1, 0, 3, 2))" if spin == "uhf" else None, # FIXME - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_lmatvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) -#with Stopwatch("1TDM (GS -> EE)"): -# # Get the R0 contractions in pdaggerq format -# pq.clear() -# pq.set_left_operators_type("EE") -# pq.set_left_operators([["1"]]) -# pq.set_right_operators_type("EE") -# pq.set_right_operators([["r1"], ["r2"]]) -# pq.add_st_operator(1.0, ["f"], ["t1", "t2"]) -# pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) -# pq.simplify() -# terms_r0 = pq.fully_contracted_strings() -# terms_r0 = remove_e0_eom(terms_r0) -# -# # Get the R1 contractions in pdaggerq format -# expr_r0 = spin_integrate(import_from_pdaggerq(terms_r0), spin) -# output_r0 = returns = (Tensor(name="r0"),) -# -# # Get the 1TDM contractions in pdaggerq format -# terms = {} -# for sectors, indices in [("oo", "ij"), ("ov", "ia"), ("vo", "ai"), ("vv", "ab")]: -# pq.clear() -# pq.set_left_operators_type("EE") -# pq.set_left_operators([["1"], ["l1"], ["l2"]]) -# pq.set_right_operators_type("EE") -# pq.set_right_operators([["r0"], ["r1"], ["r2"]]) -# pq.add_st_operator(1.0, [f"e1({','.join(indices)})"], ["t1", "t2"]) -# pq.simplify() -# terms[sectors, indices] = pq.fully_contracted_strings() -# -# # Get the 1TDM in albert format -# expr = [] -# output = [] -# returns = [] -# for sectors, indices in [("oo", "ij"), ("ov", "ia"), ("vo", "ai"), ("vv", "ab")]: -# for index_spins in get_density_spins(1, spin, indices): -# expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) -# expr_n = spin_integrate(expr_n, spin) -# if spin == "rhf": -# expr_n = tuple(e * 2 for e in expr_n) -# output_n = get_density_outputs(expr_n, f"d", indices) -# returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"d"),) -# expr.extend(expr_n) -# output.extend(output_n) -# returns.extend(returns_n) -# #output, expr = optimise_trans_dm(output, expr, spin, strategy="exhaust") -# output = tuple(output_r0) + tuple(output) -# expr = tuple(expr_r0) + tuple(expr) -# -# # Generate the 1TDM code -# for name, codegen in code_generators.items(): -# if name == "einsum": -# kwargs = { -# "preamble": get_density_einsum_preamble(1, spin), -# "postamble": get_density_einsum_postamble(1, spin), -# } -# else: -# kwargs = {} -# codegen( -# "make_trans_gs_to_ee_rdm1_f", -# returns, -# output, -# expr, -# **kwargs, -# ) - for codegen in code_generators.values(): codegen.postamble() codegen.stdout.close() diff --git a/ebcc/codegen/bootstrap_CCSDT.py b/ebcc/codegen/bootstrap_CCSDT.py index fb48193b..9fee0e05 100644 --- a/ebcc/codegen/bootstrap_CCSDT.py +++ b/ebcc/codegen/bootstrap_CCSDT.py @@ -3,23 +3,26 @@ """ import sys -sys.setrecursionlimit(2147483647) import pdaggerq -from albert.qc._pdaggerq import import_from_pdaggerq +from albert.qc._pdaggerq import remove_reference_energy, remove_reference_energy_eom +from albert.qc.spin import ghf_to_uhf, ghf_to_rhf +from albert.qc import ghf, uhf, rhf from albert.tensor import Tensor +from albert.index import Index +from albert.code._ebcc import EBCCCodeGenerator +from albert.misc import Stopwatch +from albert.opt.tools import _tensor_info -from ebcc.codegen.bootstrap_common import * +from ebcc.codegen.bootstrap_common import get_energy, get_amplitudes, get_rdm1, get_rdm2, get_eom # Get the spin case spin = sys.argv[1] # Set up the code generators code_generators = { - "einsum": EinsumCodeGen( + "einsum": EBCCCodeGenerator( stdout=open(f"{spin[0].upper()}CCSDT.py", "w"), - name_generator=name_generators[spin], - spin=spin, ), } @@ -39,22 +42,17 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms = pq.fully_contracted_strings() - terms = remove_hf_energy(terms) + terms = remove_reference_energy(terms) # Get the energy in albert format - expr = import_from_pdaggerq(terms) - expr = spin_integrate(expr, spin) - output = tuple(Tensor(name="e_cc") for _ in range(len(expr))) - output, expr = optimise(output, expr, spin, strategy="exhaust") - returns = (Tensor(name="e_cc"),) + output_expr, returns = get_energy(terms, spin) # Generate the energy code for codegen in code_generators.values(): codegen( "energy", returns, - output, - expr, + output_expr, ) with Stopwatch("T amplitudes"): @@ -83,37 +81,15 @@ terms_t3 = pq.fully_contracted_strings() # Get the T amplitudes in albert format - terms = [terms_t1, terms_t2, terms_t3] - expr = [] - output = [] - returns = [] - for n in range(3): - for index_spins in get_amplitude_spins(n + 1, spin): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"t{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"t{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust" if spin != "uhf" else "opt") + output_expr, returns = get_amplitudes([terms_t1, terms_t2, terms_t3], spin, strategy="greedy") # Generate the T amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "t1new = Namespace()\nt2new = Namespace()\nt3new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_amps", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) with Stopwatch("L amplitudes"): @@ -160,37 +136,15 @@ terms_l3 = pq.fully_contracted_strings() # Get the L amplitudes in albert format - terms = [terms_l1, terms_l2, terms_l3] - expr = [] - output = [] - returns = [] - for n in range(3): - for index_spins in get_amplitude_spins(n + 1, spin): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_l_amplitude_outputs(expr_n, f"l{n+1}new") - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"l{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="greedy") + output_expr, returns = get_amplitudes([terms_l1, terms_l2, terms_l3], spin, strategy="greedy", which="l") # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "l1new = Namespace()\nl2new = Namespace()\nl3new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_lams", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) with Stopwatch("1RDM"): @@ -204,37 +158,40 @@ terms[sectors, indices] = pq.fully_contracted_strings() # Get the 1RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [("oo", "ij"), ("ov", "ia"), ("vo", "ai"), ("vv", "ab")]: - for index_spins in get_density_spins(1, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - if spin == "rhf": - expr_n = tuple(e * 2 for e in expr_n) - output_n = get_density_outputs(expr_n, f"d", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"d"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="trav") + output_expr, returns, deltas, deltas_sources = get_rdm1(terms, spin) # Generate the 1RDM code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": get_density_einsum_preamble(1, spin), - "postamble": get_density_einsum_postamble(1, spin), - } - else: - kwargs = {} + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing + if spin != "uhf": + codegen.write("rdm1 = np.block([[rdm1.oo, rdm1.ov], [rdm1.vo, rdm1.vv]])") + else: + codegen.write("rdm1.aa = np.block([[rdm1.aa.oo, rdm1.aa.ov], [rdm1.aa.vo, rdm1.aa.vv]])") + codegen.write("rdm1.bb = np.block([[rdm1.bb.oo, rdm1.bb.ov], [rdm1.bb.vo, rdm1.bb.vv]])") + codegen( "make_rdm1_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("2RDM"): @@ -267,50 +224,42 @@ terms[sectors, indices] = pq.fully_contracted_strings() # Get the 2RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [ - ("oooo", "ijkl"), - ("ooov", "ijka"), - ("oovo", "ijak"), - ("ovoo", "iajk"), - ("vooo", "aijk"), - ("oovv", "ijab"), - ("ovov", "iajb"), - ("ovvo", "iabj"), - ("voov", "aijb"), - ("vovo", "aibj"), - ("vvoo", "abij"), - ("ovvv", "iabc"), - ("vovv", "aibc"), - ("vvov", "abic"), - ("vvvo", "abci"), - ("vvvv", "abcd"), - ]: - for index_spins in get_density_spins(2, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_density_outputs(expr_n, f"Γ", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"Γ"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="greedy") + output_expr, returns, deltas, deltas_sources = get_rdm2(terms, spin, strategy="trav" if spin == "uhf" else "exhaust") # Generate the 2RDM code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": get_density_einsum_preamble(2, spin), - "postamble": get_density_einsum_postamble(2, spin), - } + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing, handle transpose + if spin != "uhf": + codegen.write("rdm2 = pack_2e(rdm2.oooo, rdm2.ooov, rdm2.oovo, rdm2.ovoo, rdm2.vooo, rdm2.oovv, rdm2.ovov, rdm2.ovvo, rdm2.voov, rdm2.vovo, rdm2.vvoo, rdm2.ovvv, rdm2.vovv, rdm2.vvov, rdm2.vvvo, rdm2.vvvv).transpose((0, 2, 1, 3))") + else: + codegen.write("rdm2.aaaa = pack_2e(rdm2.aaaa.oooo, rdm2.aaaa.ooov, rdm2.aaaa.oovo, rdm2.aaaa.ovoo, rdm2.aaaa.vooo, rdm2.aaaa.oovv, rdm2.aaaa.ovov, rdm2.aaaa.ovvo, rdm2.aaaa.voov, rdm2.aaaa.vovo, rdm2.aaaa.vvoo, rdm2.aaaa.ovvv, rdm2.aaaa.vovv, rdm2.aaaa.vvov, rdm2.aaaa.vvvo, rdm2.aaaa.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.aabb = pack_2e(rdm2.abab.oooo, rdm2.abab.ooov, rdm2.abab.oovo, rdm2.abab.ovoo, rdm2.abab.vooo, rdm2.abab.oovv, rdm2.abab.ovov, rdm2.abab.ovvo, rdm2.abab.voov, rdm2.abab.vovo, rdm2.abab.vvoo, rdm2.abab.ovvv, rdm2.abab.vovv, rdm2.abab.vvov, rdm2.abab.vvvo, rdm2.abab.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.bbbb = pack_2e(rdm2.bbbb.oooo, rdm2.bbbb.ooov, rdm2.bbbb.oovo, rdm2.bbbb.ovoo, rdm2.bbbb.vooo, rdm2.bbbb.oovv, rdm2.bbbb.ovov, rdm2.bbbb.ovvo, rdm2.bbbb.voov, rdm2.bbbb.vovo, rdm2.bbbb.vvoo, rdm2.bbbb.ovvv, rdm2.bbbb.vovv, rdm2.bbbb.vvov, rdm2.bbbb.vvvo, rdm2.bbbb.vvvv).transpose((0, 2, 1, 3))") + codegen.write("del rdm2.abab") + codegen( "make_rdm2_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("IP-EOM"): @@ -323,7 +272,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -334,7 +283,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R3 contractions in pdaggerq format pq.clear() @@ -345,60 +294,27 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r3 = pq.fully_contracted_strings() - terms_r3 = remove_e0_eom(terms_r3) + terms_r3 = remove_reference_energy_eom(terms_r3) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2, terms_r3] - expr = [] - output = [] - returns = [] - for n in range(3): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="greedy") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2, terms_r3], spin, strategy="greedy", which="ip") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("EA-EOM"): @@ -411,7 +327,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -422,7 +338,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R3 contractions in pdaggerq format pq.clear() @@ -433,60 +349,27 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r3 = pq.fully_contracted_strings() - terms_r3 = remove_e0_eom(terms_r3) + terms_r3 = remove_reference_energy_eom(terms_r3) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2, terms_r3] - expr = [] - output = [] - returns = [] - for n in range(3): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="greedy") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2, terms_r3], spin, strategy="greedy", which="ea") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -500,7 +383,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -511,7 +394,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R3 contractions in pdaggerq format pq.clear() @@ -522,61 +405,46 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r3 = pq.fully_contracted_strings() - terms_r3 = remove_e0_eom(terms_r3) + terms_r3 = remove_reference_energy_eom(terms_r3) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2, terms_r3] - expr = [] - output = [] - returns = [] - for n in range(3): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="greedy") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2, terms_r3], spin, strategy="greedy", which="ee") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "postamble": "r2new.baba = np.transpose(r2new.abab, (1, 0, 3, 2))" if spin == "uhf" else None, # FIXME - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_matvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) with Stopwatch("L-IP-EOM"): @@ -589,7 +457,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -600,7 +468,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L3 contractions in pdaggerq format pq.clear() @@ -611,60 +479,27 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r3 = pq.fully_contracted_strings() - terms_r3 = remove_e0_eom(terms_r3) + terms_r3 = remove_reference_energy_eom(terms_r3) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2, terms_r3] - expr = [] - output = [] - returns = [] - for n in range(3): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="greedy") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2, terms_r3], spin, strategy="greedy", which="ip") # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("L-EA-EOM"): @@ -677,7 +512,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -688,7 +523,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L3 contractions in pdaggerq format pq.clear() @@ -699,60 +534,27 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r3 = pq.fully_contracted_strings() - terms_r3 = remove_e0_eom(terms_r3) + terms_r3 = remove_reference_energy_eom(terms_r3) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2, terms_r3] - expr = [] - output = [] - returns = [] - for n in range(3): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="greedy") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2, terms_r3], spin, strategy="greedy", which="ea") # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -766,7 +568,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -777,7 +579,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L3 contractions in pdaggerq format pq.clear() @@ -788,61 +590,46 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_r3 = pq.fully_contracted_strings() - terms_r3 = remove_e0_eom(terms_r3) + terms_r3 = remove_reference_energy_eom(terms_r3) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2, terms_r3] - expr = [] - output = [] - returns = [] - for n in range(3): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="greedy") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2, terms_r3], spin, strategy="greedy", which="ee") # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "postamble": "r2new.baba = np.transpose(r2new.abab, (1, 0, 3, 2))" if spin == "uhf" else None, # FIXME - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_lmatvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) for codegen in code_generators.values(): diff --git a/ebcc/codegen/bootstrap_CCSDTQ.py b/ebcc/codegen/bootstrap_CCSDTQ.py index 243229b7..7107a8b7 100644 --- a/ebcc/codegen/bootstrap_CCSDTQ.py +++ b/ebcc/codegen/bootstrap_CCSDTQ.py @@ -3,23 +3,26 @@ """ import sys -sys.setrecursionlimit(1000000) import pdaggerq -from albert.qc._pdaggerq import import_from_pdaggerq +from albert.qc._pdaggerq import remove_reference_energy, remove_reference_energy_eom +from albert.qc.spin import ghf_to_uhf, ghf_to_rhf +from albert.qc import ghf, uhf, rhf from albert.tensor import Tensor +from albert.index import Index +from albert.code._ebcc import EBCCCodeGenerator +from albert.misc import Stopwatch +from albert.opt.tools import _tensor_info -from ebcc.codegen.bootstrap_common import * +from ebcc.codegen.bootstrap_common import get_energy, get_amplitudes, get_rdm1, get_rdm2, get_eom # Get the spin case spin = sys.argv[1] # Set up the code generators code_generators = { - "einsum": EinsumCodeGen( + "einsum": EBCCCodeGenerator( stdout=open(f"{spin[0].upper()}CCSDTQ.py", "w"), - name_generator=name_generators[spin], - spin=spin, ), } @@ -39,89 +42,62 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms = pq.fully_contracted_strings() - terms = remove_hf_energy(terms) + terms = remove_reference_energy(terms) # Get the energy in albert format - expr = import_from_pdaggerq(terms) - expr = spin_integrate(expr, spin) - output = tuple(Tensor(name="e_cc") for _ in range(len(expr))) - output, expr = optimise(output, expr, spin, strategy="exhaust") - returns = (Tensor(name="e_cc"),) + output_expr, returns = get_energy(terms, spin) # Generate the energy code for codegen in code_generators.values(): codegen( "energy", returns, - output, - expr, + output_expr, ) with Stopwatch("T amplitudes"): # Get the T1 contractions in pdaggerq format pq.clear() pq.set_left_operators([["e1(i,a)"]]) - pq.add_st_operator(1.0, ["f"], ["t1", "t2", "t3", "t4"]) - pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3", "t4"]) + pq.add_st_operator(1.0, ["f"], ["t1", "t2", "t3"]) + pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_t1 = pq.fully_contracted_strings() # Get the T2 contractions in pdaggerq format pq.clear() pq.set_left_operators([["e2(i,j,b,a)"]]) - pq.add_st_operator(1.0, ["f"], ["t1", "t2", "t3", "t4"]) - pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3", "t4"]) + pq.add_st_operator(1.0, ["f"], ["t1", "t2", "t3"]) + pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_t2 = pq.fully_contracted_strings() - # Get the T2 contractions in pdaggerq format + # Get the T3 contractions in pdaggerq format pq.clear() - pq.set_left_operators([["e2(i,j,k,c,b,a)"]]) - pq.add_st_operator(1.0, ["f"], ["t1", "t2", "t3", "t4"]) - pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3", "t4"]) + pq.set_left_operators([["e3(i,j,k,c,b,a)"]]) + pq.add_st_operator(1.0, ["f"], ["t1", "t2", "t3"]) + pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3"]) pq.simplify() terms_t3 = pq.fully_contracted_strings() # Get the T3 contractions in pdaggerq format pq.clear() - pq.set_left_operators([["e3(i,j,k,l,d,c,b,a)"]]) + pq.set_left_operators([["e4(i,j,k,l,d,c,b,a)"]]) pq.add_st_operator(1.0, ["f"], ["t1", "t2", "t3", "t4"]) pq.add_st_operator(1.0, ["v"], ["t1", "t2", "t3", "t4"]) pq.simplify() terms_t4 = pq.fully_contracted_strings() # Get the T amplitudes in albert format - terms = [terms_t1, terms_t2, terms_t3, terms_t4] - expr = [] - output = [] - returns = [] - for n in range(4): - for index_spins in get_amplitude_spins(n + 1, spin): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"t{n+1}new") - returns_n = (Tensor(*indices, name=f"t{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="greedy") + output_expr, returns = get_amplitudes([terms_t1, terms_t2, terms_t3, terms_t4], spin, strategy="greedy") # Generate the T amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "t1new = Namespace()\nt2new = Namespace()\nt3new = Namespace()\nt4new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_amps", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) for codegen in code_generators.values(): diff --git a/ebcc/codegen/bootstrap_DCD.py b/ebcc/codegen/bootstrap_DCD.py index 48133592..91e7ca8b 100644 --- a/ebcc/codegen/bootstrap_DCD.py +++ b/ebcc/codegen/bootstrap_DCD.py @@ -5,21 +5,24 @@ import sys import pdaggerq -from albert.qc._pdaggerq import import_from_pdaggerq -from albert.qc.index import Index +from albert.qc._pdaggerq import remove_reference_energy, remove_reference_energy_eom +from albert.qc.spin import ghf_to_uhf, ghf_to_rhf +from albert.qc import ghf, uhf, rhf from albert.tensor import Tensor +from albert.index import Index +from albert.code._ebcc import EBCCCodeGenerator +from albert.misc import Stopwatch +from albert.opt.tools import _tensor_info -from ebcc.codegen.bootstrap_common import * +from ebcc.codegen.bootstrap_common import get_energy, get_amplitudes, get_rdm1, get_rdm2, get_eom # Get the spin case spin = sys.argv[1] # Set up the code generators code_generators = { - "einsum": EinsumCodeGen( + "einsum": EBCCCodeGenerator( stdout=open(f"{spin[0].upper()}DCD.py", "w"), - name_generator=name_generators[spin], - spin=spin, ), } @@ -35,26 +38,21 @@ # Get the energy contractions in pdaggerq format pq.clear() pq.set_left_operators([["1"]]) - pq.add_st_operator(1.0, ["f"], ["t2"]) - pq.add_st_operator(1.0, ["v"], ["t2"]) + pq.add_st_operator(1.0, ["f"], ["t1", "t2"]) + pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms = pq.fully_contracted_strings() - terms = remove_hf_energy(terms) + terms = remove_reference_energy(terms) # Get the energy in albert format - expr = import_from_pdaggerq(terms) - expr = spin_integrate(expr, spin) - output = tuple(Tensor(name="e_cc") for _ in range(len(expr))) - output, expr = optimise(output, expr, spin, strategy="exhaust") - returns = (Tensor(name="e_cc"),) + output_expr, returns = get_energy(terms, spin) # Generate the energy code for codegen in code_generators.values(): codegen( "energy", returns, - output, - expr, + output_expr, ) with Stopwatch("T amplitudes"): @@ -73,36 +71,15 @@ ] # Get the T amplitudes in albert format - expr = [] - output = [] - returns = [] - for index_spins in get_amplitude_spins(2, spin): - indices = default_indices["o"][:2] + default_indices["v"][:2] - indices = tuple(Index(i, index_spins[i]) for i in indices) - expr_n = import_from_pdaggerq(terms, index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"t2new") - returns_n = (Tensor(*indices, name=f"t2new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns = get_amplitudes([terms], spin, orders=[2]) # Generate the T amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "t1new = Namespace()\nt2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_amps", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) for codegen in code_generators.values(): diff --git a/ebcc/codegen/bootstrap_DCSD.py b/ebcc/codegen/bootstrap_DCSD.py index 30f50d9a..7a4c8495 100644 --- a/ebcc/codegen/bootstrap_DCSD.py +++ b/ebcc/codegen/bootstrap_DCSD.py @@ -5,21 +5,24 @@ import sys import pdaggerq -from albert.qc._pdaggerq import import_from_pdaggerq -from albert.qc.index import Index +from albert.qc._pdaggerq import remove_reference_energy, remove_reference_energy_eom +from albert.qc.spin import ghf_to_uhf, ghf_to_rhf +from albert.qc import ghf, uhf, rhf from albert.tensor import Tensor +from albert.index import Index +from albert.code._ebcc import EBCCCodeGenerator +from albert.misc import Stopwatch +from albert.opt.tools import _tensor_info -from ebcc.codegen.bootstrap_common import * +from ebcc.codegen.bootstrap_common import get_energy, get_amplitudes, get_rdm1, get_rdm2, get_eom # Get the spin case spin = sys.argv[1] # Set up the code generators code_generators = { - "einsum": EinsumCodeGen( + "einsum": EBCCCodeGenerator( stdout=open(f"{spin[0].upper()}DCSD.py", "w"), - name_generator=name_generators[spin], - spin=spin, ), } @@ -39,22 +42,17 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms = pq.fully_contracted_strings() - terms = remove_hf_energy(terms) + terms = remove_reference_energy(terms) # Get the energy in albert format - expr = import_from_pdaggerq(terms) - expr = spin_integrate(expr, spin) - output = tuple(Tensor(name="e_cc") for _ in range(len(expr))) - output, expr = optimise(output, expr, spin, strategy="exhaust") - returns = (Tensor(name="e_cc"),) + output_expr, returns = get_energy(terms, spin) # Generate the energy code for codegen in code_generators.values(): codegen( "energy", returns, - output, - expr, + output_expr, ) with Stopwatch("T amplitudes"): @@ -83,40 +81,18 @@ ] # Get the T amplitudes in albert format - terms = [terms_t1, terms_t2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - indices = tuple(Index(i, index_spins[i]) for i in indices) - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"t{n+1}new") - returns_n = (Tensor(*indices, name=f"t{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns = get_amplitudes([terms], spin, orders=[2]) # Generate the T amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "t1new = Namespace()\nt2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_amps", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) for codegen in code_generators.values(): codegen.postamble() codegen.stdout.close() + diff --git a/ebcc/codegen/bootstrap_DFCC2.py b/ebcc/codegen/bootstrap_DFCC2.py index 1726e81a..d0e230bb 100644 --- a/ebcc/codegen/bootstrap_DFCC2.py +++ b/ebcc/codegen/bootstrap_DFCC2.py @@ -5,23 +5,24 @@ import sys import pdaggerq -from albert.qc._pdaggerq import import_from_pdaggerq -from albert.qc.index import Index +from albert.qc._pdaggerq import remove_reference_energy, remove_reference_energy_eom +from albert.qc.spin import ghf_to_uhf, ghf_to_rhf +from albert.qc import ghf, uhf, rhf from albert.tensor import Tensor +from albert.index import Index +from albert.code._ebcc import EBCCCodeGenerator +from albert.misc import Stopwatch +from albert.opt.tools import _tensor_info -from ebcc.codegen.bootstrap_common import * +from ebcc.codegen.bootstrap_common import get_energy, get_amplitudes, get_rdm1, get_rdm2, get_eom # Get the spin case spin = sys.argv[1] -if spin == "ghf": - raise NotImplementedError # Set up the code generators code_generators = { - "einsum": EinsumCodeGen( + "einsum": EBCCCodeGenerator( stdout=open(f"{spin[0].upper()}DFCC2.py", "w"), - name_generator=name_generators[spin], - spin=spin, ), } @@ -41,23 +42,17 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms = pq.fully_contracted_strings() - terms = remove_hf_energy(terms) + terms = remove_reference_energy(terms) # Get the energy in albert format - expr = import_from_pdaggerq(terms) - expr = spin_integrate(expr, spin) - expr = tuple(e.apply(get_density_fit()) for e in expr) - output = tuple(Tensor(name="e_cc") for _ in range(len(expr))) - output, expr = optimise(output, expr, spin, strategy="exhaust") - returns = (Tensor(name="e_cc"),) + output_expr, returns = get_energy(terms, spin, density_fit=True) # Generate the energy code for codegen in code_generators.values(): codegen( "energy", returns, - output, - expr, + output_expr, ) with Stopwatch("T amplitudes"): @@ -78,38 +73,15 @@ terms_t2 = pq.fully_contracted_strings() # Get the T amplitudes in albert format - terms = [terms_t1, terms_t2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"t{n+1}new") - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"t{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns = get_amplitudes([terms_t1, terms_t2], spin, density_fit=True) # Generate the T amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "t1new = Namespace()\nt2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_amps", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) with Stopwatch("L amplitudes"): @@ -142,38 +114,15 @@ terms_l2 = pq.fully_contracted_strings() # Get the L amplitudes in albert format - terms = [terms_l1, terms_l2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_l_amplitude_outputs(expr_n, f"l{n+1}new") - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"l{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="opt") + output_expr, returns = get_amplitudes([terms_l1, terms_l2], spin, strategy="opt" if spin == "uhf" else "exhaust", which="l", density_fit=True) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "l1new = Namespace()\nl2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_lams", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) with Stopwatch("1RDM"): @@ -187,38 +136,40 @@ terms[sectors, indices] = pq.fully_contracted_strings() # Get the 1RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [("oo", "ij"), ("ov", "ia"), ("vo", "ai"), ("vv", "ab")]: - for index_spins in get_density_spins(1, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - if spin == "rhf": - expr_n = tuple(e * 2 for e in expr_n) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_density_outputs(expr_n, f"d", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"d"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns, deltas, deltas_sources = get_rdm1(terms, spin, density_fit=True) # Generate the 1RDM code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": get_density_einsum_preamble(1, spin), - "postamble": get_density_einsum_postamble(1, spin), - } - else: - kwargs = {} + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing + if spin != "uhf": + codegen.write("rdm1 = np.block([[rdm1.oo, rdm1.ov], [rdm1.vo, rdm1.vv]])") + else: + codegen.write("rdm1.aa = np.block([[rdm1.aa.oo, rdm1.aa.ov], [rdm1.aa.vo, rdm1.aa.vv]])") + codegen.write("rdm1.bb = np.block([[rdm1.bb.oo, rdm1.bb.ov], [rdm1.bb.vo, rdm1.bb.vv]])") + codegen( "make_rdm1_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("2RDM"): @@ -251,51 +202,42 @@ terms[sectors, indices] = pq.fully_contracted_strings() # Get the 2RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [ - ("oooo", "ijkl"), - ("ooov", "ijka"), - ("oovo", "ijak"), - ("ovoo", "iajk"), - ("vooo", "aijk"), - ("oovv", "ijab"), - ("ovov", "iajb"), - ("ovvo", "iabj"), - ("voov", "aijb"), - ("vovo", "aibj"), - ("vvoo", "abij"), - ("ovvv", "iabc"), - ("vovv", "aibc"), - ("vvov", "abic"), - ("vvvo", "abci"), - ("vvvv", "abcd"), - ]: - for index_spins in get_density_spins(2, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_density_outputs(expr_n, f"Γ", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"Γ"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="trav") + output_expr, returns, deltas, deltas_sources = get_rdm2(terms, spin, strategy="trav" if spin == "uhf" else "exhaust", density_fit=True) # Generate the 2RDM code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": get_density_einsum_preamble(2, spin), - "postamble": get_density_einsum_postamble(2, spin), - } + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing, handle transpose + if spin != "uhf": + codegen.write("rdm2 = pack_2e(rdm2.oooo, rdm2.ooov, rdm2.oovo, rdm2.ovoo, rdm2.vooo, rdm2.oovv, rdm2.ovov, rdm2.ovvo, rdm2.voov, rdm2.vovo, rdm2.vvoo, rdm2.ovvv, rdm2.vovv, rdm2.vvov, rdm2.vvvo, rdm2.vvvv).transpose((0, 2, 1, 3))") + else: + codegen.write("rdm2.aaaa = pack_2e(rdm2.aaaa.oooo, rdm2.aaaa.ooov, rdm2.aaaa.oovo, rdm2.aaaa.ovoo, rdm2.aaaa.vooo, rdm2.aaaa.oovv, rdm2.aaaa.ovov, rdm2.aaaa.ovvo, rdm2.aaaa.voov, rdm2.aaaa.vovo, rdm2.aaaa.vvoo, rdm2.aaaa.ovvv, rdm2.aaaa.vovv, rdm2.aaaa.vvov, rdm2.aaaa.vvvo, rdm2.aaaa.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.aabb = pack_2e(rdm2.abab.oooo, rdm2.abab.ooov, rdm2.abab.oovo, rdm2.abab.ovoo, rdm2.abab.vooo, rdm2.abab.oovv, rdm2.abab.ovov, rdm2.abab.ovvo, rdm2.abab.voov, rdm2.abab.vovo, rdm2.abab.vvoo, rdm2.abab.ovvv, rdm2.abab.vovv, rdm2.abab.vvov, rdm2.abab.vvvo, rdm2.abab.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.bbbb = pack_2e(rdm2.bbbb.oooo, rdm2.bbbb.ooov, rdm2.bbbb.oovo, rdm2.bbbb.ovoo, rdm2.bbbb.vooo, rdm2.bbbb.oovv, rdm2.bbbb.ovov, rdm2.bbbb.ovvo, rdm2.bbbb.voov, rdm2.bbbb.vovo, rdm2.bbbb.vvoo, rdm2.bbbb.ovvv, rdm2.bbbb.vovv, rdm2.bbbb.vvov, rdm2.bbbb.vvvo, rdm2.bbbb.vvvv).transpose((0, 2, 1, 3))") + codegen.write("del rdm2.abab") + codegen( "make_rdm2_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("IP-EOM"): @@ -308,7 +250,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -319,61 +261,27 @@ pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip", density_fit=True) # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("EA-EOM"): @@ -386,7 +294,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -397,61 +305,27 @@ pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea", density_fit=True) # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -465,7 +339,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -476,61 +350,46 @@ pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee", density_fit=True) # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_matvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) with Stopwatch("L-IP-EOM"): @@ -543,7 +402,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -554,61 +413,27 @@ pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip", density_fit=True) # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("L-EA-EOM"): @@ -621,7 +446,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -632,61 +457,27 @@ pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea", density_fit=True) # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -700,7 +491,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -711,61 +502,46 @@ pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee", density_fit=True) # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_lmatvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) for codegen in code_generators.values(): diff --git a/ebcc/codegen/bootstrap_DFCCD.py b/ebcc/codegen/bootstrap_DFCCD.py index 7423eb34..468d832a 100644 --- a/ebcc/codegen/bootstrap_DFCCD.py +++ b/ebcc/codegen/bootstrap_DFCCD.py @@ -5,22 +5,24 @@ import sys import pdaggerq -from albert.qc._pdaggerq import import_from_pdaggerq +from albert.qc._pdaggerq import remove_reference_energy, remove_reference_energy_eom +from albert.qc.spin import ghf_to_uhf, ghf_to_rhf +from albert.qc import ghf, uhf, rhf from albert.tensor import Tensor +from albert.index import Index +from albert.code._ebcc import EBCCCodeGenerator +from albert.misc import Stopwatch +from albert.opt.tools import _tensor_info -from ebcc.codegen.bootstrap_common import * +from ebcc.codegen.bootstrap_common import get_energy, get_amplitudes, get_rdm1, get_rdm2, get_eom # Get the spin case spin = sys.argv[1] -if spin == "ghf": - raise NotImplementedError # Set up the code generators code_generators = { - "einsum": EinsumCodeGen( + "einsum": EBCCCodeGenerator( stdout=open(f"{spin[0].upper()}DFCCD.py", "w"), - name_generator=name_generators[spin], - spin=spin, ), } @@ -40,23 +42,17 @@ pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms = pq.fully_contracted_strings() - terms = remove_hf_energy(terms) + terms = remove_reference_energy(terms) # Get the energy in albert format - expr = import_from_pdaggerq(terms) - expr = spin_integrate(expr, spin) - expr = tuple(e.apply(get_density_fit()) for e in expr) - output = tuple(Tensor(name="e_cc") for _ in range(len(expr))) - output, expr = optimise(output, expr, spin, strategy="exhaust") - returns = (Tensor(name="e_cc"),) + output_expr, returns = get_energy(terms, spin, density_fit=True) # Generate the energy code for codegen in code_generators.values(): codegen( "energy", returns, - output, - expr, + output_expr, ) with Stopwatch("T amplitudes"): @@ -69,36 +65,15 @@ terms = pq.fully_contracted_strings() # Get the T amplitudes in albert format - expr = [] - output = [] - returns = [] - for index_spins in get_amplitude_spins(2, spin): - indices = default_indices["o"][: 2] + default_indices["v"][: 2] - expr_n = import_from_pdaggerq(terms, index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"t2new") - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"t2new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns = get_amplitudes([terms], spin, orders=[2], density_fit=True) # Generate the T amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "t1new = Namespace()\nt2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_amps", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) with Stopwatch("L amplitudes"): @@ -109,44 +84,23 @@ pq.add_st_operator(1.0, ["f", "e2(a,b,j,i)"], ["t2"]) pq.add_st_operator(1.0, ["v", "e2(a,b,j,i)"], ["t2"]) pq.set_left_operators([["l2"]]) - pq.add_st_operator(1.0, ["f", "e2(a,b,j,i)"], ["t2"]) - pq.add_st_operator(1.0, ["v", "e2(a,b,j,i)"], ["t2"]) + pq.add_st_operator( 1.0, ["f", "e2(a,b,j,i)"], ["t2"]) + pq.add_st_operator( 1.0, ["v", "e2(a,b,j,i)"], ["t2"]) pq.add_st_operator(-1.0, ["e2(a,b,j,i)", "f"], ["t2"]) pq.add_st_operator(-1.0, ["e2(a,b,j,i)", "v"], ["t2"]) pq.simplify() terms = pq.fully_contracted_strings() # Get the L amplitudes in albert format - expr = [] - output = [] - returns = [] - for index_spins in get_amplitude_spins(2, spin): - indices = default_indices["v"][: 2] + default_indices["o"][: 2] - expr_n = import_from_pdaggerq(terms, index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_l_amplitude_outputs(expr_n, f"l2new") - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"l2new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="opt") + get_amplitudes([terms], spin, strategy="opt" if spin == "uhf" else "exhaust", which="l", orders=[2], density_fit=True) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "l1new = Namespace()\nl2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_lams", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) with Stopwatch("1RDM"): @@ -160,52 +114,40 @@ terms[sectors, indices] = pq.fully_contracted_strings() # Get the 1RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [("oo", "ij"), ("ov", "ia"), ("vo", "ai"), ("vv", "ab")]: - for index_spins in get_density_spins(1, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - if not (isinstance(expr_n, int) and expr_n == 0): - expr_n = spin_integrate(expr_n, spin) - if spin == "rhf": - expr_n = tuple(e * 2 for e in expr_n) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_density_outputs(expr_n, f"d", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"d"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns, deltas, deltas_sources = get_rdm1(terms, spin, density_fit=True) # Generate the 1RDM code for name, codegen in code_generators.items(): - if name == "einsum": - def get_postamble(n, spin, name="rdm{n}"): - nm = name.format(n=n) - postamble = "" - if spin != "uhf": - for occ in ("ov", "vo"): - shape = ", ".join(f"t2.shape[{'0' if o == 'o' else '-1'}]" for o in occ) - postamble += f"{nm}.{occ} = np.zeros(({shape}))\n" - else: - for s in "ab": - for occ in ("oo", "vv"): - shape = ", ".join(f"t2.{ss}{ss}.shape[{'0' if o == 'o' else '-1'}]" for ss, o in zip(s+s, occ)) - postamble += f"{nm}.{s}{s}.{occ} = np.zeros(({shape}))\n" - return postamble + get_density_einsum_postamble(n, spin) - kwargs = { - "preamble": get_density_einsum_preamble(1, spin), - "postamble": get_postamble(1, spin), - } - else: - kwargs = {} + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing + if spin != "uhf": + codegen.write("rdm1 = np.block([[rdm1.oo, rdm1.ov], [rdm1.vo, rdm1.vv]])") + else: + codegen.write("rdm1.aa = np.block([[rdm1.aa.oo, rdm1.aa.ov], [rdm1.aa.vo, rdm1.aa.vv]])") + codegen.write("rdm1.bb = np.block([[rdm1.bb.oo, rdm1.bb.ov], [rdm1.bb.vo, rdm1.bb.vv]])") + codegen( "make_rdm1_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("2RDM"): @@ -238,66 +180,42 @@ def get_postamble(n, spin, name="rdm{n}"): terms[sectors, indices] = pq.fully_contracted_strings() # Get the 2RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [ - ("oooo", "ijkl"), - ("ooov", "ijka"), - ("oovo", "ijak"), - ("ovoo", "iajk"), - ("vooo", "aijk"), - ("oovv", "ijab"), - ("ovov", "iajb"), - ("ovvo", "iabj"), - ("voov", "aijb"), - ("vovo", "aibj"), - ("vvoo", "abij"), - ("ovvv", "iabc"), - ("vovv", "aibc"), - ("vvov", "abic"), - ("vvvo", "abci"), - ("vvvv", "abcd"), - ]: - for index_spins in get_density_spins(2, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - if not (isinstance(expr_n, int) and expr_n == 0): - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_density_outputs(expr_n, f"Γ", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"Γ"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="trav") + output_expr, returns, deltas, deltas_sources = get_rdm2(terms, spin, strategy="trav" if spin == "uhf" else "exhaust", density_fit=True) # Generate the 2RDM code for name, codegen in code_generators.items(): - if name == "einsum": - def get_postamble(n, spin, name="rdm{n}"): - nm = name.format(n=n) - postamble = "" - if spin != "uhf": - for occ in ("ooov", "oovo", "ovoo", "vooo", "ovvv", "vovv", "vvov", "vvvo"): - shape = ", ".join(f"t2.shape[{'0' if o == 'o' else '-1'}]" for o in occ) - postamble += f"{nm}.{occ} = np.zeros(({shape}))\n" - else: - for s1 in "ab": - for s2 in "ab": - for occ in ("ooov", "oovo", "ovoo", "vooo", "ovvv", "vovv", "vvov", "vvvo"): - shape = ", ".join(f"t2.{s}{s}{s}{s}.shape[{'0' if o == 'o' else '-1'}]" for o, s in zip(occ, s1+s1+s2+s2)) - postamble += f"{nm}.{s1}{s1}{s2}{s2}.{occ} = np.zeros(({shape}))\n" - return postamble + get_density_einsum_postamble(n, spin) - kwargs = { - "preamble": get_density_einsum_preamble(2, spin), - "postamble": get_postamble(2, spin), - } + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing, handle transpose + if spin != "uhf": + codegen.write("rdm2 = pack_2e(rdm2.oooo, rdm2.ooov, rdm2.oovo, rdm2.ovoo, rdm2.vooo, rdm2.oovv, rdm2.ovov, rdm2.ovvo, rdm2.voov, rdm2.vovo, rdm2.vvoo, rdm2.ovvv, rdm2.vovv, rdm2.vvov, rdm2.vvvo, rdm2.vvvv).transpose((0, 2, 1, 3))") + else: + codegen.write("rdm2.aaaa = pack_2e(rdm2.aaaa.oooo, rdm2.aaaa.ooov, rdm2.aaaa.oovo, rdm2.aaaa.ovoo, rdm2.aaaa.vooo, rdm2.aaaa.oovv, rdm2.aaaa.ovov, rdm2.aaaa.ovvo, rdm2.aaaa.voov, rdm2.aaaa.vovo, rdm2.aaaa.vvoo, rdm2.aaaa.ovvv, rdm2.aaaa.vovv, rdm2.aaaa.vvov, rdm2.aaaa.vvvo, rdm2.aaaa.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.aabb = pack_2e(rdm2.abab.oooo, rdm2.abab.ooov, rdm2.abab.oovo, rdm2.abab.ovoo, rdm2.abab.vooo, rdm2.abab.oovv, rdm2.abab.ovov, rdm2.abab.ovvo, rdm2.abab.voov, rdm2.abab.vovo, rdm2.abab.vvoo, rdm2.abab.ovvv, rdm2.abab.vovv, rdm2.abab.vvov, rdm2.abab.vvvo, rdm2.abab.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.bbbb = pack_2e(rdm2.bbbb.oooo, rdm2.bbbb.ooov, rdm2.bbbb.oovo, rdm2.bbbb.ovoo, rdm2.bbbb.vooo, rdm2.bbbb.oovv, rdm2.bbbb.ovov, rdm2.bbbb.ovvo, rdm2.bbbb.voov, rdm2.bbbb.vovo, rdm2.bbbb.vvoo, rdm2.bbbb.ovvv, rdm2.bbbb.vovv, rdm2.bbbb.vvov, rdm2.bbbb.vvvo, rdm2.bbbb.vvvv).transpose((0, 2, 1, 3))") + codegen.write("del rdm2.abab") + codegen( "make_rdm2_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("IP-EOM"): @@ -310,7 +228,7 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -321,61 +239,27 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip", density_fit=True) # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("EA-EOM"): @@ -388,7 +272,7 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -399,61 +283,27 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea", density_fit=True) # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -467,7 +317,7 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -478,61 +328,46 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee", density_fit=True) # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_matvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) with Stopwatch("L-IP-EOM"): @@ -545,7 +380,7 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -556,61 +391,27 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip", density_fit=True) # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("L-EA-EOM"): @@ -623,7 +424,7 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -634,61 +435,27 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea", density_fit=True) # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -702,7 +469,7 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -713,61 +480,46 @@ def get_postamble(n, spin, name="rdm{n}"): pq.add_st_operator(1.0, ["v"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee", density_fit=True) # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_lmatvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) for codegen in code_generators.values(): diff --git a/ebcc/codegen/bootstrap_DFCCSD.py b/ebcc/codegen/bootstrap_DFCCSD.py index 17fbf09d..dd61510b 100644 --- a/ebcc/codegen/bootstrap_DFCCSD.py +++ b/ebcc/codegen/bootstrap_DFCCSD.py @@ -5,22 +5,24 @@ import sys import pdaggerq -from albert.qc._pdaggerq import import_from_pdaggerq +from albert.qc._pdaggerq import remove_reference_energy, remove_reference_energy_eom +from albert.qc.spin import ghf_to_uhf, ghf_to_rhf +from albert.qc import ghf, uhf, rhf from albert.tensor import Tensor +from albert.index import Index +from albert.code._ebcc import EBCCCodeGenerator +from albert.misc import Stopwatch +from albert.opt.tools import _tensor_info -from ebcc.codegen.bootstrap_common import * +from ebcc.codegen.bootstrap_common import get_energy, get_amplitudes, get_rdm1, get_rdm2, get_eom # Get the spin case spin = sys.argv[1] -if spin == "ghf": - raise NotImplementedError # Set up the code generators code_generators = { - "einsum": EinsumCodeGen( + "einsum": EBCCCodeGenerator( stdout=open(f"{spin[0].upper()}DFCCSD.py", "w"), - name_generator=name_generators[spin], - spin=spin, ), } @@ -40,23 +42,17 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms = pq.fully_contracted_strings() - terms = remove_hf_energy(terms) + terms = remove_reference_energy(terms) # Get the energy in albert format - expr = import_from_pdaggerq(terms) - expr = spin_integrate(expr, spin) - expr = tuple(e.apply(get_density_fit()) for e in expr) - output = tuple(Tensor(name="e_cc") for _ in range(len(expr))) - output, expr = optimise(output, expr, spin, strategy="exhaust") - returns = (Tensor(name="e_cc"),) + output_expr, returns = get_energy(terms, spin, density_fit=True) # Generate the energy code for codegen in code_generators.values(): codegen( "energy", returns, - output, - expr, + output_expr, ) with Stopwatch("T amplitudes"): @@ -77,38 +73,15 @@ terms_t2 = pq.fully_contracted_strings() # Get the T amplitudes in albert format - terms = [terms_t1, terms_t2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"t{n+1}new") - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"t{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns = get_amplitudes([terms_t1, terms_t2], spin, density_fit=True) # Generate the T amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "t1new = Namespace()\nt2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_amps", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) with Stopwatch("L amplitudes"): @@ -141,38 +114,15 @@ terms_l2 = pq.fully_contracted_strings() # Get the L amplitudes in albert format - terms = [terms_l1, terms_l2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_l_amplitude_outputs(expr_n, f"l{n+1}new") - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"l{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="opt") + output_expr, returns = get_amplitudes([terms_l1, terms_l2], spin, strategy="opt" if spin == "uhf" else "exhaust", which="l", density_fit=True) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "l1new = Namespace()\nl2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_lams", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) with Stopwatch("1RDM"): @@ -186,38 +136,40 @@ terms[sectors, indices] = pq.fully_contracted_strings() # Get the 1RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [("oo", "ij"), ("ov", "ia"), ("vo", "ai"), ("vv", "ab")]: - for index_spins in get_density_spins(1, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - if spin == "rhf": - expr_n = tuple(e * 2 for e in expr_n) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_density_outputs(expr_n, f"d", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"d"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns, deltas, deltas_sources = get_rdm1(terms, spin, density_fit=True) # Generate the 1RDM code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": get_density_einsum_preamble(1, spin), - "postamble": get_density_einsum_postamble(1, spin), - } - else: - kwargs = {} + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing + if spin != "uhf": + codegen.write("rdm1 = np.block([[rdm1.oo, rdm1.ov], [rdm1.vo, rdm1.vv]])") + else: + codegen.write("rdm1.aa = np.block([[rdm1.aa.oo, rdm1.aa.ov], [rdm1.aa.vo, rdm1.aa.vv]])") + codegen.write("rdm1.bb = np.block([[rdm1.bb.oo, rdm1.bb.ov], [rdm1.bb.vo, rdm1.bb.vv]])") + codegen( "make_rdm1_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("2RDM"): @@ -250,51 +202,42 @@ terms[sectors, indices] = pq.fully_contracted_strings() # Get the 2RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [ - ("oooo", "ijkl"), - ("ooov", "ijka"), - ("oovo", "ijak"), - ("ovoo", "iajk"), - ("vooo", "aijk"), - ("oovv", "ijab"), - ("ovov", "iajb"), - ("ovvo", "iabj"), - ("voov", "aijb"), - ("vovo", "aibj"), - ("vvoo", "abij"), - ("ovvv", "iabc"), - ("vovv", "aibc"), - ("vvov", "abic"), - ("vvvo", "abci"), - ("vvvv", "abcd"), - ]: - for index_spins in get_density_spins(2, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_density_outputs(expr_n, f"Γ", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"Γ"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="trav") + output_expr, returns, deltas, deltas_sources = get_rdm2(terms, spin, strategy="trav" if spin == "uhf" else "exhaust", density_fit=True) # Generate the 2RDM code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": get_density_einsum_preamble(2, spin), - "postamble": get_density_einsum_postamble(2, spin), - } + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing, handle transpose + if spin != "uhf": + codegen.write("rdm2 = pack_2e(rdm2.oooo, rdm2.ooov, rdm2.oovo, rdm2.ovoo, rdm2.vooo, rdm2.oovv, rdm2.ovov, rdm2.ovvo, rdm2.voov, rdm2.vovo, rdm2.vvoo, rdm2.ovvv, rdm2.vovv, rdm2.vvov, rdm2.vvvo, rdm2.vvvv).transpose((0, 2, 1, 3))") + else: + codegen.write("rdm2.aaaa = pack_2e(rdm2.aaaa.oooo, rdm2.aaaa.ooov, rdm2.aaaa.oovo, rdm2.aaaa.ovoo, rdm2.aaaa.vooo, rdm2.aaaa.oovv, rdm2.aaaa.ovov, rdm2.aaaa.ovvo, rdm2.aaaa.voov, rdm2.aaaa.vovo, rdm2.aaaa.vvoo, rdm2.aaaa.ovvv, rdm2.aaaa.vovv, rdm2.aaaa.vvov, rdm2.aaaa.vvvo, rdm2.aaaa.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.aabb = pack_2e(rdm2.abab.oooo, rdm2.abab.ooov, rdm2.abab.oovo, rdm2.abab.ovoo, rdm2.abab.vooo, rdm2.abab.oovv, rdm2.abab.ovov, rdm2.abab.ovvo, rdm2.abab.voov, rdm2.abab.vovo, rdm2.abab.vvoo, rdm2.abab.ovvv, rdm2.abab.vovv, rdm2.abab.vvov, rdm2.abab.vvvo, rdm2.abab.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.bbbb = pack_2e(rdm2.bbbb.oooo, rdm2.bbbb.ooov, rdm2.bbbb.oovo, rdm2.bbbb.ovoo, rdm2.bbbb.vooo, rdm2.bbbb.oovv, rdm2.bbbb.ovov, rdm2.bbbb.ovvo, rdm2.bbbb.voov, rdm2.bbbb.vovo, rdm2.bbbb.vvoo, rdm2.bbbb.ovvv, rdm2.bbbb.vovv, rdm2.bbbb.vvov, rdm2.bbbb.vvvo, rdm2.bbbb.vvvv).transpose((0, 2, 1, 3))") + codegen.write("del rdm2.abab") + codegen( "make_rdm2_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("IP-EOM"): @@ -307,7 +250,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -318,61 +261,27 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip", density_fit=True) # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("EA-EOM"): @@ -385,7 +294,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -396,61 +305,27 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea", density_fit=True) # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -464,7 +339,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -475,62 +350,46 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee", density_fit=True) # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "postamble": "r2new.baba = np.transpose(r2new.abab, (1, 0, 3, 2))" if spin == "uhf" else None, # FIXME - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_matvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) with Stopwatch("L-IP-EOM"): @@ -543,7 +402,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -554,61 +413,27 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip", density_fit=True) # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("L-EA-EOM"): @@ -621,7 +446,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -632,61 +457,27 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea", density_fit=True) # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -700,7 +491,7 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -711,62 +502,46 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee", density_fit=True) # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "postamble": "r2new.baba = np.transpose(r2new.abab, (1, 0, 3, 2))" if spin == "uhf" else None, # FIXME - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_lmatvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) for codegen in code_generators.values(): diff --git a/ebcc/codegen/bootstrap_DFDCD.py b/ebcc/codegen/bootstrap_DFDCD.py index af154b9f..4e49f1cd 100644 --- a/ebcc/codegen/bootstrap_DFDCD.py +++ b/ebcc/codegen/bootstrap_DFDCD.py @@ -5,21 +5,24 @@ import sys import pdaggerq -from albert.qc._pdaggerq import import_from_pdaggerq -from albert.qc.index import Index +from albert.qc._pdaggerq import remove_reference_energy, remove_reference_energy_eom +from albert.qc.spin import ghf_to_uhf, ghf_to_rhf +from albert.qc import ghf, uhf, rhf from albert.tensor import Tensor +from albert.index import Index +from albert.code._ebcc import EBCCCodeGenerator +from albert.misc import Stopwatch +from albert.opt.tools import _tensor_info -from ebcc.codegen.bootstrap_common import * +from ebcc.codegen.bootstrap_common import get_energy, get_amplitudes, get_rdm1, get_rdm2, get_eom # Get the spin case spin = sys.argv[1] # Set up the code generators code_generators = { - "einsum": EinsumCodeGen( + "einsum": EBCCCodeGenerator( stdout=open(f"{spin[0].upper()}DFDCD.py", "w"), - name_generator=name_generators[spin], - spin=spin, ), } @@ -35,27 +38,21 @@ # Get the energy contractions in pdaggerq format pq.clear() pq.set_left_operators([["1"]]) - pq.add_st_operator(1.0, ["f"], ["t2"]) - pq.add_st_operator(1.0, ["v"], ["t2"]) + pq.add_st_operator(1.0, ["f"], ["t1", "t2"]) + pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms = pq.fully_contracted_strings() - terms = remove_hf_energy(terms) + terms = remove_reference_energy(terms) # Get the energy in albert format - expr = import_from_pdaggerq(terms) - expr = spin_integrate(expr, spin) - expr = tuple(e.apply(get_density_fit()) for e in expr) - output = tuple(Tensor(name="e_cc") for _ in range(len(expr))) - output, expr = optimise(output, expr, spin, strategy="exhaust") - returns = (Tensor(name="e_cc"),) + output_expr, returns = get_energy(terms, spin, density_fit=True) # Generate the energy code for codegen in code_generators.values(): codegen( "energy", returns, - output, - expr, + output_expr, ) with Stopwatch("T amplitudes"): @@ -74,36 +71,15 @@ ] # Get the T amplitudes in albert format - expr = [] - output = [] - returns = [] - for index_spins in get_amplitude_spins(2, spin): - indices = default_indices["o"][:2] + default_indices["v"][:2] - expr_n = import_from_pdaggerq(terms, index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"t2new") - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"t2new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns = get_amplitudes([terms], spin, orders=[2], density_fit=True) # Generate the T amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "t1new = Namespace()\nt2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_amps", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) for codegen in code_generators.values(): diff --git a/ebcc/codegen/bootstrap_DFDCSD.py b/ebcc/codegen/bootstrap_DFDCSD.py index 4f40f72b..7cca030a 100644 --- a/ebcc/codegen/bootstrap_DFDCSD.py +++ b/ebcc/codegen/bootstrap_DFDCSD.py @@ -5,21 +5,24 @@ import sys import pdaggerq -from albert.qc._pdaggerq import import_from_pdaggerq -from albert.qc.index import Index +from albert.qc._pdaggerq import remove_reference_energy, remove_reference_energy_eom +from albert.qc.spin import ghf_to_uhf, ghf_to_rhf +from albert.qc import ghf, uhf, rhf from albert.tensor import Tensor +from albert.index import Index +from albert.code._ebcc import EBCCCodeGenerator +from albert.misc import Stopwatch +from albert.opt.tools import _tensor_info -from ebcc.codegen.bootstrap_common import * +from ebcc.codegen.bootstrap_common import get_energy, get_amplitudes, get_rdm1, get_rdm2, get_eom # Get the spin case spin = sys.argv[1] # Set up the code generators code_generators = { - "einsum": EinsumCodeGen( + "einsum": EBCCCodeGenerator( stdout=open(f"{spin[0].upper()}DFDCSD.py", "w"), - name_generator=name_generators[spin], - spin=spin, ), } @@ -39,23 +42,17 @@ pq.add_st_operator(1.0, ["v"], ["t1", "t2"]) pq.simplify() terms = pq.fully_contracted_strings() - terms = remove_hf_energy(terms) + terms = remove_reference_energy(terms) # Get the energy in albert format - expr = import_from_pdaggerq(terms) - expr = spin_integrate(expr, spin) - expr = tuple(e.apply(get_density_fit()) for e in expr) - output = tuple(Tensor(name="e_cc") for _ in range(len(expr))) - output, expr = optimise(output, expr, spin, strategy="exhaust") - returns = (Tensor(name="e_cc"),) + output_expr, returns = get_energy(terms, spin, density_fit=True) # Generate the energy code for codegen in code_generators.values(): codegen( "energy", returns, - output, - expr, + output_expr, ) with Stopwatch("T amplitudes"): @@ -84,40 +81,18 @@ ] # Get the T amplitudes in albert format - terms = [terms_t1, terms_t2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"t{n+1}new") - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"t{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns = get_amplitudes([terms], spin, orders=[2], density_fit=True) # Generate the T amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "t1new = Namespace()\nt2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_amps", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) for codegen in code_generators.values(): codegen.postamble() codegen.stdout.close() + diff --git a/ebcc/codegen/bootstrap_DFQCISD.py b/ebcc/codegen/bootstrap_DFQCISD.py index cd7fa95c..c1cf85f0 100644 --- a/ebcc/codegen/bootstrap_DFQCISD.py +++ b/ebcc/codegen/bootstrap_DFQCISD.py @@ -5,21 +5,24 @@ import sys import pdaggerq -from albert.qc._pdaggerq import import_from_pdaggerq -from albert.qc.index import Index +from albert.qc._pdaggerq import remove_reference_energy, remove_reference_energy_eom +from albert.qc.spin import ghf_to_uhf, ghf_to_rhf +from albert.qc import ghf, uhf, rhf from albert.tensor import Tensor +from albert.index import Index +from albert.code._ebcc import EBCCCodeGenerator +from albert.misc import Stopwatch +from albert.opt.tools import _tensor_info -from ebcc.codegen.bootstrap_common import * +from ebcc.codegen.bootstrap_common import get_energy, get_amplitudes, get_rdm1, get_rdm2, get_eom # Get the spin case spin = sys.argv[1] # Set up the code generators code_generators = { - "einsum": EinsumCodeGen( + "einsum": EBCCCodeGenerator( stdout=open(f"{spin[0].upper()}DFQCISD.py", "w"), - name_generator=name_generators[spin], - spin=spin, ), } @@ -43,23 +46,17 @@ pq.add_commutator(1.0, ["v"], ["t2"]) pq.simplify() terms = pq.fully_contracted_strings() - terms = remove_hf_energy(terms) + terms = remove_reference_energy(terms) # Get the energy in albert format - expr = import_from_pdaggerq(terms) - expr = spin_integrate(expr, spin) - expr = tuple(e.apply(get_density_fit()) for e in expr) - output = tuple(Tensor(name="e_cc") for _ in range(len(expr))) - output, expr = optimise(output, expr, spin, strategy="exhaust") - returns = (Tensor(name="e_cc"),) + output_expr, returns = get_energy(terms, spin, density_fit=True) # Generate the energy code for codegen in code_generators.values(): codegen( "energy", returns, - output, - expr, + output_expr, ) with Stopwatch("T amplitudes"): @@ -90,38 +87,15 @@ terms_t2 = pq.fully_contracted_strings() # Get the T amplitudes in albert format - terms = [terms_t1, terms_t2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"t{n+1}new") - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"t{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns = get_amplitudes([terms_t1, terms_t2], spin, density_fit=True) # Generate the T amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "t1new = Namespace()\nt2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_amps", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) with Stopwatch("L amplitudes"): @@ -178,38 +152,15 @@ terms_l2 = pq.fully_contracted_strings() # Get the L amplitudes in albert format - terms = [terms_l1, terms_l2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_l_amplitude_outputs(expr_n, f"l{n+1}new") - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"l{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="opt") + output_expr, returns = get_amplitudes([terms_l1, terms_l2], spin, strategy="opt" if spin == "uhf" else "exhaust", which="l", density_fit=True) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "l1new = Namespace()\nl2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_lams", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) with Stopwatch("1RDM"): @@ -223,38 +174,40 @@ terms[sectors, indices] = pq.fully_contracted_strings() # Get the 1RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [("oo", "ij"), ("ov", "ia"), ("vo", "ai"), ("vv", "ab")]: - for index_spins in get_density_spins(1, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - if spin == "rhf": - expr_n = tuple(e * 2 for e in expr_n) - output_n = get_density_outputs(expr_n, f"d", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"d"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns, deltas, deltas_sources = get_rdm1(terms, spin, density_fit=True) # Generate the 1RDM code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": get_density_einsum_preamble(1, spin), - "postamble": get_density_einsum_postamble(1, spin), - } - else: - kwargs = {} + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing + if spin != "uhf": + codegen.write("rdm1 = np.block([[rdm1.oo, rdm1.ov], [rdm1.vo, rdm1.vv]])") + else: + codegen.write("rdm1.aa = np.block([[rdm1.aa.oo, rdm1.aa.ov], [rdm1.aa.vo, rdm1.aa.vv]])") + codegen.write("rdm1.bb = np.block([[rdm1.bb.oo, rdm1.bb.ov], [rdm1.bb.vo, rdm1.bb.vv]])") + codegen( "make_rdm1_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("2RDM"): @@ -287,51 +240,42 @@ terms[sectors, indices] = pq.fully_contracted_strings() # Get the 2RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [ - ("oooo", "ijkl"), - ("ooov", "ijka"), - ("oovo", "ijak"), - ("ovoo", "iajk"), - ("vooo", "aijk"), - ("oovv", "ijab"), - ("ovov", "iajb"), - ("ovvo", "iabj"), - ("voov", "aijb"), - ("vovo", "aibj"), - ("vvoo", "abij"), - ("ovvv", "iabc"), - ("vovv", "aibc"), - ("vvov", "abic"), - ("vvvo", "abci"), - ("vvvv", "abcd"), - ]: - for index_spins in get_density_spins(2, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_density_outputs(expr_n, f"Γ", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"Γ"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="trav") + output_expr, returns, deltas, deltas_sources = get_rdm2(terms, spin, strategy="trav" if spin == "uhf" else "exhaust", density_fit=True) # Generate the 2RDM code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": get_density_einsum_preamble(2, spin), - "postamble": get_density_einsum_postamble(2, spin), - } + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing, handle transpose + if spin != "uhf": + codegen.write("rdm2 = pack_2e(rdm2.oooo, rdm2.ooov, rdm2.oovo, rdm2.ovoo, rdm2.vooo, rdm2.oovv, rdm2.ovov, rdm2.ovvo, rdm2.voov, rdm2.vovo, rdm2.vvoo, rdm2.ovvv, rdm2.vovv, rdm2.vvov, rdm2.vvvo, rdm2.vvvv).transpose((0, 2, 1, 3))") + else: + codegen.write("rdm2.aaaa = pack_2e(rdm2.aaaa.oooo, rdm2.aaaa.ooov, rdm2.aaaa.oovo, rdm2.aaaa.ovoo, rdm2.aaaa.vooo, rdm2.aaaa.oovv, rdm2.aaaa.ovov, rdm2.aaaa.ovvo, rdm2.aaaa.voov, rdm2.aaaa.vovo, rdm2.aaaa.vvoo, rdm2.aaaa.ovvv, rdm2.aaaa.vovv, rdm2.aaaa.vvov, rdm2.aaaa.vvvo, rdm2.aaaa.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.aabb = pack_2e(rdm2.abab.oooo, rdm2.abab.ooov, rdm2.abab.oovo, rdm2.abab.ovoo, rdm2.abab.vooo, rdm2.abab.oovv, rdm2.abab.ovov, rdm2.abab.ovvo, rdm2.abab.voov, rdm2.abab.vovo, rdm2.abab.vvoo, rdm2.abab.ovvv, rdm2.abab.vovv, rdm2.abab.vvov, rdm2.abab.vvvo, rdm2.abab.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.bbbb = pack_2e(rdm2.bbbb.oooo, rdm2.bbbb.ooov, rdm2.bbbb.oovo, rdm2.bbbb.ovoo, rdm2.bbbb.vooo, rdm2.bbbb.oovv, rdm2.bbbb.ovov, rdm2.bbbb.ovvo, rdm2.bbbb.voov, rdm2.bbbb.vovo, rdm2.bbbb.vvoo, rdm2.bbbb.ovvv, rdm2.bbbb.vovv, rdm2.bbbb.vvov, rdm2.bbbb.vvvo, rdm2.bbbb.vvvv).transpose((0, 2, 1, 3))") + codegen.write("del rdm2.abab") + codegen( "make_rdm2_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("IP-EOM"): @@ -348,7 +292,7 @@ pq.add_double_commutator(1.0, ["v"], ["t1"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -365,61 +309,27 @@ pq.add_double_commutator(0.5, ["v"], ["t2"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="opt") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip", density_fit=True) # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("EA-EOM"): @@ -436,7 +346,7 @@ pq.add_double_commutator(1.0, ["v"], ["t1"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -453,61 +363,27 @@ pq.add_double_commutator(0.5, ["v"], ["t2"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="opt") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea", density_fit=True) # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -525,7 +401,7 @@ pq.add_double_commutator(1.0, ["v"], ["t1"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -542,62 +418,46 @@ pq.add_double_commutator(0.5, ["v"], ["t2"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee", density_fit=True) # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "postamble": "r2new.baba = np.transpose(r2new.abab, (1, 0, 3, 2))" if spin == "uhf" else None, # FIXME - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_matvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) with Stopwatch("L-IP-EOM"): @@ -614,7 +474,7 @@ pq.add_double_commutator(1.0, ["v"], ["t1"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -631,61 +491,27 @@ pq.add_double_commutator(0.5, ["v"], ["t2"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="opt") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip", density_fit=True) # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("L-EA-EOM"): @@ -702,7 +528,7 @@ pq.add_double_commutator(1.0, ["v"], ["t1"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -719,61 +545,27 @@ pq.add_double_commutator(0.5, ["v"], ["t2"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="opt") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea", density_fit=True) # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -791,7 +583,7 @@ pq.add_double_commutator(1.0, ["v"], ["t1"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -808,62 +600,46 @@ pq.add_double_commutator(0.5, ["v"], ["t2"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - expr_n = tuple(e.apply(get_density_fit()) for e in expr_n) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee", density_fit=True) # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "postamble": "r2new.baba = np.transpose(r2new.abab, (1, 0, 3, 2))" if spin == "uhf" else None, # FIXME - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_lmatvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) for codegen in code_generators.values(): diff --git a/ebcc/codegen/bootstrap_MPn.py b/ebcc/codegen/bootstrap_MPn.py index c8116734..874f382a 100644 --- a/ebcc/codegen/bootstrap_MPn.py +++ b/ebcc/codegen/bootstrap_MPn.py @@ -2,14 +2,21 @@ Generate the MPn code. """ +import itertools import sys import pdaggerq -from albert.qc._pdaggerq import import_from_pdaggerq +from albert.qc._pdaggerq import remove_reference_energy, remove_reference_energy_eom +from albert.qc.spin import ghf_to_uhf, ghf_to_rhf +from albert.qc import ghf, uhf, rhf from albert.tensor import Tensor +from albert.index import Index +from albert.code._ebcc import EBCCCodeGenerator +from albert.misc import Stopwatch +from albert.opt.tools import _tensor_info +from ebcc.codegen.bootstrap_common import get_energy, get_amplitudes, get_rdm1, get_rdm2, get_eom from ebcc.codegen import bootstrap_hugenholtz as hugenholtz -from ebcc.codegen.bootstrap_common import * # Get the spin case spin = sys.argv[1] @@ -17,10 +24,8 @@ # Set up the code generators code_generators = { - "einsum": EinsumCodeGen( - stdout=open(f"{spin[0].upper()}MP{order}.py", "w"), - name_generator=name_generators[spin], - spin=spin, + "einsum": EBCCCodeGenerator( + stdout=open(f"{spin[0].upper()}CCSD.py", "w"), ), } @@ -36,55 +41,14 @@ ], []) # Get the energy in albert format - expr = import_from_pdaggerq(terms) - expr = spin_integrate(expr, spin) - output = tuple(Tensor(name="e_mp") for _ in range(len(expr))) - output, expr = optimise(output, expr, spin, strategy="exhaust") - returns = (Tensor(name="e_mp"),) - - # Get the denominators for higher orders - if order > 3: - if spin != "uhf": - lines = [] - for n in ([2] if order in (2, 3) else list(range(1, order+1))): - subscripts = ["ia", "jb", "kc", "ld", "me", "nf"][:n] - lines += [ - " denom%d = 1 / direct_sum(" % n, - " \"%s->%s%s\"," % ("+".join(subscripts), "".join([s[0] for s in subscripts]), "".join([s[1] for s in subscripts])), - ] - for subscript in subscripts: - lines.append( - " direct_sum(\"%s->%s\", np.diag(f.oo), np.diag(f.vv))," - % ("-".join(list(subscript)), subscript) - ) - lines.append(" )") - else: - lines = [ - " denom3 = Namespace()", - ] - for n in ([2] if order in (2, 3) else list(range(1, order+1))): - spins_list = [list(y) for y in sorted(set("".join(sorted(x)) for x in itertools.product("ab", repeat=n)))] - for spins in spins_list: - subscripts = ["ia", "jb", "kc", "ld", "me", "nf"][:n] - lines += [ - " denom%d.%s%s = 1 / direct_sum(" % (n, "".join(spins), "".join(spins)), - " \"%s->%s%s\"," % ("+".join(subscripts), "".join([s[0] for s in subscripts]), "".join([s[1] for s in subscripts])), - ] - for subscript, spin in zip(subscripts, spins): - lines.append( - " direct_sum(\"%s->%s\", np.diag(f.%s%s.oo), np.diag(f.%s%s.vv))," - % ("-".join(list(subscript)), subscript, spin, spin, spin, spin) - ) - lines.append(" )") - function_printer.write_python("\n".join(lines)+"\n") + output_expr, returns = get_energy(terms, spin) # Generate the energy code for codegen in code_generators.values(): codegen( "energy", returns, - output, - expr, + output_expr, ) if order == 2: @@ -99,53 +63,46 @@ } # Get the 1RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [("oo", "ij"), ("ov", "ia"), ("vo", "ai"), ("vv", "ab")]: - if (sectors, indices) not in terms: - continue - for index_spins in get_density_spins(1, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - if not (isinstance(expr_n, int) and expr_n == 0): - expr_n = spin_integrate(expr_n, spin) - if spin == "rhf": - expr_n = tuple(e * 2 for e in expr_n) - output_n = get_density_outputs(expr_n, f"d", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"d"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns, deltas, deltas_sources = get_rdm1(terms, spin) # Generate the 1RDM code for name, codegen in code_generators.items(): - if name == "einsum": - def get_postamble(n, spin, name="rdm{n}"): - nm = name.format(n=n) - postamble = "" - if spin != "uhf": - for occ in ("ov", "vo"): - shape = ", ".join(f"t2.shape[{'0' if o == 'o' else '-1'}]" for o in occ) - postamble += f"{nm}.{occ} = np.zeros(({shape}))\n" - else: - for s in "ab": - for occ in ("ov", "vo"): - shape = ", ".join(f"t2.{s}{s}{s}{s}.shape[{'0' if o == 'o' else '-1'}]" for o in occ) - postamble += f"{nm}.{s}{s}.{occ} = np.zeros(({shape}))\n" - return postamble + get_density_einsum_postamble(n, spin) - kwargs = { - "preamble": get_density_einsum_preamble(1, spin), - "postamble": get_postamble(1, spin), - } - else: - kwargs = {} + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing + if spin != "uhf": + codegen.write("rdm1.ov = np.zeros((t1.shape[0], t1.shape[1]))") + codegen.write("rdm1.vo = np.zeros((t1.shape[1], t1.shape[0]))") + codegen.write("rdm1 = np.block([[rdm1.oo, rdm1.ov], [rdm1.vo, rdm1.vv]])") + else: + codegen.write("rdm1.aa.ov = np.zeros((t1.aa.shape[0], t1.aa.shape[1]))") + codegen.write("rdm1.aa.vo = np.zeros((t1.aa.shape[1], t1.aa.shape[0]))") + codegen.write("rdm1.bb.ov = np.zeros((t1.bb.shape[0], t1.bb.shape[1]))") + codegen.write("rdm1.bb.vo = np.zeros((t1.bb.shape[1], t1.bb.shape[0]))") + codegen.write("rdm1.aa = np.block([[rdm1.aa.oo, rdm1.aa.ov], [rdm1.aa.vo, rdm1.aa.vv]])") + codegen.write("rdm1.bb = np.block([[rdm1.bb.oo, rdm1.bb.ov], [rdm1.bb.vo, rdm1.bb.vv]])") + codegen( "make_rdm1_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("2RDM"): @@ -156,124 +113,95 @@ def get_postamble(n, spin, name="rdm{n}"): } # Get the 2RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [ - ("oooo", "ijkl"), - ("ooov", "ijka"), - ("oovo", "ijak"), - ("ovoo", "iajk"), - ("vooo", "aijk"), - ("oovv", "ijab"), - ("ovov", "iajb"), - ("ovvo", "iabj"), - ("voov", "aijb"), - ("vovo", "aibj"), - ("vvoo", "abij"), - ("ovvv", "iabc"), - ("vovv", "aibc"), - ("vvov", "abic"), - ("vvvo", "abci"), - ("vvvv", "abcd"), - ]: - if (sectors, indices) not in terms: - terms[sectors, indices] = [] - for index_spins in get_density_spins(2, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - if not (isinstance(expr_n, int) and expr_n == 0): - expr_n = spin_integrate(expr_n, spin) - output_n = get_density_outputs(expr_n, f"Γ", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"Γ"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns, deltas, deltas_sources = get_rdm2(terms, spin, strategy="trav" if spin == "uhf" else "exhaust") # Generate the 2RDM code for name, codegen in code_generators.items(): - if name == "einsum": - def get_postamble(n, spin, name="rdm{n}"): - nm = name.format(n=n) - postamble = "" - if spin != "uhf": - for occ in [k[0] for k, v in terms.items() if not v]: - shape = ", ".join(f"t2.shape[{'0' if o == 'o' else '-1'}]" for o in occ) - postamble += f"{nm}.{occ} = np.zeros(({shape}))\n" - else: - for s1, s2 in [("a", "a"), ("a", "b"), ("b", "b")]: - for occ in [k[0] for k, v in terms.items() if not v]: - shape = ", ".join(f"t2.{s}{s}{s}{s}.shape[{'0' if o == 'o' else '-1'}]" for o, s in zip(occ, s1+s2+s1+s2)) - postamble += f"{nm}.{s1}{s2}{s1}{s2}.{occ} = np.zeros(({shape}))\n" - - # Pack - postamble += get_density_einsum_postamble(n, spin) - - # Add the one-body terms -- adapted from pyscf - # Can be done succinctly with tensor expressions but requires - # a second Norb^4 tensor - postamble += "\nrdm1 = make_rdm1_f(t2=t2, l2=l2)\n" - if spin == "uhf": - postamble += "delta = Namespace(\n" - postamble += " aa=np.diag(np.concatenate([np.ones(t2.aaaa.shape[0]), np.zeros(t2.aaaa.shape[-1])])),\n" - postamble += " bb=np.diag(np.concatenate([np.ones(t2.bbbb.shape[0]), np.zeros(t2.bbbb.shape[-1])])),\n" - postamble += ")\n" - postamble += "rdm1.aa -= delta.aa\n" - postamble += "rdm1.bb -= delta.bb\n" - postamble += "rdm2.aaaa += einsum(delta.aa, (0, 1), rdm1.aa, (3, 2), (0, 1, 2, 3))\n" - postamble += "rdm2.aaaa += einsum(rdm1.aa, (1, 0), delta.aa, (2, 3), (0, 1, 2, 3))\n" - postamble += "rdm2.aaaa -= einsum(delta.aa, (0, 3), rdm1.aa, (2, 1), (0, 1, 2, 3))\n" - postamble += "rdm2.aaaa -= einsum(rdm1.aa, (0, 3), delta.aa, (1, 2), (0, 1, 2, 3))\n" - postamble += "rdm2.aaaa += einsum(delta.aa, (0, 1), delta.aa, (2, 3), (0, 1, 2, 3))\n" - postamble += "rdm2.aaaa -= einsum(delta.aa, (0, 3), delta.aa, (1, 2), (0, 1, 2, 3))\n" - postamble += "rdm2.bbbb += einsum(delta.bb, (0, 1), rdm1.bb, (3, 2), (0, 1, 2, 3))\n" - postamble += "rdm2.bbbb += einsum(rdm1.bb, (1, 0), delta.bb, (2, 3), (0, 1, 2, 3))\n" - postamble += "rdm2.bbbb -= einsum(delta.bb, (0, 3), rdm1.bb, (2, 1), (0, 1, 2, 3))\n" - postamble += "rdm2.bbbb -= einsum(rdm1.bb, (0, 3), delta.bb, (1, 2), (0, 1, 2, 3))\n" - postamble += "rdm2.bbbb += einsum(delta.bb, (0, 1), delta.bb, (2, 3), (0, 1, 2, 3))\n" - postamble += "rdm2.bbbb -= einsum(delta.bb, (0, 3), delta.bb, (1, 2), (0, 1, 2, 3))\n" - postamble += "rdm2.aabb += einsum(delta.aa, (0, 1), rdm1.bb, (3, 2), (0, 1, 2, 3))\n" - postamble += "rdm2.aabb += einsum(rdm1.aa, (1, 0), delta.bb, (2, 3), (0, 1, 2, 3))\n" - postamble += "rdm2.aabb += einsum(delta.aa, (0, 1), delta.bb, (2, 3), (0, 1, 2, 3))" - elif spin == "ghf": - postamble += "delta = np.diag(np.concatenate([np.ones(t2.shape[0]), np.zeros(t2.shape[-1])]))\n" - postamble += "rdm1 -= delta\n" - postamble += "rdm2 += einsum(delta, (0, 1), rdm1, (3, 2), (0, 1, 2, 3))\n" - postamble += "rdm2 += einsum(rdm1, (1, 0), delta, (2, 3), (0, 1, 2, 3))\n" - postamble += "rdm2 -= einsum(delta, (0, 3), rdm1, (2, 1), (0, 1, 2, 3))\n" - postamble += "rdm2 -= einsum(rdm1, (0, 3), delta, (1, 2), (0, 1, 2, 3))\n" - postamble += "rdm2 += einsum(delta, (0, 1), delta, (2, 3), (0, 1, 2, 3))\n" - postamble += "rdm2 -= einsum(delta, (0, 3), delta, (1, 2), (0, 1, 2, 3))" - elif spin == "rhf": - postamble += "delta = np.diag(np.concatenate([np.ones(t2.shape[0]), np.zeros(t2.shape[-1])]))\n" - postamble += "rdm1 -= delta * 2\n" - postamble += "rdm2 += einsum(delta, (0, 1), rdm1, (3, 2), (0, 1, 2, 3)) * 2\n" - postamble += "rdm2 += einsum(rdm1, (1, 0), delta, (2, 3), (0, 1, 2, 3)) * 2\n" - postamble += "rdm2 -= einsum(delta, (0, 3), rdm1, (2, 1), (0, 1, 2, 3))\n" - postamble += "rdm2 -= einsum(rdm1, (0, 3), delta, (1, 2), (0, 1, 2, 3))\n" - postamble += "rdm2 += einsum(delta, (0, 1), delta, (2, 3), (0, 1, 2, 3)) * 4\n" - postamble += "rdm2 -= einsum(delta, (0, 3), delta, (1, 2), (0, 1, 2, 3)) * 2" - - return postamble - - def get_preamble(n, spin, name="rdm{n}"): - name = name.format(n=n) - preamble = f"{name} = Namespace()" - if spin == "uhf": - for spins in itertools.combinations_with_replacement(["a", "b"], n): - preamble += f"\n{name}.{''.join(spins+spins)} = Namespace()" - return preamble - - kwargs = { - "preamble": get_preamble(2, spin), - "postamble": get_postamble(2, spin), - } + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing, handle transpose + if spin != "uhf": + for key in itertools.product("ov", repeat=4): + if tuple(key) in {("o", "o", "v", "v") , ("v", "v", "o", "o")}: + continue + i, j, k, l = ["ov".index(k) for k in key] + codegen.write(f"rdm2.{key} = np.zeros((t1.shape[{i}], t1.shape[{j}], t1.shape[{k}], t1.shape[{l}]))") + codegen.write("rdm2 = pack_2e(rdm2.oooo, rdm2.ooov, rdm2.oovo, rdm2.ovoo, rdm2.vooo, rdm2.oovv, rdm2.ovov, rdm2.ovvo, rdm2.voov, rdm2.vovo, rdm2.vvoo, rdm2.ovvv, rdm2.vovv, rdm2.vvov, rdm2.vvvo, rdm2.vvvv).transpose((0, 2, 1, 3))") + else: + for spin in ("aaaa", "aabb", "bbbb"): + for key in itertools.product("ov", repeat=4): + if tuple(key) in {("o", "o", "v", "v") , ("v", "v", "o", "o")}: + continue + i, j, k, l = ["ov".index(k) for k in key] + si, sj, sk, sl = spin + codegen.write(f"rdm2.{spin}.{key} = np.zeros((t1.{si}.shape[{i}], t1.{sj}.shape[{j}], t1.{sk}.shape[{k}], t1.{sl}.shape[{l}]))") + codegen.write("rdm2.aaaa = pack_2e(rdm2.aaaa.oooo, rdm2.aaaa.ooov, rdm2.aaaa.oovo, rdm2.aaaa.ovoo, rdm2.aaaa.vooo, rdm2.aaaa.oovv, rdm2.aaaa.ovov, rdm2.aaaa.ovvo, rdm2.aaaa.voov, rdm2.aaaa.vovo, rdm2.aaaa.vvoo, rdm2.aaaa.ovvv, rdm2.aaaa.vovv, rdm2.aaaa.vvov, rdm2.aaaa.vvvo, rdm2.aaaa.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.aabb = pack_2e(rdm2.abab.oooo, rdm2.abab.ooov, rdm2.abab.oovo, rdm2.abab.ovoo, rdm2.abab.vooo, rdm2.abab.oovv, rdm2.abab.ovov, rdm2.abab.ovvo, rdm2.abab.voov, rdm2.abab.vovo, rdm2.abab.vvoo, rdm2.abab.ovvv, rdm2.abab.vovv, rdm2.abab.vvov, rdm2.abab.vvvo, rdm2.abab.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.bbbb = pack_2e(rdm2.bbbb.oooo, rdm2.bbbb.ooov, rdm2.bbbb.oovo, rdm2.bbbb.ovoo, rdm2.bbbb.vooo, rdm2.bbbb.oovv, rdm2.bbbb.ovov, rdm2.bbbb.ovvo, rdm2.bbbb.voov, rdm2.bbbb.vovo, rdm2.bbbb.vvoo, rdm2.bbbb.ovvv, rdm2.bbbb.vovv, rdm2.bbbb.vvov, rdm2.bbbb.vvvo, rdm2.bbbb.vvvv).transpose((0, 2, 1, 3))") + codegen.write("del rdm2.abab") + codegen.write("rdm1 = make_rdm1_f(t2=t2, l2=l2)") + if spin == "uhf": + codegen.write("delta = Namespace(") + codegen.write(" aa=np.diag(np.concatenate([np.ones(t2.aaaa.shape[0]), np.zeros(t2.aaaa.shape[-1])])),") + codegen.write(" bb=np.diag(np.concatenate([np.ones(t2.bbbb.shape[0]), np.zeros(t2.bbbb.shape[-1])])),") + codegen.write(")") + codegen.write("rdm1.aa -= delta.aa") + codegen.write("rdm1.bb -= delta.bb") + codegen.write("rdm2.aaaa += einsum(delta.aa, (0, 1), rdm1.aa, (3, 2), (0, 1, 2, 3))") + codegen.write("rdm2.aaaa += einsum(rdm1.aa, (1, 0), delta.aa, (2, 3), (0, 1, 2, 3))") + codegen.write("rdm2.aaaa -= einsum(delta.aa, (0, 3), rdm1.aa, (2, 1), (0, 1, 2, 3))") + codegen.write("rdm2.aaaa -= einsum(rdm1.aa, (0, 3), delta.aa, (1, 2), (0, 1, 2, 3))") + codegen.write("rdm2.aaaa += einsum(delta.aa, (0, 1), delta.aa, (2, 3), (0, 1, 2, 3))") + codegen.write("rdm2.aaaa -= einsum(delta.aa, (0, 3), delta.aa, (1, 2), (0, 1, 2, 3))") + codegen.write("rdm2.bbbb += einsum(delta.bb, (0, 1), rdm1.bb, (3, 2), (0, 1, 2, 3))") + codegen.write("rdm2.bbbb += einsum(rdm1.bb, (1, 0), delta.bb, (2, 3), (0, 1, 2, 3))") + codegen.write("rdm2.bbbb -= einsum(delta.bb, (0, 3), rdm1.bb, (2, 1), (0, 1, 2, 3))") + codegen.write("rdm2.bbbb -= einsum(rdm1.bb, (0, 3), delta.bb, (1, 2), (0, 1, 2, 3))") + codegen.write("rdm2.bbbb += einsum(delta.bb, (0, 1), delta.bb, (2, 3), (0, 1, 2, 3))") + codegen.write("rdm2.bbbb -= einsum(delta.bb, (0, 3), delta.bb, (1, 2), (0, 1, 2, 3))") + codegen.write("rdm2.aabb += einsum(delta.aa, (0, 1), rdm1.bb, (3, 2), (0, 1, 2, 3))") + codegen.write("rdm2.aabb += einsum(rdm1.aa, (1, 0), delta.bb, (2, 3), (0, 1, 2, 3))") + codegen.write("rdm2.aabb += einsum(delta.aa, (0, 1), delta.bb, (2, 3), (0, 1, 2, 3))" + elif spin == "ghf": + codegen.write("delta = np.diag(np.concatenate([np.ones(t2.shape[0]), np.zeros(t2.shape[-1])]))") + codegen.write("rdm1 -= delta") + codegen.write("rdm2 += einsum(delta, (0, 1), rdm1, (3, 2), (0, 1, 2, 3))") + codegen.write("rdm2 += einsum(rdm1, (1, 0), delta, (2, 3), (0, 1, 2, 3))") + codegen.write("rdm2 -= einsum(delta, (0, 3), rdm1, (2, 1), (0, 1, 2, 3))") + codegen.write("rdm2 -= einsum(rdm1, (0, 3), delta, (1, 2), (0, 1, 2, 3))") + codegen.write("rdm2 += einsum(delta, (0, 1), delta, (2, 3), (0, 1, 2, 3))") + codegen.write("rdm2 -= einsum(delta, (0, 3), delta, (1, 2), (0, 1, 2, 3))" + elif spin == "rhf": + codegen.write("delta = np.diag(np.concatenate([np.ones(t2.shape[0]), np.zeros(t2.shape[-1])]))") + codegen.write("rdm1 -= delta * 2") + codegen.write("rdm2 += einsum(delta, (0, 1), rdm1, (3, 2), (0, 1, 2, 3)) * 2") + codegen.write("rdm2 += einsum(rdm1, (1, 0), delta, (2, 3), (0, 1, 2, 3)) * 2") + codegen.write("rdm2 -= einsum(delta, (0, 3), rdm1, (2, 1), (0, 1, 2, 3))") + codegen.write("rdm2 -= einsum(rdm1, (0, 3), delta, (1, 2), (0, 1, 2, 3))") + codegen.write("rdm2 += einsum(delta, (0, 1), delta, (2, 3), (0, 1, 2, 3)) * 4") + codegen.write("rdm2 -= einsum(delta, (0, 3), delta, (1, 2), (0, 1, 2, 3)) * 2") + codegen( "make_rdm2_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("IP-EOM"): @@ -293,57 +221,24 @@ def get_preamble(n, spin, name="rdm{n}"): ] # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("EA-EOM"): @@ -363,57 +258,24 @@ def get_preamble(n, spin, name="rdm{n}"): ] # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="exhaust") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -439,58 +301,43 @@ def get_preamble(n, spin, name="rdm{n}"): ] # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "postamble": "r2new.baba = np.transpose(r2new.abab, (1, 0, 3, 2))" if spin == "uhf" else None, # FIXME - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_matvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) for codegen in code_generators.values(): diff --git a/ebcc/codegen/bootstrap_QCISD.py b/ebcc/codegen/bootstrap_QCISD.py index 4e30547c..c508e783 100644 --- a/ebcc/codegen/bootstrap_QCISD.py +++ b/ebcc/codegen/bootstrap_QCISD.py @@ -5,21 +5,24 @@ import sys import pdaggerq -from albert.qc._pdaggerq import import_from_pdaggerq -from albert.qc.index import Index +from albert.qc._pdaggerq import remove_reference_energy, remove_reference_energy_eom +from albert.qc.spin import ghf_to_uhf, ghf_to_rhf +from albert.qc import ghf, uhf, rhf from albert.tensor import Tensor +from albert.index import Index +from albert.code._ebcc import EBCCCodeGenerator +from albert.misc import Stopwatch +from albert.opt.tools import _tensor_info -from ebcc.codegen.bootstrap_common import * +from ebcc.codegen.bootstrap_common import get_energy, get_amplitudes, get_rdm1, get_rdm2, get_eom # Get the spin case spin = sys.argv[1] # Set up the code generators code_generators = { - "einsum": EinsumCodeGen( + "einsum": EBCCCodeGenerator( stdout=open(f"{spin[0].upper()}QCISD.py", "w"), - name_generator=name_generators[spin], - spin=spin, ), } @@ -43,22 +46,17 @@ pq.add_commutator(1.0, ["v"], ["t2"]) pq.simplify() terms = pq.fully_contracted_strings() - terms = remove_hf_energy(terms) + terms = remove_reference_energy(terms) # Get the energy in albert format - expr = import_from_pdaggerq(terms) - expr = spin_integrate(expr, spin) - output = tuple(Tensor(name="e_cc") for _ in range(len(expr))) - output, expr = optimise(output, expr, spin, strategy="exhaust") - returns = (Tensor(name="e_cc"),) + output_expr, returns = get_energy(terms, spin) # Generate the energy code for codegen in code_generators.values(): codegen( "energy", returns, - output, - expr, + output_expr, ) with Stopwatch("T amplitudes"): @@ -89,38 +87,15 @@ terms_t2 = pq.fully_contracted_strings() # Get the T amplitudes in albert format - terms = [terms_t1, terms_t2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - indices = tuple(Index(i, index_spins[i]) for i in indices) - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"t{n+1}new") - returns_n = (Tensor(*indices, name=f"t{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns = get_amplitudes([terms_t1, terms_t2], spin) # Generate the T amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "t1new = Namespace()\nt2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_amps", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) with Stopwatch("L amplitudes"): @@ -177,37 +152,15 @@ terms_l2 = pq.fully_contracted_strings() # Get the L amplitudes in albert format - terms = [terms_l1, terms_l2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_l_amplitude_outputs(expr_n, f"l{n+1}new") - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"l{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="opt") + output_expr, returns = get_amplitudes([terms_l1, terms_l2], spin, strategy="opt" if spin == "uhf" else "exhaust", which="l") # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": "l1new = Namespace()\nl2new = Namespace()" if spin == "uhf" else None, - "as_dict": True, - } - else: - kwargs = {} codegen( "update_lams", returns, - output, - expr, - **kwargs, + output_expr, + as_dict=True, ) with Stopwatch("1RDM"): @@ -221,37 +174,40 @@ terms[sectors, indices] = pq.fully_contracted_strings() # Get the 1RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [("oo", "ij"), ("ov", "ia"), ("vo", "ai"), ("vv", "ab")]: - for index_spins in get_density_spins(1, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - if spin == "rhf": - expr_n = tuple(e * 2 for e in expr_n) - output_n = get_density_outputs(expr_n, f"d", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"d"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="exhaust") + output_expr, returns, deltas, deltas_sources = get_rdm1(terms, spin) # Generate the 1RDM code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": get_density_einsum_preamble(1, spin), - "postamble": get_density_einsum_postamble(1, spin), - } - else: - kwargs = {} + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing + if spin != "uhf": + codegen.write("rdm1 = np.block([[rdm1.oo, rdm1.ov], [rdm1.vo, rdm1.vv]])") + else: + codegen.write("rdm1.aa = np.block([[rdm1.aa.oo, rdm1.aa.ov], [rdm1.aa.vo, rdm1.aa.vv]])") + codegen.write("rdm1.bb = np.block([[rdm1.bb.oo, rdm1.bb.ov], [rdm1.bb.vo, rdm1.bb.vv]])") + codegen( "make_rdm1_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("2RDM"): @@ -284,50 +240,42 @@ terms[sectors, indices] = pq.fully_contracted_strings() # Get the 2RDM in albert format - expr = [] - output = [] - returns = [] - for sectors, indices in [ - ("oooo", "ijkl"), - ("ooov", "ijka"), - ("oovo", "ijak"), - ("ovoo", "iajk"), - ("vooo", "aijk"), - ("oovv", "ijab"), - ("ovov", "iajb"), - ("ovvo", "iabj"), - ("voov", "aijb"), - ("vovo", "aibj"), - ("vvoo", "abij"), - ("ovvv", "iabc"), - ("vovv", "aibc"), - ("vvov", "abic"), - ("vvvo", "abci"), - ("vvvv", "abcd"), - ]: - for index_spins in get_density_spins(2, spin, indices): - expr_n = import_from_pdaggerq(terms[sectors, indices], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_density_outputs(expr_n, f"Γ", indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i], space=s) for i, s in zip(indices, sectors)), name=f"Γ"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - output, expr = optimise(output, expr, spin, strategy="trav") + output_expr, returns, deltas, deltas_sources = get_rdm2(terms, spin, strategy="trav" if spin == "uhf" else "exhaust") # Generate the 2RDM code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "preamble": get_density_einsum_preamble(2, spin), - "postamble": get_density_einsum_postamble(2, spin), - } + def preamble(): + done = set() + for delta, delta_source in zip(deltas, deltas_sources): + if delta in done: + continue + shape_source_index = 0 if delta.external_indices[0].space == "o" else 1 + codegen.tensor_declaration( + delta, + is_identity=True, + shape_source=delta_source, + shape_source_index=shape_source_index, + ) + codegen._tensor_declared.add(_tensor_info(delta)) + done.add(delta) + + def postamble(): + if name != "einsum": + raise NotImplementedError # FIXME remove packing, handle transpose + if spin != "uhf": + codegen.write("rdm2 = pack_2e(rdm2.oooo, rdm2.ooov, rdm2.oovo, rdm2.ovoo, rdm2.vooo, rdm2.oovv, rdm2.ovov, rdm2.ovvo, rdm2.voov, rdm2.vovo, rdm2.vvoo, rdm2.ovvv, rdm2.vovv, rdm2.vvov, rdm2.vvvo, rdm2.vvvv).transpose((0, 2, 1, 3))") + else: + codegen.write("rdm2.aaaa = pack_2e(rdm2.aaaa.oooo, rdm2.aaaa.ooov, rdm2.aaaa.oovo, rdm2.aaaa.ovoo, rdm2.aaaa.vooo, rdm2.aaaa.oovv, rdm2.aaaa.ovov, rdm2.aaaa.ovvo, rdm2.aaaa.voov, rdm2.aaaa.vovo, rdm2.aaaa.vvoo, rdm2.aaaa.ovvv, rdm2.aaaa.vovv, rdm2.aaaa.vvov, rdm2.aaaa.vvvo, rdm2.aaaa.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.aabb = pack_2e(rdm2.abab.oooo, rdm2.abab.ooov, rdm2.abab.oovo, rdm2.abab.ovoo, rdm2.abab.vooo, rdm2.abab.oovv, rdm2.abab.ovov, rdm2.abab.ovvo, rdm2.abab.voov, rdm2.abab.vovo, rdm2.abab.vvoo, rdm2.abab.ovvv, rdm2.abab.vovv, rdm2.abab.vvov, rdm2.abab.vvvo, rdm2.abab.vvvv).transpose((0, 2, 1, 3))") + codegen.write("rdm2.bbbb = pack_2e(rdm2.bbbb.oooo, rdm2.bbbb.ooov, rdm2.bbbb.oovo, rdm2.bbbb.ovoo, rdm2.bbbb.vooo, rdm2.bbbb.oovv, rdm2.bbbb.ovov, rdm2.bbbb.ovvo, rdm2.bbbb.voov, rdm2.bbbb.vovo, rdm2.bbbb.vvoo, rdm2.bbbb.ovvv, rdm2.bbbb.vovv, rdm2.bbbb.vvov, rdm2.bbbb.vvvo, rdm2.bbbb.vvvv).transpose((0, 2, 1, 3))") + codegen.write("del rdm2.abab") + codegen( "make_rdm2_f", returns, - output, - expr, - **kwargs, + output_expr, + preamble=preamble, + postamble=postamble, ) with Stopwatch("IP-EOM"): @@ -344,7 +292,7 @@ pq.add_double_commutator(1.0, ["v"], ["t1"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -361,60 +309,27 @@ pq.add_double_commutator(0.5, ["v"], ["t2"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="opt") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("EA-EOM"): @@ -431,7 +346,7 @@ pq.add_double_commutator(1.0, ["v"], ["t1"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -448,60 +363,27 @@ pq.add_double_commutator(0.5, ["v"], ["t2"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="opt") + output_expr_nr, returns_nr, output_expr_r, returns_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -519,7 +401,7 @@ pq.add_double_commutator(1.0, ["v"], ["t1"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the R2 contractions in pdaggerq format pq.clear() @@ -536,61 +418,46 @@ pq.add_double_commutator(0.5, ["v"], ["t2"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the R amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee") # Generate the R amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_matvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the R amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "postamble": "r2new.baba = np.transpose(r2new.abab, (1, 0, 3, 2))" if spin == "uhf" else None, # FIXME - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_matvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) with Stopwatch("L-IP-EOM"): @@ -607,7 +474,7 @@ pq.add_double_commutator(1.0, ["v"], ["t1"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -624,60 +491,27 @@ pq.add_double_commutator(0.5, ["v"], ["t2"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ip"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="opt") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ip") # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ip", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) with Stopwatch("L-EA-EOM"): @@ -694,7 +528,7 @@ pq.add_double_commutator(1.0, ["v"], ["t1"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -711,60 +545,27 @@ pq.add_double_commutator(0.5, ["v"], ["t2"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ea"): - indices = default_indices["v"][: n + 1] + default_indices["o"][: n] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="opt") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="exhaust", which="ea") # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" - if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ea", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + as_dict=True, ) if spin == "ghf": # FIXME @@ -782,7 +583,7 @@ pq.add_double_commutator(1.0, ["v"], ["t1"], ["t2"]) pq.simplify() terms_r1 = pq.fully_contracted_strings() - terms_r1 = remove_e0_eom(terms_r1) + terms_r1 = remove_reference_energy_eom(terms_r1) # Get the L2 contractions in pdaggerq format pq.clear() @@ -799,61 +600,46 @@ pq.add_double_commutator(0.5, ["v"], ["t2"], ["t2"]) pq.simplify() terms_r2 = pq.fully_contracted_strings() - terms_r2 = remove_e0_eom(terms_r2) + terms_r2 = remove_reference_energy_eom(terms_r2) # Get the L amplitudes in albert format - terms = [terms_r1, terms_r2] - expr = [] - output = [] - returns = [] - for n in range(2): - for index_spins in get_amplitude_spins(n + 1, spin, which="ee"): - indices = default_indices["o"][: n + 1] + default_indices["v"][: n + 1] - expr_n = import_from_pdaggerq(terms[n], index_spins=index_spins, l_is_lambda=False) - expr_n = spin_integrate(expr_n, spin) - output_n = get_t_amplitude_outputs(expr_n, f"r{n+1}new", indices=indices) - returns_n = (Tensor(*tuple(Index(i, index_spins[i]) for i in indices), name=f"r{n+1}new"),) - expr.extend(expr_n) - output.extend(output_n) - returns.extend(returns_n) - - (returns_nr, output_nr, expr_nr), (returns_r, output_r, expr_r) = optimise_eom(returns, output, expr, spin, strategy="trav") + returns_nr, output_expr_nr, returns_r, output_expr_r = get_eom([terms_r1, terms_r2], spin, strategy="trav", which="ee") # Generate the L amplitude intermediates code for name, codegen in code_generators.items(): - if name == "einsum": - kwargs = { - "as_dict": True, - } - else: - kwargs = {} codegen( "hbar_lmatvec_ee_intermediates", returns_nr, - output_nr, - expr_nr, - **kwargs, + output_expr_nr, + as_dict=True, ) # Generate the L amplitude code for name, codegen in code_generators.items(): - if name == "einsum": - preamble = "ints = kwargs[\"ints\"]" + def postamble(): if spin == "uhf": - preamble += "\nr1new = Namespace()\nr2new = Namespace()" - kwargs = { - "preamble": preamble, - "postamble": "r2new.baba = np.transpose(r2new.abab, (1, 0, 3, 2))" if spin == "uhf" else None, # FIXME - "as_dict": True, - } - else: - kwargs = {} + r2_abab = Tensor( + Index("i", space="o", spin="a"), + Index("j", space="v", spin="b"), + Index("a", space="o", spin="a"), + Index("b", space="v", spin="b"), + name="r2new", + ) + r2_baba = Tensor( + Index("i", space="o", spin="b"), + Index("j", space="v", spin="a"), + Index("a", space="o", spin="b"), + Index("b", space="v", spin="a"), + name="r2new", + ) + codegen.tensor_expression(r2_baba, r2_abab, is_return=True) + codegen( "hbar_lmatvec_ee", returns_r, - output_r, - expr_r, - **kwargs, + output_expr_r, + postamble=postamble, + as_dict=True, ) for codegen in code_generators.values(): diff --git a/ebcc/codegen/bootstrap_common.py b/ebcc/codegen/bootstrap_common.py index 0a8a22bf..71b4c7d8 100644 --- a/ebcc/codegen/bootstrap_common.py +++ b/ebcc/codegen/bootstrap_common.py @@ -1,847 +1,175 @@ """Common functionality for the bootstrap scripts for `ebcc`. """ -import itertools -import time -import re -from collections import defaultdict - -from albert.codegen.einsum import EinsumCodeGen as _EinsumCodeGen -from albert.codegen.base import sort_exprs -from albert.optim._gristmill import optimise as _optimise -from albert.qc.spin import generalised_to_restricted, generalised_to_unrestricted -from albert.qc.index import Index -from albert.qc.rhf import ERI as RERI, CDERI as RCDERI, R0 as RR0, L0 as RL0 -from albert.qc.uhf import ERI as UERI, CDERI as UCDERI, R0 as UR0, L0 as UL0 -from albert.qc.ghf import R0 as GR0, L0 as GL0 from albert.qc._pdaggerq import import_from_pdaggerq +from albert.qc.spin import ghf_to_rhf, ghf_to_uhf, get_amplitude_spins, get_density_spins +from albert.qc.decomp import density_fit +from albert.opt import optimise +from albert.opt.tools import optimise_eom from albert.tensor import Tensor -from albert.symmetry import Symmetry, Permutation -from albert.algebra import Mul, Add -from pdaggerq.config import OCC_INDICES, VIRT_INDICES - -ov_2e = [ - "oooo", - "ooov", - "oovo", - "ovoo", - "vooo", - "oovv", - "ovov", - "ovvo", - "voov", - "vovo", - "vvoo", - "ovvv", - "vovv", - "vvov", - "vvvo", - "vvvv", -] -ov_1e = ["oo", "ov", "vo", "vv"] - - -default_indices = { - "1": ["1"], - "o": OCC_INDICES, - "v": VIRT_INDICES, - "O": [i.upper() for i in OCC_INDICES], - "V": [i.upper() for i in VIRT_INDICES], - "i": [f"{i}_" for i in OCC_INDICES], - "a": [f"{i}_" for i in VIRT_INDICES], - "b": ["x", "y", "z", "b0", "b1", "b2", "b3"], - "x": ["P", "Q", "R", "S", "x0", "x1", "x2", "x3", "x4", "x5", "x7"], - "d": ["DUMMY1", "DUMMY2", "DUMMY3", "DUMMY4"], -} - - -default_sectors = {i: k for k, v in default_indices.items() for i in v} - -default_sizes = { - "1": 1, - "o": 200, - "v": 1000, - "O": 8, - "V": 16, - "i": 192, - "a": 984, - "b": 10, - "x": 3000, - "d": 100000, -} - -class EinsumCodeGen(_EinsumCodeGen): - """Code generator for the bootstrap scripts for `ebcc`.""" - - def __init__( - self, - einsum_func="einsum", - einsum_kwargs=None, - transpose_func="np.transpose({arg}, {transpose})", - name_generator=None, - spin="ghf", - **kwargs, - ): - if einsum_kwargs is None: - einsum_kwargs = {} - super().__init__( - einsum_func=einsum_func, - einsum_kwargs=einsum_kwargs, - transpose_func=transpose_func, - name_generator=name_generator, - spin=spin, - **kwargs, - ) - - def preamble(self, imports=None): - if imports is None: - imports = "from ebcc import numpy as np\n" - imports += "from ebcc.util import pack_2e, einsum, dirsum, Namespace" - super().preamble(imports=imports) - - def ignore_argument(self, arg): - """ - Return `True` if a potential function argument should be - ignored. - """ - return "tmp" in arg.name or arg.name in ("δ", "gc") - - def function_docstring(self, docstring): - """Write the function docstring.""" - if self.spin in ("rhf", "ghf"): - array_type = "array" - else: - array_type = "Namespace of arrays" - amplitude_names = [ - "t1", "t2", "t3", "l1", "l2", "l3", "s1", "s2", "ls1", "ls2", "u11", "u12", "lu11", "lu12", "r1", "r2", "r3" - ] - descriptions = { - "f": "Fock matrix.", - "v": "Electron repulsion integrals.", - "G": "One-boson Hamiltonian.", - "w": "Two-boson Hamiltonian.", - "g": "Electron-boson coupling.", - "e_cc": "Coupled cluster energy.", - "e_pert": "Perturbation energy.", - "rdm1": "One-particle reduced density matrix.", - "rdm2": "Two-particle reduced density matrix.", - "rdm1_b": "One-body reduced density matrix.", - "rdm_eb_cre": "Electron-boson coupling reduced density matrix, creation part.", - "rdm_eb_des": "Electron-boson coupling reduced density matrix, annihilation part.", - "dm_cre": "Single boson density matrix, creation part.", - "dm_des": "Single boson density matrix, annihilation part.", - **{f"{name}": f"{name.upper()} amplitudes." for name in amplitude_names}, - **{f"{name}new": f"Updated {name.upper()} residuals." for name in amplitude_names}, - } - types = { - "f": array_type, - "v": array_type, - "G": "array", - "w": "array", - "g": array_type, - "e_cc": "float", - "e_pert": "float", - "rdm1": array_type, - "rdm2": array_type, - "rdm1_b": "array", - "rdm_eb_cre": array_type, - "rdm_eb_des": array_type, - "dm_cre": "array", - "dm_des": "array", - **{f"{name}": array_type if not name.startswith("s") else "array" for name in amplitude_names}, - **{f"{name}new": array_type if not name.startswith("s") else "array" for name in amplitude_names}, - } - docstring = docstring.split("\n") - new_docstring = [] - for line in docstring: - if len(line.split()) and line.split()[0] in descriptions: - name = line.split()[0] - line = f"{name} : {types[name]}\n {descriptions[name]}" - new_docstring.append(line) - docstring = "\n".join(new_docstring) - super().function_docstring(docstring) - - -class Stopwatch: - def __init__(self, name=None): - self.name = name - - def __enter__(self): - self.start = time.time() - return self - - def __exit__(self, *args): - self.end = time.time() - self.elapsed = self.end - self.start - if self.name: - print(f"{self.name}: {self.elapsed:.2f}s") - - -def name_generator_rhf(tensor, add_spaces=True): - """Generate names for the RHF case.""" - if tensor.name in ("f", "v", "d", "Γ", "δ", "g", "gc", "rdm_eb_cre", "rdm_eb_des"): - if tensor.name == "d": - name = "rdm1" - elif tensor.name == "Γ": - name = "rdm2" - elif tensor.name == "δ": - name = "delta" - else: - name = tensor.name - if add_spaces: - spaces = [i.space for i in tensor.indices] - return f"{name}.{''.join(spaces)}" - else: - return name +def spin_integrate(exprs, spin): + """Perform the spin integration on the given expressions.""" + if spin == "ghf": + return exprs + elif spin == "uhf": + return ghf_to_uhf(exprs) else: - return tensor.name - - -def name_generator_uhf(tensor, add_spaces=True): - """Generate names for the UHF case.""" - if tensor.name in ("f", "v", "d", "Γ", "δ", "g", "gc", "rdm_eb_cre", "rdm_eb_des"): - if tensor.name == "d": - name = "rdm1" - elif tensor.name == "Γ": - name = "rdm2" - elif tensor.name == "δ": - name = "delta" - else: - name = tensor.name - if add_spaces: - spins = ["a" if i.spin == "α" else "b" for i in tensor.indices if i.spin is not None] - spaces = [i.space for i in tensor.indices] - return f"{name}.{''.join(spins)}.{''.join(spaces)}" + return (ghf_to_rhf(exprs),) + + +def get_energy(terms, spin, strategy="exhaust", density_fit=False): + """Get the energy expressions from `pdaggerq` terms.""" + expr = import_from_pdaggerq(terms) + expr = spin_integrate(expr, spin) + if density_fit: + expr = tuple(density_fit(e) for e in expr) + output = tuple(Tensor(name="e_cc") for _ in range(len(expr))) + output_expr = optimise(output, expr, strategy=strategy) + returns = (Tensor(name="e_cc"),) + return output_expr, returns + + +def get_amplitudes(terms_grouped, spin, strategy="exhaust", which="t", orders=None, density_fit=False): + """Get the amplitude expressions from `pdaggerq` terms.""" + expr = [] + output = [] + returns = [] + if orders is None: + orders = range(1, len(terms_grouped) + 1) + for order, terms in zip(orders, terms_grouped): + n = order - 1 + if which == "t": + indices = "ijkl"[: n + 1] + "abcd"[: n + 1] else: - return name - elif tensor.name in ("t1", "t2", "t3", "t1new", "t2new", "t3new", "l1", "l2", "l3", "l1new", "l2new", "l3new", "u11", "u11new", "u12", "u12new", "r1", "r2", "r3", "r1new", "r2new", "r3new"): - if add_spaces: - spins = ["a" if i.spin == "α" else "b" for i in tensor.indices if i.spin is not None] - return f"{tensor.name}.{''.join(spins)}" - else: - return tensor.name + indices = "abcd"[: n + 1] + "ijkl"[: n + 1] + for index_spins in get_amplitude_spins(indices[: n + 1], indices[n + 1 :], spin): + expr_n = import_from_pdaggerq(terms, index_spins=index_spins) + expr_n = spin_integrate(expr_n, spin) + output_n = [Tensor(*tuple(sorted(e.external_indices, key=lambda i: indices.index(i.name))), name=f"{which}{order}new") for e in expr_n] + returns_n = (output_n[0],) + expr.extend(expr_n) + output.extend(output_n) + returns.extend(returns_n) + if strategy is not None: + output_expr = optimise(output, expr, strategy="exhaust") else: - return tensor.name - - -name_generators = { - "rhf": name_generator_rhf, - "uhf": name_generator_uhf, - "ghf": name_generator_rhf, -} - - -def spin_integrate(expr, spin): - """Perform the spin integration.""" - if isinstance(expr, Add): - args = expr.args - else: - args = [expr] - out = tuple() - for p0 in range(0, len(args), 1000): - # Batch expressions to avoid deep recursion - p1 = min(p0 + 1000, len(args)) - if spin == "rhf": - out += (generalised_to_restricted(Add(*args[p0:p1])),) - elif spin == "uhf": - out += generalised_to_unrestricted(Add(*args[p0:p1])) - else: - out += (Add(*args[p0:p1]),) - return out - - -def remove_hf_energy(terms): - """Remove the HF energy from the terms.""" - new_terms = [] - for term in terms: - if len(term) == 2: - factor, tensor = term - if tensor.startswith("f_"): - tensor = "f" + tensor[4:] - if tensor.startswith("<") and "_" in tensor: - tensor = tensor[:-5] - if (factor, tensor) == ("+1.00000000000000", "f(i,i)"): - continue - if (factor, tensor) == ("-0.50000000000000", ""): - continue - if (factor, tensor) == ("-0.50000000000000", ""): - continue - new_terms.append(term) - return new_terms - - -def remove_e0_eom(terms): - """Remove the EOM terms to transform H v -> (H - E0) v.""" - new_terms = [] - for term in terms: - # Find if the term is disconnected - r = None - rest = [] - for t in term[1:]: - if t.startswith("r") or t.startswith("l"): - r = t - elif not t.startswith("P("): - rest.append(t) - r_inds = set(r.split("(")[1].split(")")[0].split(",")) - rest_inds = set() - for r in rest: - if "<" in r: - r = r.replace("<", "(").replace(">", ")").replace("||", ",") - rest_inds.update(r.split("(")[1].split(")")[0].split(",")) - connected = r_inds & rest_inds - if connected: - new_terms.append(term) - continue - - # We only want to remove the E0 terms: - # f(i,i) r - # f(i,a) t(a,i) r - # r - # t2(a,b,i,j) r - # t1(a,i) t1(b,j) r - if len(term) == 3: - tensor = [t for t in term[1:] if not (t.startswith("r") or t.startswith("l"))][0] - if tensor.startswith("f") and tensor[2] == tensor[4]: - continue - if tensor.startswith("<") and tensor[1] == tensor[6] and tensor[3] == tensor[8]: - continue - else: - tensors = sorted([t for t in term[1:] if not (t.startswith("r") or t.startswith("l"))]) - if tensors[0].startswith("f") and tensors[1].startswith("t"): - continue - if tensors[0].startswith("<") and all(t.startswith("t") for t in tensors[1:]): - continue + output_expr = list(zip(output, expr)) + return output_expr, returns - new_terms.append(term) - - return new_terms - - -def optimise(outputs, exprs, spin, strategy="greedy", sizes=None, **kwargs): - """Optimise the expressions.""" - - if sizes is None: - sizes = default_sizes - - index_sizes = {} - index_groups = [] - for sector, indices in default_indices.items(): - spins = ("α", "β") if spin == "uhf" and sector not in ("b", "x", "d") else (None,) - for s in spins: - index_sizes.update({Index(index, space=sector, spin=s): sizes[sector] for index in indices}) - index_groups.append([Index(index, space=sector, spin=s) for index in indices]) - - opt = _optimise( - *zip(outputs, exprs), - index_groups=index_groups, - sizes=index_sizes, - strategy=strategy, - **kwargs, - ) - - # Add expressions with the same output -- gristmill avoids this sometimes - # if you set shallow intermediates to False - output_exprs = {} - for o, e in opt: - if o not in output_exprs: - output_exprs[o] = e - elif isinstance(output_exprs[o], Add): - output_exprs[o] = Add(*output_exprs[o].args, e) - else: - output_exprs[o] += e - outputs, exprs = zip(*output_exprs.items()) - - return outputs, exprs - - -def get_t_amplitude_outputs(exprs, name, indices=None): - """Get the outputs for the T amplitude code.""" - if indices is not None: - key = lambda i: indices.index(i.name) - else: - key = lambda i: i - return [Tensor(*sorted(e.external_indices, key=key), name=name) for e in exprs] - -def get_l_amplitude_outputs(exprs, name, indices=None): - """Get the outputs for the L amplitude code.""" - def key(i): - if indices is not None: - return indices.index(i.name) - return (" bvo".index(i.space), i.spin, i.name) - return [Tensor(*sorted(e.external_indices, key=key), name=name) for e in exprs] - - -def get_density_outputs(exprs, name, indices): - """Get the outputs for the density code.""" - tensors = [] - for expr in exprs: - external_indices = sorted(expr.external_indices, key=lambda i: indices.index(i.name)) - tensors.append(Tensor(*external_indices, name=name)) - return tensors - - -def get_amplitude_spins(n, spin, which="t"): - """Get the spin cases for the amplitudes.""" - - if which in ("t", "l", "ee"): - no = nv = n - elif which == "ip": - no = n - nv = n - 1 - elif which == "ea": - no = n - 1 - nv = n - - if spin == "rhf": - case = {} - for i in range(no): - case[default_indices["o"][i]] = ["α", "β"][i % 2] - for i in range(nv): - case[default_indices["v"][i]] = ["α", "β"][i % 2] - cases = [case] +def get_rdm1(terms_sectors, spin, strategy="exhaust"): + """Get the one-body reduced density matrix expressions from `pdaggerq` terms.""" + if spin == "ghf": + from albert.qc.ghf import Delta, T1 elif spin == "uhf": - cases = [] - if which in ("ip", "ea"): - it = itertools.product("αβ", repeat=max(no, nv)) - else: - it = itertools.combinations_with_replacement("αβ", max(no, nv)) - for spins in it: - if which not in ("ip", "ea"): - # Canonicalise the spin order -- never for IP/EA...? - best = [None, 1e10] - for s in itertools.permutations(spins): - penalty = 0 - for i in range(len(s) - 1): - penalty += int(s[i] == s[i + 1]) * 2 - if s[0] != min(s): - penalty += 1 - if penalty < best[1]: - best = [s, penalty] - spins = best[0] - case = {} - for i, s in enumerate(spins): - if i < no: - case[default_indices["o"][i]] = s - if i < nv: - case[default_indices["v"][i]] = s - cases.append(case) - elif spin == "ghf": - case = {} - for i in range(no): - case[default_indices["o"][i]] = None - for i in range(nv): - case[default_indices["v"][i]] = None - cases = [case] - - return cases - - -def get_density_spins(n, spin, indices): - """Get the spin cases for the density.""" - - assert n in (1, 2) # hardcoded, TODO - - if spin == "rhf": - if n == 1: - cases = [("α", "α")] - elif n == 2: - cases = [ - ("α", "α", "α", "α"), - ("α", "β", "α", "β"), - ("β", "α", "β", "α"), - ("β", "β", "β", "β"), - ] - cases = [dict(zip(indices, case)) for case in cases] + from albert.qc.uhf import Delta, T1 + else: + from albert.qc.rhf import Delta, T1 + expr = [] + output = [] + returns = [] + deltas = [] + deltas_sources = [] + for sectors, indices in [("oo", "ij"), ("ov", "ia"), ("vo", "ai"), ("vv", "ab")]: + for index_spins in get_density_spins(indices, spin): + expr_n = import_from_pdaggerq(terms_sectors[sectors, indices], index_spins=index_spins) + expr_n = spin_integrate(expr_n, spin) + if spin == "rhf": + expr_n = tuple(e * 2 for e in expr_n) + if density_fit: + expr_n = tuple(density_fit(e) for e in expr_n) + output_n = [Tensor(*tuple(sorted(e.external_indices, key=lambda i: indices.index(i.name))), name=f"rdm1") for e in expr_n] + returns_n = (output_n[0],) + expr.extend(expr_n) + output.extend(output_n) + returns.extend(returns_n) + if len(set(sectors)) == 1: + delta = Delta(*tuple(sorted(expr_n[0].external_indices, key=lambda i: indices.index(i.name)))) + deltas.append(delta) + deltas_sources.append(next(expr_n[0].search_leaves(T1))) + output_expr = optimise(output, expr, strategy=strategy) + return output_expr, returns, deltas, deltas_sources + + +def get_rdm2(terms_sectors, spin, strategy="exhaust", density_fit=False): + """Get the two-body reduced density matrix expressions from `pdaggerq` terms.""" + if spin == "ghf": + from albert.qc.ghf import Delta, T1 elif spin == "uhf": - if n == 1: - cases = [("α", "α"), ("β", "β")] - elif n == 2: - cases = [("α", "α", "α", "α"), ("α", "β", "α", "β"), ("β", "β", "β", "β")] - cases = [dict(zip(indices, case)) for case in cases] - elif spin == "ghf": - if n == 1: - cases = [(None, None)] - elif n == 2: - cases = [(None, None, None, None)] - cases = [dict(zip(indices, case)) for case in cases] - - return cases - - -def get_density_einsum_preamble(n, spin, name="rdm{n}"): - """Get the einsum preamble for the density.""" - name = name.format(n=n) - preamble = f"{name} = Namespace()" - if spin == "uhf": - for spins in itertools.combinations_with_replacement(["a", "b"], n): - preamble += f"\n{name}.{''.join(spins+spins)} = Namespace()" - preamble += "\ndelta = Namespace(" - preamble += "\n aa=Namespace(oo=np.eye(t2.aaaa.shape[0]), vv=np.eye(t2.aaaa.shape[-1]))," - preamble += "\n bb=Namespace(oo=np.eye(t2.bbbb.shape[0]), vv=np.eye(t2.bbbb.shape[-1]))," - preamble += "\n)" + from albert.qc.uhf import Delta, T1 else: - preamble += "\ndelta = Namespace(" - preamble += "\n oo=np.eye(t2.shape[0])," - preamble += "\n vv=np.eye(t2.shape[-1])," - preamble += "\n)" - return preamble - - -def get_density_einsum_postamble(n, spin, name="rdm{n}", spaces=None): - """Get the einsum postamble for the density.""" - # TODO hardcoded - name = name.format(n=n) - if spaces is None: - if n == 1: - spaces = ["oo", "ov", "vo", "vv"] - elif n == 2: - spaces = ov_2e - if n == 1: - if spin == "uhf": - postamble = f"{name}.aa = np.block([[{name}.aa.{spaces[0]}, {name}.aa.{spaces[1]}], [{name}.aa.{spaces[2]}, {name}.aa.{spaces[3]}]])" - postamble += f"\n{name}.bb = np.block([[{name}.bb.{spaces[0]}, {name}.bb.{spaces[1]}], [{name}.bb.{spaces[2]}, {name}.bb.{spaces[3]}]])" - else: - postamble = f"{name} = np.block([[{name}.{spaces[0]}, {name}.{spaces[1]}], [{name}.{spaces[2]}, {name}.{spaces[3]}]])" - elif n == 2: - if spin == "uhf": - postamble = f"{name}.aaaa = pack_2e(%s)" % ", ".join(f"{name}.aaaa.{perm}" for perm in spaces) - postamble += f"\n{name}.abab = pack_2e(%s)" % ", ".join( - f"{name}.abab.{perm}" for perm in spaces - ) - postamble += f"\n{name}.bbbb = pack_2e(%s)" % ", ".join( - f"{name}.bbbb.{perm}" for perm in spaces - ) - postamble += f"\n{name} = Namespace(" - postamble += f"\n aaaa=np.transpose({name}.aaaa, (0, 2, 1, 3))," - postamble += f"\n aabb=np.transpose({name}.abab, (0, 2, 1, 3))," - postamble += f"\n bbbb=np.transpose({name}.bbbb, (0, 2, 1, 3))," - postamble += f"\n)" + from albert.qc.rhf import Delta, T1 + expr = [] + output = [] + returns = [] + deltas = [] + deltas_sources = [] + for sectors, indices in [ + ("oooo", "ijkl"), + ("ooov", "ijka"), + ("oovo", "ijak"), + ("ovoo", "iajk"), + ("vooo", "aijk"), + ("oovv", "ijab"), + ("ovov", "iajb"), + ("ovvo", "iabj"), + ("voov", "aijb"), + ("vovo", "aibj"), + ("vvoo", "abij"), + ("ovvv", "iabc"), + ("vovv", "aibc"), + ("vvov", "abic"), + ("vvvo", "abci"), + ("vvvv", "abcd"), + ]: + for index_spins in get_density_spins(indices, spin): + expr_n = import_from_pdaggerq(terms_sectors[sectors, indices], index_spins=index_spins) + expr_n = spin_integrate(expr_n, spin) + if density_fit: + expr_n = tuple(density_fit(e) for e in expr_n) + output_n = [Tensor(*tuple(sorted(e.external_indices, key=lambda i: indices.index(i.name))), name=f"rdm2") for e in expr_n] + returns_n = (output_n[0],) + expr.extend(expr_n) + output.extend(output_n) + returns.extend(returns_n) + if len(set(sectors)) == 1: + delta = Delta(*tuple(sorted(expr_n[0].external_indices, key=lambda i: indices.index(i.name))[:2])) + deltas.append(delta) + deltas_sources.append(next(expr_n[0].search_leaves(T1))) + output_expr = optimise(output, expr, strategy=strategy) + return output_expr, returns, deltas, deltas_sources + + +def get_eom(terms_grouped, spin, strategy="exhaust", which="ip", orders=None, density_fit=False): + """Get the EOM expressions from `pdaggerq` terms.""" + expr = [] + output = [] + returns = [] + if orders is None: + orders = range(1, len(terms_grouped) + 1) + for order, terms in zip(orders, terms_grouped): + n = order - 1 + if which == "ip": + indices = "ijkl"[: n + 1] + "abcd"[: n] + elif which == "ea": + indices = "abcd"[: n + 1] + "ijkl"[: n] else: - postamble = f"{name} = pack_2e(%s)" % ", ".join(f"{name}.{perm}" for perm in spaces) - postamble += f"\n{name} = np.transpose({name}, (0, 2, 1, 3))" - return postamble - - -def get_boson_einsum_preamble(spin): - """Get the einsum preamble for the density.""" + indices = "ijkl"[: n + 1] + "abcd"[: n + 1] + for index_spins in get_amplitude_spins(indices[: n + 1], indices[n + 1 :], spin): + expr_n = import_from_pdaggerq(terms, index_spins=index_spins, l_is_lambda=False) + expr_n = spin_integrate(expr_n, spin) + if density_fit: + expr_n = tuple(density_fit(e) for e in expr_n) + output_n = [Tensor(*sorted(e.external_indices, key=lambda i: indices.index(i.name)), name=f"r{order}new") for e in expr_n] + returns_n = (output_n[0],) + expr.extend(expr_n) + output.extend(output_n) + returns.extend(returns_n) + (returns_nr, output_expr_nr), (returns_r, output_expr_r) = optimise_eom(returns, output, expr, strategy=strategy) if spin == "uhf": - preamble = "gc = Namespace(" - preamble += "\n aa=Namespace(" - preamble += "\n boo=np.transpose(g.aa.boo, (0, 2, 1))," - preamble += "\n bov=np.transpose(g.aa.bvo, (0, 2, 1))," - preamble += "\n bvo=np.transpose(g.aa.bov, (0, 2, 1))," - preamble += "\n bvv=np.transpose(g.aa.bvv, (0, 2, 1))," - preamble += "\n )," - preamble += "\n bb=Namespace(" - preamble += "\n boo=np.transpose(g.bb.boo, (0, 2, 1))," - preamble += "\n bov=np.transpose(g.bb.bvo, (0, 2, 1))," - preamble += "\n bvo=np.transpose(g.bb.bov, (0, 2, 1))," - preamble += "\n bvv=np.transpose(g.bb.bvv, (0, 2, 1))," - preamble += "\n )," - preamble += "\n)" - else: - preamble = "gc = Namespace(" - preamble += "\n boo=np.transpose(g.boo, (0, 2, 1))," - preamble += "\n bov=np.transpose(g.bvo, (0, 2, 1))," - preamble += "\n bvo=np.transpose(g.bov, (0, 2, 1))," - preamble += "\n bvv=np.transpose(g.bvv, (0, 2, 1))," - preamble += "\n)" - return preamble - - -def get_density_fit(): - indices = {"": 0} - - def density_fit(obj): - """Get the function to apply over an expression to density fit the ERIs.""" - i = indices[""] % len(default_indices["x"]) # TODO this will break for huge ansatzes - aux_index = Index(f"{default_indices['x'][i]}", space="x") - indices[""] += 1 - if hasattr(obj, "_symbol") and obj._symbol == RERI: - left = RCDERI[(aux_index,) + obj.indices[:2]] - right = RCDERI[(aux_index,) + obj.indices[2:]] - elif hasattr(obj, "_symbol") and obj._symbol == UERI: - left = UCDERI[(aux_index,) + obj.indices[:2]] - right = UCDERI[(aux_index,) + obj.indices[2:]] - else: - return obj - return left * right - - return density_fit - - -def optimise_eom(returns, output, expr, spin, strategy="exhaust"): - """Optimise EOM expressions into intermediates and final output.""" - # Hack the expressions to make the optimiser more likely to optimise out - # intermediates that are constant between EOM iterations - new_output = [] - new_expr = [] - tensor_types = {} - for o, e in zip(output, expr): - for a in e.nested_view(): - for i in a: - if isinstance(i, Tensor): - tensor_types[i.name] = i._symbol - a = [Tensor(*x.indices, symmetry=x.symmetry, name=x.name) if isinstance(x, Tensor) else x for x in a] - for i in range(len(a)): - if isinstance(a[i], Tensor) and (a[i].name.startswith("r") or a[i].name.startswith("l")): - a[i] = a[i].copy( - Index("DUMMY1", space="d"), - *a[i].indices, - symmetry=Symmetry(*[Permutation((0,), 1) + p for p in a[i].symmetry.permutations]), - ) - oa = o.copy( - Index("DUMMY1", space="d"), - *o.indices, - ) - new_output.append(oa) - new_expr.append(Mul(*a)) - output = new_output - expr = new_expr - - output, expr = optimise(output, expr, spin, strategy=strategy) - - # Unhack - new_output = [] - new_expr = [] - for o, e in zip(output, expr): - for a in e.nested_view(): - for i in range(len(a)): - if isinstance(a[i], Tensor): - remove = [ind.space == "d" for ind in a[i].indices] - a[i] = a[i].copy( - *tuple(ind for ind, r in zip(a[i].indices, remove) if not r), - symmetry=Symmetry( - *[ - Permutation( - tuple(x-sum(remove) for j, x in enumerate(p.permutation) if not remove[j]), - p.sign, - ) - for p in a[i].symmetry.permutations - ] - ) if a[i].symmetry else None - ) - remove = [ind.space == "d" for ind in o.indices] - oa = o.copy( - *tuple(ind for ind, r in zip(o.indices, remove) if not r), - symmetry=Symmetry( - *[ - Permutation( - tuple(x-sum(remove) for j, x in enumerate(p.permutation) if not remove[j]), - p.sign, - ) - for p in o.symmetry.permutations - ] - ) if o.symmetry else None - ) - new_output.append(oa) - new_expr.append(Mul(*a)) - output = new_output - expr = new_expr - - # Extract the intermediates that don't depend on R - output_r = [] - expr_r = [] - output_nr = [] - expr_nr = [] - cache = set() - for o, e in zip(output, expr): - depends_on_r = o.name.startswith("r") or o.name.startswith("l") - if not depends_on_r: - for a in e.nested_view(): - for i in a: - if isinstance(i, Tensor): - if i.name.startswith("r") or i.name.startswith("l") or i.name in cache: - depends_on_r = True - if depends_on_r: - output_r.append(o) - expr_r.append(e) - cache.add(o.name) - else: - output_nr.append(o) - expr_nr.append(e) - - # Get the tmps needed to return from the intermediates function - returns_r = returns - returns_nr = [] - initialised_here = set() - for o, e in zip(output_r, expr_r): - if o.name.startswith("tmp"): - initialised_here.add(o.name) - for a in e.nested_view(): - for i in a: - if isinstance(i, Tensor) and i.name.startswith("tmp"): - if i.name not in initialised_here: - returns_nr.append(i) - - # Transform the names of the intermediates - new_expr_r = [] - for o, e in zip(output_r, expr_r): - add_args = [] - for args in e.nested_view(): - args = list(args) - for i, a in enumerate(args): - if isinstance(a, Tensor) and a.name.startswith("tmp") and a.name not in initialised_here: - args[i] = Tensor(*a.indices, name=f"ints.{a.name}") - add_args.append(Mul(*args)) - new_expr_r.append(Add(*add_args)) - expr_r = new_expr_r - - # Re-optimise the output part -- first resub the tmps in - while True: - oi, ei = output_r[0], expr_r[0] - ei = ei.expand() - if oi.name.startswith("r") or oi.name.startswith("l"): - break - for j, (oj, ej) in enumerate(zip(output_r[1:], expr_r[1:])): - new_muls = [] - for mul_args in ej.nested_view(): - new_args = [] - for arg in mul_args: - if isinstance(arg, Tensor) and arg.name == oi.name: - index_map = dict(zip(oi.external_indices, arg.external_indices)) - for index in ei.dummy_indices: - # avoid collision - i = 0 - new_index = None - while new_index is None or new_index in ei.external_indices or new_index in ei.dummy_indices or new_index in index_map.values(): - new_index = Index(name=default_indices[index.space][-(i + 1)], space=index.space, spin=index.spin) - i += 1 - index_map[index] = new_index - new_arg = ei.map_indices(index_map) - new_args.append(new_arg) - else: - new_args.append(arg) - new_muls.append(Mul(*new_args)) - expr_r[j + 1] = Add(*new_muls) - output_r = output_r[1:] - expr_r = expr_r[1:] - # Sum the R terms, else factorisation is not complete after the optimisation - new_output_expr_r = {} - new_output_r = [] - key = lambda o: (o.name, tuple(i.spin for i in o.indices), tuple(i.space for i in o.indices)) - for o, e in zip(output_r, expr_r): - okey = key(o) - if okey not in new_output_expr_r: - new_output_expr_r[okey] = Add(e) - new_output_r.append(o) - elif isinstance(new_output_expr_r[okey], Add): - new_output_expr_r[okey] = Add(*new_output_expr_r[okey].args, e) - else: - new_output_expr_r[okey] = Add(new_output_expr_r[okey], e) - output_r, expr_r = optimise(new_output_r, [new_output_expr_r[key(o)] for o in new_output_r], spin, strategy=strategy) - - # Replace the tensor types so the canonicalisation works - new_output_nr = [] - new_expr_nr = [] - for o, e in zip(output_nr, expr_nr): - new_args = [] - for args in e.nested_view(): - args = list(args) - for i, a in enumerate(args): - if isinstance(a, Tensor) and a.name in tensor_types: - args[i] = tensor_types[a.name][a.indices].canonicalise() - new_args.append(Mul(*args)) - new_output_nr.append(o) - new_expr_nr.append(Add(*new_args)) - output_nr = new_output_nr - expr_nr = new_expr_nr - new_output_r = [] - new_expr_r = [] - for o, e in zip(output_r, expr_r): - new_args = [] - for args in e.nested_view(): - args = list(args) - for i, a in enumerate(args): - if isinstance(a, Tensor) and a.name in tensor_types: - args[i] = tensor_types[a.name][a.indices].canonicalise() - new_args.append(Mul(*args)) - new_output_r.append(o) - new_expr_r.append(Add(*new_args)) - output_r = new_output_r - expr_r = new_expr_r - - return (returns_nr, output_nr, expr_nr), (returns, output_r, expr_r) - - -def optimise_trans_dm(output, expr, spin, strategy="exhaust"): - """Optimise TDM expressions into intermediates and final output.""" - #raise NotImplementedError("broken") - - # Split the expressions into those that depend on R0 or L0 and those that don't - output_r0 = [] - expr_r0 = [] - output_l0 = [] - expr_l0 = [] - output_r0_l0 = [] - expr_r0_l0 = [] - output_rest = [] - expr_rest = [] - for o, e in zip(output, expr): - for args in e.nested_view(): - depends_on_r0 = False - depends_on_l0 = False - mul_args = [] - for i, a in enumerate(args): - if isinstance(a, Tensor) and a.name.startswith("r0"): - depends_on_r0 = True - if isinstance(a, Tensor) and a.name.startswith("l0"): - depends_on_l0 = True - mul_args.append(a) - if depends_on_r0 and depends_on_l0: - output_r0_l0.append(o) - expr_r0_l0.append(Mul(*mul_args)) - elif depends_on_r0: - output_r0.append(o) - expr_r0.append(Mul(*mul_args)) - elif depends_on_l0: - output_l0.append(o) - expr_l0.append(Mul(*mul_args)) - else: - output_rest.append(o) - expr_rest.append(Mul(*mul_args)) - - # Store the original outputs - original_output_r0 = set(o.name for o in output_r0) - original_output_l0 = set(o.name for o in output_l0) - original_output_r0_l0 = set(o.name for o in output_r0_l0) - - # Optimise the expressions - if output_r0: - output_r0, expr_r0 = optimise(output_r0, expr_r0, spin, strategy=strategy, interm_fmt="tmp_r0_{}") - if output_l0: - output_l0, expr_l0 = optimise(output_l0, expr_l0, spin, strategy=strategy, interm_fmt="tmp_l0_{}") - if output_r0_l0: - output_r0_l0, expr_r0_l0 = optimise(output_r0_l0, expr_r0_l0, spin, strategy=strategy, interm_fmt="tmp_r0_l0_{}") - if output_rest: - output_rest, expr_rest = optimise(output_rest, expr_rest, spin, strategy=strategy) - - # Add the factors back in - new_expr_r0 = [] - for i, (o, e) in enumerate(zip(output_r0, expr_r0)): - if o.name in original_output_r0: - e = e * {"rhf": RR0, "uhf": UR0, "ghf": GR0}[spin][tuple()] - new_expr_r0.append(e) - expr_r0 = new_expr_r0 - new_expr_l0 = [] - for i, (o, e) in enumerate(zip(output_l0, expr_l0)): - if o.name in original_output_l0: - e = e * {"rhf": RL0, "uhf": UL0, "ghf": GL0}[spin][tuple()] - new_expr_l0.append(e) - expr_l0 = new_expr_l0 - new_expr_r0_l0 = [] - for i, (o, e) in enumerate(zip(output_r0_l0, expr_r0_l0)): - if o.name in original_output_r0_l0: - e = e * {"rhf": RR0, "uhf": UR0, "ghf": GR0}[spin][tuple()] - e = e * {"rhf": RL0, "uhf": UL0, "ghf": GL0}[spin][tuple()] - new_expr_r0_l0.append(e) - expr_r0_l0 = new_expr_r0_l0 - - # Combine the results - output = tuple(output_rest) + tuple(output_r0) + tuple(output_l0) + tuple(output_r0_l0) - expr = tuple(expr_rest) + tuple(expr_r0) + tuple(expr_l0) + tuple(expr_r0_l0) - - return output, expr + # R amplitudes may get wrong spins + output_expr_nr = [(output, expr.canonicalise()) for output, expr in output_expr_nr] + output_expr_r = [(output, expr.canonicalise()) for output, expr in output_expr_r] + return output_expr_nr, returns_nr, output_expr_r, returns_r