Introduction to Array Math with NumPy Arrays

In [1]:
from IPython.display import HTML
In [2]:
HTML('<img src="http://teaching.software-carpentry.org/wp-content/uploads/2012/11/array-math-cmap.png" height="500px">')
Out[2]:
In [3]:
import numpy as np

We'll need a couple of arrays for demo purposes...

In [4]:
a = np.arange(5)
b = np.arange(5, 10)
print a
print b
[0 1 2 3 4]
[5 6 7 8 9]

Array Math with Scalars

Now we can use binary operators like +, -, *, /, and ** on these arrays, which will return new arrays. We'll start with combining arrays and scalars and then look at what happens with arrays on both sides of the operator.

In [5]:
a + 6
Out[5]:
array([ 6,  7,  8,  9, 10])

You could also do this with a loop, list comprehension, or the map function, but I think you'll agree a + 6 is much easier to read and write:

In [6]:
new_a = []
for i in xrange(a.size):
    new_a.append(a[i] + 6)
np.array(new_a)
Out[6]:
array([ 6,  7,  8,  9, 10])

Array Math with Two Arrays

You can also do math with two or more arrays using the binary operators:

In [7]:
a * b
Out[7]:
array([ 0,  6, 14, 24, 36])

The new arrays was constructed by multiplying the first element of a by the first element of b, and so on.

What if the arrays aren't the same size?

In [8]:
a = np.arange(5)
b = np.arange(6)
a * b
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-8-8c541a828af2> in <module>()
      1 a = np.arange(5)
      2 b = np.arange(6)
----> 3 a * b

ValueError: operands could not be broadcast together with shapes (5) (6) 

Well that didn't work. What if the arrays are the same in at least one dimension?

In [9]:
a = np.ones((3, 2)) # 3 x 2
b = np.arange(4, 6) # 1 x 2
a * b
Out[9]:
array([[ 4.,  5.],
       [ 4.,  5.],
       [ 4.,  5.]])

NumPy saw that b could be repeated to match a's shape and did that automatically. It will also work if the number of rows matches:

In [10]:
b = np.array([[4], [5], [6]]) # 3 x 1
a * b
Out[10]:
array([[ 4.,  4.],
       [ 5.,  5.],
       [ 6.,  6.]])

Will two 2d arrays work if one can be repeated to match the other?

In [11]:
a = np.ones((4, 2)) # 4 rows, 2 columns
b = np.array([[4, 5], [6, 7]]) # 2 rows, 2 columns
a * b
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-11-14ab02e48779> in <module>()
      1 a = np.ones((4, 2)) # 4 rows, 2 columns
      2 b = np.array([[4, 5], [6, 7]]) # 2 rows, 2 columns
----> 3 a * b

ValueError: operands could not be broadcast together with shapes (4,2) (2,2) 

In this case NumPy couldn't figure out what to do and raised an exception. But let's look at what happens when one array matches the dimensionality of a subarray in the other:

In [12]:
a = np.ones((4, 3, 2))
b = np.array([[4, 5], [6, 7], [8, 9]]) # 3 x 2
a * b
Out[12]:
array([[[ 4.,  5.],
        [ 6.,  7.],
        [ 8.,  9.]],

       [[ 4.,  5.],
        [ 6.,  7.],
        [ 8.,  9.]],

       [[ 4.,  5.],
        [ 6.,  7.],
        [ 8.,  9.]],

       [[ 4.,  5.],
        [ 6.,  7.],
        [ 8.,  9.]]])

Here were are getting into arrays of higher dimensionality so it becomes harder to visualize.

When combining two arrays the operation must fit into one these categories:

  1. The arrays are the same shape.
  2. The array with smaller number of dimensions matches the last dimensions of the other array.
    • For example, an array with dimensions (4, 3, 2) and an array with dimensions (3, 2) can be combined.

Teaching Discussion

  1. I prefer to do live coding when working in front of a class so I would print this out and use it as lecture notes, but the class would see me working in an empty notebook.
  2. On their own learners can step through this, read my annotations, and change things to see what happens.
  3. Having a small amount of code followed by the result of that code makes it easy to see the result of changes, as shown in Bret Victor's demo. Thinking of the book, I think demonstrating situations that work and situations that error helps reinforce underlying mechanics for students.
In []:
 
Back to top