tequila-basic 1.9.8__py3-none-any.whl → 1.9.10__py3-none-any.whl

tequila/grouping/fermionic_methods.py

@@ -13,6 +13,7 @@ from tequila import TequilaException
 from shutil import rmtree
 import logging
 
+
 def get_psi(h_ferm, mol_name, n_elec, n_qubit, trunc, trunc_perc, get_fci):
     if get_fci:
         _, psis_fci = get_wavefunction(jordan_wigner(h_ferm), "fci", mol_name, n_elec)
@@ -25,8 +26,9 @@ def get_psi(h_ferm, mol_name, n_elec, n_qubit, trunc, trunc_perc, get_fci):
     else:
         return psi_appr
 
+
 def do_fff(h_ferm, n_elec, options=None, restart=False, metric_estim=True):
-    '''
+    """
     Main function for Fluid Fermionic Fragments methods.
     Parameters
     ---------
@@ -55,28 +57,44 @@ def do_fff(h_ferm, n_elec, options=None, restart=False, metric_estim=True):
     -------
     Optimized fragment operators in FermionOperator type.
     Unitary operators (orbtial rotations) diagonalizing each fragment.
-    '''
+    """
+
     def process_options(options):
-        mol_name, fff_method, n_iter, calc_type, trunc_perc, mix, fff_thresh = 'null', "r2", 5, 'lr', 100., 0.0, 1e-4
+        mol_name, fff_method, n_iter, calc_type, trunc_perc, mix, fff_thresh = "null", "r2", 5, "lr", 100.0, 0.0, 1e-4
         if options is not None:
-            if "mol_name" in options: mol_name=options["mol_name"]
-            if "fff_method" in options: fff_method=options["fff_method"]
-            if "n_iter" in options: n_iter=options["n_iter"]
-            if "calc_type" in options: calc_type=options["calc_type"]
-            if "trunc_perc" in options: trunc_perc=options["trunc_perc"]
-            if "mix" in options: mix=options["mix"]
-            if "fff_thresh" in options: fff_thresh=options["fff_thresh"]
+            if "mol_name" in options:
+                mol_name = options["mol_name"]
+            if "fff_method" in options:
+                fff_method = options["fff_method"]
+            if "n_iter" in options:
+                n_iter = options["n_iter"]
+            if "calc_type" in options:
+                calc_type = options["calc_type"]
+            if "trunc_perc" in options:
+                trunc_perc = options["trunc_perc"]
+            if "mix" in options:
+                mix = options["mix"]
+            if "fff_thresh" in options:
+                fff_thresh = options["fff_thresh"]
         return mol_name, fff_method, n_iter, calc_type, trunc_perc, mix, fff_thresh
 
     def check_restart(restart, mol_name):
         if mol_name == "null":
-            logging.warning("Saving restart files to SAVE/null/ specify mol_name=\'desired_path_name\' to save in SAVE/desired_path_name.")
-            logging.warning("To use these restart files, run mv SAVE/null SAVE/desired_path_name, then run with options={mol_name:\'desired_path_name\'}.")
-            print("Warning: Saving restart files to SAVE/null/ specify mol_name=\'desired_path_name\' to save in SAVE/desired_path_name.")
-            print("Warning: To use these restart files, run mv SAVE/null SAVE/desired_path_name, then run with options={mol_name:\'desired_path_name\'}.")
+            logging.warning(
+                "Saving restart files to SAVE/null/ specify mol_name='desired_path_name' to save in SAVE/desired_path_name."
+            )
+            logging.warning(
+                "To use these restart files, run mv SAVE/null SAVE/desired_path_name, then run with options={mol_name:'desired_path_name'}."
+            )
+            print(
+                "Warning: Saving restart files to SAVE/null/ specify mol_name='desired_path_name' to save in SAVE/desired_path_name."
+            )
+            print(
+                "Warning: To use these restart files, run mv SAVE/null SAVE/desired_path_name, then run with options={mol_name:'desired_path_name'}."
+            )
             rmtree("SAVE/" + mol_name.lower() + "/", ignore_errors=True)
         else:
-            if not(restart):
+            if not (restart):
                 try:
                     os.remove("SAVE/" + mol_name.lower() + "/ev_dict.pkl")
                 except OSError:
@@ -103,26 +121,29 @@ def do_fff(h_ferm, n_elec, options=None, restart=False, metric_estim=True):
                 var_new.append(None)
                 new_tbts.append(None)
             else:
-                new_all_ops.append(all_OPS[i+1])
-                var_new.append(vars_appr[i+1])
+                new_all_ops.append(all_OPS[i + 1])
+                var_new.append(vars_appr[i + 1])
                 new_tbts.append(tbts[i])
 
         for idx, i in enumerate(apply_fff_to):
             new_all_ops[i + 1] = new_ops_fff[idx]
             new_tbts[i] = new_tbts_fff[idx]
-            var_new[i + 1] = var_new_fff[idx+1]
+            var_new[i + 1] = var_new_fff[idx + 1]
 
         return new_all_ops, np.array(new_tbts), np.array(var_new)
 
     mol_name, fff_method, n_iter, calc_type, trunc_perc, mix, fff_thresh = process_options(options)
     check_restart(restart, mol_name)
 
-    if trunc_perc < 100.:
+    if trunc_perc < 100.0:
         trunc = True
-        if mix < 1e-8: mix = 1e-3
+        if mix < 1e-8:
+            mix = 1e-3
     else:
         trunc = False
-    h_ferm, obt, tbt, n_qubit, all_OPS, U_OPS, tbts, cartan_tbts = get_init_ops(h_ferm, mol_name, calc_type, spin_orb = False)
+    h_ferm, obt, tbt, n_qubit, all_OPS, U_OPS, tbts, cartan_tbts = get_init_ops(
+        h_ferm, mol_name, calc_type, spin_orb=False
+    )
     if metric_estim:
         psi_fci, psi_appr = get_psi(h_ferm, mol_name, n_elec, n_qubit, trunc, trunc_perc, metric_estim)
         print("===================================================")
@@ -132,12 +153,13 @@ def do_fff(h_ferm, n_elec, options=None, restart=False, metric_estim=True):
     else:
         psi_appr = get_psi(h_ferm, mol_name, n_elec, n_qubit, trunc, trunc_perc, False)
 
-
     print("===================================================")
     print("Getting approximate variances")
     _, vars_appr = compute_ev_var_all_ops(psi_appr, n_qubit, all_OPS, trunc=trunc)
     apply_fff_to = np.where(vars_appr[1:] > fff_thresh)[0]
-    all_ops_fff, uops_fff, cartan_tbts_fff, tbts_fff, vars_appr_fff = get_fff_obj(all_OPS, U_OPS, cartan_tbts, tbts, vars_appr, apply_fff_to)
+    all_ops_fff, uops_fff, cartan_tbts_fff, tbts_fff, vars_appr_fff = get_fff_obj(
+        all_OPS, U_OPS, cartan_tbts, tbts, vars_appr, apply_fff_to
+    )
     print("Applying F3 to {} fragments out of {}".format(len(vars_appr_fff), len(vars_appr)))
     print("===================================================")
 
@@ -157,7 +179,9 @@ def do_fff(h_ferm, n_elec, options=None, restart=False, metric_estim=True):
 
     fff_var = fff_aux(n_iter, n_qubit, O_t.shape[0], O_t.shape[1], O_t, CovOO, Cov0, Covk, uops_fff, mix)
 
-    new_obt, new_tbts_fff, meas_alloc, var_new_fff = fff_multi_iter(obt, tbts_fff, psi_appr, vars_appr_fff, fff_var, fff_method)
+    new_obt, new_tbts_fff, meas_alloc, var_new_fff = fff_multi_iter(
+        obt, tbts_fff, psi_appr, vars_appr_fff, fff_var, fff_method
+    )
     new_all_ops, new_tbts, var_new = reorganize_fff_obj(new_obt, new_tbts_fff, var_new_fff, all_OPS, tbts, vars_appr)
 
     print("Allocating measurements")
@@ -171,26 +195,27 @@ def do_fff(h_ferm, n_elec, options=None, restart=False, metric_estim=True):
     all_uops = [uop_oe] + [U_OPS[i] for i in range(len(U_OPS))]
 
     new_c_obt = np.einsum("pa, qb, pq", uop_oe, uop_oe, new_obt)
-    new_c_tbts = np.einsum('ipa, iqb, irc, isd, ipqrs -> iabcd', U_OPS, U_OPS, U_OPS, U_OPS, new_tbts)
+    new_c_tbts = np.einsum("ipa, iqb, irc, isd, ipqrs -> iabcd", U_OPS, U_OPS, U_OPS, U_OPS, new_tbts)
     return new_all_ops, np.array(all_uops), new_c_obt, new_c_tbts, meas_alloc
 
+
 def do_svd(h_ferm, n_elec):
     spin_orb = False
-    obtb = ferm.get_obt_tbt(h_ferm, spin_orb = spin_orb)
+    obtb = ferm.get_obt_tbt(h_ferm, spin_orb=spin_orb)
     obt = obtb[0]
     tbt_ham_opt = obtb[1]
     n_qubit = ferm.qubit_number(h_ferm)
-    if spin_orb != True:
+    if not spin_orb:
         obt = ferm.obt_orb_to_so(obt)
 
     n = obt.shape[0]
     CARTAN_TBTS_tmp, TBTS_tmp, OPS_tmp, U_OPS_tmp = ferm.lr_decomp(tbt_ham_opt, spin_orb=spin_orb)
-    if spin_orb != True:
+    if not spin_orb:
         U_OPS = ferm.convert_u_to_so(U_OPS_tmp)
         tbts = np.zeros([len(OPS_tmp), n, n, n, n])
         cartan_tbts = np.zeros([len(OPS_tmp), n, n, n, n])
         for i in range(len(OPS_tmp)):
-            tbts[i,:,:,:,:] = ferm.tbt_orb_to_so(TBTS_tmp[i,:,:,:,:])
+            tbts[i, :, :, :, :] = ferm.tbt_orb_to_so(TBTS_tmp[i, :, :, :, :])
             cartan_tbts[i, :, :, :, :] = ferm.tbt_orb_to_so(CARTAN_TBTS_tmp[i, :, :, :, :])
     else:
         U_OPS = U_OPS_tmp
@@ -209,24 +234,25 @@ def do_svd(h_ferm, n_elec):
     print("Allocating measurements")
     _, vars_appr = compute_ev_var_all_ops(psis_appr[0], n_qubit, all_OPS, trunc=False)
 
-    meas_alloc = ferm.compute_meas_alloc(vars_appr, obt, tbts, n_qubit, mix = 0.0)
+    meas_alloc = ferm.compute_meas_alloc(vars_appr, obt, tbts, n_qubit, mix=0.0)
 
     return all_uops, cartan_obt, cartan_tbts, meas_alloc
 
 
-def get_fermion_wise(H, U, qubit_list = []):
-    '''
+def get_fermion_wise(H, U, qubit_list=[]):
+    """
     Return z_form and orbital rotations over qubits at qubit_list
