Mercurial > hg > GlobalNeighbors

comparison globalneighbors/distance.py @ 9:638fad06e556

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
use bisect function; it has been tested faster
author Jeff Hammel <k0scist@gmail.com>
date Sun, 25 Jun 2017 11:37:52 -0700
parents e3d6919130ca
children 27925261c137
comparison
equal deleted inserted replaced
8:e3d6919130ca 9:638fad06e556
50 distances.pop() 50 distances.pop()
51 break 51 break
52 else: 52 else:
53 distances.append((i, new_distance)) 53 distances.append((i, new_distance))
54 54
55 class KeyWrapper: 55
56 class KeyWrapper(object):
57 """wrapper for python's `bisect` methods"""
56 def __init__(self, iterable, key): 58 def __init__(self, iterable, key):
57 self.it = iterable 59 self.it = iterable
58 self.key = key 60 self.key = key
59 61
60 def __getitem__(self, i): 62 def __getitem__(self, i):
61 return self.key(self.it[i]) 63 return self.key(self.it[i])
62 64
63 def __len__(self): 65 def __len__(self):
64 return len(self.it) 66 return len(self.it)
67
65 68
66 def insert_distance_bisect(distances, i, new_distance, k): 69 def insert_distance_bisect(distances, i, new_distance, k):
67 70
68 if not distances: 71 if not distances:
69 distances.append((i, new_distance)) 72 distances.append((i, new_distance))
145 148
146 # insert in order 149 # insert in order
147 for i in (id1, id2): 150 for i in (id1, id2):
148 distances = neighbors.setdefault(i, []) 151 distances = neighbors.setdefault(i, [])
149 152
150 insert_distance(distances, i, new_distance, k) 153 insert_distance_bisect(distances, i, new_distance, k)
151 154
152 return neighbors 155 return neighbors
153 156
154 157
155 def main(args=sys.argv[1:]): 158 def main(args=sys.argv[1:]):
158 # parse command line arguments 161 # parse command line arguments
159 description = """write nearest neighborfiles""" 162 description = """write nearest neighborfiles"""
160 parser = CitiesParser(description=description) 163 parser = CitiesParser(description=description)
161 parser.add_argument('output', type=argparse.FileType('w'), 164 parser.add_argument('output', type=argparse.FileType('w'),
162 help="output file to dump JSON to") 165 help="output file to dump JSON to")
166 parser.add_argument('--latlon', dest='latlon', type=float,
167 nargs=2, metavar=("LAT", "LON"),
168 default=(1., 1.),
169 help="tolerance of latitude and longitude in degrees [DEFAULT: %(default)s]")
163 parser.add_argument('--counter', '--output-counter', 170 parser.add_argument('--counter', '--output-counter',
164 dest='output_counter', 171 dest='output_counter',
165 type=int, default=100, 172 type=int, default=100,
166 help="how often to output progress updates [DEFAULT: %(default)s]") 173 help="how often to output progress updates [DEFAULT: %(default)s]")
167 parser.add_argument('-k', dest='k', 174 parser.add_argument('-k', dest='k',
176 city_locations = locations(cities) 183 city_locations = locations(cities)
177 184
178 # calculate neighbors 185 # calculate neighbors
179 neighbors = calculate_neighbors(city_locations, 186 neighbors = calculate_neighbors(city_locations,
180 k=options.k, 187 k=options.k,
188 lat_tol=options.latlon[0],
189 lon_tol=options.latlon[-1],
181 output=options.output_counter) 190 output=options.output_counter)
182 191
183 # output 192 # output
184 options.output.write(json.dumps(neighbors)) 193 options.output.write(json.dumps(neighbors))
185 194