New scripts for new albert

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():