-    '''
+    """
 
-    H = ferm.cartan_tbt_to_ferm(H, spin_orb = True)
+    H = ferm.cartan_tbt_to_ferm(H, spin_orb=True)
     z_form = QubitHamiltonian(jordan_wigner(H))
 
-    circuit = ferm.get_orb_rot(U, qubit_list=qubit_list, tol = 1e-12)
+    circuit = ferm.get_orb_rot(U, qubit_list=qubit_list, tol=1e-12)
     return [z_form, circuit]
 
+
 def get_init_ops(h_ferm, mol_name, calc_type, spin_orb, save=True):
-    '''
+    """
     Parameters
      ----------
     mol_name -
@@ -246,30 +272,30 @@ def get_init_ops(h_ferm, mol_name, calc_type, spin_orb, save=True):
     U_OPS Orbital rotations
     tbts LR Decomposition of two body integrals
     cartan_tbts Polynomial functions of Pauli Z under qubit fermion mappings
-    '''
+    """
     path = "SAVE/" + mol_name + "/"
     if os.path.isfile(path + "tensor_terms.pkl"):
         print("Using saved Hamiltonian from {}. Run with a different mol_name if this is not desired.".format(path))
-        with open(path + "tensor_terms.pkl", 'rb') as file:
+        with open(path + "tensor_terms.pkl", "rb") as file:
             INIT = pickle.load(file)
         obt = INIT[0]
         tbt_ham_opt = INIT[1]
