can now return the labels of the pois rather than just the distances

pandana/network.py

@@ -112,6 +112,7 @@ def __init__(self, node_x, node_y, edge_from, edge_to, edge_weights,
         self.impedance_names = list(edge_weights.columns)
         self.variable_names = []
         self.poi_category_names = []
+        self.poi_category_indexes = {}
         self.num_poi_categories = -1
 
         # this maps ids to indexes which are used internally
@@ -511,11 +512,13 @@ def set_pois(self, category, x_col, y_col):
 
         xys = pd.DataFrame({'x': x_col, 'y': y_col}).dropna(how='any')
 
+        self.poi_category_indexes[category] = xys.index
+
         _pyaccess.initialize_category(self.poi_category_names.index(category),
                                       xys.astype('float32'))
 
     def nearest_pois(self, distance, category, num_pois=1, max_distance=None,
-                     imp_name=None):
+                     imp_name=None, include_poi_ids=False):
         """
         Find the distance to the nearest pois from each source node.  The
         bigger values in this case mean less accessibility.
@@ -537,6 +540,13 @@ def nearest_pois(self, distance, category, num_pois=1, max_distance=None,
             Must be one of the impedance names passed in the constructor of
             this object.  If not specified, there must be only one impedance
             passed in the constructor, which will be used.
+        include_poi_ids : bool, optional
+            If this flag is set to true, the call will add columns to the
+            return DataFrame - instead of just returning the distance for
+            the nth POI, it will also return the id of that POI.  The names
+            of the columns with the poi ids will be poi1, poi2, etc - it
+            will take roughly twice as long to include these ids as to not
+            include them
 
         Returns
         -------
@@ -568,11 +578,36 @@ def nearest_pois(self, distance, category, num_pois=1, max_distance=None,
                                             self.poi_category_names.index(
                                                 category),
                                             self.graph_no,
-                                            imp_num)
-
+                                            imp_num,
+                                            0)
         a[a == -1] = max_distance
         df = pd.DataFrame(a, index=self.node_ids)
         df.columns = range(1, num_pois+1)
+
+        if include_poi_ids:
+            b = _pyaccess.find_all_nearest_pois(distance,
+                                                num_pois,
+                                                self.poi_category_names.index(
+                                                    category),
+                                                self.graph_no,
+                                                imp_num,
+                                                1)
+            df2 = pd.DataFrame(b, index=self.node_ids)
+            df2.columns = ["poi%d" % i for i in range(1, num_pois+1)]
+            for col in df2.columns:
+                # if this is still all working according to plan at this point
+                # the great magic trick is now to turn the integer position of
+                # the poi, which is painstakingly returned from the c++ code,
+                # and turn it into the actual index that was used when it was
+                # initialized as a pandas series - this really is pandas-like
+                # thinking.  it's complicated on the inside, but quite
+                # intuitive to the user I think
+                s = df2[col]
+                df2[col] = self.poi_category_indexes[category].values[s]
+                df2[col][s == -1] = np.nan
+
+            df = pd.concat([df, df2], axis=1)
+
         return df
 
     def low_connectivity_nodes(self, impedance, count, imp_name=None):

pandana/tests/test_pandana.py

@@ -211,3 +211,15 @@ def test_pois(sample_osm):
 
     with pytest.raises(AssertionError):
         net.nearest_pois(2000, "restaurants", num_pois=11)
+
+
+def test_pois_indexes(sample_osm):
+    net = sample_osm
+    x, y = random_x_y(sample_osm, 100)
+    x.index = ['lab%d' % i for i in range(len(x))]
+    y.index = x.index
+
+    net.set_pois("restaurants", x, y)
+
+    d = net.nearest_pois(2000, "restaurants", num_pois=10,
+                         include_poi_ids=True)

src/accessibility.cpp

@@ -68,9 +68,14 @@ namespace MTC {
 			accessibilityVars[category] = vars;
 		}
 
+        /* the return_nodeids parameter determines whether to
+           return the nodeids where the poi was found rather than
+           the distances - you can call this twice - once for the
+           distances and then again for the node ids */
 		std::vector<float>
 		Accessibility::findNearestPOIs(int srcnode, float maxradius,
-		    unsigned number, unsigned cat, int gno) {
+		    unsigned number, unsigned cat, int gno, 
+		    bool return_nodeids) {
 
 			assert(cat >= 0 && cat < POI_MAXVAL);
 
@@ -91,23 +96,31 @@ namespace MTC {
 				for(int i = 0 ; i < vars[nodeid].size() ; i++) {
 
 					if(vars[nodeid][i] == 0) continue;
-					ret.push_back((float)distance);
+
+					if(return_nodeids) {
+						ret.push_back((float)vars[nodeid][i]);
+					} else {
+						ret.push_back((float)distance);
+					}
 				}
 			}
 			std::sort(ret.begin(),ret.end());
 
 			return ret;
 		}
 
+
+		/* the return_nodeds param is described above */
 		std::vector<std::vector<float> >
 		Accessibility::findAllNearestPOIs(float maxradius,
-		    unsigned number, unsigned cat, int gno) {
+		    unsigned number, unsigned cat, int gno,
+		    bool return_nodeids) {
 			std::vector<std::vector<float> > dists(numnodes,
 			            std::vector<float> ( number ));
 			#pragma omp parallel for
 			for(int i = 0 ; i < numnodes ; i++) {
-				std::vector<float> d = findNearestPOIs(i, maxradius,
-					number, cat, gno);
+				std::vector<float> d = findNearestPOIs(i, maxradius, number,
+				                                       cat, gno, return_nodeids);
 				for(int j = 0 ; j < number ; j++) {
                     			if(j < d.size()) dists[i][j] = d[j];
                     			else dists[i][j] = -1;

src/accessibility.h

@@ -91,10 +91,11 @@ namespace MTC {
 
 			std::vector<float> findNearestPOIs(int srcnode,
 			    float maxradius, unsigned maxnumber, unsigned cat,
-			    int graphno=0);
+			    int graphno=0, bool return_nodeids=false);
 			std::vector<std::vector<float> >
 			    findAllNearestPOIs(float maxradius, unsigned maxnumber,
-			            unsigned cat, int graphno=0);
+			            unsigned cat, int graphno=0, 
+			            bool return_nodeids=false);
 
             DistanceVec Range(int srcnode, float radius, int graphno=0);
 

src/pyaccesswrap.cpp

@@ -120,7 +120,7 @@ initialize_category(PyObject *self, PyObject *args)
 		// conversion function below
         int nodeid = sa->ga[0]->NearestNode(pois[i*2+0],pois[i*2+1],NULL);
 		//assert(nodeid < sa->ga[0].numpois);
-        av[nodeid].push_back(nodeid);
+        av[nodeid].push_back(i);
     }
 
     sa->initializeCategory(id,av);
@@ -133,15 +133,15 @@ static PyObject *
 find_all_nearest_pois(PyObject *self, PyObject *args)
 {
     double radius;
-	int varind, num, gno, impno;
-	if (!PyArg_ParseTuple(args, "diiii", &radius, &num, &varind,
-	                        &gno, &impno))
+	int varind, num, gno, impno, return_nodeids;
+	if (!PyArg_ParseTuple(args, "diiiii", &radius, &num, &varind,
+	                        &gno, &impno, &return_nodeids))
 	    return NULL;
 
     std::shared_ptr<MTC::accessibility::Accessibility> sa = sas[gno];
 
     std::vector<std::vector<float> > nodes =
-        sa->findAllNearestPOIs(radius, num, varind, impno);
+        sa->findAllNearestPOIs(radius, num, varind, impno, return_nodeids);
 
 	npy_intp dims[2];
 	dims[0] = nodes.size();