Python:Indexing

This page just collects some notes/hints about how indexing is done in Python and numpy.

General

?

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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45

#!/usr/bin/python # -*- coding: utf-8 -*-   import numpy as np   # Test data x = np.linspace( 10.0, 20.0, 20 ) def gaussian( x, xc=0.0, sigma=1.0 ):     y = np.exp( -( np.square(x-xc) )/( (2*sigma)**2 ) )     return y y = gaussian( x, xc=15.0, sigma=2.0 ) for i, (xcur, ycur) in enumerate( zip(x,y) ):     print i, xcur, ycur print '\n\n'   # The maximum value idx_max = np.argmax( y ) x_max = x[idx_max] y_max = y[idx_max] print 'max: ', idx_max, x_max, y_max   # Find a specific value along x x_search = 15.0 idx = np.where( x>x_search )[0][0] xcur = x[idx] ycur = y[idx] print 'x search: ', x_search, idx, xcur, ycur     # Find closest to specific value idx = (np.abs(x-x_search)).argmin()     # Sort an x and y series x = [ 4.0, 3.0, 1.0, 5.0, 6.0, 8.0, 2.0, 7.0, 9.0, 10.0 ] y = [ 40.0, 30.0, 10.0, 50.0, 60.0, 80.0, 20.0, 70.0, 90.0, 100.0 ]   x = np.asarray(x) y = np.asarray(y) indices = np.argsort( x ) print indices x_sorted = x[indices] y_sorted = y[indices] print x_sorted print y_sorted

Remove a particular element from x and y

?

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17

#!/usr/bin/python # -*- coding: utf-8 -*-   import numpy as np   # Remove NaNs print '\n\n' data = np.array([[1, 2, 3, 4],            [10, 40, np.nan, 160]]) data = np.transpose(data) x = data[:,0] y = data[:,1]   idx = np.isnan( y ) xcur = x[~idx] ycur = y[~idx] print xcur, ycur

See also numpy.nan_to_num.

Replace items according to a rule

?

1 2 3 4 5 6 7 8 9 10 11 12 13

#!/usr/bin/python # -*- coding: utf-8 -*-   import numpy as np   # Replace zero elements print '\n\n' x = np.array( [-1,0,0.01,0.1,1.0,10.0] ) print x   idx = np.nonzero(x==0) x[idx] = 1e-9 print x

Use indexing to fill a matrix with a given shape/function

?

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 29 30 31 32 33 34 35 36 37 38 39

#!/usr/bin/python import numpy as np   # 2D example size = 10 extent = 1.0 Mzeros = np.zeros( (size,size) ) Mones = np.ones( (size,size) )   axis_x = np.linspace( -extent, +extent, size ) axis_y = np.linspace( -extent, +extent, size ) X, Y = np.meshgrid( axis_x, axis_y )     # Circle radius = 1.0 T = np.where( (X**2+Y**2)<radius , Mones, Mzeros ) print T     # 3D example size = 5 extent = 1.0 Mzeros = np.zeros( (size,size,size) ) Mones = np.ones( (size,size,size) )   X, Y, Z = np.mgrid[ -extent:+extent:size*1j , -extent:+extent:size*1j , -extent:+extent:size*1j ]   # Sphere radius = 1.0 T = np.where( (X**2+Y**2+Z**2)<radius , Mones, Mzeros ) print T   # Superball radius = 1.0 p = 2.1 threshold = abs(radius)**(2.0*p) T = np.where( ( np.power(np.abs(X),2.0*p) + np.power(np.abs(Y),2.0*p) + np.power(np.abs(Z),2.0*p) )<threshold, Mones, Mzeros ) print T

Use indexing to put 2D/3D array values into a 1D array

Example: Azimuthal average: each pixel in the 2D array has a particular bin in the 1D array that it's supposed to be added to.

?

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 29 30 31 32

#!/usr/bin/python import numpy as np       # 3D, sphere size = 5 extent = 1.0 Mzeros = np.zeros( (size,size,size) ) Mones = np.ones( (size,size,size) )   X, Y, Z = np.mgrid[ -extent:+extent:size*1j , -extent:+extent:size*1j , -extent:+extent:size*1j ]   radius = 1.0 vals = np.where( (X**2+Y**2+Z**2)<radius , 2.0, 0 )   # Vals are the values we want to put into the array "answer" print vals   # Indices are the index values for each 'val', telling us which bin of the "answer" array to put the value in indices = np.ones( (size,size,size) )*np.arange(0,size) print indices   answer = np.zeros( (size) ) for idx_cur in range(len(answer)):     mask = np.where( indices==idx_cur, 1, 0 )       print idx_cur, vals*mask           answer[idx_cur] = (vals*mask).sum()   print answer