Python:Indexing

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

General

#!/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

#!/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

#!/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

#!/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.

#!/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