-        with open(path + "ham.pkl", 'rb') as file:
+        with open(path + "ham.pkl", "rb") as file:
             h_ferm = pickle.load(file)
     else:
-        obtb = ferm.get_obt_tbt(h_ferm, spin_orb = spin_orb)
+        obtb = ferm.get_obt_tbt(h_ferm, spin_orb=spin_orb)
         obt = obtb[0]
         tbt_ham_opt = obtb[1]
         if save:
             Path(path).mkdir(exist_ok=True, parents=True)
-            with open(path + "tensor_terms.pkl", 'wb') as file:
+            with open(path + "tensor_terms.pkl", "wb") as file:
                 pickle.dump([obt, tbt_ham_opt], file)
-            with open(path + "ham.pkl", 'wb') as file:
+            with open(path + "ham.pkl", "wb") as file:
                 pickle.dump(h_ferm, file)
 
     n_qubit = ferm.qubit_number(h_ferm)
 
-    if spin_orb != True:
+    if not spin_orb:
         obt = ferm.obt_orb_to_so(obt)
         tbt = ferm.tbt_orb_to_so(tbt_ham_opt)
     else:
@@ -278,8 +304,12 @@ def get_init_ops(h_ferm, mol_name, calc_type, spin_orb, save=True):
 
     if calc_type.lower() == "lr":
         if os.path.isfile(path + "lr.pkl"):
