Merge pull request #112 from mattia-robbiano/main

Quimb backend: refactor and implementation of expectation function

examples/quimb_intro/quimb_introduction.ipynb

@@ -0,0 +1,543 @@
+{
+ "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": "a009a5e0-cfd4-4a49-9f7c-e82f252c6147",
+   "metadata": {},
+   "source": [
+    "#### Some hyper parameters"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "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"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "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(qimb_backend=\"jax\")"
+   ]
+  },
+  {
+   "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/mattia/main_env/lib/python3.12/site-packages/quimb/tensor/circuit.py:215: SyntaxWarning: Unsupported operation ignored: creg\n",
+      "  warnings.warn(\n",
+      "/home/mattia/main_env/lib/python3.12/site-packages/quimb/tensor/circuit.py:215: SyntaxWarning: Unsupported operation ignored: measure\n",
+      "  warnings.warn(\n",
+      "/home/mattia/main_env/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "{'nqubits': 4,\n",
+       " 'backend': qibotn (quimb),\n",
+       " 'measures': Counter({'1010': 9,\n",
+       "          '0100': 8,\n",
+       "          '1101': 15,\n",
+       "          '1011': 4,\n",
+       "          '1111': 12,\n",
+       "          '1000': 13,\n",
+       "          '0000': 8,\n",
+       "          '0010': 6,\n",
+       "          '0011': 6,\n",
+       "          '0101': 8,\n",
+       "          '1110': 5,\n",
+       "          '0110': 5,\n",
+       "          '0111': 1}),\n",
+       " 'measured_probabilities': {'1101': np.float64(0.12331159869893256),\n",
+       "  '1000': np.float64(0.11330883548333587),\n",
+       "  '1111': np.float64(0.10184806171791962),\n",
+       "  '1010': np.float64(0.03872758515126756),\n",
+       "  '0100': np.float64(0.07142939529687138),\n",
+       "  '0000': np.float64(0.08390937969317269),\n",
+       "  '0101': np.float64(0.05622305772698622),\n",
+       "  '0010': np.float64(0.09466860481989385),\n",
+       "  '0011': np.float64(0.07571277233522114),\n",
+       "  '1110': np.float64(0.07174919872959985),\n",
+       "  '0110': np.float64(0.05146064807369214),\n",
+       "  '1011': np.float64(0.053499396925872744),\n",
+       "  '0111': np.float64(0.04029185074729259)},\n",
+       " 'prob_type': 'default',\n",
+       " 'statevector': Array([[ 0.08809624-0.27594998j],\n",
+       "        [-0.05174781+0.04471217j],\n",
+       "        [ 0.00470147+0.30764672j],\n",
+       "        [-0.27208942+0.0409893j ],\n",
+       "        [ 0.18807822+0.18988408j],\n",
+       "        [ 0.2237706 +0.07842042j],\n",
+       "        [-0.18900308+0.12545314j],\n",
+       "        [ 0.17105256-0.10503749j],\n",
+       "        [ 0.24859734-0.22695419j],\n",
+       "        [-0.0411739 -0.06230037j],\n",
+       "        [ 0.17371392-0.09247189j],\n",
+       "        [-0.22748128+0.0418529j ],\n",
+       "        [ 0.09444095+0.06846087j],\n",
+       "        [-0.21784972-0.2754144j ],\n",
+       "        [-0.17359753+0.20399286j],\n",
+       "        [-0.01729754-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.12331159869893256), '1000': np.float64(0.11330883548333587), '1111': np.float64(0.10184806171791962), '1010': np.float64(0.03872758515126756), '0100': np.float64(0.07142939529687138), '0000': np.float64(0.08390937969317269), '0101': np.float64(0.05622305772698622), '0010': np.float64(0.09466860481989385), '0011': np.float64(0.07571277233522114), '1110': np.float64(0.07174919872959985), '0110': np.float64(0.05146064807369214), '1011': np.float64(0.053499396925872744), '0111': np.float64(0.04029185074729259)}\n",
+      "\n",
+      "State:\n",
+      " [[ 0.08809624-0.27594998j]\n",
+      " [-0.05174781+0.04471217j]\n",
+      " [ 0.00470147+0.30764672j]\n",
+      " [-0.27208942+0.0409893j ]\n",
+      " [ 0.18807822+0.18988408j]\n",
+      " [ 0.2237706 +0.07842042j]\n",
+      " [-0.18900308+0.12545314j]\n",
+      " [ 0.17105256-0.10503749j]\n",
+      " [ 0.24859734-0.22695419j]\n",
+      " [-0.0411739 -0.06230037j]\n",
+      " [ 0.17371392-0.09247189j]\n",
+      " [-0.22748128+0.0418529j ]\n",
+      " [ 0.09444095+0.06846087j]\n",
+      " [-0.21784972-0.2754144j ]\n",
+      " [-0.17359753+0.20399286j]\n",
+      " [-0.01729754-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",
+    "    qimb_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": 19,
+   "id": "b2a0decb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# define Hamiltonian\n",
+    "operators = [\"xzy\", \"yxzy\", \"zy\"]\n",
+    "qubits = [\"011\", \"0112\", \"01\"]\n",
+    "coefficients = [\"1\", \"2\", \"j\"]\n",
+    "hamiltonian = (operators, qubits, coefficients)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "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: 0.0\n",
+      "Elapsed time: 0.1071 seconds\n"
+     ]
+    }
+   ],
+   "source": [
+    "start = time.time()\n",
+    "expval = quimb_backend.expectation(\n",
+    "        circuit=circuit,\n",
+    "        operators_list=hamiltonian[0],\n",
+    "        sites_list=hamiltonian[1],\n",
+    "        coeffs_list=hamiltonian[2]\n",
+    "    )\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": 22,
+   "id": "fb1436c8",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Expectation value: 1.5\n",
+      "Elapsed time: 0.0501 seconds\n"
+     ]
+    }
+   ],
+   "source": [
+    "from qibo.symbols import Z, X, I\n",
+    "# We can create a symbolic Hamiltonian\n",
+    "form = 0.5 * Z(0) * Z(1) +- 1.5 *  X(0) * Z(2) + Z(3)\n",
+    "sym_hamiltonian = hamiltonians.SymbolicHamiltonian(form)\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": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[ 8.19939339e-10 -3.14190913e-08 -2.99498648e-09 -1.03641796e-07\n",
+      "  8.48652704e-10  1.00297093e-07 -6.75429277e-08 -9.78565140e-09\n",
+      " -5.11915417e-08  1.29225235e-08 -7.44280655e-08 -3.49115048e-08\n",
+      " -4.98508879e-09  6.80729357e-08 -3.29755920e-08  4.20008526e-08\n",
+      " -2.89742630e-08  1.18602941e-07 -2.88252178e-08  5.57985391e-09\n",
+      " -3.17434115e-08 -1.03342952e-08  1.34079716e-08 -7.05437886e-09\n",
+      " -4.34059650e-08 -2.18019203e-08 -5.36932561e-08 -6.38544009e-08\n",
+      "  5.85312279e-08  8.45709067e-08 -1.12777876e-09 -6.41545981e-08\n",
+      "  7.25317406e-08  4.10035668e-08 -1.29046382e-08  6.07501676e-08]\n"
+     ]
+    }
+   ],
+   "source": [
+    "def f(circuit, hamiltonian, params):\n",
+    "    circuit.set_parameters(params)\n",
+    "    return quimb_backend.expectation(\n",
+    "        circuit=circuit,\n",
+    "        operators_list=hamiltonian[0],\n",
+    "        sites_list=hamiltonian[1],\n",
+    "        coeffs_list=hamiltonian[2]\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"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "main_env",
+   "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.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

pyproject.toml

@@ -21,16 +21,18 @@ packages = [{ include = "qibotn", from = "src" }]
 
 [tool.poetry.dependencies]
 python = ">=3.11,<3.14"
-qibo = "^0.2.17"
+qibo = { git="https://github.com/qiboteam/qibo", branch="expectation"}
 quimb = { version = "^1.10.0", extras = ["tensor"] }
 cupy-cuda11x = { version = "^13.1.0", optional = true }
 cuquantum-python-cu11 = { version = "^24.1.0", optional = true }
 qmatchatea = { version = "^1.4.3", optional = true }
+qiskit = { version = "^1.4.0", optional = true }
+qtealeaves = { version = "^1.5.20", optional = true }
 
 
 [tool.poetry.extras]
 cuda = ["cupy-cuda11x", "cuquantum-python-cu11", "mpi4py"]
-qmatchatea = ["qmatchatea"]
+qmatchatea = ["qmatchatea", "qtealeaves", "qiskit"]
 
 [tool.poetry.group.docs]
 optional = true

src/qibotn/backends/__init__.py

@@ -4,7 +4,7 @@ from qibo.config import raise_error
 
 from qibotn.backends.abstract import QibotnBackend
 from qibotn.backends.cutensornet import CuTensorNet  # pylint: disable=E0401
-from qibotn.backends.quimb import QuimbBackend  # pylint: disable=E0401
+from qibotn.backends.quimb import QuimbBackend
 
 PLATFORMS = ("cutensornet", "qutensornet", "qmatchatea")
 
@@ -13,7 +13,7 @@ class MetaBackend:
     """Meta-backend class which takes care of loading the qibotn backends."""
 
     @staticmethod
-    def load(platform: str, runcard: dict = None) -> QibotnBackend:
+    def load(platform: str, runcard: dict = None, **kwargs) -> QibotnBackend:
         """Loads the backend.
 
         Args:
@@ -26,7 +26,11 @@ class MetaBackend:
         if platform == "cutensornet":  # pragma: no cover
             return CuTensorNet(runcard)
         elif platform == "quimb":  # pragma: no cover
-            return QuimbBackend(runcard)
+            quimb_backend = kwargs.get("quimb_backend", "numpy")
+            contraction_optimizer = kwargs.get("contraction_optimizer", "auto-hq")
+            return QuimbBackend(
+                quimb_backend=quimb_backend, contraction_optimizer=contraction_optimizer
+            )
         elif platform == "qmatchatea":  # pragma: no cover
             from qibotn.backends.qmatchatea import QMatchaTeaBackend
 

src/qibotn/backends/qmatchatea.py

@@ -23,6 +23,9 @@ class QMatchaTeaBackend(QibotnBackend, NumpyBackend):
         self.name = "qibotn"
         self.platform = "qmatchatea"
 
+        # Default precision
+        self.precision = "double"
+
         # Set default configurations
         self.configure_tn_simulation()
         self._setup_backend_specifics()
@@ -87,7 +90,6 @@ class QMatchaTeaBackend(QibotnBackend, NumpyBackend):
 
         # TODO: once MPI is available for Python, integrate it here
         self.qmatchatea_backend = qmatchatea.QCBackend(
-            backend="PY",  # The only alternative is Fortran, but we use Python here
             precision=qmatchatea_precision,
             device=qmatchatea_device,
             ansatz=self.ansatz,

src/qibotn/backends/quimb.py

@@ -1,29 +1,67 @@
 from collections import Counter
+from typing import Optional
 
+import quimb as qu
 import quimb.tensor as qtn
-from qibo.backends import NumpyBackend
 from qibo.config import raise_error
-from qibo.result import QuantumState
+from qibo.gates.abstract import ParametrizedGate
+from qibo.models import Circuit
 
 from qibotn.backends.abstract import QibotnBackend
 from qibotn.result import TensorNetworkResult
 
+GATE_MAP = {
+    "h": "H",
+    "x": "X",
+    "y": "Y",
+    "z": "Z",
+    "s": "S",
+    "t": "T",
+    "rx": "RX",
+    "ry": "RY",
+    "rz": "RZ",
+    "u3": "U3",  # TODO: check
+    "cx": "CX",
+    "cnot": "CNOT",
+    "cy": "CY",
+    "cz": "CZ",
+    "iswap": "ISWAP",
+    "swap": "SWAP",
+    "ccx": "CCX",
+    "ccy": "CCY",
+    "ccz": "CCZ",
+    "toffoli": "TOFFOLI",
+    "cswap": "CSWAP",
+    "fredkin": "FREDKIN",
+    "fsim": "fsim",
+    "measure": "measure",
+}
 
-class QuimbBackend(QibotnBackend, NumpyBackend):
 
-    def __init__(self):
+if not __name__ == "__main__":
-        super().__init__()
+
+    def __init__(self, quimb_backend="numpy", contraction_optimizer="auto-hq"):
+        super(self.__class__, self).__init__()
 
         self.name = "qibotn"
         self.platform = "quimb"
+        self.backend = quimb_backend
+
+        self.ansatz = None
+        self.max_bond_dimension = None
+        self.svd_cutoff = None
+        self.n_most_frequent_states = None
 
         self.configure_tn_simulation()
-        self.setup_backend_specifics()
+        self.setup_backend_specifics(
+            quimb_backend=quimb_backend, contractions_optimizer=contraction_optimizer
+        )
 
     def configure_tn_simulation(
         self,
-        ansatz: str = "MPS",
+        ansatz: str = "mps",
-        max_bond_dimension: int = 10,
+        max_bond_dimension: Optional[int] = None,
+        svd_cutoff: Optional[float] = 1e-10,
         n_most_frequent_states: int = 100,
     ):
         """
@@ -31,7 +69,8 @@ class QuimbBackend(QibotnBackend, NumpyBackend):
 
         Args:
             ansatz : str, optional
-                The tensor network ansatz to use. Currently, only "MPS" is supported. Default is "MPS".
+                The tensor network ansatz to use. Default is `None` and, in this case, a
+                generic Circuit Quimb class is used.
                 max_bond_dimension : int, optional
                 The maximum bond dimension for the MPS ansatz. Default is 10.
 
@@ -41,23 +80,50 @@ class QuimbBackend(QibotnBackend, NumpyBackend):
         """
         self.ansatz = ansatz
         self.max_bond_dimension = max_bond_dimension
+        self.svd_cutoff = svd_cutoff
         self.n_most_frequent_states = n_most_frequent_states
 
-    def setup_backend_specifics(self, qimb_backend="numpy"):
+    @property
+    def circuit_ansatz(self):
+        if self.ansatz == "mps":
+            return qtn.CircuitMPS
+        return qtn.Circuit
+
+    def setup_backend_specifics(
+        self, quimb_backend="numpy", contractions_optimizer="auto-hq"
+    ):
         """Setup backend specifics.
         Args:
             qimb_backend: str
                 The backend to use for the quimb tensor network simulation.
+            contractions_optimizer: str, optional
+                The contractions_optimizer to use for the quimb tensor network simulation.
         """
-        self.backend = qimb_backend
+        # this is not really working because it does not change the inheritance
+        if quimb_backend == "jax":
+            import jax.numpy as jnp
+
+            self.np = jnp
+        elif quimb_backend == "numpy":
+            import numpy as np
+
+            self.np = np
+        elif quimb_backend == "torch":
+            import torch
+
+            self.np = torch
+        else:
+            raise_error(ValueError, f"Unsupported quimb backend: {quimb_backend}")
+
+        self.backend = quimb_backend
+        self.contractions_optimizer = contractions_optimizer
 
     def execute_circuit(
         self,
-        circuit,
+        circuit: Circuit,
         initial_state=None,
         nshots=None,
         return_array=False,
-        **prob_kwargs,
     ):
         """
         Execute a quantum circuit using the specified tensor network ansatz and initial state.
@@ -71,10 +137,6 @@ class QuimbBackend(QibotnBackend, NumpyBackend):
                 The number of shots for sampling the circuit. If None, no sampling is performed, and the full statevector is used.
             return_array : bool, optional
                 If True, returns the statevector as a dense array. Default is False.
-            n_most_frequent_states : int, optional
-                The number of most frequent computational basis states to return. Default is 100.
-            **prob_kwargs : dict, optional
-                Additional keyword arguments for probability computation (currently unused).
 
         Returns:
             TensorNetworkResult
@@ -95,7 +157,6 @@ class QuimbBackend(QibotnBackend, NumpyBackend):
             - If `initial_state` is provided, it must be compatible with the MPS ansatz.
             - The `nshots` parameter enables sampling from the circuit's output distribution. If not specified, the full statevector is computed.
         """
-
         if initial_state is not None and self.ansatz == "MPS":
             initial_state = qtn.tensor_1d.MatrixProductState.from_dense(
                 initial_state, 2
@@ -105,27 +166,34 @@ class QuimbBackend(QibotnBackend, NumpyBackend):
                 ValueError, "Initial state not None supported only for MPS ansatz."
             )
 
-        circ_ansatz = (
+        circ_quimb = self.circuit_ansatz.from_openqasm2_str(
-            qtn.circuit.CircuitMPS if self.ansatz == "MPS" else qtn.circuit.Circuit
-        )
-        circ_quimb = circ_ansatz.from_openqasm2_str(
             circuit.to_qasm(), psi0=initial_state
         )
 
-        frequencies = Counter(circ_quimb.sample(nshots)) if nshots is not None else None
+        if nshots:
+            frequencies = Counter(circ_quimb.sample(nshots))
             main_frequencies = {
                 state: count
                 for state, count in frequencies.most_common(self.n_most_frequent_states)
             }
-        computational_states = [state for state in main_frequencies.keys()]
+            computational_states = list(main_frequencies.keys())
             amplitudes = {
                 state: circ_quimb.amplitude(state) for state in computational_states
             }
             measured_probabilities = {
                 state: abs(amplitude) ** 2 for state, amplitude in amplitudes.items()
             }
+        else:
+            frequencies = None
+            measured_probabilities = None
 
-        statevector = circ_quimb.to_dense() if return_array else None
+        statevector = (
+            circ_quimb.to_dense(
+                backend=self.backend, optimize=self.contractions_optimizer
+            )
+            if return_array
+            else None
+        )
         return TensorNetworkResult(
             nqubits=circuit.nqubits,
             backend=self,
@@ -134,3 +202,151 @@ class QuimbBackend(QibotnBackend, NumpyBackend):
             prob_type="default",
             statevector=statevector,
         )
+
+    def expectation_observable_symbolic_from_state(
+        self, circuit, operators_list, sites_list, coeffs_list, nqubits
+    ):
+        """
+        Compute the expectation value of a symbolic Hamiltonian on a quantum circuit using tensor network contraction.
+        This method takes a Qibo circuit, converts it to a Quimb tensor network circuit, and evaluates the expectation value
+        of a Hamiltonian specified by three lists of strings: operators, sites, and coefficients.
+        The expectation value is computed by summing the contributions from each term in the Hamiltonian, where each term's
+        expectation is calculated using Quimb's `local_expectation` function.
+
+        Parameters
+        ----------
+        circuit : qibo.models.Circuit
+            The quantum circuit to evaluate, provided as a Qibo circuit object.
+        operators_list : list of str
+            List of operator strings representing the symbolic Hamiltonian terms.
+        sites_list : list of str
+            List of strings, each specifying the qubits (sites) the corresponding operator acts on.
+        coeffs_list : list of str
+            List of strings representing the coefficients for each Hamiltonian term.
+        Returns
+        -------
+        float
+            The real part of the expectation value of the Hamiltonian on the given circuit state.
+        """
+        quimb_circuit = self._qibo_circuit_to_quimb(
+            circuit,
+            quimb_circuit_type=self.circuit_ansatz,
+            gate_opts={"max_bond": self.max_bond_dimension, "cutoff": self.svd_cutoff},
+        )
+
+        expectation_value = 0.0
+        for opstr, sites, coeff in zip(operators_list, sites_list, coeffs_list):
+
+            ops = self._string_to_quimb_operator(opstr)
+            coeff = coeff.real
+
+            exp_values = quimb_circuit.local_expectation(
+                ops,
+                where=sites,
+                backend=self.backend,
+                optimize=self.contractions_optimizer,
+                simplify_sequence="R",
+            )
+
+            expectation_value = expectation_value + coeff * exp_values
+
+        return self.np.real(expectation_value)
+
+    def _qibo_circuit_to_quimb(
+        self, qibo_circ, quimb_circuit_type=qtn.Circuit, **circuit_kwargs
+    ):
+        """
+        Convert a Qibo Circuit to a Quimb Circuit. Measurement gates are ignored. If are given gates not supported by Quimb, an error is raised.
+
+        Parameters
+        ----------
+        qibo_circ : qibo.models.circuit.Circuit
+            The circuit to convert.
+        quimb_circuit_type : type
+            The Quimb circuit class to use (Circuit, CircuitMPS, etc).
+        circuit_kwargs : dict
+            Extra arguments to pass to the Quimb circuit constructor.
+
+        Returns
+        -------
+        circ : quimb.tensor.circuit.Circuit
+            The converted circuit.
+        """
+        nqubits = qibo_circ.nqubits
+        circ = quimb_circuit_type(nqubits, **circuit_kwargs)
+
+        for gate in qibo_circ.queue:
+            gname = getattr(gate, "name", None)
+            qname = GATE_MAP.get(gname, None)
+            if qname == "measure":
+                continue
+            if qname is None:
+                raise_error(ValueError, f"Gate {gname} not supported in Quimb backend.")
+
+            params = getattr(gate, "parameters", ())
+            qubits = getattr(gate, "qubits", ())
+
+            is_parametrized = isinstance(gate, ParametrizedGate) and getattr(
+                gate, "trainable", True
+            )
+            if is_parametrized:
+                circ.apply_gate(qname, *params, *qubits, parametrized=is_parametrized)
+            else:
+                circ.apply_gate(
+                    qname,
+                    *params,
+                    *qubits,
+                )
+        return circ
+
+    def _string_to_quimb_operator(self, op_str):
+        """
+        Convert a Pauli string (e.g. 'xzy') to a Quimb operator using '&' chaining.
+
+        Parameters
+        ----------
+        op_str : str
+            A string like 'xzy', where each character is one of 'x', 'y', 'z', 'i'.
+
+        Returns
+        -------
+        qu_op : quimb.Qarray
+            The corresponding Quimb operator.
+        """
+        op_str = op_str.lower()
+        op = qu.pauli(op_str[0])
+        for c in op_str[1:]:
+            op = op & qu.pauli(c)
+        return op
+
+
+def QuimbBackend(
+    quimb_backend: str = "numpy", contraction_optimizer="auto-hq"
+) -> QibotnBackend:
+    bases = (QibotnBackend,)
+    methods = {
+        "__init__": __init__,
+        "configure_tn_simulation": configure_tn_simulation,
+        "setup_backend_specifics": setup_backend_specifics,
+        "execute_circuit": execute_circuit,
+        "expectation_observable_symbolic_from_state": expectation_observable_symbolic_from_state,
+        "_qibo_circuit_to_quimb": _qibo_circuit_to_quimb,
+        "_string_to_quimb_operator": _string_to_quimb_operator,
+        "circuit_ansatz": circuit_ansatz,
+    }
+    if quimb_backend == "numpy":
+        from qibo.backends import NumpyBackend
+
+        bases += (NumpyBackend,)
+    elif quimb_backend == "torch":
+        from qiboml.backends import PyTorchBackend
+
+        bases += (PyTorchBackend,)
+    elif quimb_backend == "jax":
+        from qiboml.backends import JaxBackend
+
+        bases += (JaxBackend,)
+    else:
+        raise_error(ValueError, f"Unsupported quimb backend: {quimb_backend}")
+
+    return type("QuimbBackend", bases, methods)(quimb_backend, contraction_optimizer)