docs/source/whitepaper/5-density-matrix.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f08fed11",
   "metadata": {},
   "source": [
    "# Density Matrix and Mixed State Evolution"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ca168b6",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "TensorCircuit provides two methods of simulating noisy, mixed state quantum evolution.  Full density matrix simulation of $n$ qubits is provided by using ``tc.DMCircuit(n)``, and then adding quantum operations -- both unitary gates as well as general quantum operations specified by Kraus operators -- to the circuit.  Relative to pure state simulation of $n$ qubits via ``tc.Circuit``, full density matrix simulation is twice as memory-intensive, and thus the maximum system size simulatable will be half of what can be simulated in the pure state case.  A less memory intensive option is to use the standard ``tc.Circuit(n)`` object and stochastically simulate open system evolution via Monte Carlo trajectory methods."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36cfaa32",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a7c8f51c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorcircuit as tc\n",
    "\n",
    "K = tc.set_backend(\"tensorflow\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0afa9c53",
   "metadata": {},
   "source": [
    "## Density Matrix Simulation with ``tc.DMCircuit``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68872b6a",
   "metadata": {},
   "source": [
    "We illustrate this method below, by considering a simple circuit on a single qubit, which takes as input the mixed state corresponding to a probabilistic mixture of the $\\vert{0}\\rangle$ state and the maximally mixed state\n",
    "$\\rho(\\alpha) = \\alpha\\vert 0\\rangle \\langle 0\\vert + (1-\\alpha)I/2.$\n",
    "\n",
    "This state is then passed through a circuit that applies an $X$ gate, followed by a  quantum operation corresponding to an amplitude damping channel $\\mathcal{E}_\\gamma$ with parameter $\\gamma$. This has Kraus operators\n",
    "$K_0 = \\begin{pmatrix}\n",
    "1 & 0 \\\\ 0 & \\sqrt{1-\\gamma}\n",
    "\\end{pmatrix}, \\quad K_1 = \\begin{pmatrix}\n",
    "0 & \\sqrt{\\gamma} \\\\ 0 & 0\n",
    "\\end{pmatrix}$\n",
    "This circuit thus causes the evolution\n",
    "$\\rho(\\alpha) \\xrightarrow[]{X} X\\rho(\\alpha)X\\xrightarrow[]{\\mathcal{E}_\\gamma}\\sum_{i=0}^1 K_i X\\rho(\\alpha)X K_i^\\dagger$\n",
    "\n",
    "\n",
    "To simulate this in TensorCircuit, we first create a ``tc.DMCircuit`` (density matrix circuit) object and set the input state using the ``dminputs`` optional argument (note that if a pure state input is provided to ``tc.DMCircuit``, this should be done via the ``inputs`` optional argument).\n",
    "\n",
    "$\\rho(\\alpha)$ has matrix form\n",
    "$\\rho(\\alpha) = \\begin{pmatrix}\n",
    "\\frac{1+\\alpha}{2} & \\\\ & \\frac{1-\\alpha}{2}\n",
    "\\end{pmatrix},$\n",
    "and thus (taking $\\alpha=0.6$) we initialize the density matrix circuit as follows.\n",
    "\n",
    "To implement a general quantum operation such as the amplitude damping channel, we use ``general_kraus``, supplied with the corresponding list of Kraus operators."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "332fd15c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rho(alpha):\n",
    "    return np.array([[(1 + alpha) / 2, 0], [0, (1 - alpha) / 2]])\n",
    "\n",
    "\n",
    "input_state = rho(0.6)\n",
    "dmc = tc.DMCircuit(1, dminputs=input_state)\n",
    "\n",
    "dmc.x(0)\n",
    "\n",
    "\n",
    "def amp_damp_kraus(gamma):\n",
    "    K0 = np.array([[1, 0], [0, np.sqrt(1 - gamma)]])\n",
    "    K1 = np.array([[0, np.sqrt(gamma)], [0, 0]])\n",
    "    return K0, K1\n",
    "\n",
    "\n",
    "K0, K1 = amp_damp_kraus(0.3)\n",
    "dmc.general_kraus([K0, K1], 0)  # apply channel with Kraus operators [K0,K1] to qubit 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "17aa5310",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2), dtype=complex64, numpy=\n",
       "array([[0.44+0.j, 0.  +0.j],\n",
       "       [0.  +0.j, 0.56+0.j]], dtype=complex64)>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get the output density matrix\n",
    "dmc.state()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e37cf64d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor((-0.11999999+0j), shape=(), dtype=complex64) (<tf.Tensor: shape=(1,), dtype=float32, numpy=array([1.], dtype=float32)>, -1.0)\n"
     ]
    }
   ],
   "source": [
    "# evaluate the expectation as a circuit object\n",
    "print(dmc.expectation_ps(z=[0]), dmc.measure(0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fa9e7c6",
   "metadata": {},
   "source": [
    "In the above example, we input the Kraus operators for the amplitude damping channel manually, in order to illustrate the general approach to implementing quantum channels. In fact, TensorCircuit includes built-in methods for returning the Kraus operators for a number of common channels, including the amplitude damping, depolarizing, phase damping, and reset channels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9205dac6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "amplitudedampingchannel\n",
      "depolarizingchannel\n",
      "phasedampingchannel\n",
      "resetchannel\n"
     ]
    }
   ],
   "source": [
    "# a set of built-in quantum channels\n",
    "\n",
    "for k in dir(tc.channels):\n",
    "    if k.endswith(\"channel\"):\n",
    "        print(k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "58c04997",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(4, 4), dtype=complex64, numpy=\n",
       "array([[0.49999997+0.j, 0.        +0.j, 0.4472136 +0.j, 0.        +0.j],\n",
       "       [0.        +0.j, 0.        +0.j, 0.        +0.j, 0.        +0.j],\n",
       "       [0.4472136 +0.j, 0.        +0.j, 0.49999994+0.j, 0.        +0.j],\n",
       "       [0.        +0.j, 0.        +0.j, 0.        +0.j, 0.        +0.j]],\n",
       "      dtype=complex64)>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dmc = tc.DMCircuit(2)\n",
    "dmc.h(0)\n",
    "gamma = 0.2\n",
    "K0, K1 = tc.channels.phasedampingchannel(gamma)\n",
    "dmc.general_kraus([K0, K1], 0)\n",
    "dmc.state()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "251224a1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(4, 4), dtype=complex64, numpy=\n",
       "array([[0.49999997+0.j, 0.        +0.j, 0.4472136 +0.j, 0.        +0.j],\n",
       "       [0.        +0.j, 0.        +0.j, 0.        +0.j, 0.        +0.j],\n",
       "       [0.4472136 +0.j, 0.        +0.j, 0.49999994+0.j, 0.        +0.j],\n",
       "       [0.        +0.j, 0.        +0.j, 0.        +0.j, 0.        +0.j]],\n",
       "      dtype=complex64)>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# or we can directly use the following API for shorthand\n",
    "\n",
    "dmc = tc.DMCircuit(2)\n",
    "dmc.h(0)\n",
    "gamma = 0.2\n",
    "dmc.phasedamping(0, gamma=0.2)\n",
    "dmc.state()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6cf564a",
   "metadata": {},
   "source": [
    "### AD and JIT Compatibility\n",
    "\n",
    "``tc.DMCircuit``, like ``tc.Circuit`` is also compatible with ML paradigms such as AD, jit, and vmap. See the example below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bdd81210",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 3\n",
    "nbatch = 2\n",
    "\n",
    "\n",
    "def loss(params, noisep):\n",
    "    c = tc.DMCircuit(n)\n",
    "    for i in range(n):\n",
    "        c.rx(i, theta=params[i])\n",
    "    for i in range(n):\n",
    "        c.depolarizing(i, px=noisep, py=noisep, pz=noisep)\n",
    "    return K.real(K.sum([c.expectation_ps(z=[i]) for i in range(n)]))\n",
    "\n",
    "\n",
    "loss_vvg = K.jit(\n",
    "    K.vectorized_value_and_grad(loss, argnums=(0, 1), vectorized_argnums=(0))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "db09f435",
   "metadata": {},
   "outputs": [],
   "source": [
    "vs, (gparams, gnoisep) = loss_vvg(0.1 * K.ones([nbatch, n]), 0.1 * K.ones([]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c0514b0d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorShape([2]), TensorShape([2, 3]), TensorShape([]))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vs.shape, gparams.shape, gnoisep.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88a993f3",
   "metadata": {},
   "source": [
    "Note how the noise parameter can also be differentiated and jitted!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d37b6c4f",
   "metadata": {},
   "source": [
    "## Monte Carlo Noise Simulation with ``tc.Circuit``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb05308d",
   "metadata": {},
   "source": [
    "For pure state inputs, Monte Carlo methods can be used to sample noisy quantum evolution using ``tc.Circuit`` instead of ``tc.DMCircuit`` where the mixed state is effectively simulated with an ensemble of pure states.\n",
    " \n",
    "As for density matrix simulation, quantum channels $\\mathcal{E}$ can be added to a circuit object by providing a list of their associated Kraus operators $\\{K_i\\}$.  The API is the same as for the full density matrix simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "772e634c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2,), dtype=complex64, numpy=array([0.+0.j, 1.+0.j], dtype=complex64)>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_state = np.array([1, 1] / np.sqrt(2))\n",
    "c = tc.Circuit(1, inputs=input_state)\n",
    "c.general_kraus(tc.channels.phasedampingchannel(0.5), 0)\n",
    "c.state()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9de97bcc",
   "metadata": {},
   "source": [
    "In this framework though, the output of a channel acting on $\\vert{\\psi}\\rangle$ , i.e.\n",
    "$\\mathcal{E} ( \\vert{\\psi}\\rangle\\langle{\\psi}\\vert) = \\sum_i K_i \\vert{\\psi}\\rangle\\langle{\\psi}\\vert K_i^ \\dagger$\n",
    "is viewed as an ensemble of states $\\frac{K_i\\vert{\\psi}\\rangle}{\\sqrt{\\langle{\\psi}\\vert K_i^\\dagger K_i \\vert{\\psi}\\rangle}}$ that each occur with probability $p_i = \\langle{\\psi}\\vert K_i^\\dagger K_i \\vert{\\psi}\\rangle$.  Thus, the code above stochastically produces the output of a single qubit initialized in state $\\vert{\\psi}\\rangle=\\frac{\\vert{0}\\rangle+\\vert{1}\\rangle}{\\sqrt{2}}$ being passed through a phase damping channel with parameter $\\gamma=0.5$.  \n",
    "\n",
    "The Monte Carlo simulation of channels where the Kraus operators are all unitary matrices (up to a constant factor) can be handled with additional efficiency by using ``unitary_kraus`` instead of ``general_kraus``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f414e0b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=int32, numpy=3>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "px, py, pz = 0.1, 0.2, 0.3\n",
    "c.unitary_kraus(tc.channels.depolarizingchannel(px, py, pz), 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84a47c16",
   "metadata": {},
   "source": [
    "Note the int tensor returned above indicates in this trajectory, which operator is applied on the circuit."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "330b5d24",
   "metadata": {},
   "source": [
    "### Externalizing the Randomness\n",
    "\n",
    "The ``general_kraus`` and ``unitary\\_kraus`` examples above both handle randomness generation from inside the respective methods. That is, when the list $[K_0, K_1, \\ldots, K_{m-1}]$ of Kraus operators is supplied to ``general_kraus`` or ``unitary_kraus``, the method  partitions the interval $[0,1]$ into $m$ contiguous intervals $[0,1] = I_0 \\cup I_1 \\cup \\ldots I_{m-1}$ where the length of $I_i$ is equal to the relative probability of obtaining outcome $i$. Then a uniformly random variable $x$ in $[0,1]$ is generated from within the method, and outcome $i$ selected based on which interval $x$ lies in. \n",
    "\n",
    "In TensorCircuit, we have full backend agnostic infrastructure for random number generation and management. However, the interplay between jit, random number, and backend switch is often subtle if we rely on the random number generation inside these methods. See [advance.html\\#randoms-jit-backend-agnostic-and-their-interplay](../advance.html#randoms-jit-backend-agnostic-and-their-interplay) for details.\n",
    "\n",
    "In some situations, it may be preferable to first generate the random variable from outside the method, and then pass the value generated into ``general_kraus`` or ``unitary_kraus``.  This can be done via the optional ``status`` argument:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4ce0ccf9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(1, shape=(), dtype=int32)\n",
      "tf.Tensor(3, shape=(), dtype=int32)\n"
     ]
    }
   ],
   "source": [
    "px, py, pz = 0.1, 0.2, 0.3\n",
    "x = 0.5\n",
    "print(c.unitary_kraus(tc.channels.depolarizingchannel(px, py, pz), 0, status=x))\n",
    "x = 0.8\n",
    "print(c.unitary_kraus(tc.channels.depolarizingchannel(px, py, pz), 0, status=x))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bdf395b",
   "metadata": {},
   "source": [
    "This is useful, for instance, when one wishes to use ``vmap`` to batch compute multiple runs of a Monte Carlo simulation. This is illustrated in the example below, where ``vmap`` is used to compute 10 runs of the simulation in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c9ccd608",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(10,), dtype=complex64, numpy=\n",
       "array([ 0.99999994+0.j,  0.99999994+0.j,  0.99999994+0.j, -0.99999994+0.j,\n",
       "        0.99999994+0.j,  0.99999994+0.j,  0.99999994+0.j,  0.99999994+0.j,\n",
       "       -0.99999994+0.j,  0.99999994+0.j], dtype=complex64)>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def f(x):\n",
    "    c = tc.Circuit(1)\n",
    "    c.h(0)\n",
    "    c.unitary_kraus(tc.channels.depolarizingchannel(0.1, 0.2, 0.3), 0, status=x)\n",
    "    return c.expectation_ps(x=[0])\n",
    "\n",
    "\n",
    "f_vmap = K.vmap(f, vectorized_argnums=0)\n",
    "X = K.implicit_randn(10)\n",
    "f_vmap(X)"
   ]
  }
 ],
 "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
}