docs/source/whitepaper/6-5-custom-contraction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bc07c2b3",
   "metadata": {},
   "source": [
    "# Customized Contraction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "694cc0e0",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "If the simulated circuit has large qubit counts, we recommend users try a customized contraction setup instead of the default one, which is greedy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "918c848d",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Please refer to the [installation documentation](https://cotengra.readthedocs.io/en/latest/installation.html) for cotengra, which cannot simply be obtained by pip install since it is not uploaded to PyPI. The easiest way for installation is ``pip install -U git+https://github.com/jcmgray/cotengra.git``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2ae27e57",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorcircuit as tc\n",
    "import numpy as np\n",
    "import cotengra as ctg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63d88675",
   "metadata": {},
   "source": [
    "We use the following example as a testbed for the contraction, the real contraction is invoked for ``Circuit.expectation`` API,\n",
    "and there are two stages for the contraction. The first one is contraction path searching which is used to find a better contraction path in terms of space and time. The second stage is the real contraction, where matrix multiplication is called using ML backend API. In this note, we focus on the performance of the first stage. And the contraction path solver can be customized with any type of [opt-einsum compatible path solver](https://optimized-einsum.readthedocs.io/en/stable/custom_paths.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bc3f5c65",
   "metadata": {},
   "outputs": [],
   "source": [
    "def testbed():\n",
    "    n = 40\n",
    "    d = 6\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_ps(z=[n // 2], reuse=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17e9c8bc",
   "metadata": {},
   "source": [
    "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 a greedy optimizer works, other optimizers come with exponential scaling and fail in circuit simulation scenarios.\n",
    "\n",
    "We always set ``contraction_info=True`` (default is ``False``) for the contractor system in this note, which will print contraction information summary including contraction size, flops, and write. For the definition of these metrics, also refer to cotengra docs and [the corresponding paper](https://quantum-journal.org/papers/q-2021-03-15-410/).\n",
    "\n",
    "Metrics that measure the quality of a contraction path include \n",
    "   \n",
    "   * **FLOPs**: the total number of computational operations required for all matrix multiplications involved when contracting the tensor network via the given path. This metric characterizes the total simulation time.\n",
    "    \n",
    "   * **WRITE**: the total size (the number of elements) of all tensors -- including intermediate tensors -- computed during the contraction. \n",
    "    \n",
    "   * **SIZE**: the size of the largest intermediate tensor stored in memory.\n",
    "\n",
    "Since simulations in TensorCircuit are AD-enabled, where all intermediate results need to be cached and traced, the more relevant spatial cost metric is writes instead of size.\n",
    "\n",
    "Also, we will enable ``debug_level=2`` in ``set_contractor`` (never use this option in real computation!) By enabling this, the second stage of the contraction, i.e. the real contraction, will not happen. We can focus on the contraction path information, which demonstrates the difference between different customized contractors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "64647063",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------ contraction cost summary ------\n",
      "log10[FLOPs]:  12.393  log2[SIZE]:  30  log2[WRITE]:  35.125\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=complex64, numpy=0j>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tc.set_contractor(\"greedy\", debug_level=2, contraction_info=True)\n",
    "# the default contractor\n",
    "testbed()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10d590ec",
   "metadata": {},
   "source": [
    "**cotengra optimizer**: for hyperparameters tuning, see the [documentation](https://cotengra.readthedocs.io/en/latest/advanced.html).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a0075260",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "log2[SIZE]: 15.00 log10[FLOPs]: 7.56:  45%|██████████████████▊                       | 458/1024 [02:03<02:32,  3.70it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------ contraction cost summary ------\n",
      "log10[FLOPs]:  7.565  log2[SIZE]:  15  log2[WRITE]:  19.192\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=complex64, numpy=0j>"
      ]
     },
     "execution_count": 7,
     "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",
    "# Caution: for now, parallel only works for \"ray\" in newer versions of python\n",
    "tc.set_contractor(\n",
    "    \"custom\", optimizer=opt, preprocessing=True, contraction_info=True, debug_level=2\n",
    ")\n",
    "# the opt-einsum compatible function interface is passed as the argument of optimizer\\\n",
    "# Also note how preprocessing=True merges the single qubits gate into the neighbor two-qubit gate\n",
    "testbed()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "871115c1",
   "metadata": {},
   "source": [
    "We can even include contraction reconfigure after path searching, which further greatly boosts the space efficiency for the contraction path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7c625596",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "log2[SIZE]: 15.00 log10[FLOPs]: 7.46:  32%|█████████████▍                            | 329/1024 [02:00<04:13,  2.74it/s]\n",
      "log2[SIZE]: 14.00 log10[FLOPs]: 7.02: 100%|█████████████████████████████████████████████| 20/20 [01:05<00:00,  3.30s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------ contraction cost summary ------\n",
      "log10[FLOPs]:  7.021  log2[SIZE]:  14  log2[WRITE]:  19.953\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=complex64, numpy=0j>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "opt = ctg.ReusableHyperOptimizer(\n",
    "    minimize=\"combo\",\n",
    "    max_repeats=1024,\n",
    "    max_time=120,\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",
    "# there is also a default parallel=True option for subtree_reconfigure_forest,\n",
    "# this can only be set as \"ray\" for newer version python as above\n",
    "# note how different versions of cotengra have breaking APIs in the last line: get_path or path\n",
    "# the user may need to change the API to make the example work\n",
    "\n",
    "tc.set_contractor(\n",
    "    \"custom\",\n",
    "    optimizer=opt_reconf,\n",
    "    contraction_info=True,\n",
    "    preprocessing=True,\n",
    "    debug_level=2,\n",
    ")\n",
    "testbed()"
   ]
  }
 ],
 "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
}