qite.cpp

/*******************************************************************************
 * Copyright (c) 2019 UT-Battelle, LLC.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * and Eclipse Distribution License v1.0 which accompanies this
 * distribution. The Eclipse Public License is available at
 * http://www.eclipse.org/legal/epl-v10.html and the Eclipse Distribution
 *License is available at https://eclipse.org/org/documents/edl-v10.php
 *
 * Contributors:
 *   Thien Nguyen - initial API and implementation
 *******************************************************************************/
#include "qite.hpp"

#include "Observable.hpp"
#include "xacc.hpp"
#include "xacc_service.hpp"
#include "PauliOperator.hpp"
#include "Circuit.hpp"
#include <memory>
#include <armadillo>
#include <cassert>

namespace {
const std::complex<double> I{ 0.0, 1.0};
int findMatchingPauliIndex(const std::vector<std::string>& in_OpList, const std::string& in_obsStr)
{
  // Returns true if the *operators* of the two terms are identical
  // e.g. a Z0X1 and b Z0X1 -> true 
  const auto comparePauliString = [](const std::string& in_a, const std::string& in_b) -> bool {
    // Strip the coefficient part
    auto opA = in_a.substr(in_a.find_last_of(")") + 1);
    auto opB = in_b.substr(in_b.find_last_of(")") + 1);
    opA.erase(std::remove(opA.begin(), opA.end(), ' '), opA.end()); 
    opB.erase(std::remove(opB.begin(), opB.end(), ' '), opB.end()); 
    return opA == opB;
  };

  for (int i = 0; i < in_OpList.size(); ++i)
  {
    std::shared_ptr<xacc::Observable> obs = std::make_shared<xacc::quantum::PauliOperator>();
    const std::string pauliObsStr = "1.0 " + in_OpList[i];
    obs->fromString(pauliObsStr);
    
    if (comparePauliString(obs->toString(), in_obsStr))
    {
      return i;
    }
  }
  // Failed!
  return -1;
}

// Project/flatten the target observable into the full list of all
// possible Pauli operator combinations.
// e.g. H = a X + b Z (1 qubit)
// -> { 0.0, a, 0.0, b } (the ordering is I, X, Y, Z)
std::vector<double> observableToVec(std::shared_ptr<xacc::Observable> in_observable, const std::vector<std::string>& in_pauliObsList) 
{
  std::vector<double> obsProjCoeffs(in_pauliObsList.size(), 0.0);
  for (const auto& term: in_observable->getNonIdentitySubTerms())
  {
    const auto index = findMatchingPauliIndex(in_pauliObsList, term->toString());
    assert(index >= 0);
    obsProjCoeffs[index] = term->coefficient().real();
  }
  return obsProjCoeffs;
};

// Helper to generate all permutation of Pauli observables:
// e.g.
// 1-qubit: I, X, Y, Z
// 2-qubit: II, IX, IY, IZ, XI, XX, XY, XZ, YI, YX, YY, YZ, ZI, ZX, ZY, ZZ
std::vector<std::string> generatePauliPermutation(int in_nbQubits)
{
  assert(in_nbQubits > 0);
  const int nbPermutations = std::pow(4, in_nbQubits);
  std::vector<std::string> opsList;
  opsList.reserve(nbPermutations);
  
  const std::vector<std::string> pauliOps { "X", "Y", "Z" };
  const auto addQubitPauli = [&opsList, &pauliOps](int in_qubitIdx){
    const auto currentOpListSize = opsList.size();
    for (int i = 0; i < currentOpListSize; ++i)
    {
      auto& currentOp = opsList[i];
      for (const auto& pauliOp : pauliOps)
      {
        const auto newOp = currentOp + pauliOp + std::to_string(in_qubitIdx);
        opsList.emplace_back(newOp);
      }
    }
  };
  
  opsList = { "", "X0", "Y0", "Z0" };
  for (int i = 1; i < in_nbQubits; ++i) 
  {
    addQubitPauli(i);
  }

  assert(opsList.size() == nbPermutations);
  std::sort(opsList.begin(), opsList.end());

  return opsList;
};

arma::cx_mat createSMatrix(const std::vector<std::string>& in_pauliOps, const std::vector<double>& in_tomoExp) 
{
  const auto sMatDim = in_pauliOps.size();
  arma::cx_mat S_Mat(sMatDim, sMatDim, arma::fill::zeros);
  arma::cx_vec b_Vec(sMatDim, arma::fill::zeros); 
  
  const auto calcSmatEntry = [&](int in_row, int in_col) -> std::complex<double> {
    // Map the tomography expectation to the S matrix
    // S(i, j) = <psi|sigma_dagger(i)sigma(j)|psi>
    // sigma_dagger(i)sigma(j) will produce another Pauli operator with an additional coefficient.
    // e.g. sigma_x * sigma_y = i*sigma_z
    const auto leftOp = "1.0 " + in_pauliOps[in_row];
    const auto rightOp = "1.0 " + in_pauliOps[in_col];
    xacc::quantum::PauliOperator left(leftOp);
    xacc::quantum::PauliOperator right(rightOp);
    auto product = left * right;
    // std::cout << left.toString() << " * " << right.toString() << " = " << product.toString() << "\n";
    const auto index = findMatchingPauliIndex(in_pauliOps, product.toString());
    return in_tomoExp[index]*product.coefficient();
  };

  // S matrix:
  for (int i = 0; i < sMatDim; ++i)
  {
    for (int j = 0; j < sMatDim; ++j)
    {
      S_Mat(i, j) = calcSmatEntry(i, j);
    }
  }
  return S_Mat;
}
}

