# HG changeset patch # User Jeff Hammel # Date 1498342570 25200 # Node ID 7e27e874655beaeb6c489d778844cc660f2051c5 # Parent 8e130b7bfed9b423c17203968922f6c468ae17bf test a larger grid + move distance insertion to its own function diff -r 8e130b7bfed9 -r 7e27e874655b globalneighbors/distance.py --- 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() diff -r 8e130b7bfed9 -r 7e27e874655b globalneighbors/grid.py --- 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') diff -r 8e130b7bfed9 -r 7e27e874655b globalneighbors/read.py --- 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() diff -r 8e130b7bfed9 -r 7e27e874655b tests/test_grid.py --- 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()