algorithms/qml/quantum_autoencoder/quantum_autoencoder.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "88da6df0",
   "metadata": {},
   "source": [
    "# Quantum Autoencoder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "201ce387",
   "metadata": {},
   "source": [
    "## Encoder Types\n",
    "\n",
    "### Classical Encoders\n",
    "Encode or compress classical data into smaller-sized data via a deterministic algorithm. For example, JPEG is essentially an algorithm that compresses images into smaller-sized images.\n",
    "\n",
    "### Classical Autoencoders\n",
    "Use machine-learning techniques and train a variational network for compressing data. In general, an autoencoder 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 autoencoder layout (from Wikipedia) </figcaption>\n",
    "</center>\n",
    "\n",
    "The network has three main parts: \n",
    "\n",
    "1. The encoder compresses the data into a smaller, coded layer. \n",
    "2. The latter is the input to a decoder part. \n",
    "3. Typically, training is done against the comparison between the input and the output of this network.\n",
    "\n",
    "**Classical autoencoders can also be used for anomaly detection (see below).**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5be9dbfd",
   "metadata": {},
   "source": [
    "### Quantum Autoencoders\n",
    "In a similar fashion to the classical counterpart, a quantum autoencoder compresses quantum data stored initially on $n$ qubits into a smaller quantum register of $m<n$ qubits via a variational circuit. However, quantum computing is reversible; therefore, qubits cannot be \"erased\". Alternatively, a quantum autoencoder tries to achieve the following transformation from an 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 that is in the zero state. The former is usually called the *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 autoencoder layout: uncoded data of size 5 transforms into two outputs: a coded register of size 3 and trash output of size 2 at state $|00\\rangle$ </figcaption>\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e40c41f0",
   "metadata": {},
   "source": [
    "## Training Quantum Autoencoders\n",
    "\n",
    "To train a quantum autoencoder, we define a proper cost function. Below are two common approaches, one using a swap test and the other using 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 that 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 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 Networks for Quantum Autoencoders\n",
    "\n",
    "The quantum autoencoder can be built as a quantum neural network with these parts:\n",
    "\n",
    "1. A data loading block that loads classical data on $n$ qubits.\n",
    "2. An encoder block, which is a variational quantum ansatz with 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 autoencoder 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": [
    "## Predefined Functions That Construct the Quantum Layer\n",
    "\n",
    "In the first step we build user-defined functions that allow flexible modeling:\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 the description in 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": [
    "## Example: Autoencoder for Domain Wall Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "428691bc",
   "metadata": {},
   "source": [
    "In the following example we 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 encode this data of size 4 on 2 qubits. Let us build the corresponding quantum layer based on the predefined 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": [
    "We 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, so what remains is to define the execution preferences and the classical postprocess. The classical output is defined as $1-\\alpha^2$, with $\\alpha$ being the probability of the test qubit being 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 the Dataset\n",
    "\n",
    "The cost function 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": [
    "### Defining 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 Hyperparameters\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 initialize the network with trained parameters and run only one epoch for demonstration purposes. Reasonable training with the above hyperparameters 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": [
    "### Verifying\n",
    "\n",
    "Once we have trained the network, we can build a new network with the trained variables. We verify our encoder by taking only the encoding block, changing the postprocess, etc. \n",
    "Below, we verify our quantum autoencoder by comparing the input with the output of an encoder-decoder network.\n",
    "\n",
    "We create the following network containing 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 weights are allocated together 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 the quantum autoencoder 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": [
    "### 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. We have to wrap the encoder function because the `invert` quantum function accepts a single in/out 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 to model 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": [
    "### Defining the Classical Output of the Network \n",
    "\n",
    "For the validator postprocessing, we take the output with the maximum 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": [
    "### Creating the Network and Assigning 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 the input with the output of the validator for 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": [
    "### Detecting Anomalies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2479f917",
   "metadata": {},
   "source": [
    "We can use our trained network for anomaly detection. Let's see what happens to the trash qubits when we insert an anomaly; namely, 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 predefined 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 Autoencoder\n",
    "\n",
    "Another way to introduce a cost function is by estimating 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 to define the corresponding Qlayer: the quantum program and postprocessing."
   ]
  },
  {
   "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_alt = create_model(main)\n",
    "\n",
    "qprog_ae_alt = synthesize(qmod_ae_alt)\n",
    "show(qprog_ae_alt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1160bc2e-84ea-42de-b48c-9a2113bee98c",
   "metadata": {},
   "source": [
    "### Executing and Postprocessing\n",
    "\n",
    "The size of the trash register is 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_alt,\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
}