vqite_quimb.py

"""
This module performs VQITE simulations using tensor-network (TN) methods.
Quimb is the library for performing TN contractions.
Cotengra library is used to find optimal contraction paths.
mpi4py is used for parallelization.
Examples of calculations are given in the accompanying notebooks
vqite_timing_test.ipynb and one_step_timing_test_and_MV_avqite_comparison.ipynb.

Packages information:
---------------------
NumPy version = 1.24.4
SciPy version = 1.12.0
mpi4py version = 3.1.4
Quimb version = 1.8.4 (errors can occur for the version 1.9.0)
Cotengra version = 0.6.2
Autoray version = 0.7.0
"""

import numpy as np
import scipy
import pickle
import time
from tqdm import tqdm
from typing import (
    List,
    Optional,
    Tuple,
    Union
)
import random
from mpi4py import MPI

import quimb as qu
import quimb.tensor as qtn

import cotengra as ctg


class model_H:
    """
    Class for Hamiltonians constructed using incar_file.

    Attributes:
    -----------
    incar_file : str
        Incar file (of the AVQITE format).
    paulis : List[str]
        List of the Pauli strings comprising the Hamiltonian.
    coefs : List[float]
        List of the coefficients in the Hamiltonian corresponding to the Pauli
        strings in paulis.
    """
    def __init__(
        self,
        incar_file: str
    ):
        self.incar_file = incar_file
        with open(self.incar_file) as fp:
            incar_content = fp.read()
        h_pos = incar_content.find("h")
        pool_pos = incar_content.find("pool")
        h_string = incar_content[h_pos+14:pool_pos-14]

        self.paulis = "".join([el for el in h_string if el=='I' or el=='X'
                            or el=='Y' or el=='Z' or el == '\n']).split('\n')
        coefs_str = "".join([el for el in h_string if el.isdigit() or el=="-"
                            or el=="." or el == "*"]).split('*')
        self.coefs = [float(el) for el in coefs_str[0:-1]]