-            print("Using saved LR decomposition saved in {}. Run with a different mol_name if this is not desired.".format(path))
-            with open(path + "lr.pkl", 'rb') as file:
+            print(
+                "Using saved LR decomposition saved in {}. Run with a different mol_name if this is not desired.".format(
+                    path
+                )
+            )
+            with open(path + "lr.pkl", "rb") as file:
                 INIT = pickle.load(file)
             CARTAN_TBTS_tmp = INIT[0]
             TBTS_tmp = INIT[1]
@@ -287,16 +317,16 @@ def get_init_ops(h_ferm, mol_name, calc_type, spin_orb, save=True):
         else:
             CARTAN_TBTS_tmp, TBTS_tmp, OPS_tmp, U_OPS_tmp = ferm.lr_decomp(tbt_ham_opt, spin_orb=spin_orb)
             if save:
-                with open(path + "lr.pkl", 'wb') as file:
+                with open(path + "lr.pkl", "wb") as file:
                     pickle.dump([CARTAN_TBTS_tmp, TBTS_tmp, U_OPS_tmp], file)
         OPS = ferm.convert_tbts_to_frags(TBTS_tmp, spin_orb)
 
-    if spin_orb != True:
+    if not spin_orb:
         U_OPS = ferm.convert_u_to_so(U_OPS_tmp)
         tbts = np.zeros([len(OPS), n, n, n, n])
         cartan_tbts = np.zeros([len(OPS), n, n, n, n])
         for i in range(len(OPS)):
-            tbts[i,:,:,:,:] = ferm.tbt_orb_to_so(TBTS_tmp[i,:,:,:,:])
+            tbts[i, :, :, :, :] = ferm.tbt_orb_to_so(TBTS_tmp[i, :, :, :, :])
             cartan_tbts[i, :, :, :, :] = ferm.tbt_orb_to_so(CARTAN_TBTS_tmp[i, :, :, :, :])
 
     else:
@@ -309,8 +339,9 @@ def get_init_ops(h_ferm, mol_name, calc_type, spin_orb, save=True):
         all_OPS.append(OPS[i])
     return h_ferm, obt, tbt, n_qubit, all_OPS, U_OPS, tbts, cartan_tbts
 
+
 def get_wavefunction(Hq, wf_type, mol_name, n_elec, N=1, save=True):
-    '''
+    """
     Parameters
      ----------
     h_ferm -
@@ -326,15 +357,16 @@ def get_wavefunction(Hq, wf_type, mol_name, n_elec, N=1, save=True):
     -------
     energies Eigenenergies
     psi Eigenstates
-    '''
+    """
     n_qubits = count_qubits(Hq)
     if wf_type.lower() == "fci":
-       return get_fci_states(Hq, mol_name, n_elec, n_qubits, N=N, save=save)
+        return get_fci_states(Hq, mol_name, n_elec, n_qubits, N=N, save=save)
     elif wf_type.lower() == "cisd":
-       return get_cisd_states(Hq, mol_name, n_elec, n_qubits, N=N, save=save)
+        return get_cisd_states(Hq, mol_name, n_elec, n_qubits, N=N, save=save)
+
 
 def get_fci_states(Hq, mol_name, n_elec, n_qubits, N=1, save=True):
-    '''
+    """
     Parameters
      ----------
     Hq -
@@ -350,12 +382,12 @@ def get_fci_states(Hq, mol_name, n_elec, n_qubits, N=1, save=True):
     -------
     e_fci Energy of FCI ground state
     psi_fci Wavefunction of FCI ground state
