{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%matplotlib inline" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\nArg-K-Min reduction\n===================\n\nUsing the :mod:`pykeops.numpy` API, we define a dataset of N points in \$\\mathbb R^D\$ and compute for each\npoint the indices of its K nearest neighbours (including itself).\n\n \n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Setup\n----------\n\nStandard imports:\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import time\n\nimport matplotlib.pyplot as plt\nimport numpy as np\n\nfrom pykeops.numpy import Genred" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Define our dataset:\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "N = 1000 # Number of points\nD = 2 # Dimension of the ambient space\nK = 3 # Number of neighbors to look for\n\ndtype = 'float32' # May be 'float32' or 'float64'\n\nx = np.random.rand(N,D).astype(dtype)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "KeOps Kernel\n-------------\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "formula = 'SqDist(x,y)' # Use a simple Euclidean (squared) norm\nvariables = ['x = Vi('+str(D)+')', # First arg : i-variable, of size D\n 'y = Vj('+str(D)+')'] # Second arg: j-variable, of size D\n\n# N.B.: The number K is specified as an optional argument `opt_arg`\nmy_routine = Genred(formula, variables, reduction_op='ArgKMin', axis=1, \n dtype=dtype, opt_arg=K)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Using our new :class:`pykeops.numpy.Genred` routine, \nwe perform a K-nearest neighbor search ( **reduction_op** = ``\"ArgKMin\"`` )\nover the \$j\$ variable \$y_j\$ ( **axis** = 1):\n\n

#### Note

If CUDA is available and **backend** is ``\"auto\"`` or not specified,\n KeOps will:\n\n 1. Load the data on the GPU\n 2. Perform the computation on the device \n 3. Unload the result back to the CPU\n\n as it is assumed to be most efficient for large-scale problems.\n By specifying **backend** = ``\"CPU\"`` in the call to ``my_routine``, \n you can bypass this procedure and use a simple C++ ``for`` loop instead.

\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# Dummy first call to warm-up the GPU and thus get an accurate timing:\nmy_routine( np.random.rand(10,D).astype(dtype),\n np.random.rand(10,D).astype(dtype) )\n\n# Actually perform our K-nn search:\nstart = time.time()\nind = my_routine(x, x, backend=\"auto\")\nprint(\"Time to perform the K-nn search: \",round(time.time()-start,5),\"s\")\n\n# The result is now an (N,K) array of integers:\nprint(\"Output values :\")\nprint(ind)\n\nplt.figure(figsize=(8,8))\nplt.scatter(x[:,0], x[:,1], s= 25*500 / len(x))\n\nfor k in range(K): # Highlight some points and their nearest neighbors\n plt.scatter(x[ ind[:4,k], 0],x[ ind[:4,k], 1], s= 100)\n \n\nplt.axis(\"equal\") ; plt.axis([0,1,0,1])\nplt.tight_layout(); plt.show()" ] } ], "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": 0 }