class Quimb_vqite:
    """
    Class for performing VQITE simualtions using Quimb.
    Form of the ansatz is read out from a file generated by AVQITE.

    Attributes:
    -----------
    _incar_file : str
        Path to the incar file.
        Incar file is used to read out the Hamiltonian and the reference state.
    _ansatz_file : str
        Path to the ansatz file.
    _output_file : str
        Path to the output file.
    _init_params : str
        What initial parameters to use for the ansatz.
        Possible options: 'random', 'zeros', 'avqite' or a list of parameters.
    _comm : MPI.COMM_WORLD
    _size : int
        Total number of MPI processes.
    _rank : int
        Rank on an MPI process.
    _num_qubits : int
        Number of qubits in the system. Determined from the incar file.
    _ansatz : List[str]
        Ansatz form. Obtained from the file generated by AVQITE.
    _params_solution : List[float]
        Parameters of the ansatz calculated by AVQITE. These are not updated
        throughout this calculation.
    params : List[float]
        Parameters of the ansatz. These are updated throughout this calculation.
    _m : numpy.ndarray
        M matrix used in VQITE.
    _m_width : numpy.ndarray
        Contraction width matrix for the M matrix.
    _m_cost : numpy.ndarray
        Contraction cost matrix for the M matrix.
    _v : numpy.ndarray
        V vector used in VQITE.
    _ref_state : str
        Reference state. Read out from the incar file.
    _init_qc : quimb.tensor.circuit.Circuit
        Initial quantum circuit that incorporates the possible reference state
        gates.
    _pauli_rot_gates_list : List[quimb.tensor.circuit.Gate]
        List of Quimb gates representing Pauli rotations in the ansatz.
    _pauli_rot_dag_gates_list : List[quimb.tensor.circuit.Gate]
        List of Quimb gates representing inverse Pauli rotations in the ansatz.
    _base_circuits : List[quimb.tensor.circuit.Circuit]
        List of quantum circuits representing the ansatz up to ith rotation,
        where i is the index in the list.
    h_terms_reh_dict : dict
        Dictionary of TN "reh" dictionaries for calculating expectation value
        of each pauli for the ansatz state.
    optimize_dict : dict
        Dictionary of TN contraction paths for calculating expectation value
        of each pauli for the ansatz state.
    """
    def __init__(
        self,
        incar_file: str,
        ansatz_file: str,
        output_file: str,
        init_params = "random"
    ):
        self._incar_file = incar_file
        self._ansatz_file = ansatz_file
        self._output_file = output_file
        self._init_params = init_params

        self._comm = MPI.COMM_WORLD
        self._size = self._comm.Get_size()
        self._rank = self._comm.Get_rank()

        #Reads out the Hamiltonian from the incar file.
        #The number of qubits is determined from there.
        self._H = model_H(self._incar_file)
        self._num_qubits = len(self._H.paulis[0])

        #Reads out the form of the ansatz and the parameters of the ansatz from
        #the ansatz file.
        #The ansatz file should be in the AVQITE format.
        (self._ansatz,
         self._params_solution) = read_adaptvqite_ansatz(self._ansatz_file)
        #For the purposes of VQITE, we might want to set the initial parameters
        #to be random.
        if self._init_params == "random":
            self.params = [self._params_solution[i]+random.uniform(-0.05, 0.05)
                                              for i in range(len(self._ansatz))]
        elif self._init_params == "zeros":
            self.params = [0.0 for i in range(len(self._ansatz))]
        elif self._init_params == "avqite":
            self.params = self._params_solution.copy()
        elif (type(self._init_params)==list and
                                    len(self._init_params)==len(self._ansatz)):
            self.params = self._init_params.copy()
        else:
            raise NotImplementedError(
                "self._init_params has to be either random, avqite, or a list"
            )

        #Matrix M and cvctor V used in VQITE.
        self._m = np.zeros((len(self._ansatz),len(self._ansatz)))
        self._m_width = np.zeros((len(self._ansatz),len(self._ansatz)))
        self._m_cost = np.zeros((len(self._ansatz),len(self._ansatz)))

        self._v = np.zeros(len(self._ansatz))

        #Reads out the incar file.
        with open(self._incar_file) as fp:
            incar_content = fp.read()
        ref_st_r_pos = incar_content.find("ref_state")
        #Reads out the reference state from the incar file.
        self._ref_state = incar_content[
                            ref_st_r_pos+13:ref_st_r_pos+13+self._num_qubits
                            ]

        #Initializes a quantum circuit.
        self._init_qc = qtn.Circuit(N=self._num_qubits)
        #If the reference state contains "1"s, adds corresponding bit-flips.
        if all([(el=='0') or (el=='1') for el in self._ref_state]):
            [self._init_qc.apply_gate('X',i)
                            for i,el in enumerate(self._ref_state) if el=='1']
        else:
            raise ValueError(
                "Reference state is supposed to be a string of 0s and 1s"
            )
        #Creates and saves a set of gates in Quimb corresponding to Pauli
        #rotations (and their inverse) from the ansatz.
        #This will be used throghout the calculations.
        #Here we create separate lists for circuits and for gates because Quimb
        #does not have functionality to reparameterize gates, only circuits.
        self._pauli_rot_circuits_list = [
            add_pauli_rotation_gate(
                qc=qtn.Circuit(N=self._num_qubits),
                pauli_string=self._ansatz[i],
                theta=self.params[i],
                decompose_rzz=False
            ) for i in range(len(self._ansatz))
        ]
        self._pauli_rot_gates_list = [
            self._pauli_rot_circuits_list[i].gates
                                    for i in range(len(self._ansatz))
        ]
        self._pauli_rot_dag_circuits_list = [
            add_pauli_rotation_gate(
                qc=qtn.Circuit(N=self._num_qubits),
                pauli_string=self._ansatz[i],
                theta=-self.params[i],
                decompose_rzz=False
            ) for i in range(len(self._ansatz))
        ]
        self._pauli_rot_dag_gates_list = [
           self._pauli_rot_dag_circuits_list[i].gates
                                    for i in range(len(self._ansatz))
        ]
        #Creates and saves circuits in Quimb correspondning to the product of
        #Pauli rotations up to mu'th rotation in the ansatz list, where
        #mu is the index.
        #This will be used throghout the calculations.
        self._base_circuits = [self.circuit_2(mu)
                                    for mu in range(len(self._ansatz)+1)]


    def update_params(self):
        """
        Update self._pauli_rot_circuits_list, self._pauli_rot_gates_list,
        self._pauli_rot_dag_circuits_list, self._pauli_rot_dag_gates_list,
        and self._base_circuits for the current values of self.params.
        """
        for k in range(len(self._ansatz)):
            old_params_dict = self._pauli_rot_circuits_list[k].get_params()
            new_params_dict = dict()
            for i,key in enumerate(old_params_dict.keys()):
                new_params_dict[key]= np.array([self.params[k]])
            self._pauli_rot_circuits_list[k].set_params(new_params_dict)
        self._pauli_rot_gates_list = [
            self._pauli_rot_circuits_list[k].gates
                                            for k in range(len(self._ansatz))
        ]
        for k in range(len(self._ansatz)):
            old_params_dict = self._pauli_rot_dag_circuits_list[k].get_params()
            new_params_dict = dict()
            for i,key in enumerate(old_params_dict.keys()):
                new_params_dict[key]= np.array([-self.params[k]])
            self._pauli_rot_dag_circuits_list[k].set_params(new_params_dict)
        self._pauli_rot_dag_gates_list = [
           self._pauli_rot_dag_circuits_list[k].gates
                                            for k in range(len(self._ansatz))
        ]
        for mu in range(len(self._ansatz)+1):
            old_params_dict = self._base_circuits[mu].get_params()
            new_params_dict = dict()
            for i,key in enumerate(old_params_dict.keys()):
                new_params_dict[key]= np.array([self.params[i]])
            self._base_circuits[mu].set_params(new_params_dict)


    def compute_m(self, which_nonzero=None, **kwargs):
        """
        Computes matrix M in VQITE in parallel.
        Each parallel process calculates a different part of the matrix.

        Parameters:
        -----------
        which_nonzero : List[int]
            Indices of matrix M that are known to be nonzero.
            If None, the entire matrix is calculated.
        **kwargs
            Arguments used in Quimb methods for tensor contraction
            evaluations, such as:
                optimize : str
                    Optimizer to use when looking for contraction paths.
                simplify_sequence : str
                    TN simplifications to use when looking for contraction paths.
                backend : str
                    Backend to use when performing the contractions.
                    Usually specified if GPU acceleration is needed.
                ...
        """
        if which_nonzero==None:
            ind_list = [(mu,nu)
                        for nu in range(len(self._ansatz))
                        for mu in range(nu+1)]
        else:
            ind_list = which_nonzero

        bins_sizes = [int(len(ind_list)/self._size) for i in range(self._size)]
        for i in range(len(ind_list) - int(len(ind_list)/self._size)*self._size):
            bins_sizes[i] = bins_sizes[i]+1
        start = sum(bins_sizes[:self._rank])
        end = start + bins_sizes[self._rank]

        m_interm = np.zeros(end-start)
        m_interm_cost = np.zeros(end-start)
        m_interm_width = np.zeros(end-start)

        m_nonzero = np.zeros(len(ind_list))
        m_nonzero_cost = np.zeros(len(ind_list))
        m_nonzero_width = np.zeros(len(ind_list))

        for i,(mu,nu) in enumerate(ind_list[start:end]):
            contr_mu_nu=self.contr1_est(mu=mu, nu=nu, **kwargs)
            m_interm[i] = (
                contr_mu_nu[-1] +
                self.contr2_est(mu = mu, **kwargs)[-1]*
                self.contr2_est(mu = nu, **kwargs)[-1]
            )
            (m_interm_width[i],
            m_interm_cost[i]) = (contr_mu_nu[0],contr_mu_nu[1])

        sendcountes=tuple(bins_sizes)
        displacements=tuple([sum(bins_sizes[:i]) for i in range(self._size)])

        self._comm.Allgatherv(
            [m_interm,  MPI.DOUBLE],
            [m_nonzero, sendcountes, displacements, MPI.DOUBLE]
        )
        self._comm.Allgatherv(
            [m_interm_cost,  MPI.DOUBLE],
            [m_nonzero_cost, sendcountes, displacements, MPI.DOUBLE]
        )
        self._comm.Allgatherv(
            [m_interm_width,  MPI.DOUBLE],
            [m_nonzero_width, sendcountes, displacements, MPI.DOUBLE]
        )
        self._m = np.zeros((len(self._ansatz),len(self._ansatz)))
        for i in range(len(ind_list)):
            self._m[ind_list[i]] = m_nonzero[i]
            self._m[ind_list[i][::-1]] = m_nonzero[i]
        self._m_width = np.zeros((len(self._ansatz),len(self._ansatz)))
        for i in range(len(ind_list)):
            self._m_width[ind_list[i]] = m_nonzero_width[i]
            self._m_width[ind_list[i][::-1]] = m_nonzero_width[i]
        self._m_cost = np.zeros((len(self._ansatz),len(self._ansatz)))
        for i in range(len(ind_list)):
            self._m_cost[ind_list[i]] = m_nonzero_cost[i]
            self._m_cost[ind_list[i][::-1]] = m_nonzero_cost[i]


    def compute_v(self, optimize = 'greedy', **kwargs):
        """
        Computes vector V in VQITE in parallel using parameter shift rule.
        Each parallel process calculates the expectationvalue of a particular
        Pauli string in the Hamiltonian.

        Parameters:
        -----------
        optimize : str or dict
            Optimizer to use when looking for contraction paths.
        **kwargs
            Arguments used in Quimb methods for tensor contraction
            evaluations, such as:
                simplify_sequence : str
                    TN simplifications to use when looking for contraction paths.
                backend : str
                    Backend to use when performing the contractions.
                    Usually specified if GPU acceleration is needed.
                ...
        """

        n_of_exp_vals = len(self.params)*2*len(self._H.paulis)

        self._exp_vals = np.zeros(n_of_exp_vals)

        bins_sizes = [int(n_of_exp_vals/self._size)
                                        for i in range(self._size)]
        for i in range(n_of_exp_vals -
                            int(n_of_exp_vals/self._size)*self._size):
            bins_sizes[i] = bins_sizes[i]+1
        start = sum(bins_sizes[:self._rank])
        end = start + bins_sizes[self._rank]

        exp_vals_iterm = np.zeros(end-start)

        for i,ind in enumerate(range(start, end)):
            #parameter index for this process
            mu = int(int(ind/len(self._H.paulis))/2)

            params = self.params.copy()
            #parameter shift rule
            if int(ind/len(self._H.paulis)) % 2 == 0:
                params[mu] = params[mu]+np.pi/2
            else:
                params[mu] = params[mu]-np.pi/2

            qc = self._base_circuits[-1].copy()
            #update parameters
            old_params_dict = qc.get_params()
            new_params_dict = dict()
            for j,key in enumerate(old_params_dict.keys()):
                new_params_dict[key]= np.array([params[j]])
            qc.set_params(new_params_dict)

            pauli_str_ind = ind % len(self._H.paulis)
            pauli_str = self._H.paulis[pauli_str_ind]

            if type(optimize) == dict:
                exp_vals_iterm[i] = np.real(p_str_exp_eval(
                    qc=qc,
                    pauli_str=pauli_str,
                    optimize = optimize[pauli_str],
                    **kwargs
                ))
            else:
                exp_vals_iterm[i] = np.real(p_str_exp_eval(
                    qc=qc,
                    pauli_str=pauli_str,
                    optimize = optimize,
                    **kwargs
                ))

        sendcountes=tuple(bins_sizes)
        displacements=tuple([sum(bins_sizes[:i]) for i in range(self._size)])
        #collecting an array of the expectation values for all Pauli strings
        self._comm.Allgatherv(
            [exp_vals_iterm,  MPI.DOUBLE],
            [self._exp_vals, sendcountes, displacements, MPI.DOUBLE]
        )
        #computing Hamiltonian expectation values for different parameters
        H_exp_vals = [
            sum([self._exp_vals[param_ind*len(self._H.coefs)+i]*self._H.coefs[i]
                for i in range(len(self._H.coefs))]
            )
            for param_ind in range(len(self.params)*2)
        ]
        self._v = [
            np.real(
                -1/2*(H_exp_vals[param_ind*2] - H_exp_vals[param_ind*2+1]) / 2
            )
            for param_ind in range(len(self.params))
        ]


    def get_dthdt(self, delta, m, v):
        """
        Computes parameter vector gradient.

        Parameters:
        -----------
        delta : float
            Tikhonov regularization parameter for computing inverse of a matrix.
        m : numpy.ndarray
            Matrix M.
        v : numpy.ndarray
            Vector V.
        """
        a = m + delta*np.eye(m.shape[0])
        ainv = np.linalg.inv(a)
        dthdt = ainv.dot(v)
        return dthdt


    def vqite(
        self,
        delta = 1e-4,
        dt = 0.02,
        optimize_m='greedy',
        optimize_v='greedy',
        **kwargs
    ):
        """
        Performs VQITE routine.

        Parameters:
        -----------
        delta : float
            Tikhonov regularization parameter for computing inverse of a matrix.
        dt : float
            Imaginary time step.
        optimize_m : string
            Optimizer to use when looking for contraction paths for M matrix.
        optimize_v : string or dict[cotengra.core.ContractionTree]
            Optimizer to use when looking for contraction paths for V vector.
            If dict, then entries should correspond to a contraction tree for
            each Pauli in the Hamiltonian (this is to reuse contraction paths).
        **kwargs
            Arguments used in Quimb methods for tensor contraction
            evaluations, such as (note that optimize parameter is specified
            separately):
                simplify_sequence : str
                    TN simplifications to use when looking for contraction paths.
                backend : str
                    Backend to use when performing the contractions.
                    Usually specified if GPU acceleration is needed.
                ...
        """
        _iter=0
        if self._rank==0:
            with open(self._output_file, "a") as f:
                print("Starting VQITE calculation...", file=f)
        while True:
            t1 = MPI.Wtime()
            if _iter==0:
                #For the first iteration, need to compute the entire matrix
                #since it is not known a priori which elements are zero.
                self.compute_m(
                    optimize=optimize_m,
                    which_nonzero=None,
                    **kwargs
                )
                #Save locations of nonzero elements.
                non_zero_els = np.where((np.abs(self._m)>1e-14) == True)
                self.which_nonzero = [(non_zero_els[0][i],non_zero_els[1][i])
                                    for i in range(len(non_zero_els[0]))
                                    if non_zero_els[0][i]<=non_zero_els[1][i]]
                if self._rank==0:
                    with open(self._output_file, "a") as f:
                        print("# of nonzero elements of M: ",
                              len(self.which_nonzero), file=f)
            else:
                #For iterations after the first one, need to compute only
                #nonzero elements
                self.compute_m(
                    optimize=optimize_m,
                    which_nonzero=self.which_nonzero,
                    **kwargs
                )
            t2 = MPI.Wtime()
            self.compute_v(optimize=optimize_v,**kwargs)
            t3 = MPI.Wtime()
            dthdt = self.get_dthdt(delta = delta, m = self._m, v= self._v)
            params_new = [p + pp*dt for p, pp in zip(self.params, dthdt)]
            self.params = params_new
            self.update_params()
            self._e = self.h_exp_val(
                params = self.params,
                optimize = optimize_v,
                **kwargs
            )
            if self._rank==0:
                with open(self._output_file, "a") as f:
                    print(
                        "iter: ",_iter,
                        ", M matrix time: ", t2-t1,
                        ", V vector time: ", t3-t2,
                        ", Energy: ", self._e,
                        file=f
                    )
            self._vmax = np.max(np.abs(self._v))
            #Convergence condition.
            if self._vmax < 1e-4:
                break
            _iter+=1


    def h_terms_find_contractions(
        self,
        **kwargs
    ):
        """
        Finds TN contractions for computing the expectation value of each
        Pauli string in the Hamiltonian.
        The obtained contractions are saved in dictionaries
        self.h_terms_reh_dict (for reh) and self.optimize_dict (for reh trees).
        """
        self.h_terms_reh_dict = dict()
        self.optimize_dict = dict()
        qc = self._base_circuits[-1].copy()
        for pauli_str in self._H.paulis:
            self.h_terms_reh_dict[pauli_str] = p_str_exp_contr_path(
                qc = qc,
                pauli_str = pauli_str,
                **kwargs
            )
            self.optimize_dict[pauli_str] = self.h_terms_reh_dict[
                pauli_str
            ]['tree']


    def h_exp_val(
        self,
        params = None,
        optimize = 'greedy',
        **kwargs
    ):
        """
        Computes expectation value of the Hamiltonian using Quimb.

        Parameters:
        ----------
        params : List[float] or None
            List of parameters to be used in the ansatz state.
            If None, current parameters within the object are used.
        optimize : str ot dict
            Optimizer to use when looking for contraction paths.
            If str, then provide Quimb value.
            If dict, then provide a dictionary with a rehearsal tree for each
            Pauli string in the Hamiltonian.
        **kwargs
            Arguments used in Quimb methods for tensor contraction
            evaluations, such as (note that optimize parameter is specified
            separately):
                simplify_sequence : str
                    TN simplifications to use when looking for contraction paths.
                backend : str
                    Backend to use when performing the contractions.
                    Usually specified if GPU acceleration is needed.
                ...
        """
        qc = self._base_circuits[-1].copy()

        if params != None:
            old_params_dict = qc.get_params()
            new_params_dict = dict()
            for i,key in enumerate(old_params_dict.keys()):
                new_params_dict[key]= np.array([params[i]])
            qc.set_params(new_params_dict)

        h_exp_vals = []
        for pauli_str in self._H.paulis:
            if type(optimize) == dict:
                exp_val = p_str_exp_eval(
                    qc=qc,
                    pauli_str=pauli_str,
                    optimize = optimize[pauli_str],
                    **kwargs
                )
                h_exp_vals.append(exp_val)
            else:
                h_exp_vals.append( p_str_exp_eval(
                    qc=qc,
                    pauli_str=pauli_str,
                    optimize = optimize,
                    **kwargs
                ) )
        exp_value = sum([h_exp_vals[i]*self._H.coefs[i]
                                    for i in range(len(self._H.coefs))])
        return exp_value


    def circuit_1(
        self,
        mu: int,
        nu: int,
        A_mu: str,
        A_nu: str
    ):
        """
        Constructs the following quantum circuit (see AVQITE paper for details):
        U^{\dag}_{0,\nu-1} A_{\nu} U_{\mu,\nu-1} A_{\mu} U_{0,\mu-1}|ref>,
        where |ref> is the reference state.

        Parameters:
        ----------
        mu : int
            Index where Pauli A_mu is placed.
        nu : int
            Index where Pauli A_mu is placed.
        A_mu : str
            Pauli string A_{\mu}.
        A_nu : str
            Pauli string A_{\nu}.
        """
        if mu >= nu:
            raise ValueError("Here mu<nu is required.")
        if mu > len(self._ansatz) or nu > len(self._ansatz):
            raise ValueError("mu, nu has to be smaller than "
                                "the number of operators in the ansatz")
        qc = self._base_circuits[mu].copy()
        qc.apply_gates(
            pauli_string_to_quimb_gates(pauli_string=A_mu),
            contract=False
        )
        for i in range(mu,nu):
             qc.apply_gates(self._pauli_rot_gates_list[i], contract=False)
        qc.apply_gates(
            pauli_string_to_quimb_gates(pauli_string=A_nu),
            contract=False
        )
        for i in reversed(range(nu)):
            qc.apply_gates(self._pauli_rot_dag_gates_list[i], contract=False)

        [self._init_qc.apply_gate('X',i)
                            for i,el in enumerate(self._ref_state) if el=='1']

        return qc


    def circuit_2(
        self,
        mu: int
    ):
        """
        Constructs the following quantum circuit (see AVQITE paper for details):
        U_{0,\mu-1}|ref>, where |ref> is the reference state.

        Parameters:
        ----------
        mu : int
            Index up to which Pauli rotations from the ansatz are used.
        """
        qc = self._init_qc.copy()
        for i in range(mu):
            qc.apply_gates(self._pauli_rot_gates_list[i], contract=False)
        return qc


    def contr1_est(
        self,
        mu: int,
        nu: int,
        backend = None,
        **kwargs
    ):
        """
        Calculates contraction width, cost, and value for the following tensor
        (see AVQITE paper for details):
        <ref|U^{\dag}_{0,\nu-1} A_{\nu} U_{\mu,\nu-1} A_{\mu} U_{0,\mu-1}|ref>.
        The tensor and the contraction are obtained as the evaluation of the
        overlap between state
        U^{\dag}_{0,\nu-1} A_{\nu} U_{\mu,\nu-1} A_{\mu} U_{0,\mu-1}|ref> and
        |ref>.

        Parameters:
        -----------
        mu : int
            Index where Pauli A_mu is placed.
            Pauli A_mu is the mu'th Pauli in the ansatz.
        nu : int
            Index where Pauli A_nu is placed.
            Pauli A_nu is the nu'th Pauli in the ansatz.
        backend : str
            Backend to use when performing the contractions.
            Usually specified if GPU acceleration is needed.
        **kwargs
            Arguments used in Quimb methods for tensor contraction
            evaluations, such as:
                optimize : str
                    Optimizer to use when looking for contraction paths.
                simplify_sequence : str
                    TN simplifications to use when looking for contraction paths.
                ...

        Returns:
        --------
        width : float64
            Contraction width.
        cost : float64
            Contraction cost.
        contraction : complex128
            Contraction value.
        """
        if mu > len(self._ansatz) or nu > len(self._ansatz):
            raise ValueError("mu, nu has to be smaller than "
                                "the number of operators in the ansatz")
        if mu>nu:
            raise ValueError("it is assumed here that mu<=nu")
        if mu<nu:
            qc = self.circuit_1(
                mu,
                nu,
                A_mu = self._ansatz[mu],
                A_nu = self._ansatz[nu]
            )
            reh = qc.amplitude_rehearse(
                '0'*self._num_qubits,
                **kwargs
            )
            width, cost = reh['W'], reh['C']
            contraction = reh['tn'].contract(
                all,
                optimize=reh['tree'],
                output_inds=(),
                backend=backend
            )
        if mu==nu:
            width, cost, contraction = (1,0,1)
        contraction = np.real(contraction)/4
        return width, cost, contraction


    def contr2_est(
        self,
        mu: int,
        backend = None,
        **kwargs
    ):
        """
        Calculates contraction width, cost, and value for the following tensor
        (see AVQITE paper for details):
        <ref|U^{\dag}_{0,\mu-1} A_{\mu} U_{0,\mu-1}|ref>.
        The tensor and the contraction are obtained as the evaluation of the
        expectation value of operator A_{\mu} (\mu'th operator from the ansatz).

        Parameters:
        -----------
        mu : int
            Index where Pauli A_mu is placed.
            Pauli A_mu is the mu'th Pauli in the ansatz.
        backend : str
            Backend to use when performing the contractions.
            Usually specified if GPU acceleration is needed.
        **kwargs
            Arguments used in Quimb methods for tensor contraction
            evaluations, such as:
                optimize : str
                    Optimizer to use when looking for contraction paths.
                simplify_sequence : str
                    TN simplifications to use when looking for contraction paths.
                ...

        Returns:
        --------
        reh['W'] : float64
            Contraction width.
        reh['C'] : float64
            Contraction cost.
        contraction : complex128
            Contraction value.
        """
        if mu > len(self._ansatz):
            raise ValueError("mu has to be smaller than "
                                "the number of operators in the ansatz")
        qc = self._base_circuits[mu]
        reh = p_str_exp_contr_path(
            qc = qc,
            pauli_str = self._ansatz[mu],
            **kwargs
        )
        contraction = reh['tn'].contract(
            all,
            optimize=reh['tree'],
            output_inds=(),
            backend=backend
        )
        contraction = np.real(1j*contraction/2)
        return reh['W'], reh['C'], contraction