-    '''
+    """
 
     path = "SAVE/" + mol_name.lower() + "/"
     if os.path.isfile(path + "psi_fci.pkl"):
         print("Using saved psi_fci in {}. Run with a different mol_name if this is not desired.".format(path))
-        with open(path + "psi_fci.pkl", 'rb') as file:
+        with open(path + "psi_fci.pkl", "rb") as file:
             INIT = pickle.load(file)
         e_fci = INIT[0]
         psi_fci = INIT[1]
@@ -371,26 +403,27 @@ def get_fci_states(Hq, mol_name, n_elec, n_qubits, N=1, save=True):
     size_H = sparse_H.get_shape()[0]
     if M >= size_H - 1:
         M = size_H - 2
-    w,v = sp.sparse.linalg.eigsh(sparse_H, k = max(10, M), which = "SA")
+    w, v = sp.sparse.linalg.eigsh(sparse_H, k=max(10, M), which="SA")
     srt_arg = np.argsort(w)
     w = w[srt_arg]
     v = v[:, srt_arg]
     values = []
     vectors = []
     for i in range(len(w)):
-        Nel = expectation(get_sparse_operator(Nop, n_qubits), v[:,i])
+        Nel = expectation(get_sparse_operator(Nop, n_qubits), v[:, i])
         if np.abs(Nel - n_elec) < 1e-6:
             values.append(w[i])
-            vectors.append(v[:,i])
-        if len(values) == N or i == len(w)-1:
+            vectors.append(v[:, i])
+        if len(values) == N or i == len(w) - 1:
             if save:
                 Path(path).mkdir(exist_ok=True)
-                with open(path + "psi_fci.pkl", 'wb') as file:
+                with open(path + "psi_fci.pkl", "wb") as file:
                     pickle.dump([values, vectors], file)
             return values, vectors
 
+
 def get_cisd_states(Hq, mol_name, n_elec, n_qubits, N=1, save=True):
-    '''
+    """
     Parameters
      ----------
     Hq -
@@ -406,11 +439,11 @@ def get_cisd_states(Hq, mol_name, n_elec, n_qubits, N=1, save=True):
     -------
     e_cisd Energy of CISD ground state
     psi_cisd Wavefunction of CISD ground state
-    '''
+    """
     path = "SAVE/" + mol_name.lower() + "/"
     if os.path.isfile(path + "psi_cisd.pkl"):
         print("Using saved psi_cisd in {}. Run with a different mol_name if this is not desired.".format(path))
-        with open(path + "psi_cisd.pkl", 'rb') as file:
+        with open(path + "psi_cisd.pkl", "rb") as file:
             INIT = pickle.load(file)
         e_cisd = INIT[0]
         psi_cisd = INIT[1]
@@ -423,7 +456,7 @@ def get_cisd_states(Hq, mol_name, n_elec, n_qubits, N=1, save=True):
         M = size_H - 1
     else:
         M = N
-    w,v = sp.sparse.linalg.eigsh(H_mat_cisd, k = M, which = "SA")
+    w, v = sp.sparse.linalg.eigsh(H_mat_cisd, k=M, which="SA")
     order = np.argsort(w)
     values = w[order].tolist()
     vectors = []
@@ -431,17 +464,18 @@ def get_cisd_states(Hq, mol_name, n_elec, n_qubits, N=1, save=True):
         wfs = np.zeros(2**n_qubits)
         for iidx, iindx in enumerate(indices):
             wfs[iindx] = v[iidx, i]
-        wfs = wfs/np.linalg.norm(wfs)
+        wfs = wfs / np.linalg.norm(wfs)
         vectors.append(wfs)
     if save:
         Path(path).mkdir(exist_ok=True)
-        with open(path + "psi_cisd.pkl", 'wb') as file:
+        with open(path + "psi_cisd.pkl", "wb") as file:
             pickle.dump([values, vectors], file)
 
     return values, vectors
 
+
 def compute_and_print_ev_var(psi, h_ferm, all_OPS, meas_alloc=None):
-    '''
+    """
     Parameters
      ----------
     psi Wavefunction
@@ -451,12 +485,13 @@ def compute_and_print_ev_var(psi, h_ferm, all_OPS, meas_alloc=None):
     it is computed according to the variances [see Quantum 5, 385 (2021)].
 
     Prints out the variances and expectation of each fragment over psi.
-    '''
+    """
     n_qubit = ferm.qubit_number(h_ferm)
     h_const = h_ferm.constant
     exps, variances = compute_ev_var_all_ops(psi, n_qubit, all_OPS)
-    if meas_alloc is None: meas_alloc = ferm.compute_meas_alloc(variances)
-    scaled_variances = np.divide(variances,meas_alloc)
+    if meas_alloc is None:
+        meas_alloc = ferm.compute_meas_alloc(variances)
+    scaled_variances = np.divide(variances, meas_alloc)
 
     scaled_variances_sum = np.sum(scaled_variances)
     print("Full variances:")
@@ -466,8 +501,9 @@ def compute_and_print_ev_var(psi, h_ferm, all_OPS, meas_alloc=None):
     print("Variance metric value is {}".format(scaled_variances_sum))
     print("Exp value is {}".format(np.sum(exps) + h_const))
 
+
 def compute_ev_var_all_ops(psi, n_qubit, all_OPS, trunc=False):
-    '''
+    """
     Parameters
      ----------
     psi Wavefunction
