Imagine that we want to add two vectors called a
and b
. The word vector is used here in the mathematical sense, which means a one-dimensional array. We will learn about specialized NumPy arrays that represent matrices in Chapter 4, Statistics and Linear Algebra. The vector a
holds the squares of integers 0
to n; for instance, if n
is equal to 3
, a
contains 0
, 1
, or 4
. The vector b
holds the cubes of integers 0
to n, so if n
is equal to 3
, then the vector b
is equal to 0
, 1
, or 8
. How would you do that using plain Python? After we come up with a solution, we will compare it to the NumPy equivalent.
The following function solves the vector addition problem using pure Python without NumPy:
def pythonsum(n):
a = list(range(n))
b = list(range(n))
c = []
for i in range(len(a)):
a[i] = i ** 2
b[i] = i ** 3
c.append(a[i] + b[i])
return c
The following is a function that solves the vector addition problem with NumPy:
def numpysum(n):
a = numpy.arange(n) ** 2
b = numpy.arange(n) ** 3
c = a + b
return c
Note that numpysum()
does not need a for
loop. We also used the arange()
function from NumPy, which creates a NumPy array for us with integers from 0
to n. The arange()
function was imported; that is why it is prefixed with numpy
.
Now comes the fun part. We mentioned earlier that NumPy is faster when it comes to array operations. How much faster is Numpy, though? The following program will show us by measuring the elapsed time in microseconds for the numpysum()
and pythonsum()
functions. It also prints the last two elements of the vector sum. Let's check that we get the same answers using Python and NumPy:
#!/usr/bin/env/python
import sys
from datetime import datetime
import numpy as np
"""
This program demonstrates vector addition the Python way.
Run the following from the command line:
python vectorsum.py n
Here, n is an integer that specifies the size of the vectors.
The first vector to be added contains the squares of 0 up to n.
The second vector contains the cubes of 0 up to n.
The program prints the last 2 elements of the sum and the elapsed time:
"""
def numpysum(n):
a = np.arange(n) ** 2
b = np.arange(n) ** 3
c = a + b
return c
def pythonsum(n):
a = list(range(n))
b = list(range(n))
c = []
for i in range(len(a)):
a[i] = i ** 2
b[i] = i ** 3
c.append(a[i] + b[i])
return c
size = int(sys.argv[1])
start = datetime.now()
c = pythonsum(size)
delta = datetime.now() - start
print("The last 2 elements of the sum", c[-2:])
print("PythonSum elapsed time in microseconds", delta.microseconds)
start = datetime.now()
c = numpysum(size)
delta = datetime.now() - start
print("The last 2 elements of the sum", c[-2:])
print("NumPySum elapsed time in microseconds", delta.microseconds)
The output of the program for 1000
, 2000
, and 3000
vector elements is as follows:
$ python3 vectorsum.py 1000
The last 2 elements of the sum [995007996, 998001000]
PythonSum elapsed time in microseconds 976
The last 2 elements of the sum [995007996 998001000]
NumPySum elapsed time in microseconds 87
$ python3 vectorsum.py 2000
The last 2 elements of the sum [7980015996, 7992002000]
PythonSum elapsed time in microseconds 1623
The last 2 elements of the sum [7980015996 7992002000]
NumPySum elapsed time in microseconds 143
$ python3 vectorsum.py 4000
The last 2 elements of the sum [63920031996, 63968004000]
PythonSum elapsed time in microseconds 3417
The last 2 elements of the sum [63920031996 63968004000]
NumPySum elapsed time in microseconds 237
Clearly, NumPy is much faster than the equivalent normal Python code. One thing is certain; we get the same results whether we are using NumPy or not. However, the result that is printed differs in representation. Note that the result from the numpysum()
function does not have any commas. How come? Obviously, we are not dealing with a Python list, but with a NumPy array. We will learn more about NumPy arrays in the Chapter 2, NumPy Arrays.