algorithms/qml/quantum_autoencoder/quantum_autoencoder.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "88da6df0",
   "metadata": {},
   "source": [
    "# Quantum Autoencoder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "201ce387",
   "metadata": {},
   "source": [
    "### Classical encoders:\n",
    "Refer to encoding/compressing classical data into a smaller sized data via deterministic algorithm. For example, JPEG is essentially an algorithm which compresses images into a smaller sized images.\n",
    "\n",
    "### Classical auto-encoders:\n",
    "One can use machine-learning technics and train a variational network for compressing data. In general, an auto-encoder network looks as follows:\n",
    "\n",
    "<center>\n",
    "<img src=\"https://docs.classiq.io/resources/Autoencoder_structure.png\" style=\"width:50%\">\n",
    "<figcaption align = \"middle\"> Classical auto encoder layout (taken from Wikipedia) </figcaption>\n",
    "</center>\n",
    "\n",
    "\n",
    "The network has three main parts: (1) The encoder part that compresses the data into a smaller, coded layer (2). The latter is the input to a decoder part (3). Typically, training is done against the comparison between the input and the output of this network.\n",
    "\n",
    "**Classical auto encoders can also be used for anomaly detection (see below).**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5be9dbfd",
   "metadata": {},
   "source": [
    "### Quantum auto-encoders:\n",
    "In a similar fashion to the classical counterpart, quantum auto-encoder refers to \"compressing\" quantum data stored initially on $n$ qubits into a smaller quantum register of $m<n$ qubits, via variational circuit. However, quantum computing is reversibale, and thus qubits cannot be \"erased\". Therefore, alternatively, a quantum autoencoder tries to acheive the following transformation from uncoded quantum register of size $n$ to a coded one of size $m$:\n",
    "$$\n",
    "|\\psi\\rangle_n \\rightarrow |\\psi'\\rangle_m|0\\rangle_{n-m}\n",
    "$$\n",
    "Namely, we try to decouple the initial state to a product state of a smaller register of size $m$ and a register which is in the zero state. The former is usually called *coded* state and the latter the *trash* state.\n",
    "\n",
    "<center>\n",
    "<img src=\"https://docs.classiq.io/resources/q_auto_encoder.png\" style=\"width:100%\">\n",
    "<figcaption align = \"middle\"> Quantum auto encoder layout: uncoded data of size 5 transforms into two outputs, a coded register of size 3 and trash outputs of size 2 at state $|00\\rangle$ </figcaption>\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e40c41f0",
   "metadata": {},
   "source": [
    "# Training of quantum auto encoders\n",
    "\n",
    "To train a quantum auto encoder one should define a proper cost function. Below we propose two common approaches, one using a swap test and the other uses Hamiltonian measurements. We focus on the swap test case, and comment on the other approach at the end of this notebook.\n",
    "\n",
    "## The swap test\n",
    "\n",
    "The swap test is a quantum function which checks the overlap between two quantum states: the inputs of the function are two quantum registers of the same size, $|\\psi_1\\rangle, \\,|\\psi_2\\rangle$, and it returns as an output a single \"test\" qubit whose state encodes the overlap between the two inputs: $|q\\rangle_{\\rm test} = \\alpha|0\\rangle + \\sqrt{1-\\alpha^2}|1\\rangle$, with\n",
    "$$\n",
    "\\alpha^2 = \\frac{1}{2}\\left(1+|\\langle \\psi_1 |\\psi_2 \\rangle |^2\\right).\n",
    "$$\n",
    "Thus, the probability to measure the test qubit at state $0$ is 1 if the states are identical (up to a global phase), and 0 if the states are orthogonal to each other.\n",
    "\n",
    "The quantum model starts with an H gate on the test qubit, followed by swapping between the two states controlled on the test qubit, and a final H gate on the test qubit:\n",
    "\n",
    "<center>\n",
    "<table><tr>\n",
    "<td> <img src=\"https://docs.classiq.io/resources/swap_test_closed.png\" style=\"width:100%\">\n",
    "<td> <img src=\"https://docs.classiq.io/resources/swap_test_opened.png\" style=\"width:100%\">\n",
    "</tr></table>\n",
    "<figcaption align = \"middle\"> Closed (left panel) and opened (right panel) visualization of the swap test algorithm </figcaption>\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "443577d1",
   "metadata": {},
   "source": [
    "## Quantum neural network for quantum auto encoder\n",
    "\n",
    "The quantum auto encoder can be built as a quantum neural network, having the following three parts:\n",
    "\n",
    "1. A data loading block that loads classical data on $n$ qubits.\n",
    "2. An encoder block which is some variational quantum ansatz; input port of size $n$ and output ports of size $m$ and $n-m$.\n",
    "3. A swap test block between the $n-m$ trash output of the encoder and new $n-m$ zero registers.\n",
    "\n",
    "We train the network such that the test qubit of the swap test is at state |0⟩ with probability 1.\n",
    "<center>\n",
    "<img src=\"https://docs.classiq.io/resources/qae_qlayer.png\" style=\"width:100%\">\n",
    "<figcaption align = \"middle\">  Quantum auto encoder layout: uncoded data of size 5 transforms into two outputs, a coded register of size 3 and trash outputs of size 2 at state $|00\\rangle$\n",
    "</figcaption>\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a685c3b",
   "metadata": {},
   "source": [
    "# Pre-user-defined functions which will be used to construct the quantum layer\n",
    "\n",
    "As a first step we build some user-defined functions which allow us flexible modeling. We have three functions:\n",
    "1. `angle_encoding`: This function loads data of size `num_qubits` on `num_qubits` qubits via RY gates. It has an output port named `qpv`.\n",
    "2. `encoder_ansatz` : A simple variational ansatz for encoding `num_qubits` qubits on `num_encoding_qubits` qubits (see description within the code-block). The input port is `x`, and the output ports are `coded` and `trash`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f2b814d9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:51:55.917096Z",
     "iopub.status.busy": "2024-05-07T14:51:55.916566Z",
     "iopub.status.idle": "2024-05-07T14:51:59.024617Z",
     "shell.execute_reply": "2024-05-07T14:51:59.023600Z"
    },
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "from classiq import *\n",
    "from classiq.qmod.symbolic import pi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "75634100",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:51:59.030300Z",
     "iopub.status.busy": "2024-05-07T14:51:59.028905Z",
     "iopub.status.idle": "2024-05-07T14:51:59.037330Z",
     "shell.execute_reply": "2024-05-07T14:51:59.036640Z"
    }
   },
   "outputs": [],
   "source": [
    "@qfunc\n",
    "def angle_encoding(exe_params: CArray[CReal], qbv: Output[QArray[QBit]]) -> None:\n",
    "    allocate(exe_params.len, qbv)\n",
    "    repeat(\n",
    "        count=exe_params.len,\n",
    "        iteration=lambda index: RY(pi * exe_params[index], qbv[index]),\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7132bf7f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:51:59.041648Z",
     "iopub.status.busy": "2024-05-07T14:51:59.040513Z",
     "iopub.status.idle": "2024-05-07T14:51:59.048695Z",
     "shell.execute_reply": "2024-05-07T14:51:59.048146Z"
    }
   },
   "outputs": [],
   "source": [
    "@qfunc\n",
    "def encoder_ansatz(\n",
    "    num_qubits: CInt,\n",
    "    num_encoding_qubits: CInt,\n",
    "    exe_params: CArray[CReal],\n",
    "    x: Input[QArray[QBit, \"num_qubits\"]],\n",
    "    trash: Output[QArray[QBit, \"num_qubits-num_encoding_qubits\"]],\n",
    "    coded: Output[QArray[QBit, \"num_encoding_qubits\"]],\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    This is a parametric model which has num_trash_qubits = num_qubits-num_encoding_qubits as an output.\n",
    "    It contains num_trash_qubits layers, each composed of RY gates and CX gates with a linear connectivity,\n",
    "    and a final layer with RY gate on each of the trash qubits is applied.\n",
    "    \"\"\"\n",
    "\n",
    "    def single_layer(rep: CInt) -> None:\n",
    "        repeat(\n",
    "            count=num_qubits,\n",
    "            iteration=lambda index: RY(exe_params[rep * num_qubits + index], x[index]),\n",
    "        )\n",
    "        repeat(\n",
    "            count=num_qubits - 1,\n",
    "            iteration=lambda index: CX(x[index], x[index + 1]),\n",
    "        )\n",
    "\n",
    "    repeat(count=num_qubits - num_encoding_qubits, iteration=single_layer)\n",
    "    bind(x, [coded, trash])\n",
    "    repeat(\n",
    "        count=num_qubits - num_encoding_qubits,\n",
    "        iteration=lambda index: RY(\n",
    "            exe_params[(num_qubits - num_encoding_qubits) * num_qubits + index],\n",
    "            trash[index],\n",
    "        ),\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "493d4499",
   "metadata": {},
   "source": [
    "# An example: auto encoder for domain wall data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "428691bc",
   "metadata": {},
   "source": [
    "In the following example we will try to encode data which has a domain wall structure. Let us define the relevant data for strings of size 4."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "535d9181",
   "metadata": {},
   "source": [
    "## The data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1cdcd74c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:51:59.051331Z",
     "iopub.status.busy": "2024-05-07T14:51:59.050956Z",
     "iopub.status.idle": "2024-05-07T14:51:59.055305Z",
     "shell.execute_reply": "2024-05-07T14:51:59.054655Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "domain wall data:\n",
      " [[0 0 1 1]\n",
      " [0 0 0 1]\n",
      " [0 1 1 1]]\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "domain_wall_data = np.array([[0, 0, 1, 1], [0, 0, 0, 1], [0, 1, 1, 1]])\n",
    "print(\"domain wall data:\\n\", domain_wall_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a09e977",
   "metadata": {},
   "source": [
    "## The quantum program"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "decfc82e-fae0-4d3d-91c8-3b2a660af884",
   "metadata": {},
   "source": [
    "We will try to encode this data on size 4 on 2 qubits. Let us built the corresponding quantum layer based on the pre-defined functions above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6fb450af-1613-4284-8625-ff7f2ae12557",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:51:59.057848Z",
     "iopub.status.busy": "2024-05-07T14:51:59.057452Z",
     "iopub.status.idle": "2024-05-07T14:51:59.061978Z",
     "shell.execute_reply": "2024-05-07T14:51:59.061324Z"
    }
   },
   "outputs": [],
   "source": [
    "NUM_QUBITS = 4\n",
    "NUM_ENCODING_QUBITS = 2\n",
    "num_trash_qubits = NUM_QUBITS - NUM_ENCODING_QUBITS\n",
    "num_weights_in_encoder = NUM_QUBITS * num_trash_qubits + num_trash_qubits"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2454d248-5614-4300-a5e3-d4826469e1c8",
   "metadata": {},
   "source": [
    "Let us construct the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5f1bc0cd",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:51:59.065357Z",
     "iopub.status.busy": "2024-05-07T14:51:59.064898Z",
     "iopub.status.idle": "2024-05-07T14:51:59.109586Z",
     "shell.execute_reply": "2024-05-07T14:51:59.108752Z"
    }
   },
   "outputs": [],
   "source": [
    "@qfunc\n",
    "def main(\n",
    "    w: CArray[CReal, num_weights_in_encoder],\n",
    "    input: CArray[CReal, NUM_QUBITS],\n",
    "    trash: Output[QArray[QBit, num_trash_qubits]],\n",
    "    coded: Output[QArray[QBit, NUM_ENCODING_QUBITS]],\n",
    "    test: Output[QBit],\n",
    ") -> None:\n",
    "    x = QArray(\"x\")\n",
    "    psi2 = QArray(\"psi2\")\n",
    "    allocate(num_trash_qubits, psi2)\n",
    "    angle_encoding(exe_params=input, qbv=x)\n",
    "    encoder_ansatz(\n",
    "        num_qubits=NUM_QUBITS,\n",
    "        num_encoding_qubits=NUM_ENCODING_QUBITS,\n",
    "        exe_params=w,\n",
    "        x=x,\n",
    "        trash=trash,\n",
    "        coded=coded,\n",
    "    )\n",
    "\n",
    "    swap_test(state1=trash, state2=psi2, test=test)\n",
    "\n",
    "\n",
    "ae_qmod = create_model(main)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4597538f",
   "metadata": {},
   "source": [
    "We synthesize and visualize the quantum layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1ef859c2-02ef-46e7-a4fc-ff02ae77b28a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:51:59.114780Z",
     "iopub.status.busy": "2024-05-07T14:51:59.113594Z",
     "iopub.status.idle": "2024-05-07T14:52:01.559527Z",
     "shell.execute_reply": "2024-05-07T14:52:01.558874Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Opening: https://platform.classiq.io/circuit/dae2bad0-c23b-4a53-966e-3ea1c721051c?version=0.41.0.dev39%2B79c8fd0855\n"
     ]
    }
   ],
   "source": [
    "qprog = synthesize(ae_qmod)\n",
    "show(qprog)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a41ccaa",
   "metadata": {},
   "source": [
    "## The network\n",
    "\n",
    "The network for training contains only a quantum layer. The corresponding quantum program was already defined above, what is left is to define some execution preferences and the classical post-process. The classical output is defined as $1-\\alpha^2$, with $\\alpha$ being the probability of the test qubit to be at state 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "20c207bc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:01.562592Z",
     "iopub.status.busy": "2024-05-07T14:52:01.562150Z",
     "iopub.status.idle": "2024-05-07T14:52:02.797968Z",
     "shell.execute_reply": "2024-05-07T14:52:02.797311Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from classiq.applications.qnn import QLayer\n",
    "from classiq.applications.qnn.types import (\n",
    "    MultipleArguments,\n",
    "    ResultsCollection,\n",
    "    SavedResult,\n",
    ")\n",
    "from classiq.execution import (\n",
    "    ExecutionPreferences,\n",
    "    execute_qnn,\n",
    "    set_quantum_program_execution_preferences,\n",
    ")\n",
    "from classiq.synthesis import SerializedQuantumProgram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9f55ec56",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:02.801081Z",
     "iopub.status.busy": "2024-05-07T14:52:02.800375Z",
     "iopub.status.idle": "2024-05-07T14:52:02.804668Z",
     "shell.execute_reply": "2024-05-07T14:52:02.804076Z"
    }
   },
   "outputs": [],
   "source": [
    "num_shots = 4096\n",
    "\n",
    "\n",
    "def execute(\n",
    "    quantum_program: SerializedQuantumProgram, arguments: MultipleArguments\n",
    ") -> ResultsCollection:\n",
    "    quantum_program = set_quantum_program_execution_preferences(\n",
    "        quantum_program, preferences=ExecutionPreferences(num_shots=num_shots)\n",
    "    )\n",
    "    return execute_qnn(quantum_program, arguments)\n",
    "\n",
    "\n",
    "def post_process(result: SavedResult) -> torch.Tensor:\n",
    "    alpha_sqaured = result.value.counts_of_output(\"test\")[\"0\"] / num_shots\n",
    "    out = 1 - alpha_sqaured\n",
    "    return torch.tensor(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d7913804",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:02.807180Z",
     "iopub.status.busy": "2024-05-07T14:52:02.806812Z",
     "iopub.status.idle": "2024-05-07T14:52:02.815094Z",
     "shell.execute_reply": "2024-05-07T14:52:02.814520Z"
    }
   },
   "outputs": [],
   "source": [
    "def create_net(*args, **kwargs) -> nn.Module:\n",
    "    class Net(nn.Module):\n",
    "        def __init__(self, *args, **kwargs):\n",
    "            super().__init__()\n",
    "\n",
    "            self.qlayer = QLayer(\n",
    "                qprog,\n",
    "                execute,\n",
    "                post_process,\n",
    "                *args,\n",
    "                **kwargs,\n",
    "            )\n",
    "\n",
    "        def forward(self, x):\n",
    "            x = self.qlayer(x)\n",
    "            return x\n",
    "\n",
    "    return Net(*args, **kwargs)\n",
    "\n",
    "\n",
    "encoder_train_network = create_net()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7590c55d",
   "metadata": {},
   "source": [
    "## Creating dataset\n",
    "\n",
    "The cost function we would like to minimize is $|1-\\alpha^2|$ for all our training data. Looking at the qlayer output this means that we should define the corresponding labels as $0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "931ecd5a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:02.817442Z",
     "iopub.status.busy": "2024-05-07T14:52:02.817102Z",
     "iopub.status.idle": "2024-05-07T14:52:02.820913Z",
     "shell.execute_reply": "2024-05-07T14:52:02.820320Z"
    }
   },
   "outputs": [],
   "source": [
    "class MyDWDataset:\n",
    "    def __init__(self, data, labels) -> None:\n",
    "        self.data = torch.from_numpy(data).float()\n",
    "        self.labels = torch.unsqueeze(torch.from_numpy(labels), dim=-1).float()\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.data.shape[0]\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.data[idx], self.labels[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "775728a9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:02.823056Z",
     "iopub.status.busy": "2024-05-07T14:52:02.822880Z",
     "iopub.status.idle": "2024-05-07T14:52:02.827516Z",
     "shell.execute_reply": "2024-05-07T14:52:02.826900Z"
    }
   },
   "outputs": [],
   "source": [
    "labels = np.array([0, 0, 0])\n",
    "train_dataset = MyDWDataset(domain_wall_data, labels)\n",
    "train_data_loader = DataLoader(\n",
    "    train_dataset, batch_size=2, shuffle=True, drop_last=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "608173a0",
   "metadata": {},
   "source": [
    "## Define the training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "938eab77",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:02.829661Z",
     "iopub.status.busy": "2024-05-07T14:52:02.829493Z",
     "iopub.status.idle": "2024-05-07T14:52:02.833752Z",
     "shell.execute_reply": "2024-05-07T14:52:02.833184Z"
    }
   },
   "outputs": [],
   "source": [
    "import time as time\n",
    "\n",
    "\n",
    "def train(\n",
    "    model: nn.Module,\n",
    "    data_loader: DataLoader,\n",
    "    loss_func: nn.modules.loss._Loss,\n",
    "    optimizer: optim.Optimizer,\n",
    "    epoch: int = 40,\n",
    ") -> None:\n",
    "    for index in range(epoch):\n",
    "        start = time.time()\n",
    "        for data, label in data_loader:\n",
    "            optimizer.zero_grad()\n",
    "            output = model(data)\n",
    "            loss = loss_func(torch.squeeze(output), torch.squeeze(label))\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        print(time.time() - start)\n",
    "        print(index, f\"\\tloss = {loss.item()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee98061f",
   "metadata": {},
   "source": [
    "## Setting some hyper-parameters\n",
    "\n",
    "The L1 loss function fits the intended cost function we aim to minimize."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ee0a0551",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:02.836010Z",
     "iopub.status.busy": "2024-05-07T14:52:02.835828Z",
     "iopub.status.idle": "2024-05-07T14:52:02.838960Z",
     "shell.execute_reply": "2024-05-07T14:52:02.838372Z"
    }
   },
   "outputs": [],
   "source": [
    "_LEARNING_RATE = 0.3\n",
    "loss_func = nn.L1Loss()\n",
    "optimizer = optim.SGD(encoder_train_network.parameters(), lr=_LEARNING_RATE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e38d76f1",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "In this demo we will initialize the network with trained parameters, and run only 1 epoch for demonstration. A reasonable training with the above hyper-parameters can be achieved with $\\sim 40$ epochs. To train the network from the beginning uncomment the following code line:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "868a106c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:02.841035Z",
     "iopub.status.busy": "2024-05-07T14:52:02.840865Z",
     "iopub.status.idle": "2024-05-07T14:52:02.844053Z",
     "shell.execute_reply": "2024-05-07T14:52:02.843452Z"
    }
   },
   "outputs": [],
   "source": [
    "trained_weights = torch.nn.Parameter(\n",
    "    torch.Tensor(\n",
    "        [1.5227, 0.3588, 0.6905, 1.4777, 1.5718, 1.5615, 1.5414, 0.6021, 0.1254, 0.9903]\n",
    "    )\n",
    ")\n",
    "encoder_train_network.qlayer.weight = trained_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "4a0f11fa",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:02.846070Z",
     "iopub.status.busy": "2024-05-07T14:52:02.845891Z",
     "iopub.status.idle": "2024-05-07T14:52:09.887741Z",
     "shell.execute_reply": "2024-05-07T14:52:09.887070Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7.038687229156494\n",
      "0 \tloss = 0.00146484375\n"
     ]
    }
   ],
   "source": [
    "data_loader = train_data_loader\n",
    "\n",
    "train(encoder_train_network, data_loader, loss_func, optimizer, epoch=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fb11e18",
   "metadata": {},
   "source": [
    "## Verification\n",
    "\n",
    "Once we trained our network, we can build a new network with the trained variables. We can thus verify our encoder by taking only the encoding block, changing post_process, etc.\n",
    "\n",
    "Below we verify our quantum auto encoder by comparing between input and output of an encoder-decoder network:\n",
    "\n",
    "We will create the following network, which contains two quantum blocks:\n",
    "* The first two blocks of the previous network: a block for loading the inputs followed by our quantum encoder.\n",
    "* The inverse of the quantum encoder, where the inputs for the trash qubits are new zero inputs.\n",
    "\n",
    "**The network's weights will be allocated with the trained ones**\n",
    "\n",
    "<center>\n",
    "<img src=\"https://docs.classiq.io/resources/encoder_decoder.png\" style=\"width:100%\">\n",
    "<figcaption align = \"middle\"> Qlayer for verifying quantum auto encoder, containing three quantum blocks: (1) data encoding, (2) trained encoder, and (3) inverse of trained encoder (i.e., decoder). </figcaption>\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "812d2bc5",
   "metadata": {},
   "source": [
    "### We start with building the quantum layer for the validator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "952545a9-6e50-4a62-8538-f896220d15de",
   "metadata": {},
   "source": [
    "The validator contains the inverse of our encoder, for this, we will have to wrap our encoder function. This is because the `invert` quantum function except a single inout port."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "e2eb15db-866e-4ac1-8b7d-a3efa239f994",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:09.890580Z",
     "iopub.status.busy": "2024-05-07T14:52:09.890066Z",
     "iopub.status.idle": "2024-05-07T14:52:09.894377Z",
     "shell.execute_reply": "2024-05-07T14:52:09.893824Z"
    }
   },
   "outputs": [],
   "source": [
    "@qfunc\n",
    "def encoder_ansatz_wrapper(\n",
    "    num_qubits: CInt,\n",
    "    num_encoding_qubits: CInt,\n",
    "    exe_params: CArray[CReal],\n",
    "    qbv: QArray[QBit, \"num_qubits\"],\n",
    ") -> None:\n",
    "    coded = QArray(\"coded\")\n",
    "    trash = QArray(\"trash\")\n",
    "    encoder_ansatz(\n",
    "        num_qubits=num_qubits,\n",
    "        num_encoding_qubits=num_encoding_qubits,\n",
    "        exe_params=exe_params,\n",
    "        x=qbv,\n",
    "        trash=trash,\n",
    "        coded=coded,\n",
    "    )\n",
    "    bind([coded, trash], qbv)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e01fcd94-072b-4483-8373-404826057669",
   "metadata": {},
   "source": [
    "We are now ready for modeling our verification quantum layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e87cfbcc-cc20-4728-a2b9-e8abf18ec52d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:09.896624Z",
     "iopub.status.busy": "2024-05-07T14:52:09.896454Z",
     "iopub.status.idle": "2024-05-07T14:52:09.926569Z",
     "shell.execute_reply": "2024-05-07T14:52:09.926000Z"
    }
   },
   "outputs": [],
   "source": [
    "@qfunc\n",
    "def main(\n",
    "    w: CArray[CReal, num_weights_in_encoder],\n",
    "    input: CArray[CReal, NUM_QUBITS],\n",
    "    decoded: Output[QArray[QBit, NUM_QUBITS]],\n",
    "    trash: Output[QArray[QBit, num_trash_qubits]],\n",
    ") -> None:\n",
    "    psi2 = QArray(\"psi2\")\n",
    "    coded = QArray(\"coded\")\n",
    "    allocate(num_trash_qubits, psi2)\n",
    "    angle_encoding(exe_params=input, qbv=decoded)\n",
    "    encoder_ansatz(\n",
    "        num_qubits=NUM_QUBITS,\n",
    "        num_encoding_qubits=NUM_ENCODING_QUBITS,\n",
    "        exe_params=w,\n",
    "        x=decoded,\n",
    "        trash=trash,\n",
    "        coded=coded,\n",
    "    )\n",
    "\n",
    "    bind([coded, psi2], decoded)\n",
    "    invert(\n",
    "        stmt_block=lambda: encoder_ansatz_wrapper(\n",
    "            num_qubits=NUM_QUBITS,\n",
    "            num_encoding_qubits=NUM_ENCODING_QUBITS,\n",
    "            exe_params=w,\n",
    "            qbv=decoded,\n",
    "        ),\n",
    "    )\n",
    "\n",
    "\n",
    "qmod_validator = create_model(main)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "02ed5fc2-85a9-4473-bbbb-a443a765919d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:09.928600Z",
     "iopub.status.busy": "2024-05-07T14:52:09.928429Z",
     "iopub.status.idle": "2024-05-07T14:52:12.227982Z",
     "shell.execute_reply": "2024-05-07T14:52:12.227270Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Opening: https://platform.classiq.io/circuit/9aee699f-59eb-4d4e-b543-c0e346619889?version=0.41.0.dev39%2B79c8fd0855\n"
     ]
    }
   ],
   "source": [
    "qprog_validator = synthesize(qmod_validator)\n",
    "show(qprog_validator)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f3f9983",
   "metadata": {},
   "source": [
    "### Next, we define the classical output of the network. For the validator post-process we take the output with the maximal counts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "41266c52",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:12.232075Z",
     "iopub.status.busy": "2024-05-07T14:52:12.231653Z",
     "iopub.status.idle": "2024-05-07T14:52:12.236387Z",
     "shell.execute_reply": "2024-05-07T14:52:12.235709Z"
    }
   },
   "outputs": [],
   "source": [
    "def execute_validator(\n",
    "    quantum_program: SerializedQuantumProgram, arguments: MultipleArguments\n",
    ") -> ResultsCollection:\n",
    "    return execute_qnn(quantum_program, arguments)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "e5e8f6f2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:12.239751Z",
     "iopub.status.busy": "2024-05-07T14:52:12.239371Z",
     "iopub.status.idle": "2024-05-07T14:52:12.243655Z",
     "shell.execute_reply": "2024-05-07T14:52:12.242996Z"
    }
   },
   "outputs": [],
   "source": [
    "def post_process_validator(result: SavedResult) -> torch.Tensor:\n",
    "    counts = result.value.counts_of_output(\"decoded\")\n",
    "\n",
    "    max_key = max(counts, key=counts.get)\n",
    "\n",
    "    return torch.tensor([int(k) for k in max_key])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39a94e36",
   "metadata": {},
   "source": [
    "### We create the network and assign the trained parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "43b1a4f8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:12.247036Z",
     "iopub.status.busy": "2024-05-07T14:52:12.246591Z",
     "iopub.status.idle": "2024-05-07T14:52:12.262065Z",
     "shell.execute_reply": "2024-05-07T14:52:12.261452Z"
    }
   },
   "outputs": [],
   "source": [
    "def create_encoder_decoder_net(*args, **kwargs) -> nn.Module:\n",
    "    class Net(nn.Module):\n",
    "        def __init__(self, *args, **kwargs):\n",
    "            super().__init__()\n",
    "\n",
    "            self.qlayer = QLayer(\n",
    "                qprog_validator,\n",
    "                execute_validator,\n",
    "                post_process_validator,\n",
    "                *args,\n",
    "                **kwargs,\n",
    "            )\n",
    "\n",
    "        def forward(self, x):\n",
    "            x = self.qlayer(x)\n",
    "            return x\n",
    "\n",
    "    return Net(*args, **kwargs)\n",
    "\n",
    "\n",
    "validator_network = create_encoder_decoder_net()\n",
    "validator_network.qlayer.weight = encoder_train_network.qlayer.weight"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0db9d1a8",
   "metadata": {},
   "source": [
    "Now we can compare between the input and output of the validator, for each different data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "28c93e9e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:12.266494Z",
     "iopub.status.busy": "2024-05-07T14:52:12.265482Z",
     "iopub.status.idle": "2024-05-07T14:52:16.622434Z",
     "shell.execute_reply": "2024-05-07T14:52:16.621830Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input= [0.0, 0.0, 1.0, 1.0] ,   output= [0.0, 0.0, 1.0, 1.0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input= [0.0, 0.0, 0.0, 1.0] ,   output= [0.0, 0.0, 0.0, 1.0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input= [0.0, 1.0, 1.0, 1.0] ,   output= [0.0, 1.0, 1.0, 1.0]\n"
     ]
    }
   ],
   "source": [
    "validator_data_loader = DataLoader(\n",
    "    train_dataset, batch_size=1, shuffle=True, drop_last=False\n",
    ")\n",
    "\n",
    "for data, label in validator_data_loader:\n",
    "    output = validator_network(data)\n",
    "    print(\"input=\", data.tolist()[0], \",   output=\", output.tolist()[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dca60d3",
   "metadata": {},
   "source": [
    "# Usage for anomaly detection"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2479f917",
   "metadata": {},
   "source": [
    "We can use our trained network for anomaly detection. Let see what happens to the trash qubits when we insert anomaly, namely, a non domain-wall data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "31253024",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:16.624992Z",
     "iopub.status.busy": "2024-05-07T14:52:16.624545Z",
     "iopub.status.idle": "2024-05-07T14:52:16.628713Z",
     "shell.execute_reply": "2024-05-07T14:52:16.628149Z"
    }
   },
   "outputs": [],
   "source": [
    "input_anomaly_data = np.array(\n",
    "    [[0, 0, 1, 1], [0, 0, 0, 1], [0, 1, 1, 1], [1, 0, 1, 0], [1, 1, 1, 1]]\n",
    ")\n",
    "anomaly_labels = np.array([0, 0, 0, 0, 0])\n",
    "anomaly_dataset = MyDWDataset(input_anomaly_data, anomaly_labels)\n",
    "anomaly_data_loader = DataLoader(\n",
    "    anomaly_dataset, batch_size=1, shuffle=True, drop_last=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f45de115",
   "metadata": {},
   "source": [
    "We print all the anomaly data based on some pre-defined accuracy for the cost function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "e9ac2d78",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:16.631149Z",
     "iopub.status.busy": "2024-05-07T14:52:16.630672Z",
     "iopub.status.idle": "2024-05-07T14:52:23.745511Z",
     "shell.execute_reply": "2024-05-07T14:52:23.743996Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "anomaly: [1.0, 1.0, 1.0, 1.0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "anomaly: [1.0, 0.0, 1.0, 0.0]\n"
     ]
    }
   ],
   "source": [
    "tolerance = 1e-2\n",
    "for data, label in anomaly_data_loader:\n",
    "    output = encoder_train_network(data)\n",
    "    if abs(output.tolist()[0]) > tolerance:\n",
    "        print(\"anomaly:\", data.tolist()[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1949ec76",
   "metadata": {},
   "source": [
    "# Alternative network for training a quantum auto encoder\n",
    "\n",
    "Another way to introduce a cost function is via estimation of Hamiltonians. Measuring the Pauli $Z$ matrix on a qubit at the general state $|q\\rangle=a|0\\rangle+b|1\\rangle$ is $\\langle q |Z|q \\rangle=a^2-b^2$. Therefore, a cost function can be defined by taking expectation values on the trash output (without a swap-test) as follows:\n",
    "$$\n",
    "\\text{Cost} = \\frac{1}{2}\\sum^{\\text{num of trash qubits}}_{k=1} 1 - \\langle Z_k \\rangle.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "188687e5-95c1-4841-ac59-3d47de5e1748",
   "metadata": {},
   "source": [
    "Below we show how the define the corresponding Qlayer: the quantum program and post-processing:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cf8c081-3853-4d43-b003-5a3aa50ebf12",
   "metadata": {},
   "source": [
    "## The quantum program"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "cdcab73f-a9cd-4779-9043-ee1813d82b38",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:23.750598Z",
     "iopub.status.busy": "2024-05-07T14:52:23.749376Z",
     "iopub.status.idle": "2024-05-07T14:52:26.038135Z",
     "shell.execute_reply": "2024-05-07T14:52:26.037439Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Opening: https://platform.classiq.io/circuit/00ce49a6-7aa2-4e35-9295-d598d672b71d?version=0.41.0.dev39%2B79c8fd0855\n"
     ]
    }
   ],
   "source": [
    "@qfunc\n",
    "def main(\n",
    "    w: CArray[CReal, num_weights_in_encoder],\n",
    "    input: CArray[CReal, NUM_QUBITS],\n",
    "    trash: Output[QArray[QBit, num_trash_qubits]],\n",
    ") -> None:\n",
    "    x = QArray(\"x\")\n",
    "    coded = QArray(\"coded\")\n",
    "    angle_encoding(exe_params=input, qbv=x)\n",
    "    encoder_ansatz(\n",
    "        num_qubits=NUM_QUBITS,\n",
    "        num_encoding_qubits=NUM_ENCODING_QUBITS,\n",
    "        exe_params=w,\n",
    "        x=x,\n",
    "        trash=trash,\n",
    "        coded=coded,\n",
    "    )\n",
    "\n",
    "\n",
    "qmod_ae = create_model(main)\n",
    "\n",
    "qprog_ae = synthesize(qmod_ae)\n",
    "show(qprog_ae)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1160bc2e-84ea-42de-b48c-9a2113bee98c",
   "metadata": {},
   "source": [
    "## Execution and post-process\n",
    "\n",
    "The trash register is of size 2, we measure the Pauli $Z$ matrix on each of its qubits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "a56ea5a2-b040-4923-a8fa-1b46faeb5d1a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:26.040751Z",
     "iopub.status.busy": "2024-05-07T14:52:26.040244Z",
     "iopub.status.idle": "2024-05-07T14:52:26.044916Z",
     "shell.execute_reply": "2024-05-07T14:52:26.044072Z"
    }
   },
   "outputs": [],
   "source": [
    "from classiq.applications.chemistry import PauliOperator, PauliOperators\n",
    "\n",
    "\n",
    "def execute(\n",
    "    quantum_program: SerializedQuantumProgram, arguments: MultipleArguments\n",
    ") -> ResultsCollection:\n",
    "    return execute_qnn(\n",
    "        quantum_program,\n",
    "        arguments,\n",
    "        observable=PauliOperator(pauli_list=[(\"IZ\", 1), (\"ZI\", 1)]),\n",
    "    )\n",
    "\n",
    "\n",
    "def post_process(result: SavedResult) -> torch.Tensor:\n",
    "    out = 1 / 2 * (2 - np.real(result.value.value))\n",
    "    return torch.tensor(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "f87fef6b-36b3-4707-b1cd-7042b0ff053c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-05-07T14:52:26.048898Z",
     "iopub.status.busy": "2024-05-07T14:52:26.048111Z",
     "iopub.status.idle": "2024-05-07T14:52:26.059915Z",
     "shell.execute_reply": "2024-05-07T14:52:26.058954Z"
    }
   },
   "outputs": [],
   "source": [
    "def create_net(*args, **kwargs) -> nn.Module:\n",
    "    class Net(nn.Module):\n",
    "        def __init__(self, *args, **kwargs):\n",
    "            super().__init__()\n",
    "\n",
    "            self.qlayer = QLayer(\n",
    "                qprog_ae,\n",
    "                execute,\n",
    "                post_process,\n",
    "                *args,\n",
    "                **kwargs,\n",
    "            )\n",
    "\n",
    "        def forward(self, x):\n",
    "            x = self.qlayer(x)\n",
    "            return x\n",
    "\n",
    "    return Net(*args, **kwargs)\n",
    "\n",
    "\n",
    "encoder_train_network = create_net()"
   ]
  }
 ],
 "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}