1
2 """
3 Operations on numpy arrays not present in the standard package.
4 """
5 import numpy as np
6 import pylab as pl
7 from scipy import spatial
8 import itertools
9
10
12 """
13 Distil unique rows/cols from an array
14
15 C{axis=0}: unique rows
16 C{axis=1}: unique cols
17
18 Example with rows:
19 >>> c = np.sort(np.random.normal(size=(3,5)),axis=0)
20 >>> d = np.r_[c,c,c]
21 >>> du = np.sort(unique_arr(d,axis=0),axis=0)
22 >>> np.all(du==c)
23 True
24
25 Example with columns:
26 >>> c = np.sort(np.random.normal(size=(3,5)),axis=1)
27 >>> d = np.hstack([c,c])
28 >>> du = np.sort(unique_arr(d,axis=1),axis=1)
29 >>> np.all(du==c)
30 True
31
32 @param a: array to remove duplicate entries from
33 @type a: numpy array
34 @param axis: keep unique elements in rows (axis=0) or columns (axis=1)
35 @type axis: integer
36 @param return_index: return index array with unique elements
37 @type return_index: bool
38 @return: unique array(,indices)
39 @rtype: numpy array(, numpy array)
40 """
41
42 if axis==0:
43 a = np.ascontiguousarray(a)
44 a_,index = np.unique(a.view([('',a.dtype)]*a.shape[1]),return_index=True)
45 a = a_.view(a.dtype).reshape(-1,a.shape[1])
46 elif axis==1:
47 a = np.transpose(a)
48 a = np.ascontiguousarray(a)
49 a_,index = np.unique(a.view([('',a.dtype)]*a.shape[1]),return_index=True)
50 a = a_.view(a.dtype).reshape(-1,a.shape[1])
51 a = np.transpose(a)
52 if return_index:
53 return a,index
54 else:
55 return a
56
58 """
59 Sort an array along several axes
60
61 >>> a = np.array([[ 0., 1.],\
62 [ 1., 1.],\
63 [ 1., 0.],\
64 [ 0., 0.],\
65 [ 0., 3.],\
66 [ 1., 0.],\
67 [ 1., 3.],\
68 [ 1., 2.],\
69 [ 1., 3.],\
70 [ 0., 4.]])
71 >>> sort_order(a,[0,1])
72 array([[ 0., 0.],
73 [ 0., 1.],
74 [ 0., 3.],
75 [ 0., 4.],
76 [ 1., 0.],
77 [ 1., 0.],
78 [ 1., 1.],
79 [ 1., 2.],
80 [ 1., 3.],
81 [ 1., 3.]])
82
83 @param a: numpy array to sort
84 @type a: numpy array
85 @param order: order of the columns to sort
86 @type order: list of integers (indices)
87 @return: sorted array
88 @rtype: numpy array
89 """
90 a = np.ascontiguousarray(a.copy())+0.
91 a_ = np.sort(a.view([('',a.dtype)]*a.shape[1]), order=['f%d'%(i) for i in order], axis=0)
92 a = a_.view(a.dtype).reshape(-1,a.shape[1])
93 return a
94
96 """
97 Calculate the argmax of a 2D array.
98
99 Example usage:
100
101 >>> output = np.zeros(100)
102 >>> for i in xrange(100):
103 ... a = np.random.normal(size=(5,6))
104 ... x,y = argmax2D(a)
105 ... output[i] = (a[x,y] == a.max())
106 >>> sum(output)
107 100.0
108
109 @param a: array to calculate the position of the maximum of
110 @type a: numpy 2D array
111 @rtype: (index,index)
112 @return: x,y coordinates
113 """
114 index = np.argmax(a)
115 y,x = index%a.shape[1],index//a.shape[1]
116 return x,y
117
119 """
120 Calculate the argmin of a 2D array.
121
122 Example usage:
123
124 >>> output = np.zeros(100)
125 >>> for i in xrange(100):
126 ... a = np.random.normal(size=(5,6))
127 ... x,y = argmin2D(a)
128 ... output[i] = (a[x,y] == a.min())
129 >>> sum(output)
130 100.0
131
132 @param a: array to calculate the position of the minimum of
133 @type a: numpy 2D array
134 @rtype: (index,index)
135 @return: x,y coordinates
136 """
137 index = np.argmin(a)
138 y,x = index%a.shape[1],index//a.shape[1]
139 return x,y
140
141
143 """
144 Return closest-match indices from b in a.
145
146 Example usage:
147 >>> a = np.random.uniform(size=10)
148 >>> b = a[np.array(np.random.uniform(size=3)*10,int)]
149 >>> ind = match_arrays(a,b)
150 >>> all(a[ind] == b)
151 True
152 """
153 sa = np.argsort(a)
154 a_ = a[sa]
155 a_ = a_[:-1] + np.diff(a_)/2.
156 closest_index = np.searchsorted(a_,b)
157 indices = sa[closest_index]
158 return indices
159
160 -def stdw(data,weights=None):
161 """
162 Calculates the weighted (sample) standard deviation of a list of numbers.
163
164 @type data: list
165 @param data: input data, must be a two dimensional list in format [value, weight]
166 @rtype: float
167 @return: weighted standard deviation
168 """
169 if len(data)==1:
170 return 0
171 listMean=np.average(data,weights=weights)
172 sum=0
173 wSum=0
174 wNonZero=0
175 for el,w in zip(data,weights):
176 if w>0.0:
177 sum=sum+float((el-listMean)/w)*float((el-listMean)/w)
178 wSum=wSum+float(1.0/w)*float(1.0/w)
179
180 nFactor=float(len(data))/float(len(data)-1)
181 stdev=np.sqrt(nFactor*(sum/wSum))
182
183 return stdev
184
185
186
188 """
189 3 point Lagrangian differentiation.
190
191 Returns z = dy/dx
192
193 Example usage:
194
195 >>> X = np.array([ 0.1, 0.3, 0.4, 0.7, 0.9])
196 >>> Y = np.array([ 1.2, 2.3, 3.2, 4.4, 6.6])
197 >>> deriv(X,Y)
198 array([ 3.16666667, 7.83333333, 7.75 , 8.2 , 13.8 ])
199 """
200
201 x0_x1 = np.roll(x,1)-x
202 x1_x2 = x-np.roll(x,-1)
203 x0_x2 = np.roll(x,1)-np.roll(x,-1)
204 derivee = np.roll(y,1)*x1_x2/(x0_x1*x0_x2)
205 derivee = derivee + y*(1/x1_x2-1/x0_x1)
206 derivee = derivee - np.roll(y,-1)*x0_x1/(x0_x2*x1_x2)
207
208 derivee[0]=y[0]*(1./x0_x1[1]+1./x0_x2[1])
209 derivee[0]=derivee[0]-y[1]*x0_x2[1]/(x0_x1[1]*x1_x2[1])
210 derivee[0]=derivee[0]+y[2]*x0_x1[1]/(x0_x2[1]*x1_x2[1])
211 nm3=len(x)-3
212 nm2=len(x)-2
213 nm1=len(x)-1
214 derivee[nm1]=-y[nm3]*x1_x2[nm2]/(x0_x1[nm2]*x0_x2[nm2])
215 derivee[nm1]=derivee[nm1]+y[nm2]*x0_x2[nm2]/(x0_x1[nm2]*x1_x2[nm2])
216 derivee[nm1]=derivee[nm1]-y[nm1]*(1./x0_x2[nm2]+1./x1_x2[nm2])
217
218 return derivee
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
235 """
236 Generate random points in a non-convex continuous rectangular grid.
237
238 This routine will subdivide the grid in smaller cubicles and then populate
239 those with a number of points proportional to the relative size of the
240 cubicle. Because of this, C{size} will never be the true size of the
241 sample returned, but will always be a little bit higher or lower.
242
243 B{Warning: This function requiresevery square to be populated by at least
244 one point. If this is not the case, it will be forced that every square
245 has a point.}
246
247 >>> xs = np.array([1,2,0,1,2,3,0,1,2,3,1,2.])
248 >>> ys = np.array([4,4,3,3,3,3,2,2,2,2,1,1.])
249 >>> gridpoints = np.column_stack([xs,ys])
250 >>> sample = random_rectangular_grid(gridpoints,10000)
251
252 >>> p = pl.figure()
253 >>> p = pl.plot(xs,ys,'ro')
254 >>> p = pl.plot(sample[:,0],sample[:,1],'ko',ms=2)
255
256 ]include figure]]ivs_aux_numpy_ext_grid.png]
257
258 Gridpoints should be Ngridpoints x Ncols
259 """
260
261 tree = spatial.KDTree(gridpoints)
262 axis_values = [np.sort(np.unique(col)) for col in gridpoints.T]
263 axis_indices = [np.arange(1,len(axis)) for axis in axis_values]
264
265
266
267
268 hypercube_sides = [abs(np.diff(axis)) for axis in axis_values]
269 hypercube_volume = 0.
270 random_ranges_sizes = []
271 for combination,limit_indices in itertools.izip(itertools.product(*hypercube_sides),itertools.product(*axis_indices)):
272 limit_indices = np.array(limit_indices)
273
274
275 for corner in itertools.product(*zip(limit_indices-1,limit_indices)):
276 dist = tree.query([axis_values[i][j] for i,j in enumerate(corner)])[0]
277 if dist!=0:
278 break
279
280
281 else:
282 low = [axis_values[i][j-1] for i,j in enumerate(limit_indices)]
283 high = [axis_values[i][j] for i,j in enumerate(limit_indices)]
284 volume = np.product(combination)
285 hypercube_volume+= volume
286 random_ranges_sizes.append([low,high,volume])
287
288
289 sample = np.vstack([np.random.uniform(low=low,high=high,size=(int(max(vol/hypercube_volume*size,1)),len(low))) for low,high,vol in random_ranges_sizes])
290
291 return sample
292
293
294
295
297 """
298 Convert an array to a record array.
299 dtype = [('name1',int),('name2',float)]
300
301 >>> x = np.array([[1,1],[2,2]])
302 >>> y = recarr(x,[('ones',int),('twos',int)])
303 >>> print y['ones']
304 [1 1]
305
306 @param x: array to convert
307 @type x: numpy array
308 @param mydtype: dtype of record array
309 @type mydtype: list of tuples ('column name',type)
310 @return: convert record array
311 @rtype: numpy record array
312 """
313 y = [tuple([col[i] for col in x]) for i in range(len(x[0]))]
314 y = np.array(y,dtype=mydtype)
315 return y
316
318 """
319 Add rows to a record array
320
321 >>> x = np.array([[1,1],[2,2]])
322 >>> y = recarr(x,[('ones',int),('twos',int)])
323 >>> z = recarr_addrows(y,[[1,2]])
324 >>> print z['ones']
325 [1 1 1]
326
327 @param x: original record array
328 @type x: numpy record array
329 @param rows: list of lists/tuples or 2D-numpy array
330 @type rows: list or ndarray
331 @return: extended record array
332 @rtype: numpy record array
333 """
334 rows = [tuple(row) for row in rows]
335 x = np.core.records.fromrecords(x.tolist()+rows,dtype=x.dtype)
336 return x
337
339 """
340 Add columns to a record array
341
342 >>> x = np.array([[1,1],[2,2]])
343 >>> y = recarr(x,[('ones',int),('twos',int)])
344 >>> z = recarr_addcols(y,[[3,3],[4,4]],[('threes',int),('fours',int)])
345 >>> print z['fours']
346 [4 4]
347
348 @param x: original record array
349 @type x: numpy record array
350 @param cols: list of lists/tuples or 2D-numpy array
351 @type cols: list or ndarray
352 @param dtypes_ext: dtype of extra columns
353 @type dtypes_ext: list of tuples ('column name',type)
354 @return: extended record array
355 @rtype: numpy record array
356 """
357 names = list(x.dtype.names)
358 dtypes = [(name,x.dtype[names.index(name)].str) for name in names]
359 dtypes += dtypes_ext
360 rows = []
361 for i in range(len(x)):
362 rows.append(tuple(list(x[i]) + [col[i] for col in cols]))
363 x = np.core.records.fromrecords(rows,dtype=dtypes)
364 return x
365
367 """
368 Join to record arrays column wise.
369 """
370 arr1 = arr1.copy()
371 for field in arr2.dtype.names:
372 arr1 = pl.mlab.rec_append_fields(arr1,field,arr2[field])
373 return arr1
374
375
376
377
378
380 """
381 Find intersection of two 2D lines.
382
383 Example usage:
384
385 >>> pcoeff1 = [1,2,-2.]
386 >>> pcoeff2 = [-2.12,2.45,3.7321]
387 >>> x = np.linspace(-3,3,7)
388 >>> A = np.column_stack([x,np.polyval(pcoeff1,x)])
389 >>> B = np.column_stack([x,np.polyval(pcoeff2,x)])
390
391 We find two intersections:
392
393 >>> xs,ys = find_intersections(A,B)
394 >>> print xs,ys
395 [-1.22039755 1.34367003] [-2.77960245 2.71835017]
396
397 >>> p = pl.figure()
398 >>> p = pl.plot(x,A[:,1],'bo-')
399 >>> p = pl.plot(x,B[:,1],'go-')
400 >>> p = pl.plot(xs,ys,'rs',ms=5)
401
402 Returns empty arrays when there are no intersections.
403
404 ]include figure]]ivs_aux_numpy_ext_intersect.png]
405 """
406
407 amin = lambda x1, x2: np.where(x1<x2, x1, x2)
408 amax = lambda x1, x2: np.where(x1>x2, x1, x2)
409 aall = lambda abools: np.dstack(abools).all(axis=2)
410 slope = lambda line: (lambda d: d[:,1]/d[:,0])(np.diff(line, axis=0))
411
412 x11, x21 = np.meshgrid(A[:-1, 0], B[:-1, 0])
413 x12, x22 = np.meshgrid(A[1:, 0], B[1:, 0])
414 y11, y21 = np.meshgrid(A[:-1, 1], B[:-1, 1])
415 y12, y22 = np.meshgrid(A[1:, 1], B[1:, 1])
416
417 m1, m2 = np.meshgrid(slope(A), slope(B))
418 m1inv, m2inv = 1/m1, 1/m2
419
420 yi = (m1*(x21-x11-m2inv*y21) + y11)/(1 - m1*m2inv)
421 xi = (yi - y21)*m2inv + x21
422
423 xconds = (amin(x11, x12) < xi, xi <= amax(x11, x12),
424 amin(x21, x22) < xi, xi <= amax(x21, x22) )
425 yconds = (amin(y11, y12) < yi, yi <= amax(y11, y12),
426 amin(y21, y22) < yi, yi <= amax(y21, y22) )
427
428 return xi[aall(xconds)], yi[aall(yconds)]
429
430
431
432 if __name__=="__main__":
433 import doctest
434 doctest.testmod()
435 pl.show()
436