@@ -478,7 +514,7 @@ def compute_ev_var_all_ops(psi, n_qubit, all_OPS, trunc=False):
     -------
     exps Expectations of the fragments over psi
     variances Variances of the fragments over psi
-    '''
+    """
     num_frags = len(all_OPS)
     exps = np.zeros(num_frags)
     variances = np.zeros(num_frags)
@@ -487,8 +523,9 @@ def compute_ev_var_all_ops(psi, n_qubit, all_OPS, trunc=False):
         variances[i] = ferm.variance_value(all_OPS[i], psi, n_qubit, trunc=trunc)
     return exps, variances
 
+
 def init_ev_dict(mol_name, psi, n_qubit, trunc=False, spin_orb=True, save=True):
-    '''
+    """
     Parameters
      ----------
     psi Wavefunction
@@ -498,14 +535,14 @@ def init_ev_dict(mol_name, psi, n_qubit, trunc=False, spin_orb=True, save=True):
     Returns
     -------
     ev_dict_all Returns a dictionary of expectaion and variances of fermionic operators over psi
-    '''
+    """
     if spin_orb:
         n = n_qubit
     else:
         n = n_qubit // 2
     path = "SAVE/" + mol_name.lower() + "/"
     if os.path.isfile(path + "ev_dict.pkl"):
-        with open(path + "ev_dict.pkl", 'rb') as file:
+        with open(path + "ev_dict.pkl", "rb") as file:
             ev_dict_all = pickle.load(file)
     else:
         ev_dict_E = ferm.get_E(psi, n, n_qubit, trunc=trunc)
@@ -513,16 +550,18 @@ def init_ev_dict(mol_name, psi, n_qubit, trunc=False, spin_orb=True, save=True):
         ev_dict_all = ferm.reorganize(n, ev_dict_E, ev_dict_EE)
         if save:
             Path(path).mkdir(exist_ok=True)
-            with open(path + "ev_dict.pkl", 'wb') as file:
+            with open(path + "ev_dict.pkl", "wb") as file:
                 pickle.dump(ev_dict_all, file)
     return ev_dict_all
 
+
 def check_method(method):
     if method.lower() not in ["full", "r1", "r2"]:
-       raise TequilaException("method has to be specified as one from Full, R1 or R2")
+        raise TequilaException("method has to be specified as one from Full, R1 or R2")
+
 
 def compute_O_t(U_OPS, method, tbts, mol_name, save=True):
-    '''
+    """
     Parameters
      ----------
     U_OPS Orbital rotations
@@ -532,56 +571,57 @@ def compute_O_t(U_OPS, method, tbts, mol_name, save=True):
     Returns
     -------
     O_t O_alpha (arXiv:2208.14490v3 - Section 2.3)
-    '''
+    """
     check_method(method)
 
     path = "SAVE/" + mol_name.lower() + "/" + method.lower() + "/"
     if os.path.isfile(path + "O_t.pkl"):
-        with open(path + "O_t.pkl", 'rb') as file:
+        with open(path + "O_t.pkl", "rb") as file:
             O_t = pickle.load(file)
         return O_t
 
     num_frags = U_OPS.shape[0]
     n = U_OPS.shape[1]
     if method.lower() == "full":
-        n_param = n//2
+        n_param = n // 2
     else:
         n_param = 1
-        ratios = np.zeros([num_frags,n])
+        ratios = np.zeros([num_frags, n])
         for frag1 in range(num_frags):
             for p1 in range(n):
                 if method.lower() == "r1":
-                    ratios[frag1, p1] = tbts[frag1, p1,p1,p1,p1]
+                    ratios[frag1, p1] = tbts[frag1, p1, p1, p1, p1]
                 else:
-                    ratios[frag1, p1] = np.sum([tbts[frag1,p1,p1,r1,r1] for r1 in range(n)])
+                    ratios[frag1, p1] = np.sum([tbts[frag1, p1, p1, r1, r1] for r1 in range(n)])
     O_t_tmp = np.zeros([num_frags, n, n, n])
     O_t = np.zeros([num_frags, n_param, n, n])
-    Otmp = np.zeros([n,n])
+    Otmp = np.zeros([n, n])
     for k in range(num_frags):
         Umat = U_OPS[k, :, :]
         for p in range(n):
             for r, s in product(range(n), repeat=2):
-                Otmp[r,s] = Umat[r,p] * Umat[s,p].conjugate()
-            O_t_tmp[k,p,:,:] = Otmp
+                Otmp[r, s] = Umat[r, p] * Umat[s, p].conjugate()
+            O_t_tmp[k, p, :, :] = Otmp
     if method.lower() == "full":
         for k in range(num_frags):
             for p in range(n_param):
                 for alpha in range(2):
-                    O_t[k,p,:,:] += O_t_tmp[k,2*(p)+alpha,:,:]
+                    O_t[k, p, :, :] += O_t_tmp[k, 2 * (p) + alpha, :, :]
     else:
         for k in range(num_frags):
             for p1 in range(n):
-                O_t[k,0,:,:] += ratios[k,p1] * O_t_tmp[k,p1,:,:]
+                O_t[k, 0, :, :] += ratios[k, p1] * O_t_tmp[k, p1, :, :]
 
     if save:
         Path(path).mkdir(exist_ok=True)
-        with open(path + "O_t.pkl", 'wb') as file:
-                pickle.dump(O_t, file)
+        with open(path + "O_t.pkl", "wb") as file:
+            pickle.dump(O_t, file)
 
     return O_t
 
+
 def compute_all_covs(all_OPS, O_t, psi, n_qubit, mol_name, method, trunc=False, save=True):
-    '''
+    """
     Parameters
      ----------
     psi Wavefunction
