1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27 """
28 A collection of iteration utilities.
29 """
30
31
32 import math
33 import numpy
34 import random
35
36
37 from glue import git_version
38
39
40 __author__ = "Kipp Cannon <kipp.cannon@ligo.org>"
41 __version__ = "git id %s" % git_version.id
42 __date__ = git_version.date
43
44
45
46
47
48
49
50
51
52
53
55 """
56 A generator for iterating over the elements of multiple sequences
57 simultaneously. With N sequences given as input, the generator
58 yields all possible distinct N-tuples that contain one element from
59 each of the input sequences.
60
61 Example:
62
63 >>> x = MultiIter([0, 1, 2], [10, 11])
64 >>> list(x)
65 [(0, 10), (1, 10), (2, 10), (0, 11), (1, 11), (2, 11)]
66
67 The elements in each output tuple are in the order of the input
68 sequences, and the left-most input sequence is iterated over first.
69
70 Internally, the input sequences themselves are each iterated over
71 only once, so it is safe to pass generators as arguments. Also,
72 this generator is significantly faster if the longest input
73 sequence is given as the first argument. For example, this code
74
75 >>> lengths = range(1, 12)
76 >>> for x in MultiIter(*map(range, lengths)):
77 ... pass
78 ...
79
80 runs approximately 5 times faster if the lengths list is reversed.
81 """
82 if len(sequences) > 1:
83
84
85
86
87
88
89 head = tuple((x,) for x in sequences[0])
90 for t in MultiIter(*sequences[1:]):
91 for h in head:
92 yield h + t
93 elif sequences:
94 for t in sequences[0]:
95 yield (t,)
96
97
99 """
100 A generator for iterating over all choices of n elements from the
101 input sequence vals. In each result returned, the original order
102 of the values is preserved.
103
104 Example:
105
106 >>> x = choices(["a", "b", "c"], 2)
107 >>> list(x)
108 [('a', 'b'), ('a', 'c'), ('b', 'c')]
109
110 The order of combinations in the output sequence is always the
111 same, so if choices() is called twice with two different sequences
112 of the same length the first combination in each of the two output
113 sequences will contain elements from the same positions in the two
114 different input sequences, and so on for each subsequent pair of
115 output combinations.
116
117 Example:
118
119 >>> x = choices(["a", "b", "c"], 2)
120 >>> y = choices(["1", "2", "3"], 2)
121 >>> zip(x, y)
122 [(('a', 'b'), ('1', '2')), (('a', 'c'), ('1', '3')), (('b', 'c'), ('2', '3'))]
123
124 Furthermore, the order of combinations in the output sequence is
125 such that if the input list has n elements, and one constructs the
126 combinations choices(input, m), then each combination in
127 choices(input, n-m).reverse() contains the elements discarded in
128 forming the corresponding combination in the former.
129
130 Example:
131
132 >>> x = ["a", "b", "c", "d", "e"]
133 >>> X = list(choices(x, 2))
134 >>> Y = list(choices(x, len(x) - 2))
135 >>> Y.reverse()
136 >>> zip(X, Y)
137 [(('a', 'b'), ('c', 'd', 'e')), (('a', 'c'), ('b', 'd', 'e')), (('a', 'd'), ('b', 'c', 'e')), (('a', 'e'), ('b', 'c', 'd')), (('b', 'c'), ('a', 'd', 'e')), (('b', 'd'), ('a', 'c', 'e')), (('b', 'e'), ('a', 'c', 'd')), (('c', 'd'), ('a', 'b', 'e')), (('c', 'e'), ('a', 'b', 'd')), (('d', 'e'), ('a', 'b', 'c'))]
138 """
139 if n == len(vals):
140 yield tuple(vals)
141 elif n > 1:
142 n -= 1
143 for i, v in enumerate(vals[:-n]):
144 v = (v,)
145 for c in choices(vals[i+1:], n):
146 yield v + c
147 elif n == 1:
148 for v in vals:
149 yield (v,)
150 elif n == 0:
151 yield ()
152 else:
153
154 raise ValueError(n)
155
156
158 """
159 Yield the unique items of an iterable, preserving order.
160 http://mail.python.org/pipermail/tutor/2002-March/012930.html
161
162 Example:
163
164 >>> x = uniq([0, 0, 2, 6, 2, 0, 5])
165 >>> list(x)
166 [0, 2, 6, 5]
167 """
168 temp_dict = {}
169 for e in iterable:
170 if e not in temp_dict:
171 yield temp_dict.setdefault(e, e)
172
173
175 """
176 Yield the non-unique items of an iterable, preserving order. If an
177 item occurs N > 0 times in the input sequence, it will occur N-1
178 times in the output sequence.
179
180 Example:
181
182 >>> x = nonuniq([0, 0, 2, 6, 2, 0, 5])
183 >>> list(x)
184 [0, 2, 0]
185 """
186 temp_dict = {}
187 for e in iterable:
188 if e in temp_dict:
189 yield e
190 temp_dict.setdefault(e, e)
191
192
194 """
195 Example:
196 >>> nested = [[1,2], [[3]]]
197 >>> list(flatten(nested))
198 [1, 2, [3]]
199 """
200 if levels == 0:
201 for x in sequence:
202 yield x
203 else:
204 for x in sequence:
205 for y in flatten(x, levels - 1):
206 yield y
207
208
209
210
211
212
213
214
215
216
217
219 """
220 Like Python's filter() builtin, but modifies the sequence in place.
221
222 Example:
223
224 >>> l = range(10)
225 >>> inplace_filter(lambda x: x > 5, l)
226 >>> l
227 [6, 7, 8, 9]
228
229 Performance considerations: the function iterates over the
230 sequence, shuffling surviving members down and deleting whatever
231 top part of the sequence is left empty at the end, so sequences
232 whose surviving members are predominantly at the bottom will be
233 processed faster.
234 """
235 target = 0
236 for source in range(len(sequence)):
237 if func(sequence[source]):
238 sequence[target] = sequence[source]
239 target += 1
240 del sequence[target:]
241
242
243
244
245
246
247
248
249
250
251
253 """
254 A generator that yields the values from several ordered iterables
255 in order.
256
257 Example:
258
259 >>> x = [0, 1, 2, 3]
260 >>> y = [1.5, 2.5, 3.5, 4.5]
261 >>> z = [1.75, 2.25, 3.75, 4.25]
262 >>> list(inorder(x, y, z))
263 [0, 1, 1.5, 1.75, 2, 2.25, 2.5, 3, 3.5, 3.75, 4.25, 4.5]
264 >>> list(inorder(x, y, z, key=lambda x: x * x))
265 [0, 1, 1.5, 1.75, 2, 2.25, 2.5, 3, 3.5, 3.75, 4.25, 4.5]
266
267 >>> x.sort(key=lambda x: abs(x-3))
268 >>> y.sort(key=lambda x: abs(x-3))
269 >>> z.sort(key=lambda x: abs(x-3))
270 >>> list(inorder(x, y, z, key=lambda x: abs(x - 3)))
271 [3, 2.5, 3.5, 2.25, 3.75, 2, 1.75, 4.25, 1.5, 4.5, 1, 0]
272
273 >>> x = [3, 2, 1, 0]
274 >>> y = [4.5, 3.5, 2.5, 1.5]
275 >>> z = [4.25, 3.75, 2.25, 1.75]
276 >>> list(inorder(x, y, z, reverse = True))
277 [4.5, 4.25, 3.75, 3.5, 3, 2.5, 2.25, 2, 1.75, 1.5, 1, 0]
278 >>> list(inorder(x, y, z, key = lambda x: -x))
279 [4.5, 4.25, 3.75, 3.5, 3, 2.5, 2.25, 2, 1.75, 1.5, 1, 0]
280
281 NOTE: this function will never reverse the order of elements in
282 the input iterables. If the reverse keyword argument is False (the
283 default) then the input sequences must yield elements in increasing
284 order, likewise if the keyword argument is True then the input
285 sequences must yield elements in decreasing order. Failure to
286 adhere to this yields undefined results, and for performance
287 reasons no check is performed to validate the element order in the
288 input sequences.
289 """
290 reverse = kwargs.pop("reverse", False)
291 keyfunc = kwargs.pop("key", lambda x: x)
292 if kwargs:
293 raise TypeError("invalid keyword argument '%s'" % list(kwargs.keys())[0])
294 nextvals = {}
295 for iterable in iterables:
296 next = iter(iterable).next
297 try:
298 nextval = next()
299 nextvals[next] = keyfunc(nextval), nextval, next
300 except StopIteration:
301 pass
302 if not nextvals:
303
304 return
305 if reverse:
306 select = lambda seq: max(seq, key = lambda elem: elem[0])
307 else:
308 select = lambda seq: min(seq, key = lambda elem: elem[0])
309 values = nextvals.itervalues
310 if len(nextvals) > 1:
311 while 1:
312 _, val, next = select(values())
313 yield val
314 try:
315 nextval = next()
316 nextvals[next] = keyfunc(nextval), nextval, next
317 except StopIteration:
318 del nextvals[next]
319 if len(nextvals) < 2:
320 break
321
322 (_, val, next), = values()
323 yield val
324 while 1:
325 yield next()
326
327
328
329
330
331
332
333
334
335
336
338 """
339 Yields integers in the range [lo, hi) where 0 <= lo < hi. Each
340 return value is a two-element tuple. The first element is the
341 random integer, the second is the natural logarithm of the
342 probability with which that integer will be chosen.
343
344 The CDF for the distribution from which the integers are drawn goes
345 as [integer]^{n}, where n > 0. Specifically, it's
346
347 CDF(x) = (x^{n} - lo^{n}) / (hi^{n} - lo^{n})
348
349 n = 1 yields a uniform distribution; n > 1 favours larger
350 integers, n < 1 favours smaller integers.
351 """
352 if not 0 <= lo < hi:
353 raise ValueError("require 0 <= lo < hi: lo = %d, hi = %d" % (lo, hi))
354 if n <= 0.:
355 raise ValueError("n <= 0: %g" % n)
356 elif n == 1.:
357
358 try:
359 lnP = math.log(1. / (hi - lo))
360 except ValueError:
361 raise ValueError("[lo, hi) domain error")
362 hi -= 1
363 rnd = random.randint
364 while 1:
365 yield rnd(lo, hi), lnP
366
367
368 lnP = numpy.arange(lo, hi + 1, dtype = "double")**n
369 lnP -= lnP[0]
370 lnP /= lnP[-1]
371
372 lnP = tuple(numpy.log(lnP[1:] - lnP[:-1]))
373 if numpy.isinf(lnP).any():
374 raise ValueError("[lo, hi) domain error")
375
376 beta = lo**n / (hi**n - lo**n)
377 n = 1. / n
378 alpha = hi / (1. + beta)**n
379 flr = math.floor
380 rnd = random.random
381 while 1:
382 index = int(flr(alpha * (rnd() + beta)**n))
383
384
385 assert index >= lo
386 yield index, lnP[index - lo]
387