{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4610aeac",
   "metadata": {},
   "source": [
    "## Approximate State Preparation using MPS Sequential Encoding\n",
    "\n",
    "Author: Amir Ali Malekani Nezhad (Qualition)\n",
    "\n",
    "State Preparation is the task of compiling a quantum statevector $\\ket{\\psi}$ to a quantum circuit. There have been many approaches proposed by Shende, Mottonen, Perciano, Ran, Rudolph, et al who have proposed different schema for tackling this problem. Additionally, some like Pollmann et al have focused on specific types of states which made their algorithm a better fit compared to current state-of-the-art.\n",
    "\n",
    "This topic has been researched since the early 2000s, and the summary of which is presented below:\n",
    "- **Exact encoding of arbitrary statevectors:** Amplitude encoding was proposed as the most efficient exact encoding approach to compile arbitrary statevectors. These approaches primarily developed by Shende [[1]] and Mottonen [[2]] and Iten et al [[3]] focus on encoding a vector of $2^N$ elements over a quantum circuit with $N$ qubits and a depth of $2^N$ approximately. Such an approach as inferred from quantum information theory proved that we require an exponential number of gates to encode such an exponential data structure. Whilst these approaches are actively used in most modern packages they still bottleneck the implementation of practical quantum algorithms using this subroutine on actual hardware.\n",
    "- **Variational encoding of arbitrary statevectors:** With the issue present in Amplitude Encoding, during the 2015 to early 2020s many researchers pivoted towards using QML to approximately encode statevectors. This approach differed significantly from that of Amplitude encoding as it was stochastic (no analytical decomposition like Amplitude encoding), and focused on an iterative approach to gradually train a QNN to fit to the target state's distribution [[4]]. These approaches also had their own challenges, i.e., \"orthogonality catastrophe\". Orthogonality catastrophe in quantum machine learning refers to a situation where a small change in the parameters of a quantum model causes the quantum state to change drastically, becoming almost completely different (orthogonal) from the original state. This causes an exponential decay in the fidelity measure as the size of the system (number of qubits) increases. This caused more focus on the expressivity of the ansatzes proposed, and the training approach to avoid Barren Plateaus.\n",
    "- **MPS encoding:** First seen in 2008 [[5]] and becoming prominently known with Ran's work in 2019 [[6]], MPS became a massive success in approximate compilation of statevectors. The approach utilized MPS representation of dense tensors (aka the statevector) to reduce the number of parameters needed to represent the tensor from $2^N$ to $2N\\chi^2$ where $\\chi$ is called the bond dimension and serves as a hyperparameter controlling the fidelity of the encoding. Given a constant and modest bond dimension as we increase $N$, we in fact reduce the scaling from $O(2^N)$ to $O(N)$. This allowed the limitation imposed by exact encoding to be relaxed to $O(N)$ and provided a significant reduction in depth. However, even MPS approach yet again had challenges. MPS tackles entanglement in a 1D manner, meaning long-range entanglement (if q0 and qN are entangled, this is long-range, whereas q0 and q1 entangled would be considered short-range) had to be captured via deeper circuits. Additionally, Ran's approach required many layers to adequately encode certain types of states, and is known for being very slow to converge. To tackle this, Rudolph et al [[7]] published a work in 2022 which utilized Environment Tensor Updates to optimize the circuit produced by Ran and find the optimal parameters for the circuit without changing depth. This solved the issue with Ran, and is considered state-of-the-art as of writing this notebook.\n",
    "\n",
    "We can see there have been many different approaches to tackle this problem, each solving some problems and introducing new ones, and in pursuing perfection we have come a long long way from early 2000s works.\n",
    "\n",
    "[1]: https://arxiv.org/abs/quant-ph/0406176\n",
    "[2]: https://arxiv.org/abs/quant-ph/0407010\n",
    "[3]: https://arxiv.org/abs/1501.06911\n",
    "[4]: https://arxiv.org/abs/1612.02806\n",
    "[5]: https://arxiv.org/abs/2008.10322\n",
    "[6]: https://arxiv.org/abs/1908.07958\n",
    "[7]: https://arxiv.org/abs/2209.00595"
   ]
  },
  {
   "attachments": {
    "image.png": {
     "image/png": "[encoded data, 55828 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "fb9a73cb",
   "metadata": {},
   "source": [
    "### Ran's approach\n",
    "\n",
    "In simple terms, Ran's approach constructs an MPS from the statevector we want to compile. To exactly encode the MPS we don't need more than one layer, however, we would need $\\lceil log_2(\\chi) \\rceil + 1$ qubits for the tensors depending on their bond dimension. This leads to larger circuits and makes it harder to transfer the circuits to hardware.\n",
    "\n",
    "He proposed to truncate the bond dimension of the MPS to 2, and encode the truncated version to a circuit which would now only require 1 and 2 qubit gates. This approximately encodes the MPS. To improve the fidelity, we need to generate many layers. Below is the pseudocode borrowed from Rudolph et al. Ran generates a layer, applies its inverse to the original MPS to \"update\" it or more accurately disentangle it (push it closer to product state $\\ket{00 \\cdots 0}$), and generates a new layer from the updated original MPS. We do so until we run out of layers, or reach close enough to product state.\n",
    "\n",
    "Thus, an approximate encoding using an iterative disentangling approach.\n",
    "\n",
    "![image.png](attachment:image.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "451a3cdc",
   "metadata": {},
   "source": [
    "Freeze `numba==0.60.0` to avoid index error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ccba5c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install quimb==1.11.2 numba==0.60.0 --break-system-packages -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "90a99991",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ace07/Documents/GitHub/QMPRS/.venv/lib/python3.11/site-packages/cotengra/hyperoptimizers/hyper.py:54: UserWarning: Couldn't find `optuna`, `cmaes`, or `nevergrad` so will use completely random sampling in place of hyper-optimization.\n",
      "  warnings.warn(\n",
      "/home/ace07/Documents/GitHub/QMPRS/.venv/lib/python3.11/site-packages/cotengra/hyperoptimizers/hyper.py:36: UserWarning: Couldn't import `kahypar` - skipping from default hyper optimizer and using basic `labels` method instead.\n",
      "  warnings.warn(\n",
      "/home/ace07/Documents/GitHub/QMPRS/.venv/lib/python3.11/site-packages/cotengra/hyperoptimizers/hyper.py:73: UserWarning: Couldn't find `optuna`, `cmaes`, or `nevergrad` so will use completely random sampling in place of hyper-optimization.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from numpy.typing import NDArray\n",
    "import quimb.tensor as qtn\n",
    "\n",
    "import cudaq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "614b5b60",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gram_schmidt(matrix: NDArray[np.complex128]) -> NDArray[np.complex128]:\n",
    "    \"\"\" Perform Gram-Schmidt orthogonalization on the columns of a matrix\n",
    "    to define the unitary block to encode the MPS.\n",
    "\n",
    "    Notes\n",
    "    -----\n",
    "    If a column is (approximately) zero, it is replaced with a random vector.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    `matrix` : NDArray[np.complex128]\n",
    "        Input matrix with complex entries.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    `unitary` : NDArray[np.complex128]\n",
    "        A unitary matrix with orthonormal columns derived from the input matrix.\n",
    "        If a column is (approximately) zero, it is replaced with a random vector.\n",
    "    \"\"\"\n",
    "    num_rows, num_columns = matrix.shape\n",
    "    unitary = np.zeros((num_rows, num_columns), dtype=np.complex128)\n",
    "    orthonormal_basis: list[NDArray[np.complex128]] = []\n",
    "\n",
    "    for j in range(num_columns):\n",
    "        column = matrix[:, j]\n",
    "\n",
    "        # If column is (approximately) zero, replace with random\n",
    "        if np.allclose(column, 0):\n",
    "            column = np.random.uniform(-1, 1, num_rows)\n",
    "            if np.iscomplexobj(matrix):\n",
    "                column = column + 1j * np.random.uniform(-1, 1, num_rows)\n",
    "\n",
    "        # Gram-Schmidt orthogonalization\n",
    "        for basis_vector in orthonormal_basis:\n",
    "            column -= (basis_vector.conj().T @ column) * basis_vector\n",
    "\n",
    "        # Handle near-zero vectors (linear dependence)\n",
    "        norm = np.linalg.norm(column)\n",
    "        if norm < 1e-12:\n",
    "            is_complex = np.iscomplexobj(matrix)\n",
    "            column = np.random.uniform(-1, 1, num_rows)\n",
    "            if is_complex:\n",
    "                column += 1j * np.random.uniform(-1, 1, num_rows)\n",
    "            for basis_vector in orthonormal_basis:\n",
    "                column -= (basis_vector.conj().T @ column) * basis_vector\n",
    "\n",
    "        unitary[:, j] = column / np.linalg.norm(column)\n",
    "        orthonormal_basis.append(unitary[:, j])\n",
    "\n",
    "    return unitary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "436b6040",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sequential:\n",
    "    def __init__(self, target_fidelity) -> None:\n",
    "        \"\"\"Initialize the Sequential class.\n",
    "\n",
    "        Args:\n",
    "            target_fidelity (float): The target fidelity for the MPS encoding.\n",
    "        \"\"\"\n",
    "        self.target_fidelity = target_fidelity\n",
    "\n",
    "    def generate_layer(self, mps: qtn.MatrixProductState) -> list[tuple[list[int], NDArray[np.complex128]]]:\n",
    "        \"\"\" Convert a Matrix Product State (MPS) to a circuit representation\n",
    "        using a single unitary layer.\n",
    "\n",
    "        Args:\n",
    "            mps (qtn.MatrixProductState): The MPS to convert.\n",
    "\n",
    "        Returns:\n",
    "            unitary_layer (list[tuple[list[int], NDArray[np.complex128]]]): A list of\n",
    "            tuples representing the unitary layer of the circuit.\n",
    "                Each tuple contains:\n",
    "                - A list of qubit indices (in LSB order) that the unitary acts on.\n",
    "                - A unitary matrix (as a 2D NumPy array) that encodes the MPS.\n",
    "        \"\"\"\n",
    "        num_sites = mps.L\n",
    "\n",
    "        unitary_layer: list[tuple[list[int], NDArray[np.complex128]]] = []\n",
    "\n",
    "        for i, tensor in enumerate(reversed(mps.arrays)):\n",
    "            i = num_sites - i - 1\n",
    "\n",
    "            # MPS representation uses 1D entanglement, thus we need to define\n",
    "            # the range of the indices via the tensor shape\n",
    "            # i.e., if q0 and q3 are entangled, then regardless of q1 and q2 being\n",
    "            # entangled the entanglement range would be q0-q3\n",
    "            if i == 0:\n",
    "                d_right, d = tensor.shape\n",
    "                tensor = tensor.reshape((1, d_right, d))\n",
    "            if i == num_sites - 1:\n",
    "                d_left, d = tensor.shape\n",
    "                tensor = tensor.reshape((d_left, 1, d))\n",
    "\n",
    "            tensor = np.swapaxes(tensor, 1, 2)\n",
    "\n",
    "            # Combine the physical index and right-virtual index of the tensor to construct an isometry\n",
    "            # matrix\n",
    "            d_left, d, d_right = tensor.shape\n",
    "            isometry = tensor.reshape((d * d_left, d_right))\n",
    "\n",
    "            # Put qubit ordering in LSB (we provide this so that users can modify between LSB and MSB)\n",
    "            # To put into MSB, comment the second line below\n",
    "            qubits = reversed(range(i - int(np.ceil(np.log2(d_left))), i + 1))\n",
    "            qubits = [abs(qubit - num_sites + 1) for qubit in qubits] # type: ignore\n",
    "\n",
    "            # Create all-zero matrix and add the isometry columns\n",
    "            matrix = np.zeros((isometry.shape[0], isometry.shape[0]), dtype=isometry.dtype)\n",
    "\n",
    "            # Keep columns for which all ancillas are in the zero state\n",
    "            matrix[:, : isometry.shape[1]] = isometry\n",
    "\n",
    "            # Perform Gram-Schmidt orthogonalization to ensure the columns are orthonormal\n",
    "            unitary = gram_schmidt(matrix)\n",
    "\n",
    "            # Store the unitary layer for the circuit\n",
    "            unitary_layer.append((qubits, unitary)) # type: ignore\n",
    "\n",
    "        return unitary_layer\n",
    "\n",
    "    def mps_to_circuit_approx(\n",
    "            self,\n",
    "            statevector: NDArray[np.complex128],\n",
    "            max_num_layers: int,\n",
    "            chi_max: int\n",
    "        ) -> cudaq.Kernel:\n",
    "        r\"\"\" Approximately encodes the MPS into a circuit via multiple layers\n",
    "        of exact encoding of bond 2 truncated MPS.\n",
    "\n",
    "        Notes\n",
    "        -----\n",
    "        Whilst we can encode the MPS exactly in a single layer, we require\n",
    "        $log(chi) + 1$ qubits for each tensor, which results in larger circuits.\n",
    "        This function uses bond 2 which allows us to encode the MPS using one and\n",
    "        two qubit gates, which results in smaller circuits, and easier to run on\n",
    "        hardware.\n",
    "\n",
    "        This is the core idea of Ran's paper [1].\n",
    "\n",
    "        [1] https://arxiv.org/abs/1908.07958\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        `statevector` : NDArray[np.complex128]\n",
    "            The statevector to convert, which should be normalized.\n",
    "        `max_num_layers` : int\n",
    "            The maximum number of layers to use in the circuit.\n",
    "        `chi_max` : int\n",
    "            The maximum bond dimension for the MPS compression.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        `circuit` : cudaq.Kernel\n",
    "            The CUDA Quantum circuit that approximates the statevector.\n",
    "        \"\"\"\n",
    "        mps = qtn.MatrixProductState.from_dense(statevector)\n",
    "        mps: qtn.MatrixProductState = qtn.tensor_1d_compress.tensor_network_1d_compress(mps, max_bond=chi_max) # type: ignore\n",
    "        mps.permute_arrays()\n",
    "\n",
    "        mps.compress(form=\"left\", max_bond=chi_max)\n",
    "        mps.left_canonicalize(normalize=True)\n",
    "\n",
    "        compressed_mps = mps.copy(deep=True)\n",
    "        disentangled_mps = mps.copy(deep=True)\n",
    "\n",
    "        circuit = cudaq.make_kernel()\n",
    "        qr = circuit.qalloc(mps.L)\n",
    "\n",
    "        unitary_layers = []\n",
    "\n",
    "        # Initialize the zero state |00...0> to serve as comparison\n",
    "        # for the disentangled MPS\n",
    "        zero_state = np.zeros((2**mps.L,), dtype=np.complex128)\n",
    "        zero_state[0] = 1.0  # |00...0> state\n",
    "\n",
    "        # Ran's approach uses a iterative disentanglement of the MPS\n",
    "        # where each layer compresses the MPS to a maximum bond dimension of 2\n",
    "        # and applies the inverse of the layer to disentangle the MPS\n",
    "        # After a few layers we are adequately close to |00...0> state\n",
    "        # after which we can simply reverse the layers (no inverse) and apply them\n",
    "        # to the |00...0> state to obtain the MPS state\n",
    "        for layer_index in range(max_num_layers):\n",
    "            # Compress the MPS from the previous layer to a maximum bond dimension of 2,\n",
    "            # |ψ_k> -> |ψ'_k>\n",
    "            compressed_mps = disentangled_mps.copy(deep=True)\n",
    "\n",
    "            # Normalize improves fidelity of the encoding\n",
    "            compressed_mps.normalize()\n",
    "            compressed_mps.compress(form=\"left\", max_bond=2)\n",
    "\n",
    "            # Generate unitary layer\n",
    "            unitary_layer = self.generate_layer(compressed_mps)\n",
    "            unitary_layers.append(unitary_layer)\n",
    "\n",
    "            # To update the MPS definition, apply the inverse of U_k to disentangle |ψ_k>,\n",
    "            # |ψ_(k+1)> = inv(U_k) @ |ψ_k>\n",
    "            for i, _ in enumerate(unitary_layer):\n",
    "                inverse = unitary_layer[-(i + 1)][1].conj().T\n",
    "\n",
    "                if inverse.shape[0] == 4:\n",
    "                    disentangled_mps.gate_split_(\n",
    "                        inverse, (i - 1, i)\n",
    "                    )\n",
    "                else:\n",
    "                    disentangled_mps.gate_(\n",
    "                        inverse, (i), contract=True\n",
    "                    )\n",
    "\n",
    "            # Compress the disentangled MPS to a maximum bond dimension of chi_max\n",
    "            # This is to ensure that the disentangled MPS does not grow too large\n",
    "            # and improves the fidelity of the encoding\n",
    "            disentangled_mps: qtn.MatrixProductState = qtn.tensor_1d_compress.tensor_network_1d_compress( # type: ignore\n",
    "                disentangled_mps, max_bond=chi_max\n",
    "            )\n",
    "\n",
    "            fidelity = np.vdot(disentangled_mps.to_dense(), zero_state)\n",
    "\n",
    "            if fidelity >= self.target_fidelity:\n",
    "                # If the disentangled MPS is close enough to the zero state,\n",
    "                # we can stop the disentanglement process\n",
    "                print(f\"Reached target fidelity in {layer_index + 1} layers.\")\n",
    "                break\n",
    "\n",
    "        if layer_index == max_num_layers - 1:\n",
    "            print(f\"Reached {fidelity} with the maximum number of layers {max_num_layers}.\")\n",
    "\n",
    "        # The layers disentangle the MPS to a state close to |00...0>\n",
    "        # inv(U_k) ... inv(U_1) |ψ> = |00...0>\n",
    "        # So, we have to reverse the layers and apply them to the |00...0> state\n",
    "        # to obtain the MPS state\n",
    "        # |ψ> = U_1 ... U_k |00...0>\n",
    "        unitary_layers.reverse()\n",
    "\n",
    "        # Apply the unitary layers to the |00...0> state\n",
    "        gate_index = 0\n",
    "        for unitary_layer in unitary_layers:\n",
    "            for qubits, unitary in unitary_layer:\n",
    "                gate_name = f\"unitary_{gate_index}\"\n",
    "                cudaq.register_operation(gate_name, unitary)\n",
    "                qubits = [qr[qubit_index] for qubit_index in qubits]\n",
    "                qubits.reverse()\n",
    "                circuit.__getattr__(gate_name)(*qubits)\n",
    "                gate_index += 1\n",
    "\n",
    "        return circuit\n",
    "\n",
    "    def __call__(\n",
    "            self,\n",
    "            statevector: NDArray[np.complex128],\n",
    "            max_num_layers: int,\n",
    "            chi_max: int\n",
    "        ) -> cudaq.Kernel:\n",
    "        \"\"\" Call the instance to generate a circuit from the MPS.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        `statevector` : NDArray[np.complex128]\n",
    "            The statevector to convert.\n",
    "        `max_num_layers` : int\n",
    "            The maximum number of layers to use in the circuit.\n",
    "        `chi_max` : int\n",
    "            The maximum bond dimension for the MPS compression.\n",
    "\n",
    "        Returns:\n",
    "            cudaq.Kernel: The generated quantum circuit.\n",
    "        \"\"\"\n",
    "        return self.mps_to_circuit_approx(statevector, max_num_layers, chi_max)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "23edf921",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(1)\n",
    "\n",
    "def generate_random_state(num_qubits: int) -> NDArray[np.complex128]:\n",
    "    \"\"\" Generate a random normalized quantum state vector.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    `num_qubits` : int\n",
    "        The number of qubits in the state.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    `statevector` : NDArray[np.complex128]\n",
    "        A normalized state vector of shape (2**num_qubits,).\n",
    "    \"\"\"\n",
    "    statevector = np.random.rand(2**num_qubits) + 1j * np.random.rand(2**num_qubits)\n",
    "    return statevector / np.linalg.norm(statevector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d753b977",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reached target fidelity in 1 layers.\n",
      "(0.9930301077191865+2.657153153062808e-08j)\n"
     ]
    }
   ],
   "source": [
    "statevector = generate_random_state(4)\n",
    "\n",
    "encoder = Sequential(target_fidelity=0.99)\n",
    "qc = encoder(statevector, max_num_layers=3, chi_max=2**9)\n",
    "\n",
    "# Sanity check\n",
    "print(np.dot(\n",
    "    np.array(cudaq.get_state(qc)).conj().T, statevector\n",
    "))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a4e4472",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reached (0.9566448180001509-6.505556919443983e-18j) with the maximum number of layers 3.\n",
      "(0.956644869705398-2.680174415736758e-09j)\n"
     ]
    }
   ],
   "source": [
    "statevector = generate_random_state(6)\n",
    "\n",
    "encoder = Sequential(target_fidelity=0.99)\n",
    "qc = encoder(statevector, max_num_layers=3, chi_max=2**9)\n",
    "\n",
    "# Sanity check\n",
    "print(np.dot(\n",
    "    np.array(cudaq.get_state(qc)).conj().T, statevector\n",
    "))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
