{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quantum circuit Born machine (QCBM)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial you will implement your first QCBM and train it to generate a Gaussian distribution over integers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "We first sample a trainingset from the Gaussian distribution. \n",
    "Then, we implement an ansatz of alternating parametrized rotation layers and entangling layers. \n",
    "Finally, we optimize the parameters with respect to the squared MMD-loss and test our trained model by plotting the distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import minimize\n",
    "from time import time\n",
    "from functools import partial\n",
    "import cirq, sympy\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Ansatz hyperparameters\n",
    "n_qubits = 3\n",
    "depth = 2\n",
    "n_params = 2 * depth * n_qubits\n",
    "\n",
    "# Begin with statevector simulator\n",
    "shots = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sampling the trainingdata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gaussian_pdf(num_bit, mu, sigma):\n",
    "    '''get gaussian distribution function'''\n",
    "    x = np.arange(2**num_bit)\n",
    "    pl = 1. / np.sqrt(2 * np.pi * sigma**2) * \\\n",
    "        np.exp(-(x - mu)**2 / (2. * sigma**2))\n",
    "    return pl/pl.sum()\n",
    "\n",
    "pg = gaussian_pdf(n_qubits, mu=2**(n_qubits-1)-0.5, sigma=2**(n_qubits-2))\n",
    "plt.plot(pg, 'ro')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementing the ansatz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Layer of single qubit z rotations\n",
    "def rot_z_layer(n_qubits, parameters):\n",
    "    if n_qubits != len(parameters):\n",
    "        raise ValueError(\"Too many or few parameters, must equal n_qubits\")\n",
    "    for i in range(n_qubits):\n",
    "        yield cirq.Rz(2 * parameters[i])(cirq.GridQubit(i, 0))\n",
    "\n",
    "# Layer of single qubit y rotations\n",
    "def rot_y_layer(n_qubits, parameters):\n",
    "    if n_qubits != len(parameters):\n",
    "        raise ValueError(\"Too many of few parameters, must equal n_qubits\")\n",
    "    for i in range(n_qubits):\n",
    "        yield cirq.Ry(parameters[i])(cirq.GridQubit(i, 0))\n",
    "\n",
    "# Layer of entangling CZ(i,i+1 % n_qubits) gates.\n",
    "def entangling_layer(n_qubits):\n",
    "    if n_qubits == 2:\n",
    "        yield cirq.CZ(cirq.GridQubit(0, 0), cirq.GridQubit(1, 0))\n",
    "        return\n",
    "    for i in range(n_qubits):\n",
    "        yield cirq.CZ(cirq.GridQubit(i, 0), cirq.GridQubit((i+1) % n_qubits, 0))\n",
    "\n",
    "# Variational circuit, i.e., the ansatz.\n",
    "def variational_circuit(n_qubits, depth, theta):\n",
    "    if len(theta) != (2 * depth * n_qubits):\n",
    "        raise ValueError(\"Theta of incorrect dimension, must equal 2*depth*n_qubits\")\n",
    "    \n",
    "    # Initializing qubits and circuit\n",
    "    qubits = [cirq.GridQubit(i, 0) for i in range(n_qubits)]\n",
    "    circuit = cirq.Circuit()\n",
    "    \n",
    "    # Adding layers of rotation gates and entangling gates.\n",
    "    for d in range(depth):\n",
    "        # Adding single qubit rotations\n",
    "        circuit.append(rot_z_layer(n_qubits, theta[d * 2 * n_qubits : (d+1) * 2 * n_qubits : 2]))\n",
    "        circuit.append(rot_y_layer(n_qubits, theta[d * 2 * n_qubits + 1 : (d+1) * 2 * n_qubits + 1 : 2]))\n",
    "        # Adding entangling layer\n",
    "        circuit.append(entangling_layer(n_qubits))\n",
    "    \n",
    "    return circuit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta_entry_symbols = [sympy.Symbol('theta_' + str(i)) for i in range(2 * n_qubits * depth)]\n",
    "theta_symbol = sympy.Matrix(theta_entry_symbols)\n",
    "ansatz = variational_circuit(n_qubits, depth, theta_symbol)\n",
    "print(ansatz.to_text_diagram(transpose=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Estimating the probabilities of our parameterized quantum circuit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Estimate all probabilities of the PQCs distribution.\n",
    "def estimate_probs(circuit, theta, n_shots=shots):\n",
    "    # Creating parameter resolve dict by adding state and theta.\n",
    "    try:\n",
    "        theta_mapping = [('theta_' + str(i), theta[i]) for i in range(len(theta))]\n",
    "    except IndexError as error:\n",
    "        print(\"Could not resolve theta symbol, array of wrong size.\")\n",
    "    resolve_dict = dict(theta_mapping)\n",
    "    resolver = cirq.ParamResolver(resolve_dict)\n",
    "    resolved_circuit = cirq.resolve_parameters(circuit, resolver)\n",
    "    \n",
    "    # Use statevector simulator\n",
    "    if n_shots == 0:\n",
    "        final_state = cirq.final_wavefunction(resolved_circuit)\n",
    "        probs = np.array([abs(final_state[i])**2 for i in range(len(final_state))])\n",
    "        \n",
    "    # Run the circuit.\n",
    "    else:\n",
    "        # Adding measurement at the end.\n",
    "        resolved_circuit.append(cirq.measure(*resolved_circuit.all_qubits(), key='m'))\n",
    "        results = cirq.sample(resolved_circuit, repetitions=n_shots)\n",
    "        frequencies = results.histogram(key='m')\n",
    "        probs = np.zeros(2**n_qubits)\n",
    "        for key, value in frequencies.items():\n",
    "            probs[key] = value / n_shots\n",
    "    \n",
    "    return probs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Kernel expectation\n",
    "The function kernel_expectation(px, py, kernel_matrix) evaluates:\n",
    "\n",
    "$$ \\mathbb{E}_{x \\sim p_x, y \\sim p_y}\\left[K(x, y) \\right],$$\n",
    "\n",
    "where $K$ is the kernel_matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function that computes the kernel for the MMD loss\n",
    "def multi_rbf_kernel(x, y, sigma_list):\n",
    "    '''\n",
    "    multi-RBF kernel.\n",
    "    \n",
    "    Args:\n",
    "        x (1darray|2darray): the collection of samples A.\n",
    "        x (1darray|2darray): the collection of samples B.\n",
    "        sigma_list (list): a list of bandwidths.\n",
    "        \n",
    "    Returns:\n",
    "        2darray: kernel matrix.\n",
    "    '''\n",
    "    ndim = x.ndim\n",
    "    if ndim == 1:\n",
    "        exponent = np.abs(x[:, None] - y[None, :])**2\n",
    "    elif ndim == 2:\n",
    "        exponent = ((x[:, None, :] - y[None, :, :])**2).sum(axis=2)\n",
    "    else:\n",
    "        raise\n",
    "    K = 0.0\n",
    "    for sigma in sigma_list:\n",
    "        gamma = 1.0 / (2 * sigma)\n",
    "        K = K + np.exp(-gamma * exponent)\n",
    "    return K\n",
    "\n",
    "# Function that computes expectation of kernel in MMD loss\n",
    "def kernel_expectation(px, py, kernel_matrix):\n",
    "    return px.dot(kernel_matrix).dot(py)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Squared MMD loss\n",
    "Implement the squared_MMD_loss(p, q, kernel_matrix) function that computes:\n",
    "\n",
    "$$ L(p, q) = \\mathbb{E}_{x \\sim p, y \\sim p}\\left[K(x, y) \\right] - 2\\mathbb{E}_{x \\sim p, y \\sim q}\\left[K(x, y) \\right] + \\mathbb{E}_{x \\sim q, y \\sim q}\\left[K(x, y) \\right],$$\n",
    "\n",
    "where $K$ is the kernel_matrix. \n",
    "\n",
    "Also, implement the loss(theta, circuit, target, kernel_matrix, n_shots=shots) that computes:\n",
    "$$ L(p_{\\theta}, target),$$\n",
    "\n",
    "where $p_{\\theta}$ is the distribution induced by your parameterized quantum circuit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function that computes the squared MMD loss related to the given kernel_matrix.\n",
    "def squared_MMD_loss(probs, target, kernel_matrix):\n",
    "    # TODO: Implement function!\n",
    "    # MMD_loss = ...\n",
    "    return MMD_loss\n",
    "\n",
    "# The loss function that we aim to minimize.\n",
    "def loss(theta, circuit, target, kernel_matrix, n_shots=shots):\n",
    "    # TODO: Implement function!\n",
    "    # loss = ...\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Estimating the gradient\n",
    "Implement the gradient(theta, q, kernel_matrix, n_shots=shots) function that computes:\n",
    "\n",
    "$$ \\nabla_{\\theta}L(p_{\\theta}, q) = \\left(\\frac{\\partial L(p_\\theta, q)}{\\partial \\theta_j} \\right)_{j}.$$\n",
    "\n",
    "You should use that by the parameter-shift rule we have:\n",
    "\n",
    "$$ \\frac{\\partial L(p_\\theta, q)}{\\partial \\theta_j} = \\mathbb{E}_{x \\sim p_{\\theta^+}, y \\sim p_\\theta}\\left[K(x, y) \\right] - \\mathbb{E}_{x \\sim p_{\\theta^-}, y \\sim p_\\theta}\\left[K(x, y) \\right] - \\mathbb{E}_{x \\sim p_{\\theta^+}, y \\sim q}\\left[K(x, y) \\right] + \\mathbb{E}_{x \\sim p_{\\theta^-}, y \\sim q}\\left[K(x, y) \\right],$$\n",
    "\n",
    "where $\\theta^+ = \\theta + \\frac{\\pi}{2}e_j$ and $\\theta^- = \\theta - \\frac{\\pi}{2}e_j$ with $e_j$ denoting the $j$-th unit vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the gradient.\n",
    "def gradient(theta, target, kernel_matrix, n_shots=shots):\n",
    "    grad = []\n",
    "    for i in range(len(theta)):\n",
    "      # TODO: compute the entries of the gradient! \n",
    "      # grad[i] = ...\n",
    "    return np.array(grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training\n",
    "It is time to train our quantum circuit Borne machine! Please allow for 30secs of training time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MMD kernel\n",
    "basis = np.arange(2**n_qubits)\n",
    "sigma_list = [0.25,4]\n",
    "kernel_matrix = multi_rbf_kernel(basis, basis, sigma_list)\n",
    "\n",
    "# Initial theta\n",
    "theta0 = np.random.random(n_params)*2*np.pi\n",
    "\n",
    "# Initializing loss function with our ansatz, target and kernel matrix\n",
    "loss_ansatz = partial(loss, circuit=ansatz, target=pg, kernel_matrix=kernel_matrix)\n",
    "\n",
    "# Callback function to track status \n",
    "step = [0]\n",
    "tracking_cost = []\n",
    "def callback(x, *args, **kwargs):\n",
    "    step[0] += 1\n",
    "    tracking_cost.append(loss_ansatz(x))\n",
    "    print('step = %d, loss = %s'%(step[0], loss_ansatz(x)))\n",
    "\n",
    "# Training the QCBM.\n",
    "start_time = time()\n",
    "final_params = minimize(loss_ansatz,\n",
    "                        theta0, \n",
    "                        method=\"L-BFGS-B\", \n",
    "                        jac=partial(gradient, target=pg, kernel_matrix=kernel_matrix),\n",
    "                        tol=10**-5, \n",
    "                        options={'maxiter':50, 'disp': 0, 'gtol':1e-10, 'ftol':0}, \n",
    "                        callback=callback)\n",
    "end_time = time()\n",
    "print(end_time-start_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Looking at final parameters, its cost function and the generated distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(list(range(len(tracking_cost))), tracking_cost)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(estimate_probs(ansatz, final_params.x), 'ro')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}