Quadratic-Labs/PrivateNearestNeighbors-FHE

1{2 "cells": [3  {4   "cell_type": "markdown",5   "metadata": {},6   "source": [7    "# Nearest Example"8   ]9  },10  {11   "cell_type": "markdown",12   "metadata": {},13   "source": [14    "## Server's Data Setup\n",15    "The server owns coordinates to points of interest like restaurants and commerces. The coordinates are kept in a LookupTable"16   ]17  },18  {19   "cell_type": "code",20   "execution_count": 1,21   "metadata": {},22   "outputs": [],23   "source": [24    "from concrete import fhe\n",25    "import numpy\n",26    "\n",27    "\n",28    "# Database of Points of Interests\n",29    "points_array = numpy.array([\n",30    "    [2, 3], [1, 5], [3, 2], [5, 2], [1, 1],\n",31    "    [9, 4], [13, 2], [14, 13], [9, 8], [8, 0],\n",32    "    [2, 10], [3, 8], [8, 12], [4, 10], [7, 7],\n",33    "])\n",34    "N_PTS = points_array.shape[0]\n",35    "points = fhe.LookupTable(points_array.flatten())\n",36    "\n",37    "\n",38    "def get_point(index):\n",39    "    return (points[2*index], points[2*index + 1])\n",40    "\n",41    "\n",42    "def all_distances(x, y):\n",43    "    xs = numpy.arange(0, 2 * N_PTS, 2)\n",44    "    ys = numpy.arange(1, 2 * N_PTS, 2)\n",45    "    a = abs(points[xs] - x)\n",46    "    b = abs(points[ys] - y)\n",47    "    return a + b"48   ]49  },50  {51   "cell_type": "markdown",52   "metadata": {},53   "source": [54    "We use swap sort to find the $K$ nearest points to a given point. However, we are interested in the indices of the elements, not just their distances. We must therefore work on tuples of index and distance, effectively implementing numpy argpartition."55   ]56  },57  {58   "cell_type": "code",59   "execution_count": 2,60   "metadata": {},61   "outputs": [],62   "source": [63    "# TLUs\n",64    "relu = fhe.univariate(lambda x: x if x > 0 else 0)\n",65    "is_positive = fhe.univariate(lambda x: 1 if x > 0 else 0)\n",66    "arg_selection = fhe.univariate(lambda x: (x-1)//2 if x % 2 else 0)  # relu packed with a flag (alternating between 0 and relu)"67   ]68  },69  {70   "cell_type": "code",71   "execution_count": 3,72   "metadata": {},73   "outputs": [],74   "source": [75    "def swap(this_idx, this_dist, that_idx, that_dist):\n",76    "    \"\"\"\n",77    "    Swaps this and that if this > that. \n",78    "    We must pass both the index and the distance for both this and that.\n",79    "\n",80    "    Returns:\n",81    "      idxmin, min, idxmax, max of this and that based on distance\n",82    "    \"\"\"\n",83    "    diff = this_dist - that_dist\n",84    "    idx = arg_selection(2 * (this_idx - that_idx) + is_positive(diff))\n",85    "    dist = relu(diff)\n",86    "\n",87    "    idx_min = this_idx - idx\n",88    "    idx_max = that_idx + idx \n",89    "    dist_min = this_dist - dist\n",90    "    dist_max = that_dist + dist\n",91    "    return fhe.array([idx_min, dist_min, idx_max, dist_max])\n",92    "\n",93    "\n",94    "@fhe.compiler({\"x\": \"encrypted\", \"y\": \"encrypted\"})\n",95    "def knn(x, y):\n",96    "    dist = all_distances(x, y)\n",97    "    idx = list(range(N_PTS))\n",98    "    for k in range(2):\n",99    "        for i in range(k+1, N_PTS):\n",100    "             idx[k], dist[k], idx[i], dist[i] = swap(idx[k], dist[k], idx[i], dist[i])\n",101    "    return fhe.array([get_point(idx[j]) for j in range(2)])\n",102    "\n",103    "\n",104    "inputset = [(4, 3), (0, 0), (15, 3), (4, 15)]\n",105    "\n",106    "circuit = knn.compile(inputset)\n"107   ]108  },109  {110   "cell_type": "markdown",111   "metadata": {},112   "source": [113    "## Client\n",114    "The client simply invokes the server's nearest neighbours circuit."115   ]116  },117  {118   "cell_type": "code",119   "execution_count": 4,120   "metadata": {},121   "outputs": [122    {123     "name": "stdout",124     "output_type": "stream",125     "text": [126      "57.9 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)\n"127     ]128    }129   ],130   "source": [131    "%%timeit -r 1 -n 1\n",132    "circuit.client.keys.generate()"133   ]134  },135  {136   "cell_type": "code",137   "execution_count": 5,138   "metadata": {},139   "outputs": [],140   "source": [141    "def nearest(x, y):\n",142    "    ex, ey = circuit.encrypt(x, y)\n",143    "    res = circuit.run(ex, ey)   # Simulate request to the server\n",144    "    return circuit.decrypt(res)"145   ]146  },147  {148   "cell_type": "markdown",149   "metadata": {},150   "source": [151    "## Benchmarks"152   ]153  },154  {155   "cell_type": "code",156   "execution_count": 6,157   "metadata": {},158   "outputs": [159    {160     "name": "stdout",161     "output_type": "stream",162     "text": [163      "21.7 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)\n"164     ]165    }166   ],167   "source": [168    "%%timeit -r 1 -n 1\n",169    "nearest(4, 3)"170   ]171  }172 ],173 "metadata": {174  "kernelspec": {175   "display_name": "zama",176   "language": "python",177   "name": "python3"178  },179  "language_info": {180   "codemirror_mode": {181    "name": "ipython",182    "version": 3183   },184   "file_extension": ".py",185   "mimetype": "text/x-python",186   "name": "python",187   "nbconvert_exporter": "python",188   "pygments_lexer": "ipython3",189   "version": "3.9.5"190  }191 },192 "nbformat": 4,193 "nbformat_minor": 4194}195