15 May, 2019

NumPy by Example

This originally was in my Scientific Python 101 article, I've split it now as it was a long article and sometimes I need just to have a look at this code as a reminder of how things work.
If you're interested in a similar "by example" introduction for Mathematica, check this other one out.

Execute each section in an IPython cell (copy and paste, then shift-enter)

# Set-up the environment:
# First, we'll have to import NumPy and SciPy packages.
# IPython has "magics" (macros) to help common tasks like this:
%pylab
np.set_printoptions(suppress = True)
# Optional: change the way division works:
#from __future__ import division
# 1/2 = 0.5 instead of integer division... problem is that then 2/2 would be float as well if we enable this, so, beware.

# In IPython notebook, cell execution is shift-enter
# In IPython-QT enter evaluates and control-enter is used for multiline input
# Tab is used for completion, shift-tab after a function name shows its help
# Note that IPython will display the output of the LAST -expression- in a cell
# but variable assignments in Python are NOT expressions, so they will
# suppress output, unlike Matlab or Mathematica. You will need to add print
# statements after the assignments in the following examples to see the output
data = [[1,2],[3,4],[5,6]]
# if you want to suppress output in the rare cases it's generated, use ;
data; # try without the ; to see the difference

# There is a magic for help as well
%quickref
# Also, you can use ? and ?? for details on a symbol
get_ipython().magic(u'pinfo %pylab')

# In addition to the Python built-in functions help() and dir(symbol)
# Other important magics are: %reset, %reset_selective, %whos
# %prun (profile), %pdb (debug), %run, %edit

1*3 # evaluate this in a separate cell...

_*3 # ...now _ refers to the output of last evaluated cell
# again, not that useful because you can't refer to the previous expression
# evaluated inside a cell, just to the previous evaluation of an entire cell

# A numpy array can be created from a homogeneous list-like object
arr = np.array(data)
# Or an uninitialized one can be created by specifying its shape
arr = np.ndarray((3,3))
# There are also many other "constructors"
arr = np.identity(5)
arr = np.zeros((4,5))
arr = np.ones((4,))
arr = np.random.random((4,4))

# Arrays can also be created from functions
arr = np.fromfunction((lambda x: x*x), shape = (10,))
# ...or by parsing a string
arr = np.fromstring('1, 2', dtype=int, sep=',')

# Arrays are assigned by reference
arr2 = arr
# To create a copy, use copy
arr2 = arr.copy()

# An array shape is a descriptor, it can be changed with no copy
arr = np.zeros((4,4))
arr = arr.reshape((2,8)) # returns the same data in a new view
# numpy also supports matrices, which are arrays contrained to be 2d
mat = np.asmatrix(arr)
# Not all operations avoid copies, flatten creates a copy
arr.flatten()
# While ravel doesn't
arr = np.zeros((4,4))
arr.ravel()

# By default numpy arrays are created in C order (row-major), but
# you can arrange them in fortran order as well on creation,
# or make a fortran-order copy of an existing array
arr = np.asfortranarray(arr)
# Data is always contiguously packed in memory

# Arrays can be indexed as with python lists/tuples
arr = np.zeros((4,4))
arr[0][0]
# Or with a multidimensional index
arr[0,0]
# Negative indices start from the end
arr[-1,-1] # same as arr[3,3] for this array

# Or, like matlab, via splicing of a range: start:end
arr = arange(10)
arr[1:3] # elements 1,2
arr[:3] # 0,1,2
arr[5:] # 5,6,7,8,9
arr[5:-1] # 5,6,7,8
arr[0:4:2] # step 2: 0,3
arr = arr.reshape(2,5)
arr[0,:] # first row
arr[:,0] # first column

# Also, splicing works on a list of indices (see also choose)
arr = arr.reshape(10)
arr[[1,3,5]]

arr=np.array([1,2,3])
arr2=np.array([0,3,2])
arr[arr > arr2]

# flat returns an 1D-iterator
arr = arr.reshape(2,5)
arr.flat[3] # same as arr.reshape(arr.size)[3]

# Core operations on arrays are "ufunc"tions, element-wise
# vectorized operations
arr = arange(0,5)
arr2 = arange(5,10)
arr + arr2

arr = np.array([[0,1],[2,3]]) # shape (2,2)
arr2 = np.array([1,1]) # shape (1,)
# If we do arr+arr2 arr2.ndim
# an input can be used if shape in all dimensions sizes, or if
# the dimensions that don't match have size 1 in its shape
arr + arr2 # arr2 added to each row of arr!

# Broadcasting also works for assignment:
arr[...] = arr2 # [[1,1],[1,1]] note the [...] to access all contents
arr2 = arr # without [...] we just say arr2 refers to the same object as arr
arr[1,:] = 0 # This now is [[1,1],[0,0]]
# flat can be used with broadcasting too
arr.flat = 3 # [[3,3],[3,3]]
arr.flat[[1,3]] = 2 # [[3,2],[2,3]]

