changeset 5:7e27e874655b

test a larger grid + move distance insertion to its own function
author Jeff Hammel <k0scist@gmail.com>
date Sat, 24 Jun 2017 15:16:10 -0700
parents 8e130b7bfed9
children 316e1d54ffd4
files globalneighbors/distance.py globalneighbors/grid.py globalneighbors/read.py tests/test_grid.py
diffstat 4 files changed, 38 insertions(+), 11 deletions(-) [+]
line wrap: on
line diff
--- a/globalneighbors/distance.py	Sat Jun 24 14:48:55 2017 -0700
+++ b/globalneighbors/distance.py	Sat Jun 24 15:16:10 2017 -0700
@@ -32,6 +32,23 @@
     return degrees*DEGREESTORADIANS
 
 
+def insert_distance(distances, i, new_distance, k):
+    """
+    insert `(i, new_distance)` into `distances`
+    keeping distances in order with maximum of size `k`
+    """
+
+     # TODO: Binary Search Tree
+    for _index, (geoid, old_distance) in enumerate(distances):
+        if new_distance < old_distance:
+            distances.insert(_index, (i, new_distance))
+            if len(distances) == k+1:
+                distances.pop()
+            break
+    else:
+        distances.append((i, new_distance))
+
+
 def calculate_distances(locations, r=Rearth):
     """
     WARNING! This is an N-squared approach
@@ -101,15 +118,7 @@
                 if len(distances) == k and new_distance >= distances[-1][-1]:
                     continue
 
-                # TODO: Binary Search Tree
-                for _index, (geoid, old_distance) in enumerate(distances):
-                    if new_distance < old_distance:
-                        distances.insert(_index, (i, new_distance))
-                        if len(distances) == k+1:
-                            distances.pop()
-                        break
-                else:
-                    distances.append((i, new_distance))
+                insert_distance(distances, i, new_distance, k)
 
     return neighbors
 
@@ -143,7 +152,7 @@
                                     output=options.output_counter)
 
     # output
-    options.output.write(json.dumps(neighbors, indent=2))
+    options.output.write(json.dumps(neighbors))
 
 if __name__ == '__main__':
     main()
--- a/globalneighbors/grid.py	Sat Jun 24 14:48:55 2017 -0700
+++ b/globalneighbors/grid.py	Sat Jun 24 15:16:10 2017 -0700
@@ -35,5 +35,10 @@
         """
         return neighbors of points i, j
         """
+        imat = []
+        jmat = []
+
+        # latitude
         if i:
+            imat.append(i-1)
             raise NotImplementedError('TODO')
--- a/globalneighbors/read.py	Sat Jun 24 14:48:55 2017 -0700
+++ b/globalneighbors/read.py	Sat Jun 24 15:16:10 2017 -0700
@@ -37,7 +37,7 @@
         try:
             cast = cast_row(row, types=types)
         except AssertionError:
-            sys.stderr.write("Error, row {}".format(index))
+            sys.stderr.write("Error, row {}\n{}\n".format(index, row))
             raise
         if fields:
             cast = {key: value for key, value in cast.items()
--- a/tests/test_grid.py	Sat Jun 24 14:48:55 2017 -0700
+++ b/tests/test_grid.py	Sat Jun 24 15:16:10 2017 -0700
@@ -9,6 +9,7 @@
 from common import datafile
 from globalneighbors.grid import LatLonGrid
 from globalneighbors.locations import locations
+from globalneighbors.read import read_cities
 from globalneighbors.read import read_city_list
 
 
@@ -44,6 +45,16 @@
         city_locations = locations(read_city_list(samplefile))
         self.grid_locations(city_locations)
 
+    def test_10000(self):
+        """test 10000 cities"""
+
+        filename = datafile('10000cities.tsv')
+        assert os.path.exists(filename)
+        with open(filename) as f:
+            city_locations = locations(read_cities(f))
+        self.grid_locations(city_locations)
+
+
     ### generic (utility) functions
 
     def grid_locations(self, locations):
@@ -63,5 +74,7 @@
                 n_locations += len(grid[(i,j)])
         assert n_locations == len(locations)
 
+
+
 if __name__ == '__main__':
     unittest.main()