@@ -594,11 +634,11 @@ def compute_all_covs(all_OPS, O_t, psi, n_qubit, mol_name, method, trunc=False,
     Returns
     -------
     all_covs --------------------
-    '''
+    """
     path = "SAVE/" + mol_name.lower() + "/" + method.lower() + "/"
     check_method(method)
     if os.path.isfile(path + "all_covs.pkl"):
-        with open(path + "all_covs.pkl", 'rb') as file:
+        with open(path + "all_covs.pkl", "rb") as file:
             all_covs = pickle.load(file)
         return all_covs
 
@@ -607,17 +647,18 @@ def compute_all_covs(all_OPS, O_t, psi, n_qubit, mol_name, method, trunc=False,
         ops1 = []
         for p in range(O_t.shape[1]):
             ops1.append(ferm.obt_to_ferm(O_t[frag_idx, p, :, :], True))
-        all_covs.append(ferm.compute_covk(ops1, all_OPS[frag_idx+1], psi, n_qubit, trunc=trunc))
+        all_covs.append(ferm.compute_covk(ops1, all_OPS[frag_idx + 1], psi, n_qubit, trunc=trunc))
 
     if save:
         Path(path).mkdir(exist_ok=True)
-        with open(path + "all_covs.pkl", 'wb') as file:
+        with open(path + "all_covs.pkl", "wb") as file:
             pickle.dump(all_covs, file)
 
     return all_covs
 
+
 def compute_cov_OO(O_t, ev_dict_all, mol_name, method, save=True):
-    '''
+    """
     Parameters
      ----------
     O_t O_alpha (arXiv:2208.14490v3 - Section 2.3)
@@ -628,11 +669,11 @@ def compute_cov_OO(O_t, ev_dict_all, mol_name, method, save=True):
     Returns
     -------
     covmat Dictionary of covariances between O_alpha's
-    '''
+    """
     path = "SAVE/" + mol_name.lower() + "/" + method.lower() + "/"
     check_method(method)
     if os.path.isfile(path + "cov_OO.pkl"):
-        with open(path + "cov_OO.pkl", 'rb') as file:
+        with open(path + "cov_OO.pkl", "rb") as file:
             covmat = pickle.load(file)
         return covmat
 
@@ -646,17 +687,18 @@ def compute_cov_OO(O_t, ev_dict_all, mol_name, method, save=True):
                 for q in range(n):
                     ind1 = n * (k) + p
                     ind2 = n * (l) + q
-                    covmat[ind1, ind2] = ferm.covariance_ob_ob(O_t[k,p,:,:], O_t[l,q,:,:], ev_dict_all[0])
+                    covmat[ind1, ind2] = ferm.covariance_ob_ob(O_t[k, p, :, :], O_t[l, q, :, :], ev_dict_all[0])
 
     if save:
         Path(path).mkdir(exist_ok=True)
-        with open(path + "cov_OO.pkl", 'wb') as file:
+        with open(path + "cov_OO.pkl", "wb") as file:
             pickle.dump(covmat, file)
 
     return covmat
 
+
 def compute_cov_O(O_t, H0, ev_dict_all, mol_name, method, save=True):
-    '''
+    """
     Parameters
      ----------
     O_t O_alpha (arXiv:2208.14490v3 - Section 2.3)
@@ -668,11 +710,11 @@ def compute_cov_O(O_t, H0, ev_dict_all, mol_name, method, save=True):
     Returns
     -------
     covvec Covariance of O_alpha's and original fragments
-    '''
+    """
     path = "SAVE/" + mol_name.lower() + "/" + method.lower() + "/"
     check_method(method)
     if os.path.isfile(path + "cov_O.pkl"):
-        with open(path + "cov_O.pkl", 'rb') as file:
+        with open(path + "cov_O.pkl", "rb") as file:
             covvec = pickle.load(file)
         return covvec
 
@@ -683,21 +725,23 @@ def compute_cov_O(O_t, H0, ev_dict_all, mol_name, method, save=True):
     for k in range(nf):
         for p in range(n):
             ind = n * (k) + p
-            covl = ferm.covariance_ob_ob(O_t[k,p,:,:], H0, ev_dict_all[0])
-            covr = ferm.covariance_ob_ob(H0, O_t[k,p,:,:], ev_dict_all[0])
+            covl = ferm.covariance_ob_ob(O_t[k, p, :, :], H0, ev_dict_all[0])
+            covr = ferm.covariance_ob_ob(H0, O_t[k, p, :, :], ev_dict_all[0])
             covvec[ind] = covl + covr
 
     if save:
         Path(path).mkdir(exist_ok=True)
-        with open(path + "cov_O.pkl", 'wb') as file:
+        with open(path + "cov_O.pkl", "wb") as file:
             pickle.dump(covvec, file)
 
     return covvec
 
+
 class fff_aux:
-    '''
+    """
     Class containing all variables needed for FFF.
-    '''
+    """
+
     def __init__(self, n_iter, nq, nf, n, o_t, coo, c0, ck, uops, mix):
         self.n_iter = n_iter
         self.nq = nq
@@ -710,8 +754,9 @@ class fff_aux:
         self.uops = uops
         self.mix = mix
 
+
 def fff_multi_iter(obt, tbts, psi, varbs, fff_var, method):
-    '''
+    """
     Parameters
      ----------
     psi Wavefunction
@@ -726,24 +771,26 @@ def fff_multi_iter(obt, tbts, psi, varbs, fff_var, method):
     new_obt Repartitioned one body term
     new_tbts Repartitioned two body terms
     m0 Optimal Measurement Allocation
-    '''
+    """
     check_method(method)
 
     ntmp = obt.shape[0]
-    obt_list = np.zeros([ fff_var.n_iter+1, ntmp, ntmp])
-    tbts_list = np.zeros([fff_var.n_iter+1, fff_var.nf, ntmp, ntmp, ntmp, ntmp])
-    var_list = np.zeros([fff_var.n_iter+1, fff_var.nf+1])
-    obt_list[0,:,:] = obt
-    tbts_list[0,:,:,:,:,:] = tbts
-    var_list[0,:] = varbs
+    obt_list = np.zeros([fff_var.n_iter + 1, ntmp, ntmp])
+    tbts_list = np.zeros([fff_var.n_iter + 1, fff_var.nf, ntmp, ntmp, ntmp, ntmp])
+    var_list = np.zeros([fff_var.n_iter + 1, fff_var.nf + 1])
+    obt_list[0, :, :] = obt
+    tbts_list[0, :, :, :, :, :] = tbts
+    var_list[0, :] = varbs
     new_c0 = fff_var.c0
     new_ck = fff_var.ck
     for i in range(fff_var.n_iter):
         print("Progress: iteration #{} out of {}".format(i + 1, fff_var.n_iter))
-        new_obt, new_tbts, new_vars, new_c0, new_ck = ferm.fff_1_iter(obt_list[i,:,:], tbts_list[i,:,:,:,:,:], var_list[i,:], new_c0, new_ck, fff_var)
-        obt_list[i+1,:,:] = new_obt
-        tbts_list[i+1,:,:,:,:,:] = new_tbts
-        var_list[i+1,:] = new_vars
-    var0 = var_list[len(var_list)-1,:]
+        new_obt, new_tbts, new_vars, new_c0, new_ck = ferm.fff_1_iter(
+            obt_list[i, :, :], tbts_list[i, :, :, :, :, :], var_list[i, :], new_c0, new_ck, fff_var
+        )
+        obt_list[i + 1, :, :] = new_obt
+        tbts_list[i + 1, :, :, :, :, :] = new_tbts
+        var_list[i + 1, :] = new_vars
+    var0 = var_list[len(var_list) - 1, :]
     m0 = ferm.compute_meas_alloc(var0, obt, tbts, fff_var.nq, fff_var.mix)
-    return obt_list[len(obt_list)-1,:,:], tbts_list[len(tbts_list)-1,:,:,:,:,:], m0, var0
+    return obt_list[len(obt_list) - 1, :, :], tbts_list[len(tbts_list) - 1, :, :, :, :, :], m0, var0