{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "656bb283-ac6d-48d2-a029-3c417c9961f8",
   "metadata": {},
   "source": [
    "## Introduction to Quimb backend in QiboTN\n",
    "\n",
    "#### Some imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6722d94e-e311-48f9-b6df-c6d829bf67fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "# from scipy import stats\n",
    "\n",
    "# import qibo\n",
    "from qibo import Circuit, gates, hamiltonians\n",
    "from qibo.backends import construct_backend"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c5a8939",
   "metadata": {},
   "source": [
    "#### Some hyper parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64162116-1555-4a68-811c-01593739d622",
   "metadata": {},
   "outputs": [],
   "source": [
    "# construct qibotn backend\n",
    "quimb_backend = construct_backend(backend=\"qibotn\", platform=\"quimb\")\n",
    "\n",
    "# set number of qubits\n",
    "nqubits = 4\n",
    "\n",
    "# set numpy random seed\n",
    "np.random.seed(42)\n",
    "\n",
    "quimb_backend.setup_backend_specifics(quimb_backend=\"jax\", contractions_optimizer='auto-hq')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "926cfea5",
   "metadata": {},
   "source": [
    "Quimb accepts different methods for optimizing the way it does contractions, that we pass through \"contractions_optimizer\". \n",
    "We could also define our own cotengra contraction optimizer! \n",
    "\n",
    "cotengra is a Python library designed for **optimising contraction trees** and performing efficient contractions of large tensor‐networks.\n",
    "You can find it here: [https://github.com/jcmgray/cotengra](https://github.com/jcmgray/cotengra)\n",
    "\n",
    "For the sake of this tutorial however the default \"auto-hq\" will be fine :) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0a1da82",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cotengra as ctg\n",
    "ctg_opt = ctg.ReusableHyperOptimizer(\n",
    "    max_time=10,\n",
    "    minimize='combo',\n",
    "    slicing_opts=None,\n",
    "    parallel=True,\n",
    "    progbar=True\n",
    ")\n",
    "# quimb_backend.setup_backend_specifics(quimb_backend=\"jax\", contractions_optimizer='ctg_opt')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "252f5cd1-5932-4de6-8076-4a357d50ebad",
   "metadata": {},
   "source": [
    "#### Constructing a parametric quantum circuit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4a22a172-f50d-411d-afa3-fa61937c7b3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_circuit(nqubits, nlayers):\n",
    "    \"\"\"Construct a parametric quantum circuit.\"\"\"\n",
    "    circ = Circuit(nqubits)\n",
    "    for _ in range(nlayers):\n",
    "        for q in range(nqubits):\n",
    "            circ.add(gates.RY(q=q, theta=0.))\n",
    "            circ.add(gates.RZ(q=q, theta=0.))\n",
    "        [circ.add(gates.CNOT(q%nqubits, (q+1)%nqubits) for q in range(nqubits))]\n",
    "    circ.add(gates.M(*range(nqubits)))\n",
    "    return circ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "76f23c57-6d08-496b-9a27-52fb63bbfcb1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0: ─RY─RZ─o─────X─RY─RZ─o─────X─RY─RZ─o─────X─M─\n",
      "1: ─RY─RZ─X─o───|─RY─RZ─X─o───|─RY─RZ─X─o───|─M─\n",
      "2: ─RY─RZ───X─o─|─RY─RZ───X─o─|─RY─RZ───X─o─|─M─\n",
      "3: ─RY─RZ─────X─o─RY─RZ─────X─o─RY─RZ─────X─o─M─\n"
     ]
    }
   ],
   "source": [
    "circuit = build_circuit(nqubits=nqubits, nlayers=3)\n",
    "circuit.draw()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "07b2c097-cea2-42ec-8f1d-b4bbb5b71d98",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting random parameters\n",
    "circuit.set_parameters(\n",
    "    parameters=np.random.uniform(-np.pi, np.pi, len(circuit.get_parameters())),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd0cea52-03f5-4366-a01a-a5a84aa8ebc7",
   "metadata": {},
   "source": [
    "#### Setting up the tensor network simulator\n",
    "\n",
    "Depending on the simulator, various parameters can be set. One can customize the tensor network execution via the `backend.configure_tn_simulation` function, whose face depends on the specific backend provider."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2ee03e94-d794-4a51-9e76-01e8d8a259ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Customization of the tensor network simulation in the case of quimb backend\n",
    "# Here we use only some of the possible arguments\n",
    "quimb_backend.configure_tn_simulation(\n",
    "    #ansatz=\"MPS\",\n",
    "    max_bond_dimension=10\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "648d85b8-445d-4081-aeed-1691fbae67be",
   "metadata": {},
   "source": [
    "#### Executing through the backend\n",
    "\n",
    "The `backend.execute_circuit` method can be used then. We can simulate results in three ways:\n",
    "1. reconstruction of the final state only if `return_array` is set to `True`;\n",
    "2. computation of the relevant probabilities of the final state.\n",
    "3. reconstruction of the relevant state's frequencies (only if `nshots` is not `None`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "35a244c3-adba-4b8b-b28c-0ab592b0f7cf",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/andrea/python_envs/3.11/lib/python3.11/site-packages/quimb/tensor/circuit.py:215: SyntaxWarning: Unsupported operation ignored: creg\n",
      "  warnings.warn(\n",
      "/home/andrea/python_envs/3.11/lib/python3.11/site-packages/quimb/tensor/circuit.py:215: SyntaxWarning: Unsupported operation ignored: measure\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'nqubits': 4,\n",
       " 'backend': qibotn (quimb),\n",
       " 'measures': Counter({'1101': 14,\n",
       "          '1000': 12,\n",
       "          '0010': 11,\n",
       "          '0011': 11,\n",
       "          '0110': 9,\n",
       "          '0000': 8,\n",
       "          '1010': 7,\n",
       "          '1110': 6,\n",
       "          '0100': 5,\n",
       "          '1111': 5,\n",
       "          '1011': 5,\n",
       "          '0101': 4,\n",
       "          '0111': 1,\n",
       "          '0001': 1,\n",
       "          '1100': 1}),\n",
       " 'measured_probabilities': {'1101': np.float64(0.12331159869893284),\n",
       "  '1000': np.float64(0.11330883548333684),\n",
       "  '0010': np.float64(0.0946686048198943),\n",
       "  '0011': np.float64(0.07571277233522157),\n",
       "  '0110': np.float64(0.051460648073692314),\n",
       "  '0000': np.float64(0.08390937969317334),\n",
       "  '1010': np.float64(0.03872758515126775),\n",
       "  '1110': np.float64(0.07174919872960006),\n",
       "  '0100': np.float64(0.07142939529687146),\n",
       "  '1111': np.float64(0.10184806171791994),\n",
       "  '1011': np.float64(0.053499396925872716),\n",
       "  '0101': np.float64(0.05622305772698606),\n",
       "  '0111': np.float64(0.040291850747292815),\n",
       "  '0001': np.float64(0.004677011195208322),\n",
       "  '1100': np.float64(0.013605984872668443)},\n",
       " 'prob_type': 'default',\n",
       " 'statevector': Array([[ 0.08809626-0.27595j   ],\n",
       "        [-0.05174781+0.04471214j],\n",
       "        [ 0.00470146+0.30764672j],\n",
       "        [-0.27208942+0.04098931j],\n",
       "        [ 0.18807825+0.1898841j ],\n",
       "        [ 0.22377063+0.07842041j],\n",
       "        [-0.18900302+0.12545316j],\n",
       "        [ 0.17105258-0.10503745j],\n",
       "        [ 0.24859732-0.22695422j],\n",
       "        [-0.04117391-0.0623003j ],\n",
       "        [ 0.17371394-0.09247189j],\n",
       "        [-0.22748126+0.04185291j],\n",
       "        [ 0.09444097+0.06846087j],\n",
       "        [-0.21784975-0.2754144j ],\n",
       "        [-0.17359754+0.20399287j],\n",
       "        [-0.01729751-0.31866732j]], dtype=complex64)}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# # Simple execution (defaults)\n",
    "outcome = quimb_backend.execute_circuit(circuit=circuit, nshots=100, return_array=True)\n",
    "\n",
    "# # Print outcome\n",
    "vars(outcome)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84ec0b48-f6b4-495c-93b8-8e42d1a8b0df",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "One can access to the specific contents of the simulation outcome."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c0443efc-21ef-4ed5-9cf4-785d204a1881",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Probabilities:\n",
      " {'1101': np.float64(0.12331159869893284), '1000': np.float64(0.11330883548333684), '0010': np.float64(0.0946686048198943), '0011': np.float64(0.07571277233522157), '0110': np.float64(0.051460648073692314), '0000': np.float64(0.08390937969317334), '1010': np.float64(0.03872758515126775), '1110': np.float64(0.07174919872960006), '0100': np.float64(0.07142939529687146), '1111': np.float64(0.10184806171791994), '1011': np.float64(0.053499396925872716), '0101': np.float64(0.05622305772698606), '0111': np.float64(0.040291850747292815), '0001': np.float64(0.004677011195208322), '1100': np.float64(0.013605984872668443)}\n",
      "\n",
      "State:\n",
      " [[ 0.08809626-0.27595j   ]\n",
      " [-0.05174781+0.04471214j]\n",
      " [ 0.00470146+0.30764672j]\n",
      " [-0.27208942+0.04098931j]\n",
      " [ 0.18807825+0.1898841j ]\n",
      " [ 0.22377063+0.07842041j]\n",
      " [-0.18900302+0.12545316j]\n",
      " [ 0.17105258-0.10503745j]\n",
      " [ 0.24859732-0.22695422j]\n",
      " [-0.04117391-0.0623003j ]\n",
      " [ 0.17371394-0.09247189j]\n",
      " [-0.22748126+0.04185291j]\n",
      " [ 0.09444097+0.06846087j]\n",
      " [-0.21784975-0.2754144j ]\n",
      " [-0.17359754+0.20399287j]\n",
      " [-0.01729751-0.31866732j]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(f\"Probabilities:\\n {outcome.probabilities()}\\n\")\n",
    "print(f\"State:\\n {outcome.state()}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9531f9d6",
   "metadata": {},
   "source": [
    "### Compute expectation values\n",
    "\n",
    "Another important feature of this backend is the `expectation` function. In fact, we can compute expectation values of given observables thorugh a Qibo-friendly interface.\n",
    "\n",
    "---\n",
    "\n",
    "Let's start by importing some symbols, thanks to which we can build our observable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "647f2073",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import jax\n",
    "from qibo.backends import construct_backend\n",
    "from qibo import Circuit, gates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "74c63a41",
   "metadata": {},
   "outputs": [],
   "source": [
    "# construct qibotn backend\n",
    "quimb_backend = construct_backend(backend=\"qibotn\", platform=\"quimb\")\n",
    "\n",
    "quimb_backend.setup_backend_specifics(\n",
    "    quimb_backend    =\"jax\", \n",
    "    contractions_optimizer='auto-hq'\n",
    "    )\n",
    "\n",
    "quimb_backend.configure_tn_simulation(\n",
    "    max_bond_dimension=10\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "b2a0decb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qibo.symbols import X, Z, Y\n",
    "from qibo.hamiltonians import XXZ\n",
    "\n",
    "# define Hamiltonian\n",
    "hamiltonian = XXZ(4, dense=False, backend=quimb_backend)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "bd734be8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define circuit\n",
    "def build_circuit(nqubits, nlayers):\n",
    "    circ = Circuit(nqubits)\n",
    "    for layer in range(nlayers):\n",
    "        for q in range(nqubits):\n",
    "            circ.add(gates.RY(q=q, theta=0.))\n",
    "            circ.add(gates.RZ(q=q, theta=0.))\n",
    "            circ.add(gates.RX(q=q, theta=0.))\n",
    "        for q in range(nqubits - 1):\n",
    "            circ.add(gates.CNOT(q, q + 1))\n",
    "            circ.add(gates.SWAP(q, q + 1))\n",
    "    circ.add(gates.M(*range(nqubits)))\n",
    "    return circ\n",
    "\n",
    "def build_circuit_problematic(nqubits, nlayers):\n",
    "    circ = Circuit(nqubits)\n",
    "    for _ in range(nlayers):\n",
    "        for q in range(nqubits):\n",
    "            circ.add(gates.RY(q=q, theta=0.))\n",
    "            circ.add(gates.RZ(q=q, theta=0.))\n",
    "        [circ.add(gates.CNOT(q%nqubits, (q+1)%nqubits) for q in range(nqubits))]\n",
    "    circ.add(gates.M(*range(nqubits)))\n",
    "    return circ\n",
    "\n",
    "\n",
    "nqubits = 4\n",
    "circuit = build_circuit(nqubits=nqubits, nlayers=3)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "fe63ff24",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Expectation value: 2.0\n",
      "Elapsed time: 0.0268 seconds\n"
     ]
    }
   ],
   "source": [
    "start = time.time()\n",
    "expval = hamiltonian.expectation(circuit)\n",
    "\n",
    "elapsed = time.time() - start\n",
    "print(f\"Expectation value: {expval}\")\n",
    "print(f\"Elapsed time: {elapsed:.4f} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d976a849",
   "metadata": {},
   "source": [
    "Try with Qibo (which is by default using the Qibojit backend)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "fb1436c8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Qibo 0.2.21|INFO|2025-10-27 16:24:00]: Using numpy backend on /CPU:0\n",
      "WARNING:root:Calculation of expectation values starting from the state is deprecated, use the ``expectation_from_state`` method if you really need it, or simply pass the circuit you want to calculate the expectation value from.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Expectation value: 2.0\n",
      "Elapsed time: 0.0360 seconds\n"
     ]
    }
   ],
   "source": [
    "sym_hamiltonian = XXZ(4, dense=False, backend=None)\n",
    "\n",
    "#  Let's show it\n",
    "sym_hamiltonian.form\n",
    "\n",
    "# Compute expectation value\n",
    "start = time.time()\n",
    "result = sym_hamiltonian.expectation(circuit().state())\n",
    "elapsed = time.time() - start\n",
    "print(f\"Expectation value: {result}\")\n",
    "print(f\"Elapsed time: {elapsed:.4f} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77bef077",
   "metadata": {},
   "source": [
    "They match! 🥳"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50130ae6",
   "metadata": {},
   "source": [
    "We can also compute gradient of expectation function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "6a3b26e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/andrea/python_envs/3.11/lib/python3.11/site-packages/quimb/tensor/circuit.py:4927: UserWarning: Unsupported options for computing local_expectation with an MPS circuit supplied, ignoring: R, None, None, jax, None\n",
      "  warnings.warn(\n",
      "/home/andrea/python_envs/3.11/lib/python3.11/site-packages/quimb/tensor/circuit.py:4927: UserWarning: Unsupported options for computing local_expectation with an MPS circuit supplied, ignoring: R, None, None, jax, None\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.24630009  0.8370421  -0.11103702 -0.12855841  0.41325414 -0.0628037\n",
      "  0.51638705  0.794163   -0.27972788 -1.0718998   0.02731732  1.0153619\n",
      " -0.34494495  1.5744264   0.26920277 -0.36333832  0.12331417  0.5196531\n",
      "  1.1294655   0.29257926 -0.18237355  0.8914014  -0.9471657   0.3492473\n",
      " -0.3477673   0.24325958  0.04818404 -0.87983793  0.47196424  0.36605012\n",
      "  1.005       0.65054715 -0.94860053  0.14459445  0.36571163 -0.2550101 ]\n"
     ]
    }
   ],
   "source": [
    "def f(circuit, hamiltonian, params):\n",
    "    circuit.set_parameters(params)\n",
    "    return hamiltonian.expectation(\n",
    "        circuit=circuit,\n",
    "    )\n",
    "\n",
    "parameters = np.random.uniform(-np.pi, np.pi, size=len(circuit.get_parameters()))\n",
    "print(jax.grad(f, argnums=2)(circuit, hamiltonian, parameters))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeafa5a6-2afa-429c-a101-effa84bac1d2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}