{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries\n",
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "# Load test module for sanity check\n",
    "from test_utils import test"
   ]
  },
  {
   "cell_type": "markdown",
   "execution_count": null,
   "metadata": {},
   "source": [
    "Data Generation\n",
    "==="
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(10)\n",
    "P, Q = (np.random.rand(i, 2) for i in (4, 5))\n",
    "P_big, Q_big = (np.random.rand(i, 80) for i in (100, 120))\n",
    "\n",
    "print(P, \"\\n\\n\", Q)"
   ]
  },
  {
   "cell_type": "markdown",
   "execution_count": null,
   "metadata": {},
   "source": [
    "Solution\n",
    "==="
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def naive(P, Q):\n",
    "    \"\"\"\n",
    "    A naive solution for finding pairwise distances between poins in P and Q\n",
    "\n",
    "    Args:\n",
    "        P: numpy array of shape=(p, 2)\n",
    "        Q: numpy array of shape=(q, 2)\n",
    "    Returns:\n",
    "        D: numpy array of shape=(p, q)\n",
    "\n",
    "    >>> naive(np.array([[0, 1]]), np.array([[2, 3], [4, 5]]))\n",
    "    array([[2.82842712, 5.65685425]])\n",
    "    \"\"\"\n",
    "    # ***************************************************\n",
    "    # INSERT YOUR CODE HERE\n",
    "    # TODO: implement a naive solution\n",
    "    # ***************************************************\n",
    "    raise NotImplementedError\n",
    "\n",
    "\n",
    "test(naive)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "execution_count": null,
   "metadata": {},
   "source": [
    "### Use matching indices\n",
    "\n",
    "Instead of iterating through indices, one can use them directly to parallelize the operations with Numpy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rows, cols = np.indices((P.shape[0], Q.shape[0]))\n",
    "print(rows, end=\"\\n\\n\")\n",
    "print(cols)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(P[rows.ravel()], end=\"\\n\\n\")\n",
    "print(Q[cols.ravel()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def with_indices(P, Q):\n",
    "    \"\"\"\n",
    "    An optimized solution using matching indices\n",
    "\n",
    "    Args:\n",
    "        P: numpy array of shape=(p, 2)\n",
    "        Q: numpy array of shape=(q, 2)\n",
    "    Returns:\n",
    "        D: numpy array of shape=(p, q)\n",
    "\n",
    "    >>> with_indices(np.array([[0, 1]]), np.array([[2, 3], [4, 5]]))\n",
    "    array([[2.82842712, 5.65685425]])\n",
    "    \"\"\"\n",
    "    # ***************************************************\n",
    "    # INSERT YOUR CODE HERE\n",
    "    # TODO: implement an optimized solution\n",
    "    # ***************************************************\n",
    "    raise NotImplementedError\n",
    "\n",
    "\n",
    "test(with_indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "execution_count": null,
   "metadata": {},
   "source": [
    "### Use a library\n",
    "\n",
    "`scipy` is the equivalent of matlab toolboxes and have a lot to offer. Actually the pairwise computation is part of the library through the `spatial` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.spatial.distance import cdist\n",
    "\n",
    "\n",
    "def scipy_version(P, Q):\n",
    "    \"\"\"\n",
    "    A solution using scipy\n",
    "\n",
    "    Args:\n",
    "        P: numpy array of shape=(p, 2)\n",
    "        Q: numpy array of shape=(q, 2)\n",
    "\n",
    "    Returns:\n",
    "        D: numpy array of shape=(p, q)\n",
    "\n",
    "    >>> scipy_version(np.array([[0, 1]]), np.array([[2, 3], [4, 5]]))\n",
    "    array([[2.82842712, 5.65685425]])\n",
    "    \"\"\"\n",
    "    return cdist(P, Q)"
   ]
  },
  {
   "cell_type": "markdown",
   "execution_count": null,
   "metadata": {},
   "source": [
    "### Numpy Magic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tensor_broadcasting(P, Q):\n",
    "    \"\"\"\n",
    "    A solution using tensor broadcasting\n",
    "\n",
    "    Args:\n",
    "        P: numpy array of shape=(p, 2)\n",
    "        Q: numpy array of shape=(q, 2)\n",
    "\n",
    "    Returns:\n",
    "        D: numpy array of shape=(p, q)\n",
    "\n",
    "    >>> tensor_broadcasting(np.array([[0, 1]]), np.array([[2, 3], [4, 5]]))\n",
    "    array([[2.82842712, 5.65685425]])\n",
    "    \"\"\"\n",
    "    return np.sqrt(np.sum((P[:, np.newaxis, :] - Q[np.newaxis, :, :]) ** 2, axis=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "execution_count": null,
   "metadata": {},
   "source": [
    "# Compare methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "methods = [\n",
    "    naive,\n",
    "    naive_2,  # This is another possible solution. Feel free to comment it out if you have only one solution.\n",
    "    with_indices,\n",
    "    with_indices_2,  # This is another possible solution. Feel free to comment it out if you have only one solution.\n",
    "    scipy_version,\n",
    "    tensor_broadcasting,\n",
    "]\n",
    "timers = []\n",
    "for f in methods:\n",
    "    r = %timeit -o f(P_big, Q_big)\n",
    "    timers.append(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 6))\n",
    "plt.bar(\n",
    "    np.arange(len(methods)), [r.best * 1000 for r in timers], log=False\n",
    ")  # Set log to True for logarithmic scale\n",
    "plt.xticks(np.arange(len(methods)) + 0.2, [f.__name__ for f in methods], rotation=30)\n",
    "plt.xlabel(\"Method\")\n",
    "plt.ylabel(\"Time (ms)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}