def add_pauli_rotation_gate(
    qc: "quimb.tensor.circuit.Circuit",
    pauli_string: str,
    theta: float,
    decompose_rzz: bool = True
):
    """
    Appends a Pauli rotation gate to a Quimb Circuit.
    Convention for Pauli string ordering is opposite to the Qiskit convention.
    For example, in string "XYZ" Pauli "X" acts on the first qubit.

    Parameters
    ----------
    qc : "quimb.tensor.circuit.Circuit"
        Quimb Circuit to which the Pauli rotation gate is appended.
    pauli_string : str
        Pauli string defining the rotation.
    theta : float
        Rotation angle.
    decompose_rzz : bool
        If decompose_rzz==True, all rzz gates are decompsed into cx-rz-cx.
        Otherwise, the final circuit contains rzz gates.

    Returns
    -------
    qc: Parameterized "quimb.tensor.circuit.Circuit"
    """
    if qc.N != len(pauli_string):
        raise ValueError("Circuit and Pauli string are of different size")
    if all([pauli=='I' or pauli=='X' or pauli=='Y' or pauli=='Z'
            for pauli in pauli_string])==False:
        raise ValueError("Pauli string does not have a correct format")

    nontriv_pauli_list = [(i,pauli)
                        for i,pauli in enumerate(pauli_string) if pauli!='I']
    if len(nontriv_pauli_list)==1:
        if nontriv_pauli_list[0][1]=='X':
            qc.apply_gate(
                'RX',
                theta,
                nontriv_pauli_list[0][0],
                parametrize=True,
                gate_opts={'contract': False}
            )
        if nontriv_pauli_list[0][1]=='Y':
            qc.apply_gate(
                'RY',
                theta,
                nontriv_pauli_list[0][0],
                parametrize=True,
                gate_opts={'contract': False}
            )
        if nontriv_pauli_list[0][1]=='Z':
            qc.apply_gate(
                'RZ',
                theta,
                nontriv_pauli_list[0][0],
                parametrize=True,
                gate_opts={'contract': False}
            )
    elif (len(nontriv_pauli_list)==2 and
                    nontriv_pauli_list[0][1]+nontriv_pauli_list[1][1] == 'XX'):
            qc.apply_gate(
                'RXX',
                theta,
                nontriv_pauli_list[0][0],
                nontriv_pauli_list[1][0],
                parametrize=True,
                gate_opts={'contract': False}
            )
    elif (len(nontriv_pauli_list)==2 and
                    nontriv_pauli_list[0][1]+nontriv_pauli_list[1][1] == 'YY'):
            qc.apply_gate(
                'RYY',
                theta,
                nontriv_pauli_list[0][0],
                nontriv_pauli_list[1][0],
                parametrize=True,
                gate_opts={'contract': False}
            )
    else:
        for (i,pauli) in nontriv_pauli_list:
            if pauli=='X':
                qc.apply_gate('H',i)
            if pauli=='Y':
                qc.apply_gate('SDG',i)
                qc.apply_gate('H',i)
        for list_ind in range(len(nontriv_pauli_list)-2):
            qc.apply_gate(
                'CX',
                nontriv_pauli_list[list_ind][0],
                nontriv_pauli_list[list_ind+1][0]
            )
        if decompose_rzz==True:
            qc.apply_gate(
                'CX',
                nontriv_pauli_list[len(nontriv_pauli_list)-2][0],
                nontriv_pauli_list[len(nontriv_pauli_list)-1][0]
            )
            qc.apply_gate(
                'RZ',
                theta,
                nontriv_pauli_list[len(nontriv_pauli_list)-1][0],
                parametrize=True,
                gate_opts={'contract': False}
            )
            qc.apply_gate(
                'CX',
                nontriv_pauli_list[len(nontriv_pauli_list)-2][0],
                nontriv_pauli_list[len(nontriv_pauli_list)-1][0]
            )
        if decompose_rzz==False:
            qc.apply_gate(
                'RZZ',
                theta,
                nontriv_pauli_list[len(nontriv_pauli_list)-2][0],
                nontriv_pauli_list[len(nontriv_pauli_list)-1][0],
                parametrize=True,
                gate_opts={'contract': False}
            )
        for list_ind in reversed(range(len(nontriv_pauli_list)-2)):
            qc.apply_gate(
                'CX',
                nontriv_pauli_list[list_ind][0],
                nontriv_pauli_list[list_ind+1][0]
            )
        for (i,pauli) in nontriv_pauli_list:
            if pauli=='X':
                qc.apply_gate('H',i)
            if pauli=='Y':
                qc.apply_gate('H',i)
                qc.apply_gate('S',i)
    return qc


