{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `pyqtorch`\n",
    "\n",
    "A fast large scale emulator for quantum machine learning on a PyTorch backend.\n",
    "\n",
    "\n",
    "## Installation\n",
    "\n",
    "To install the library for development, you can go into any virtual environment of your\n",
    "choice and install it normally with `pip` (including extra dependencies for development):\n",
    "\n",
    "```\n",
    "pip install pyqtorch\n",
    "```\n",
    "\n",
    "## Contribute\n",
    "\n",
    "If you want to contribute to the package, make sure to execute tests and MyPy checks\n",
    "otherwise the automatic pipeline will not pass. To do so, the recommended way is\n",
    "to use `hatch` for managing the environments:\n",
    "\n",
    "```shell\n",
    "hatch env create tests\n",
    "hatch --env tests run python -m pytest -vvv --cov pyqtorch tests\n",
    "hatch --env tests run python -m mypy pyqtorch tests\n",
    "```\n",
    "\n",
    "If you don't want to use `hatch`, you can use the environment manager of your\n",
    "choice (e.g. Conda) and execute the following:\n",
    "\n",
    "```shell\n",
    "pip install -e .[dev]\n",
    "pytest -vvv --cov pyqtorch tests\n",
    "mypy pyqtorch tests\n",
    "```\n",
    "\n",
    "## Getting started with `pyqtorch`\n",
    "\n",
    "### Gates\n",
    "\n",
    "`pyqtorch` implements most of the commonly used gates like Pauli gates, rotation\n",
    "gates, and controlled gates. Every gate accepts a sequence of `qubits` on which\n",
    "it operates and a total number `n_qubits` of the state that it will operate on:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/niklas/Library/Application Support/hatch/env/virtual/qucint/6NOL9orC/qucint/lib/python3.9/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": [
       "tensor([[0.+0.j],\n",
       "        [1.+0.j]], dtype=torch.complex128)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import pyqtorch.modules as pyq\n",
    "\n",
    "gate = pyq.X(qubits=[0], n_qubits=1)\n",
    "z = pyq.zero_state(n_qubits=1)\n",
    "\n",
    "gate(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[1.+0.j],\n",
       "         [0.+0.j]],\n",
       "\n",
       "        [[0.+0.j],\n",
       "         [0.+0.j]]], dtype=torch.complex128)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gate = pyq.CNOT(qubits=[0,1], n_qubits=2)\n",
    "z = pyq.zero_state(n_qubits=2)\n",
    "gate(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 2, 1])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In `pyqtorch` the state is a `n_qubit+1` dimensional `Tensor`, for example a\n",
    "state with 3 qubits has the shape `(2, 2, 2, 1)` (i.e. one dimension for each\n",
    "qubit, plus one dimension for the batch size).\n",
    "\n",
    "\n",
    "_**NOTE:**_ We always work with batched state in `pyqtorch`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 2, 2, 1])\n",
      "torch.Size([2, 2, 2, 5])\n"
     ]
    }
   ],
   "source": [
    "z = pyq.zero_state(n_qubits=3)\n",
    "print(z.shape)\n",
    "z = pyq.zero_state(n_qubits=3, batch_size=5)\n",
    "print(z.shape)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `QuantumCircuit`s\n",
    "\n",
    "To compose multiple gates we use a `QuantumCircuit` which is constructed from\n",
    "a list of operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0.+0.j],\n",
       "         [0.+0.j]],\n",
       "\n",
       "        [[0.+0.j],\n",
       "         [1.+0.j]]], dtype=torch.complex128)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "circ = pyq.QuantumCircuit(\n",
    "    n_qubits=2,\n",
    "    operations=[\n",
    "        pyq.X([0], 2),\n",
    "        pyq.CNOT([0,1], 2)\n",
    "    ]\n",
    ")\n",
    "\n",
    "z = pyq.zero_state(2)\n",
    "circ(z)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Every gate and circuit in `pyqtorch` accepts a state and an optional tensor of angles.\n",
    "If the gate/circuit does not depend on any angles, the second argument is ignored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0.+0.j],\n",
       "         [0.+0.j]],\n",
       "\n",
       "        [[0.+0.j],\n",
       "         [1.+0.j]]], dtype=torch.complex128)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "theta = torch.rand(1)\n",
    "circ(z, theta)  # theta is ignored"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0.9885+0.0000j],\n",
       "         [0.0000+0.0000j]],\n",
       "\n",
       "        [[0.0000+0.0000j],\n",
       "         [0.0000-0.1511j]]], dtype=torch.complex128)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "circ = pyq.QuantumCircuit(\n",
    "    n_qubits=2,\n",
    "    operations=[\n",
    "        pyq.RX([0], 2), # rotation instead of X gate\n",
    "        pyq.CNOT([0,1], 2)\n",
    "    ]\n",
    ")\n",
    "\n",
    "circ(z, theta)  # theta is used!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The vanilla `QuantumCircuit` is always passing the same `theta` tensor to its operations, meaning\n",
    "the `forward` method of the circuit is:\n",
    "```python\n",
    "class QuantumCircuit(torch.nn.Module):\n",
    "\n",
    "    # ...\n",
    "\n",
    "    def forward(self, state: torch.Tensor, thetas: torch.Tensor = None) -> torch.Tensor:\n",
    "        for op in self.operations:\n",
    "            state = op(state, thetas)\n",
    "        return state\n",
    "```\n",
    "\n",
    "The `FeaturemapLayer` is a convenience constructor for a `QuantumCircuit` which accepts an operation\n",
    "to put on every qubit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "QuantumCircuit(\n",
      "  (operations): ModuleList(\n",
      "    (0): RX(qubits=[0], n_qubits=3)\n",
      "    (1): RX(qubits=[1], n_qubits=3)\n",
      "    (2): RX(qubits=[2], n_qubits=3)\n",
      "  )\n",
      ")\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 2, 2, 4])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "circ = pyq.FeaturemapLayer(n_qubits=3, Op=pyq.RX)\n",
    "print(circ)\n",
    "\n",
    "states = pyq.zero_state(n_qubits=3, batch_size=4)\n",
    "inputs = torch.rand(4)\n",
    "\n",
    "# the same batch of inputs are passed to the operations\n",
    "circ(states, inputs).shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Trainable `QuantumCircuit`s aka `VariationalLayer`s\n",
    "\n",
    "If you want the angles of your circuit to be trainable you can use a `VariationalLayer`.\n",
    "The `VariationalLayer` ignores the second input (because it has trainable angle parameters)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[ 0.2419+0.0000j],\n",
       "          [ 0.0000+0.2235j]],\n",
       "\n",
       "         [[ 0.0000-0.0524j],\n",
       "          [ 0.0484+0.0000j]]],\n",
       "\n",
       "\n",
       "        [[[ 0.0000+0.6759j],\n",
       "          [-0.6244+0.0000j]],\n",
       "\n",
       "         [[ 0.1464+0.0000j],\n",
       "          [ 0.0000+0.1352j]]]], dtype=torch.complex128,\n",
       "       grad_fn=<ViewBackward0>)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "circ = pyq.VariationalLayer(n_qubits=3, Op=pyq.RX)\n",
    "\n",
    "state = pyq.zero_state(3)\n",
    "this_argument_is_ignored = None\n",
    "circ(state, this_argument_is_ignored)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Composing `QuantumCircuit`s\n",
    "\n",
    "As every gate and circuit in `pyqtorch` accept the same arguments we can easily\n",
    "compose them to larger circuits, i.e. to implement a hardware efficient ansatz:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "QuantumCircuit(\n",
      "  (operations): ModuleList(\n",
      "    (0): VariationalLayer(\n",
      "      (operations): ModuleList(\n",
      "        (0): RX(qubits=[0], n_qubits=3)\n",
      "        (1): RX(qubits=[1], n_qubits=3)\n",
      "        (2): RX(qubits=[2], n_qubits=3)\n",
      "      )\n",
      "    )\n",
      "    (1): VariationalLayer(\n",
      "      (operations): ModuleList(\n",
      "        (0): RY(qubits=[0], n_qubits=3)\n",
      "        (1): RY(qubits=[1], n_qubits=3)\n",
      "        (2): RY(qubits=[2], n_qubits=3)\n",
      "      )\n",
      "    )\n",
      "    (2): VariationalLayer(\n",
      "      (operations): ModuleList(\n",
      "        (0): RX(qubits=[0], n_qubits=3)\n",
      "        (1): RX(qubits=[1], n_qubits=3)\n",
      "        (2): RX(qubits=[2], n_qubits=3)\n",
      "      )\n",
      "    )\n",
      "    (3): EntanglingLayer(\n",
      "      (operations): ModuleList(\n",
      "        (0): CNOT(qubits=[0, 1], n_qubits=3)\n",
      "        (1): CNOT(qubits=[1, 2], n_qubits=3)\n",
      "        (2): CNOT(qubits=[2, 0], n_qubits=3)\n",
      "      )\n",
      "    )\n",
      "    (4): VariationalLayer(\n",
      "      (operations): ModuleList(\n",
      "        (0): RX(qubits=[0], n_qubits=3)\n",
      "        (1): RX(qubits=[1], n_qubits=3)\n",
      "        (2): RX(qubits=[2], n_qubits=3)\n",
      "      )\n",
      "    )\n",
      "    (5): VariationalLayer(\n",
      "      (operations): ModuleList(\n",
      "        (0): RY(qubits=[0], n_qubits=3)\n",
      "        (1): RY(qubits=[1], n_qubits=3)\n",
      "        (2): RY(qubits=[2], n_qubits=3)\n",
      "      )\n",
      "    )\n",
      "    (6): VariationalLayer(\n",
      "      (operations): ModuleList(\n",
      "        (0): RX(qubits=[0], n_qubits=3)\n",
      "        (1): RX(qubits=[1], n_qubits=3)\n",
      "        (2): RX(qubits=[2], n_qubits=3)\n",
      "      )\n",
      "    )\n",
      "    (7): EntanglingLayer(\n",
      "      (operations): ModuleList(\n",
      "        (0): CNOT(qubits=[0, 1], n_qubits=3)\n",
      "        (1): CNOT(qubits=[1, 2], n_qubits=3)\n",
      "        (2): CNOT(qubits=[2, 0], n_qubits=3)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      ")\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.5790-0.0103j],\n",
       "          [ 0.1601+0.0274j]],\n",
       "\n",
       "         [[-0.1465+0.3101j],\n",
       "          [-0.0429+0.2236j]]],\n",
       "\n",
       "\n",
       "        [[[-0.1688-0.3029j],\n",
       "          [-0.2445-0.2112j]],\n",
       "\n",
       "         [[-0.4516+0.0601j],\n",
       "          [ 0.0749-0.1762j]]]], dtype=torch.complex128,\n",
       "       grad_fn=<PermuteBackward0>)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def hea(n_qubits: int, n_layers: int) -> pyq.QuantumCircuit:\n",
    "    ops = []\n",
    "    for _ in range(n_layers):\n",
    "        ops.append(pyq.VariationalLayer(n_qubits, pyq.RX))\n",
    "        ops.append(pyq.VariationalLayer(n_qubits, pyq.RY))\n",
    "        ops.append(pyq.VariationalLayer(n_qubits, pyq.RX))\n",
    "        ops.append(pyq.EntanglingLayer(n_qubits))\n",
    "    return pyq.QuantumCircuit(n_qubits, ops)\n",
    "\n",
    "circ = hea(3,2)\n",
    "print(circ)\n",
    "\n",
    "state = pyq.zero_state(3)\n",
    "circ(state)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "qucint",
   "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.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
