PANDAS spatial clustering

1. Assignments to Pandas dataframes always work better when doing entire columns at once, instead of filling a single row or a small number of rows at a time. In the code below I first completely define a NumPy array of cluster numbers and only after it is completely defined do I pass it into the Pandas DataFrame.

2. Even aside from the speed issue, your for loop where you assign cluster numbers is very confusing because of a misuse of enumerate: you have clusterNr, idx in enumerate(nngrf) but the idiom is always idx, value in enumerate(values). You have the position of the index and the value switched. In your particular case this works and is a nice trick but you should document how this works in your code if anyone besides you is ever meant to read it.

3. I changed variable names to be PEP8 compliant and to be easier to read. I also added some comments, docstrings, and formatting. Particular changes of note were defining n_points and n_samples as variables, wrapping the code you wrote into a function (so that I could line profile it), making RADIUS a variable defined in all-caps (which in Python suggests it is a constant), and defining the radius in the call to .radius_neighbors_graph to be in terms of RADIUS rather than just be another magic number hard-coded into the code. I think most of these changes improve your code's readability and make it more in line with style guidelines for Python.

4. Minor point: coercing to boolean with .astype('bool') works on SciPy sparse CSR matrices, and so doing the coercion before converting to a non-sparse NumPy array with .toarray() should be slightly faster and use less memory than doing things the other way around -- no time is wasted on converting zeroes.

   import numpy as np
   import pandas as pd
   from sklearn.neighbors import NearestNeighbors
   %load_ext line_profiler
   
   def knn_clusters(n_points, n_samples):
       """
       Finds the which of n_samples points is closest to each of n_points randomly defined points.
       """
       RADIUS = 0.01
   
       def get_random_coords(samples=1000, offsetX=52.2, offsetY=13.1, width=0.5):
           """Generate random coordinates in a pandas dataframe"""
           points = np.random.rand(samples, 2) * width
           data = pd.DataFrame(points, columns=['latitude', 'longitude'])
           data.latitude += offsetX
           data.longitude += offsetY
   
           # set spatial properties
           data.columnX = 'latitude'
           data.columnY = 'longitude'
   
           return data
   
       # some random points
       points = get_random_coords(n_points)
   
       # a subset of those points
       samples = points.sample(n_samples)
   
       # KNN
       tree = NearestNeighbors(n_neighbors=2, 
                               radius=RADIUS, 
                               leaf_size=30, 
                               algorithm="ball_tree", 
                               n_jobs=1,
                              ).fit(points)
   
       # 
       nn_graph = tree.radius_neighbors_graph(samples, 
                                              radius = 10 * RADIUS, 
                                              mode = 'connectivity',
                                             ).astype('bool').toarray()
   
       # faster assigment to pandas dataframes: entire columns at once
       cluster_number, point_number = np.where(nn_graph)
       cluster_assignments = -np.ones(n_points)
       cluster_assignments[point_number] = cluster_number
       points.loc[:, 'cluster'] = cluster_assignments
   
       # for comparison: assignment by specific rows is slow
       points['cluster_alt'] = -1
       for clusterNr, idx in enumerate(nn_graph):
           points.loc[idx, 'cluster_alt'] = clusterNr
   
       # ensure equality of both approaches
       assert np.all(np.equal(points.cluster, points.cluster_alt))
   
       return points, nn_graph 
   

The results of line profiling with %lprun -f knn_clusters knn_clusters(10000, 1000) were:

Timer unit: 1e-06 s

Total time: 0.999948 s
File: <ipython-input-76-7be6bc3ea686>
Function: knn_clusters at line 6

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     6                                           def knn_clusters(n_points, n_samples):
     7                                               """
     8                                               Finds the which of n_samples points is closest to each of n_points randomly defined points.
     9                                               """
    10         1            2      2.0      0.0      RADIUS = 0.01
    11                                               
    12         1            2      2.0      0.0      def get_random_coords(samples=1000, offsetX=52.2, offsetY=13.1, width=0.5):
    13                                                   """Generate random coordinates in a pandas dataframe"""
    14                                                   points = np.random.rand(samples, 2) * width
    15                                                   data = pd.DataFrame(points, columns=['latitude', 'longitude'])
    16                                                   data.latitude += offsetX
    17                                                   data.longitude += offsetY
    18                                           
    19                                                   # set spatial properties
    20                                                   data.columnX = 'latitude'
    21                                                   data.columnY = 'longitude'
    22                                           
    23                                                   return data
    24                                               
    25                                               # some random points
    26         1         3710   3710.0      0.4      points = get_random_coords(n_points)
    27                                               
    28                                               # a subset of those points
    29         1         7687   7687.0      0.8      samples = points.sample(n_samples)
    30                                               
    31                                               # KNN
    32         1            4      4.0      0.0      tree = NearestNeighbors(n_neighbors=2, 
    33         1            2      2.0      0.0                              radius=RADIUS, 
    34         1            2      2.0      0.0                              leaf_size=30, 
    35         1            2      2.0      0.0                              algorithm="ball_tree", 
    36         1          275    275.0      0.0                              n_jobs=1,
    37         1         6029   6029.0      0.6                             ).fit(points)
    38                                               
    39                                               # 
    40         1            7      7.0      0.0      nn_graph = tree.radius_neighbors_graph(samples, 
    41         1            4      4.0      0.0                                             radius=10*RADIUS, 
    42         1        90594  90594.0      9.1                                             mode='connectivity',
    43         1        36823  36823.0      3.7                                            ).astype('bool').toarray()
    44                                               
    45                                               # faster assigment to pandas dataframes: entire columns at once
    46         1        73476  73476.0      7.3      cluster_number, point_number = np.where(nn_graph)
    47         1           74     74.0      0.0      cluster_assignments = -np.ones(n_points)
    48         1        11675  11675.0      1.2      cluster_assignments[point_number] = cluster_number
    49         1         2548   2548.0      0.3      points.loc[:, 'cluster'] = cluster_assignments
    50                                               
    51                                               # for comparison: assignment by specific rows is slow
    52         1          889    889.0      0.1      points['cluster_alt'] = -1
    53      1001         2464      2.5      0.2      for clusterNr, idx in enumerate(nn_graph):
    54      1000       763337    763.3     76.3          points.loc[idx, 'cluster_alt'] = clusterNr
    55                                                   
    56                                               # ensure equality of both approaches
    57         1          341    341.0      0.0      assert np.all(np.equal(points.cluster, points.cluster_alt))
    58                                               
    59         1            1      1.0      0.0      return points, nn_graph

Thus assigning an entire column at once was just under 10x faster.