def read_adaptvqite_ansatz(
    filename: str
):
    """
    Reads the ansatz from a file resulting from adaptvqite calculation.

    Parameters
    ----------
    filename : str
        Name of a file containing the results of adaptvqite calculation.
        Has to be given in .pkle format.

    Returns
    -------
    ansatz_adaptvqite : List[str]
        List of Pauli strings entering the ansatz.
    params_adaptvqite : List[float64]
        Parameters (angles) of the ansatz.
    """
    if filename[-5:] != '.pkle':
        raise ImportError("Ansatz file should be given in .pkle format")

    with open(filename, 'rb') as inp:
        data_inp = pickle.load(inp)
        ansatz_adaptvqite = data_inp[0]
        params_adaptvqite = data_inp[1]

    return ansatz_adaptvqite, params_adaptvqite


def pauli_string_to_quimb_gates(pauli_string):
    """
    Converts a Pauli string into a Quimb gate.
    """
    gates = ()
    for i,el in enumerate(pauli_string):
        if el=='X':
            gates = gates + (qtn.circuit.Gate(
                label='X',
                params=[],
                qubits=(i,)
            ),)
        if el=='Y':
            gates = gates + (qtn.circuit.Gate(
                label='Y',
                params=[],
                qubits=(i,)
            ),)
        if el=='Z':
            gates = gates + (qtn.circuit.Gate(
                label='Z',
                params=[],
                qubits=(i,)
            ),)
    return gates


