{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3ae160d4",
   "metadata": {},
   "source": [
    "# The usage of contractor\n",
    "\n",
    "## Overview\n",
    "\n",
    "In this tutorial, we will demonstrate how to utilize different types of TensorNetwork contractors for the circuit simulation to achieve a better space-time tradeoff. The customization of the contractor is the main highlight for the TensorCircuit package since a better contractor can make better use of the power of the TensorNetwok simulation engine. [WIP]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e689831",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "770fe50e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorcircuit as tc\n",
    "import numpy as np\n",
    "import cotengra as ctg\n",
    "import opt_einsum as oem\n",
    "\n",
    "K = tc.set_backend(\"tensorflow\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "781a1e8d",
   "metadata": {},
   "source": [
    "## Testbed System\n",
    "\n",
    "We provide tensor networks for two circuits, and test the contraction efficiency for these two systems, the first system is small while the second one is large."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ef09ad07",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get state for small system\n",
    "def small_tn():\n",
    "    n = 10\n",
    "    d = 4\n",
    "    param = K.ones([2 * d, n])\n",
    "    c = tc.Circuit(n)\n",
    "    c = tc.templates.blocks.example_block(c, param, nlayers=d)\n",
    "    return c.state()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d80c3252",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get expectation for extra large system\n",
    "def large_tn():\n",
    "    n = 60\n",
    "    d = 8\n",
    "    param = K.ones([2 * d, n])\n",
    "    c = tc.Circuit(n)\n",
    "    c = tc.templates.blocks.example_block(c, param, nlayers=d, is_split=True)\n",
    "    # the two qubit gate is split and truncated with SVD decomposition\n",
    "    return c.expectation([tc.gates.z(), [n // 2]], reuse=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bd06739",
   "metadata": {},
   "source": [
    "## Opt-einsum\n",
    "\n",
    "There are several contractor optimizers provided by opt-einsum and shipped with the TensorNetwork package. Since TensorCircuit is built on top of TensorNetwork, we can use these simple contractor optimizers. Though for any moderate system, only greedy optimizer works, other optimizers come with exponential scaling and fail in circuit simulation scenarios.\n",
    "\n",
    "We always set ``contraction_info=True`` (default ``False``) for the contractor system, which will print contraction information summary including contraction size, flops, and writes. For the definition of these metrics, also refer to cotengra docs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f35a1974",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "functools.partial(<function custom at 0x7fd5a0a3d430>, optimizer=<function contraction_info_decorator.<locals>.new_algorithm at 0x7fd588e281f0>, memory_limit=None, debug_level=2)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# if we set nothing, the default optimizer is greedy, i.e.\n",
    "tc.set_contractor(\"greedy\", debug_level=2, contraction_info=True)\n",
    "# We set debug_level=2 to not really run the contraction computation\n",
    "# i.e. by set debug_level>0, only contraction information and the return shape is correct, the result is wrong"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "25180630",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------ contraction cost summary ------\n",
      "log10[FLOPs]:  5.132  log2[SIZE]:  11  log2[WRITE]:  13.083\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1024,), dtype=complex64, numpy=\n",
       "array([0.+0.j, 0.+0.j, 0.+0.j, ..., 0.+0.j, 0.+0.j, 0.+0.j],\n",
       "      dtype=complex64)>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "small_tn()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "aa8b26b8",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------ contraction cost summary ------\n",
      "log10[FLOPs]:  17.766  log2[SIZE]:  44  log2[WRITE]:  49.636\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=complex64, numpy=0j>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "large_tn()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "34cb17e3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "functools.partial(<function custom_stateful at 0x7fd5a0a3d4c0>, optimizer=<class 'opt_einsum.path_random.RandomGreedy'>, opt_conf=None, contraction_info=True, debug_level=2, max_time=60, max_repeats=128, minimize='size')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# we can use more fancy contractor in opt-einsum but not in tensornetwork\n",
    "# custom_stateful is used for contraction path finder which has a life cycle of one-time path solver\n",
    "tc.set_contractor(\n",
    "    \"custom_stateful\",\n",
    "    optimizer=oem.RandomGreedy,\n",
    "    max_time=60,\n",
    "    max_repeats=128,\n",
    "    minimize=\"size\",\n",
    "    debug_level=2,\n",
    "    contraction_info=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3441b3e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------ contraction cost summary ------\n",
      "log10[FLOPs]:  4.925  log2[SIZE]:  10  log2[WRITE]:  12.531\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1024,), dtype=complex64, numpy=\n",
       "array([0.+0.j, 0.+0.j, 0.+0.j, ..., 0.+0.j, 0.+0.j, 0.+0.j],\n",
       "      dtype=complex64)>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "small_tn()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ef946a8c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------ contraction cost summary ------\n",
      "log10[FLOPs]:  11.199  log2[SIZE]:  26  log2[WRITE]:  28.183\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=complex64, numpy=0j>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "large_tn()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf774c11",
   "metadata": {},
   "source": [
    "## Cotengra"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9466a596",
   "metadata": {},
   "source": [
    "for more advanced contractors, we ask help for the sota contractor optimizer on the market: cotengra"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b07b777a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "functools.partial(<function custom at 0x7fd5a0a3d430>, optimizer=<function contraction_info_decorator.<locals>.new_algorithm at 0x7fd588e28ee0>, memory_limit=None, debug_level=2, preprocessing=True)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "opt = ctg.ReusableHyperOptimizer(\n",
    "    methods=[\"greedy\", \"kahypar\"],\n",
    "    parallel=True,\n",
    "    minimize=\"write\",\n",
    "    max_time=120,\n",
    "    max_repeats=1024,\n",
    "    progbar=True,\n",
    ")\n",
    "tc.set_contractor(\n",
    "    \"custom\", optimizer=opt, preprocessing=True, contraction_info=True, debug_level=2\n",
    ")\n",
    "## for more setup on cotengra optimizers, see the reference in\n",
    "## https://cotengra.readthedocs.io/en/latest/advanced.html\n",
    "## preprocessing=True will merge all single-qubit gates into neighboring two-qubit gates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c4ffe73a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "log2[SIZE]: 10.00 log10[FLOPs]: 4.90: 100%|█████████████████████████████████████████| 1024/1024 [00:28<00:00, 35.45it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------ contraction cost summary ------\n",
      "log10[FLOPs]:  4.900  log2[SIZE]:  10  log2[WRITE]:  12.255\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1024,), dtype=complex64, numpy=\n",
       "array([0.+0.j, 0.+0.j, 0.+0.j, ..., 0.+0.j, 0.+0.j, 0.+0.j],\n",
       "      dtype=complex64)>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "small_tn()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0b8a440a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "log2[SIZE]: 20.00 log10[FLOPs]: 9.50:   4%|█▊                                         | 43/1024 [02:09<49:22,  3.02s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------ contraction cost summary ------\n",
      "log10[FLOPs]:  9.501  log2[SIZE]:  20  log2[WRITE]:  24.090\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=complex64, numpy=0j>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "large_tn()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69a5e59f",
   "metadata": {},
   "source": [
    "We can also apply subtree reconf after cotengra find the path, which in general further (and usually greatly) improve flops and write for the contraction. Indeed, the subtree reconf postprocessing is in general more important than increasing the search time or repeats for the optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "1d2594e5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "functools.partial(<function custom at 0x7fd5a0a3d430>, optimizer=<function contraction_info_decorator.<locals>.new_algorithm at 0x7fd58c832700>, memory_limit=None, debug_level=2, preprocessing=True)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "opt = ctg.ReusableHyperOptimizer(\n",
    "    minimize=\"combo\",\n",
    "    max_repeats=1024,\n",
    "    max_time=240,\n",
    "    progbar=True,\n",
    ")\n",
    "\n",
    "\n",
    "def opt_reconf(inputs, output, size, **kws):\n",
    "    tree = opt.search(inputs, output, size)\n",
    "    tree_r = tree.subtree_reconfigure_forest(\n",
    "        progbar=True, num_trees=10, num_restarts=20, subtree_weight_what=(\"size\",)\n",
    "    )\n",
    "    return tree_r.get_path()\n",
    "\n",
    "\n",
    "tc.set_contractor(\n",
    "    \"custom\",\n",
    "    optimizer=opt_reconf,\n",
    "    contraction_info=True,\n",
    "    preprocessing=True,\n",
    "    debug_level=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "5214d1a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "log2[SIZE]: 10.00 log10[FLOPs]: 4.87: 100%|█████████████████████████████████████████| 1024/1024 [02:29<00:00,  6.86it/s]\n",
      "log2[SIZE]: 10.00 log10[FLOPs]: 4.86: 100%|█████████████████████████████████████████████| 20/20 [00:31<00:00,  1.57s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------ contraction cost summary ------\n",
      "log10[FLOPs]:  4.859  log2[SIZE]:  10  log2[WRITE]:  12.583\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1024,), dtype=complex64, numpy=\n",
       "array([0.+0.j, 0.+0.j, 0.+0.j, ..., 0.+0.j, 0.+0.j, 0.+0.j],\n",
       "      dtype=complex64)>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "small_tn()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "23eabe08",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "log2[SIZE]: 21.00 log10[FLOPs]: 9.62:   9%|███▉                                       | 93/1024 [04:04<40:49,  2.63s/it]\n",
      "log2[SIZE]: 17.00 log10[FLOPs]: 8.66: 100%|█████████████████████████████████████████████| 20/20 [03:08<00:00,  9.44s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------ contraction cost summary ------\n",
      "log10[FLOPs]:  8.657  log2[SIZE]:  17  log2[WRITE]:  25.035\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=complex64, numpy=0j>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "large_tn()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8355b1dc",
   "metadata": {},
   "source": [
    "We can also extract the tensornetwork directly for the circuit or observable calculations and we can do the contraction by ourselves using any method we like. Moreover, all these contractors or our customized external contractions can still be compatible with jit, automatic differentiation. Specifically, the contraction path solver, though taking some time overhead, is only evaluated once due to the jit infrastructure (note for demonstration usage, we don't decorate our contraction function with ``K.jit`` here)."
   ]
  }
 ],
 "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}