using namespace xacc;
namespace xacc {
namespace algorithm {
bool QITE::initialize(const HeterogeneousMap &parameters) 
{
  bool initializeOk = true;
  if (!parameters.pointerLikeExists<Accelerator>("accelerator")) 
  {
    std::cout << "'accelerator' is required.\n";
    initializeOk = false;
  }

  if (!parameters.keyExists<int>("steps")) 
  {
    std::cout << "'steps' is required.\n";
    initializeOk = false;
  }

  if (!parameters.keyExists<double>("step-size")) 
  {
    std::cout << "'step-size' is required.\n";
    initializeOk = false;
  }

  if (!parameters.pointerLikeExists<Observable>("observable")) 
  {
    std::cout << "'observable' is required.\n";
    initializeOk = false;
  }
  
  if (initializeOk)
  {
    m_accelerator = xacc::as_shared_ptr(parameters.getPointerLike<Accelerator>("accelerator"));
    m_nbSteps = parameters.get<int>("steps");
    m_dBeta = parameters.get<double>("step-size");
    m_observable = xacc::as_shared_ptr(parameters.getPointerLike<Observable>("observable"));
  }

  m_analytical = true;
  if (parameters.keyExists<bool>("analytical")) 
  {
    m_analytical = parameters.get<bool>("analytical");
  }

  m_approxOps.clear();
  m_energyAtStep.clear();
  
  return initializeOk;
}

const std::vector<std::string> QITE::requiredParameters() const 
{
  return { "accelerator", "steps", "step-size", "observable" };
}

std::shared_ptr<CompositeInstruction> QITE::constructPropagateCircuit() const
{
  auto gateRegistry = xacc::getService<xacc::IRProvider>("quantum");
  auto propagateKernel = gateRegistry->createComposite("statePropCircuit");

  // Adds ansatz if provided
  if (m_ansatz)
  {
    propagateKernel->addInstructions(m_ansatz->getInstructions());
  } 

  const auto pauliTermToString = [](const std::shared_ptr<xacc::Observable>& in_term){
    std::string pauliTermStr = in_term->toString();
    std::stringstream s;
    s.precision(12);
    s << std::fixed << in_term->coefficient();
    // Find the parenthesis
    const auto startPosition = pauliTermStr.find("(");
    const auto endPosition = pauliTermStr.find(")");

    if (startPosition != std::string::npos && endPosition != std::string::npos)
    {
      const auto length = endPosition - startPosition + 1;
      pauliTermStr.replace(startPosition, length, s.str());
    }
    return pauliTermStr;
  };

  // Progagates by Trotter steps
  // Using those A operators that have been 
  // optimized up to this point.
  for (const auto& aObs : m_approxOps)
  {
    // Circuit is: exp(-idt*A),
    // i.e. regular evolution which approximates the imaginary time evolution.
    for (const auto& term : aObs->getNonIdentitySubTerms())
    {
      const auto pauliStr = pauliTermToString(term);
      auto expCirc = std::dynamic_pointer_cast<xacc::quantum::Circuit>(xacc::getService<Instruction>("exp_i_theta"));
      const bool expandOk = expCirc->expand({ std::make_pair("pauli", pauliStr) });
      assert(expandOk);
      auto evaled = expCirc->operator()({ m_dBeta });
      propagateKernel->addInstructions(evaled->getInstructions());
    }
  }

  // std::cout << "Progagated kernel:\n" << propagateKernel->toString() << "\n";
  return propagateKernel;
}

double QITE::calcCurrentEnergy(int in_nbQubits) const
{
  // Trotter kernel up to this point
  auto propagateKernel = constructPropagateCircuit();
  auto kernels = m_observable->observe(propagateKernel);
  std::vector<double> coefficients;
  std::vector<std::string> kernelNames;
  std::vector<std::shared_ptr<CompositeInstruction>> fsToExec;

  double identityCoeff = 0.0;
  for (auto &f : kernels) 
  {
    kernelNames.push_back(f->name());
    std::complex<double> coeff = f->getCoefficient();
    int nFunctionInstructions = 0;
    if (f->nInstructions() > 0 && f->getInstruction(0)->isComposite()) 
    {
      nFunctionInstructions = propagateKernel->nInstructions() + f->nInstructions() - 1;
    } 
    else 
    {
      nFunctionInstructions = f->nInstructions();
    }

    if (nFunctionInstructions > propagateKernel->nInstructions()) 
    {
      fsToExec.push_back(f);
      coefficients.push_back(std::real(coeff));
    } 
    else 
    {
      identityCoeff += std::real(coeff);
    }
  }

  auto tmpBuffer = xacc::qalloc(in_nbQubits);
  m_accelerator->execute(tmpBuffer, fsToExec);
  auto buffers = tmpBuffer->getChildren();

  double energy = identityCoeff;
  for (int i = 0; i < buffers.size(); ++i) 
  {
    auto expval = buffers[i]->getExpectationValueZ();
    energy += expval * coefficients[i];
  }
  std::cout << "Energy = " << energy << "\n";
  return energy;
}

std::shared_ptr<Observable> QITE::calcAOps(const std::shared_ptr<AcceleratorBuffer>& in_buffer, std::shared_ptr<CompositeInstruction> in_kernel, std::shared_ptr<Observable> in_hmTerm) const
{
  const auto pauliOps = generatePauliPermutation(in_buffer->size());
  assert(in_hmTerm->getSubTerms().size() == 1);
  assert(in_hmTerm->getNonIdentitySubTerms().size() == 1);

  // Step 1: Observe the kernels using the various Pauli
  // operators to calculate S and b.
  std::vector<double> sigmaExpectation(pauliOps.size());
  sigmaExpectation[0] = 1.0;
  for (int i = 1; i < pauliOps.size(); ++i)
  {
    std::shared_ptr<Observable> tomoObservable = std::make_shared<xacc::quantum::PauliOperator>();
    const std::string pauliObsStr = "1.0 " + pauliOps[i];
    tomoObservable->fromString(pauliObsStr);
    assert(tomoObservable->getSubTerms().size() == 1);
    assert(tomoObservable->getNonIdentitySubTerms().size() == 1);
    auto obsKernel = tomoObservable->observe(in_kernel).front();
    auto tmpBuffer = xacc::qalloc(in_buffer->size());
    m_accelerator->execute(tmpBuffer, obsKernel);
    const auto expval = tmpBuffer->getExpectationValueZ();
    sigmaExpectation[i] = expval;
  }

  // Step 2: Calculate S matrix and b vector
  // i.e. set up the linear equation Sa = b
  const auto sMatDim = pauliOps.size();
  arma::cx_mat S_Mat(sMatDim, sMatDim, arma::fill::zeros);
  arma::cx_vec b_Vec(sMatDim, arma::fill::zeros); 
  
  const auto calcSmatEntry = [&](const std::vector<double>& in_tomoExp, int in_row, int in_col) -> std::complex<double> {
    // Map the tomography expectation to the S matrix
    // S(i, j) = <psi|sigma_dagger(i)sigma(j)|psi>
    // sigma_dagger(i)sigma(j) will produce another Pauli operator with an additional coefficient.
    // e.g. sigma_x * sigma_y = i*sigma_z
    const auto leftOp = "1.0 " + pauliOps[in_row];
    const auto rightOp = "1.0 " + pauliOps[in_col];
    xacc::quantum::PauliOperator left(leftOp);
    xacc::quantum::PauliOperator right(rightOp);
    auto product = left * right;
    const auto index = findMatchingPauliIndex(pauliOps, product.toString());
    return in_tomoExp[index]*product.coefficient();
  };

  // S matrix:
  for (int i = 0; i < sMatDim; ++i)
  {
    for (int j = 0; j < sMatDim; ++j)
    {
      S_Mat(i, j) = calcSmatEntry(sigmaExpectation, i, j);
    }
  }

  // b vector:
  const auto obsProjCoeffs = observableToVec(in_hmTerm, pauliOps);
  // std::cout << "Observable Pauli Vec: [";
  // for (const auto& elem: obsProjCoeffs)
  // {
  //   std::cout << elem << ", ";
  // }
  // std::cout << "]\n";

  // Calculate c: Eq. 3 in https://arxiv.org/pdf/1901.07653.pdf
  double c = 1.0;
  for (int i = 0; i < obsProjCoeffs.size(); ++i)
  {
    c -= 2.0 * m_dBeta * obsProjCoeffs[i] * sigmaExpectation[i];
  }
  
  // std::cout << "c = " << c << "\n";
  for (int i = 0; i < sMatDim; ++i)
  {
    std::complex<double> b = (sigmaExpectation[i]/ std::sqrt(c) - sigmaExpectation[i])/ m_dBeta;
    for (int j = 0; j < obsProjCoeffs.size(); ++j)
    {
      // The expectation of the pauli product of the Hamiltonian term
      // and the sweeping pauli term.
      const auto expectVal = calcSmatEntry(sigmaExpectation, i, j);
      b -= obsProjCoeffs[j] * expectVal / std::sqrt(c);
    }
    b = I*b - I*std::conj(b);
    // Set b_Vec
    b_Vec(i) = b;
  }

  // std::cout << "S Matrix: \n" << S_Mat << "\n"; 
  // std::cout << "B Vector: \n" << b_Vec << "\n"; 
  // Add regularizer
  arma::cx_mat dalpha(sMatDim, sMatDim, arma::fill::eye); 
  dalpha = 0.1 * dalpha;

  auto lhs = S_Mat + S_Mat.st() + dalpha;
  auto rhs = -b_Vec;

  // std::cout << "LHS Matrix: \n" << lhs << "\n"; 
  // std::cout << "RHS Vector: \n" << rhs << "\n"; 
  arma::cx_vec a_Vec = arma::solve(lhs, rhs);
  // std::cout << "Result A Vector: \n" << a_Vec << "\n"; 

  // Now, we have the decomposition of A observable in the basis of
  // all possible Pauli combinations.
  assert(a_Vec.n_elem == pauliOps.size());
  const std::string aObsStr = [&](){
    std::stringstream s;
    s.precision(12);
    s << std::fixed << a_Vec(0);
    
    for (int i = 1; i < pauliOps.size(); ++i)
    {
      s << " + " << (-2.0 * a_Vec(i)) << " " << pauliOps[i];
    }

    return s.str();
  }(); 

  // Step 3: compute the approximate A observable/Hamiltonian.
  // This operator will drive the exp_i_theta evolution
  // which emulate the imaginary time evolution of the original observable.
  std::shared_ptr<Observable> updatedAham = std::make_shared<xacc::quantum::PauliOperator>();
  updatedAham->fromString(aObsStr);
  return updatedAham; 
}

void QITE::execute(const std::shared_ptr<AcceleratorBuffer> buffer) const 
{
  if(!m_analytical)
  {
    // Run on hardware/simulator using quantum gates/measure
    // Initial energy
    m_energyAtStep.emplace_back(calcCurrentEnergy(buffer->size()));
    
    // Time stepping:
    for (int i = 0; i < m_nbSteps; ++i)
    {
      for (const auto& hamTerm : m_observable->getNonIdentitySubTerms())
      {
        // Propagates the state via Trotter steps:
        auto kernel = constructPropagateCircuit();
        // Optimizes/calculates next A ops
        auto nextAOps = calcAOps(buffer, kernel, hamTerm);
        m_approxOps.emplace_back(nextAOps);
      }

      m_energyAtStep.emplace_back(calcCurrentEnergy(buffer->size()));
    }
    assert(m_energyAtStep.size() == m_nbSteps + 1);
    // Last energy value
    buffer->addExtraInfo("opt-val", ExtraInfo(m_energyAtStep.back()));
    // Also returns the full list of energy values 
    // at each Trotter step.
    buffer->addExtraInfo("exp-vals", ExtraInfo(m_energyAtStep));
  }
  else
  {
    // Analytical run:
    // This serves two purposes:
    // (1) Validate the convergence (e.g. Trotter step size) before running via gates.
    // (2) Derive the circuit analytically for running.
    // exp(-dtH)
    const auto expMinusHamTerm = [](const arma::cx_mat& in_hMat, const arma::cx_vec& in_psi, double in_dt) {
      assert(in_hMat.n_rows == in_hMat.n_cols);
      assert(in_hMat.n_rows == in_psi.n_elem);
      arma::cx_mat hMatExp = arma::expmat(-in_dt*in_hMat);
      arma::cx_vec result = hMatExp * in_psi;
      const double norm = arma::norm(result, 2);
      result = result / norm;
      return std::make_pair(result, norm);
    };
    
    const auto getTomographyExpVec = [](int in_nbQubits, const arma::cx_vec& in_psi, const arma::cx_vec& in_delta) {
      const auto pauliOps = generatePauliPermutation(in_nbQubits);
      std::vector<std::complex<double>> sigmaExpectation(pauliOps.size());
      std::vector<std::complex<double>> bVec(pauliOps.size());

      sigmaExpectation[0] = 1.0;
      bVec[0] = arma::cdot(in_delta, in_psi);

      for (int i = 1; i < pauliOps.size(); ++i)
      {
        auto tomoObservable = std::make_shared<xacc::quantum::PauliOperator>();
        const std::string pauliObsStr = "1.0 " + pauliOps[i];
        tomoObservable->fromString(pauliObsStr);
        assert(tomoObservable->getSubTerms().size() == 1);
        assert(tomoObservable->getNonIdentitySubTerms().size() == 1);
        arma::cx_mat hMat(1 << in_nbQubits, 1 << in_nbQubits, arma::fill::zeros);
        const auto hamMat = tomoObservable->toDenseMatrix(in_nbQubits);
        for (int i = 0; i < hMat.n_rows; ++i)
        {
          for (int j = 0; j < hMat.n_cols; ++j)
          {
            const int index = i*hMat.n_rows + j;
            hMat(i, j) = hamMat[index];
          }
        }

        arma::cx_vec pauliApplied = hMat*in_psi;
        sigmaExpectation[i] = arma::cdot(in_psi, pauliApplied);
        bVec[i] = arma::cdot(in_delta, pauliApplied);
      }

      return std::make_pair(sigmaExpectation, bVec);
    };

    // Initial state
    arma::cx_vec psiVec(1 << buffer->size(), arma::fill::zeros);
    psiVec(1) = 1.0;
    // Time stepping:
    for (int i = 0; i < m_nbSteps; ++i)
    {
      arma::cx_mat hMat(1 << buffer->size(), 1 << buffer->size(), arma::fill::zeros);
      
      auto identityTerm = m_observable->getIdentitySubTerm();
      if (identityTerm)
      {
        arma::cx_mat idTerm(1 << buffer->size(), 1 << buffer->size(), arma::fill::eye);
        hMat += identityTerm->coefficient() * idTerm;
      }

      double stateVecNorm = 1.0;
      for (const auto& hamTerm : m_observable->getNonIdentitySubTerms())
      {
        auto pauliCast = std::dynamic_pointer_cast<xacc::quantum::PauliOperator>(hamTerm);
        const auto hamMat = pauliCast->toDenseMatrix(buffer->size());
      
        for (int i = 0; i < hMat.n_rows; ++i)
        {
          for (int j = 0; j < hMat.n_cols; ++j)
          {
            const int index = i*hMat.n_rows + j;
            hMat(i, j) += hamMat[index];
          }
        }
      }
      double normAfter = 0.0;
      arma::cx_vec dPsiVec(1 << buffer->size(), arma::fill::zeros);
      std::tie(dPsiVec, normAfter) = expMinusHamTerm(hMat, psiVec, m_dBeta);
      stateVecNorm *= normAfter;
      // Eq. 8, SI of https://arxiv.org/pdf/1901.07653.pdf
      dPsiVec = dPsiVec - psiVec;
      std::vector<std::complex<double>> pauliExp;
      std::vector<std::complex<double>> bVec;
      std::tie(pauliExp, bVec) = getTomographyExpVec(buffer->size(), psiVec, dPsiVec);
      std::vector<double> pauliExpValues;
      for (const auto& val: pauliExp)
      {
        pauliExpValues.emplace_back(val.real());
      }

      arma::cx_mat sMat = createSMatrix(generatePauliPermutation(buffer->size()), pauliExpValues);
      arma::cx_vec b_Vec(bVec.size(), arma::fill::zeros); 
      for (int i = 0; i < bVec.size(); ++i)
      {
        b_Vec(i) = -I*bVec[i] + I*std::conj(bVec[i]);
      }
      
      auto lhs = sMat + sMat.st();
      auto rhs = b_Vec;
      arma::cx_vec a_Vec = arma::solve(lhs, rhs);
      const auto pauliOps = generatePauliPermutation(buffer->size());
      const std::string aObsStr = [&](){
        std::stringstream s;
        s.precision(12);
        s << std::fixed << a_Vec(0);
        
        for (int i = 1; i < pauliOps.size(); ++i)
        {
          s << " + " << a_Vec(i) << " " << pauliOps[i];
        }

        return s.str();
      }(); 

      std::shared_ptr<xacc::quantum::PauliOperator> updatedAham = std::make_shared<xacc::quantum::PauliOperator>();
      updatedAham->fromString(aObsStr);
      const auto aHamMat = updatedAham->toDenseMatrix(buffer->size());
      arma::cx_mat aMat(1 << buffer->size(), 1 << buffer->size(), arma::fill::zeros);

      for (int i = 0; i < aMat.n_rows; ++i)
      {
        for (int j = 0; j < aMat.n_cols; ++j)
        {
          const int index = i*aMat.n_rows + j;
          aMat(i, j) = aHamMat[index];
        }
      }

      // Evolve exp(-iAt)
      arma::cx_mat aMatExp = arma::expmat(-I*aMat);
      arma::cx_mat psiUpdate = aMatExp*psiVec;
      psiVec = psiUpdate;
      const std::complex<double> energyRaw = arma::cdot(psiUpdate, hMat*psiUpdate);
      std::cout << "Energy = " << energyRaw << "\n";
      m_energyAtStep.emplace_back(energyRaw.real());
    }

    m_energyAtStep.emplace_back(m_energyAtStep.back());
    assert(m_energyAtStep.size() == m_nbSteps + 1);
    // Last energy value
    buffer->addExtraInfo("opt-val", ExtraInfo(m_energyAtStep.back()));
    // Also returns the full list of energy values 
    // at each Trotter step.
    buffer->addExtraInfo("exp-vals", ExtraInfo(m_energyAtStep));
  } 
}

std::vector<double> QITE::execute(const std::shared_ptr<AcceleratorBuffer> buffer, const std::vector<double>& x) 
{
  // We don't do any parameter optimization here,
  // hence don't support this!
  xacc::error("This method is unsupported!");
  return {};
}
} // namespace algorithm
} // namespace xacc