def p_str_exp_contr_path(
    qc,
    pauli_str: str,
    **kwargs
):
    """
    Constucts a contraction path for a TN evaluating the expectation value of a
    Pauli string.

    Parameters:
    -----------
    qc : quimb.tensor.circuit.Circuit
        Quantum circuit representing a state for which the expectation value
        is computed.
    **kwargs
        Arguments used in Quimb methods for tensor contraction
        evaluations, such as:
            optimize : str
                Optimizer to use when looking for contraction paths.
            simplify_sequence : str
                TN simplifications to use when looking for contraction paths.
            backend : str
                Backend to use when performing the contractions.
                Usually specified if GPU acceleration is needed.
            ...

    Returns:
    --------
    reh : dict
        Dictionary representing the cotraction path and the TN to be contracted:
        reh['tree'] -- contraction path
        reh['tn'] -- TN
    """
    where = [i for i,p in enumerate(pauli_str) if p!= 'I']
    paulis = [p for i,p in enumerate(pauli_str) if p!= 'I']
    operator = qu.pauli(paulis[0])
    for i in range(1,len(where)):
        operator = operator & qu.pauli(paulis[i])
    reh = qc.local_expectation_rehearse(
        operator,
        where,
        **kwargs
    )
    return reh


def p_str_exp_eval(
    qc,
    pauli_str: str,
    **kwargs
):
    """
    Evaluates the expectation value of a Pauli string using TN contraction.

    Parameters:
    -----------
    qc : quimb.tensor.circuit.Circuit
        Quantum circuit representing a state for which the expectation value
        is computed.
    pauli_str : str
        Pauli string representing an observable.
    **kwargs
        Arguments used in Quimb methods for tensor contraction
        evaluations, such as:
            optimize : str
                Optimizer to use when looking for contraction paths.
            simplify_sequence : str
                TN simplifications to use when looking for contraction paths.
            backend : str
                Backend to use when performing the contractions.
                Usually specified if GPU acceleration is needed.
            ...
    """
    where = [i for i,p in enumerate(pauli_str) if p!= 'I']
    paulis = [p for i,p in enumerate(pauli_str) if p!= 'I']
    operator = qu.pauli(paulis[0])
    for i in range(1,len(where)):
        operator = operator & qu.pauli(paulis[i])
    exp_val = qc.local_expectation(
        operator,
        where,
        **kwargs
    )
    return exp_val