# It's possible to manually add ones in the shape using newaxis
print(arr[np.newaxis(),:,:].shape) # (1,2,2)
print(arr[:,np.newaxis(),:].shape) # (2,1,2)

# There are many ways to generate list of indices as well
arr = arange(5) # 0,1,2,3,4
arr[np.nonzero(arr)] += 2 # 0,3,4,5,6
arr = np.identity(3)
arr[np.diag_indices(3)] = 0 # diagonal elements, same as np.diag(arg.shape[0])
arr[np.tril_indices(3)] = 1 # lower triangle elements
arr[np.unravel_index(5,(3,3))] = 2 # returns an index given the flattened index and a shape

# Iteration over arrays can be done with for loops and indices
# Oterating over single elements is of course slower than native
# numpy operators. Prefer vector operations with splicing and masking
# Cython, Numba, Weave or Numexpr can be used when performance matters.
arr = np.arange(10)
for idx in range(arr.size):
print idx
# For multidimensiona arrays there are indices and iterators
arr = np.identity(3)
for idx in np.ndindex(arr.shape):
print arr[idx]
for idx, val in np.ndenumerate(arr):
print idx, val # assigning to val won't change arr
for val in arr.flat:
print val # assigning to val won't change arr
for val in np.nditer(arr):
print val # same as before
val[...] += 1 # this changes arr

# Vector and Matrix multiplication are done with dot
arr = np.array([1,2,3])
arr2 = np.identity(3)
np.dot(arr2, arr)
# Or by using the matrix object, note that 1d vectors
# are interpreted as rows when "seen" by a matrix object
np.asmatrix(arr2) * np.asmatrix(arr).transpose()

# Comparisons are also element-wise and generate masks
arr2 = np.array([2,0,0])
print (arr2 > arr)
# Branching can be done checking predicates for any true or all true
if (arr2 > arr).any():
print "at least one greater"

# Mapping a function over an array should NOT be done w/comprehensions
arr = np.arange(5)
[x*2 for x in arr] # this will return a list, not an array
# Instead use apply_along_axis, axis 0 is rows
np.apply_along_axis((lambda x: x*2), 0, arr)
# apply_along_axis is equivalent to a python loop, for simple expressions like the above, it's much slower than broadcasting
# It's also possible to vectorize python functions, but they won't execute faster
def test(x):
return x*2
testV = np.vectorize(test)
testV(arr)

# Scipy adds a wealth of numerical analysis functions, it's simple
# so I won't write about it.
# Matplotlib (replicates Matlab plotting) is worth having a quick look.
# IPython supports matplotlib integration and can display plots inline
%matplotlib inline
# Unfortunately if you chose the inline backend you will lose the ability
# of interacting with the plots (zooming, panning...)
# Use a separate backend like %matplotlib qt if interaction is needed

# Simple plots are simple
test = np.linspace(0, 2*pi)
plt.plot(test, np.sin(test)) # %pylab imports matplotib too
plt.show() # IPython will show automatically a plot in a cell, show() starts a new plot
plt.plot(test, np.cos(test)) # a separate plot
plt.plot(test, test) # this will be part of the cos plot
#plt.show()

# Multiple plots can also be done directly with a single plot statement
test = np.arange(0., 5., 0.1)
# Notes that ** is exponentiation.
# Styles in strings: red squares and blue triangles
plt.plot(test, test**2, 'rs', test, test**3, 'b^')
plt.show()
# It's also possible to do multiple plots in a grid with subplot
plt.subplot(211)
plt.plot(test, np.sin(test))
plt.subplot(212)
plt.plot(test, np.cos(test), 'r--')
#plt.show()

# Matplotlib plots use a hierarchy of objects you can edit to
# craft the final image. There are also global objects that are
# used if you don't specify any. This is the same as before
# using separate objects instead of the global ones
fig = plt.figure()
plt.plot(test, np.sin(test), axes = axes1)
plt.plot(test, np.cos(test), 'r--', axes = axes2)
#fig.show()

# All components that do rendering are called artists
# Figure and Axes are container artists
# Line2D, Rectangle, Text, AxesImage and so on are primitive artists
# Top-level commands like plot generate primitives to create a graphic

# Matplotlib is extensible via toolkits. Toolkits are very important.
# For example mplot3d is a toolkit that enables 3d drawing:
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
points = np.random.random((100,3))
# note that scatter wants separate x,y,z arrays
ax.scatter(points[:,0], points[:,1], points[:,2])
#fig.show()

# Matplotlib is quite extensive (and ugly) and can do much more
# It can do animations, but it's very painful, by changing the
# data artists use after having generated them. Not advised.
# What I end up doing is to generate a sequence on PNG from
# matplotlib and play them in sequence. In general MPL is very slow.

# If you want to just plot functions, instead of discrete data,
# sympy plots are a usable alternative to manual sampling:
from sympy import symbols as sym
from sympy import plotting as splt
from sympy import sin
x = sym('x')
splt.plot(sin(x), (x, 0, 2*pi))