Rick Muller, Sandia National Laboratories

version 0.6

This work is licensed under a Creative Commons Attribution-ShareAlike 3.0 Unported License.

Python is the programming language of choice for many scientists to a large degree because it offers a great deal of power to analyze and model scientific data with relatively little overhead in terms of learning, installation or development time. It is a language you can pick up in a weekend, and use for the rest of one's life.

The Python Tutorial is a great place to start getting a feel for the language. To complement this material, I taught a Python Short Course years ago to a group of computational chemists during a time that I was worried the field was moving too much in the direction of using canned software rather than developing one's own methods. I wanted to focus on what working scientists needed to be more productive: parsing output of other programs, building simple models, experimenting with object oriented programming, extending the language with C, and simple GUIs.

I'm trying to do something very similar here, to cut to the chase and focus on what scientists need. In the last year or so, the IPython Project has put together a notebook interface that I have found incredibly valuable. A large number of people have released very good IPython Notebooks that I have taken a huge amount of pleasure reading through. Some ones that I particularly like include:

- Rob Johansson's excellent notebooks, including Scientific Computing with Python and Computational Quantum Physics with QuTiP lectures;
- XKCD style graphs in matplotlib;
- A collection of Notebooks for using IPython effectively
- A gallery of interesting IPython Notebooks

I find IPython notebooks an easy way both to get important work done in my everyday job, as well as to communicate what I've done, how I've done it, and why it matters to my coworkers. I find myself endlessly sweeping the IPython subreddit hoping someone will post a new notebook. In the interest of putting more notebooks out into the wild for other people to use and enjoy, I thought I would try to recreate some of what I was trying to get across in the original Python Short Course, updated by 15 years of Python, Numpy, Scipy, Matplotlib, and IPython development, as well as my own experience in using Python almost every day of this time.

There are two branches of current releases in Python: the older-syntax Python 2, and the newer-syntax Python 3. This schizophrenia is largely intentional: when it became clear that some non-backwards-compatible changes to the language were necessary, the Python dev-team decided to go through a five-year (or so) transition, during which the new language features would be introduced and the old language was still actively maintained, to make such a transition as easy as possible. We're now (2013) past the halfway point, and, IMHO, at the first time when I'm considering making the change to Python 3.

Nonetheless, I'm going to write these notes with Python 2 in mind, since this is the version of the language that I use in my day-to-day job, and am most comfortable with. If these notes are important and are valuable to people, I'll be happy to rewrite the notes using Python 3.

With this in mind, these notes assume you have a Python distribution that includes:

- Python version 2.7;
- Numpy, the core numerical extensions for linear algebra and multidimensional arrays;
- Scipy, additional libraries for scientific programming;
- Matplotlib, excellent plotting and graphing libraries;
- IPython, with the additional libraries required for the notebook interface.

A good, easy to install option that supports Mac, Windows, and Linux, and that has all of these packages (and much more) is the Entought Python Distribution, also known as EPD, which appears to be changing its name to Enthought Canopy. Enthought is a commercial company that supports a lot of very good work in scientific Python development and application. You can either purchase a license to use EPD, or there is also a free version that you can download and install.

Here are some other alternatives, should you not want to use EPD:

**Linux** Most distributions have an installation manager. Redhat has yum, Ubuntu has apt-get. To my knowledge, all of these packages should be available through those installers.

**Mac** I use Macports, which has up-to-date versions of all of these packages.

**Windows** The PythonXY package has everything you need: install the package, then go to Start > PythonXY > Command Prompts > IPython notebook server.

**Cloud** This notebook is currently not running on the IPython notebook viewer, but will be shortly, which will allow the notebook to be viewed but not interactively. I'm keeping an eye on Wakari, from Continuum Analytics, which is a cloud-based IPython notebook. Wakari appears to support free accounts as well. Continuum is a company started by some of the core Enthought Numpy/Scipy people focusing on big data.

Continuum also supports a bundled, multiplatform Python package called Anaconda that I'll also keep an eye on.

This is a quick introduction to Python. There are lots of other places to learn the language more thoroughly. I have collected a list of useful links, including ones to other learning resources, at the end of this notebook. If you want a little more depth, Python Tutorial is a great place to start, as is Zed Shaw's Learn Python the Hard Way.

The lessons that follow make use of the IPython notebooks. There's a good introduction to notebooks in the IPython notebook documentation that even has a nice video on how to use the notebooks. You should probably also flip through the IPython tutorial in your copious free time.

Briefly, notebooks have code cells (that are generally followed by result cells) and text cells. The text cells are the stuff that you're reading now. The code cells start with "In []:" with some number generally in the brackets. If you put your cursor in the code cell and hit Shift-Enter, the code will run in the Python interpreter and the result will print out in the output cell. You can then change things around and see whether you understand what's going on. If you need to know more, see the IPython notebook documentation or the IPython tutorial.

Many of the things I used to use a calculator for, I now use Python for:

In [1]:

```
2+2
```

Out[1]:

4

In [2]:

```
(50-5*6)/4
```

Out[2]:

5

There are some gotchas compared to using a normal calculator.

In [3]:

```
7/3
```

Out[3]:

2

Python integer division, like C or Fortran integer division, truncates the remainder and returns an integer. At least it does in version 2. In version 3, Python returns a floating point number. You can get a sneak preview of this feature in Python 2 by importing the module from the future features:

```
from __future__ import division
```

In [4]:

```
7/3.
```

Out[4]:

2.3333333333333335

In [5]:

```
7/float(3)
```

Out[5]:

2.3333333333333335

In the last few lines, we have sped by a lot of things that we should stop for a moment and explore a little more fully. We've seen, however briefly, two different data types: **integers**, also known as *whole numbers* to the non-programming world, and **floating point numbers**, also known (incorrectly) as *decimal numbers* to the rest of the world.

We've also seen the first instance of an **import** statement. Python has a huge number of libraries included with the distribution. To keep things simple, most of these variables and functions are not accessible from a normal Python interactive session. Instead, you have to import the name. For example, there is a **math** module containing many useful functions. To access, say, the square root function, you can either first

```
from math import sqrt
```

and then

In [6]:

```
sqrt(81)
```

Out[6]:

9.0

or you can simply import the math library itself

In [7]:

```
import math
math.sqrt(81)
```

Out[7]:

9.0

You can define variables using the equals (=) sign:

In [8]:

```
width = 20
length = 30
area = length*width
area
```

Out[8]:

600

If you try to access a variable that you haven't yet defined, you get an error:

In [9]:

```
volume
```

and you need to define it:

In [ ]:

```
depth = 10
volume = area*depth
volume
```

You can name a variable *almost* anything you want. It needs to start with an alphabetical character or "_", can contain alphanumeric charcters plus underscores ("_"). Certain words, however, are reserved for the language:

```
and, as, assert, break, class, continue, def, del, elif, else, except,
exec, finally, for, from, global, if, import, in, is, lambda, not, or,
pass, print, raise, return, try, while, with, yield
```

Trying to define a variable using one of these will result in a syntax error:

In [10]:

```
return = 0
```

File "<ipython-input-10-2b99136d4ec6>", line 1 return = 0 ^ SyntaxError: invalid syntax

Strings are lists of printable characters, and can be defined using either single quotes

In [11]:

```
'Hello, World!'
```

Out[11]:

'Hello, World!'

or double quotes

In [12]:

```
"Hello, World!"
```

Out[12]:

'Hello, World!'

But not both at the same time, unless you want one of the symbols to be part of the string.

In [13]:

```
"He's a Rebel"
```

Out[13]:

"He's a Rebel"

In [14]:

```
'She asked, "How are you today?"'
```

Out[14]:

'She asked, "How are you today?"'

In [15]:

```
greeting = "Hello, World!"
```

The **print** statement is often used for printing character strings:

In [16]:

```
print greeting
```

Hello, World!

But it can also print data types other than strings:

In [17]:

```
print "The area is ",area
```

The area is 600

You can use the + operator to concatenate strings together:

In [18]:

```
statement = "Hello," + "World!"
print statement
```

Hello,World!

Don't forget the space between the strings, if you want one there.

In [19]:

```
statement = "Hello, " + "World!"
print statement
```

Hello, World!

You can use + to concatenate multiple strings in a single statement:

In [20]:

```
print "This " + "is " + "a " + "longer " + "statement."
```

This is a longer statement.

Very often in a programming language, one wants to keep a group of similar items together. Python does this using a data type called **lists**.

In [21]:

```
days_of_the_week = ["Sunday","Monday","Tuesday","Wednesday","Thursday","Friday","Saturday"]
```

You can access members of the list using the **index** of that item:

In [22]:

```
days_of_the_week[2]
```

Out[22]:

'Tuesday'

*n*th element from the end of the list, you can use a negative index. For example, the -1 element of a list is the last element:

In [23]:

```
days_of_the_week[-1]
```

Out[23]:

'Saturday'

You can add additional items to the list using the .append() command:

In [24]:

```
languages = ["Fortran","C","C++"]
languages.append("Python")
print languages
```

['Fortran', 'C', 'C++', 'Python']

The **range()** command is a convenient way to make sequential lists of numbers:

In [25]:

```
range(10)
```

Out[25]:

[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

In [26]:

```
range(2,8)
```

Out[26]:

[2, 3, 4, 5, 6, 7]

*step* of 1 between elements. You can also give a fixed step size via a third command:

In [27]:

```
evens = range(0,20,2)
evens
```

Out[27]:

[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]

In [28]:

```
evens[3]
```

Out[28]:

6

Lists do not have to hold the same data type. For example,

In [29]:

```
["Today",7,99.3,""]
```

Out[29]:

['Today', 7, 99.3, '']

However, it's good (but not essential) to use lists for similar objects that are somehow logically connected. If you want to group different data types together into a composite data object, it's best to use **tuples**, which we will learn about below.

You can find out how long a list is using the **len()** command:

In [30]:

```
help(len)
```

In [31]:

```
len(evens)
```

Out[31]:

10

One of the most useful things you can do with lists is to *iterate* through them, i.e. to go through each element one at a time. To do this in Python, we use the **for** statement:

In [32]:

```
for day in days_of_the_week:
print day
```

Sunday Monday Tuesday Wednesday Thursday Friday Saturday

This code snippet goes through each element of the list called **days_of_the_week** and assigns it to the variable **day**. It then executes everything in the indented block (in this case only one line of code, the print statement) using those variable assignments. When the program has gone through every element of the list, it exists the block.

(Almost) every programming language defines blocks of code in some way. In Fortran, one uses END statements (ENDDO, ENDIF, etc.) to define code blocks. In C, C++, and Perl, one uses curly braces {} to define these blocks.

Python uses a colon (":"), followed by indentation level to define code blocks. Everything at a higher level of indentation is taken to be in the same block. In the above example the block was only a single line, but we could have had longer blocks as well:

In [33]:

```
for day in days_of_the_week:
statement = "Today is " + day
print statement
```

**range()** command is particularly useful with the **for** statement to execute loops of a specified length:

In [34]:

```
for i in range(20):
print "The square of ",i," is ",i*i
```

Lists and strings have something in common that you might not suspect: they can both be treated as sequences. You already know that you can iterate through the elements of a list. You can also iterate through the letters in a string:

In [35]:

```
for letter in "Sunday":
print letter
```

S u n d a y

*slicing* operation, which you can also use on any sequence. We already know that we can use *indexing* to get the first element of a list:

In [36]:

```
days_of_the_week[0]
```

Out[36]:

'Sunday'

If we want the list containing the first two elements of a list, we can do this via

In [37]:

```
days_of_the_week[0:2]
```

Out[37]:

['Sunday', 'Monday']

or simply

In [38]:

```
days_of_the_week[:2]
```

Out[38]:

['Sunday', 'Monday']

If we want the last items of the list, we can do this with negative slicing:

In [39]:

```
days_of_the_week[-2:]
```

Out[39]:

['Friday', 'Saturday']

which is somewhat logically consistent with negative indices accessing the last elements of the list.

You can do:

In [40]:

```
workdays = days_of_the_week[1:6]
print workdays
```

['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday']

Since strings are sequences, you can also do this to them:

In [41]:

```
day = "Sunday"
abbreviation = day[:3]
print abbreviation
```

Sun

**range()** function specifies the step):

In [42]:

```
numbers = range(0,40)
evens = numbers[2::2]
evens
```

Out[42]:

[2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38]

We have now learned a few data types. We have integers and floating point numbers, strings, and lists to contain them. We have also learned about lists, a container that can hold any data type. We have learned to print things out, and to iterate over items in lists. We will now learn about **boolean** variables that can be either True or False.

We invariably need some concept of *conditions* in programming to control branching behavior, to allow a program to react differently to different situations. If it's Monday, I'll go to work, but if it's Sunday, I'll sleep in. To do this in Python, we use a combination of **boolean** variables, which evaluate to either True or False, and **if** statements, that control branching based on boolean values.

For example:

In [43]:

```
if day == "Sunday":
print "Sleep in"
else:
print "Go to work"
```

Sleep in

(Quick quiz: why did the snippet print "Go to work" here? What is the variable "day" set to?)

Let's take the snippet apart to see what happened. First, note the statement

In [44]:

```
day == "Sunday"
```

Out[44]:

True

*equality testing*. If the two items are equal, it returns True, otherwise it returns False. In this case, it is comparing two variables, the string "Sunday", and whatever is stored in the variable "day", which, in this case, is the other string "Saturday". Since the two strings are not equal to each other, the truth test has the false value.

The if statement that contains the truth test is followed by a code block (a colon followed by an indented block of code). If the boolean is true, it executes the code in that block. Since it is false in the above example, we don't see that code executed.

The first block of code is followed by an **else** statement, which is executed if nothing else in the above if statement is true. Since the value was false, this code is executed, which is why we see "Go to work".

You can compare any data types in Python:

In [45]:

```
1 == 2
```

Out[45]:

False

In [46]:

```
50 == 2*25
```

Out[46]:

True

In [47]:

```
3 < 3.14159
```

Out[47]:

True

In [48]:

```
1 == 1.0
```

Out[48]:

True

In [49]:

```
1 != 0
```

Out[49]:

True

In [50]:

```
1 <= 2
```

Out[50]:

True

In [51]:

```
1 >= 1
```

Out[51]:

True

We see a few other boolean operators here, all of which which should be self-explanatory. Less than, equality, non-equality, and so on.

Particularly interesting is the 1 == 1.0 test, which is true, since even though the two objects are different data types (integer and floating point number), they have the same *value*. There is another boolean operator **is**, that tests whether two objects are the same object:

In [52]:

```
1 is 1.0
```

Out[52]:

False

We can do boolean tests on lists as well:

In [53]:

```
[1,2,3] == [1,2,4]
```

Out[53]:

False

In [54]:

```
[1,2,3] < [1,2,4]
```

Out[54]:

True

In [55]:

```
hours = 5
0 < hours < 24
```

Out[55]:

True

If statements can have **elif** parts ("else if"), in addition to if/else parts. For example:

In [56]:

```
if day == "Sunday":
print "Sleep in"
elif day == "Saturday":
print "Do chores"
else:
print "Go to work"
```

Sleep in

Of course we can combine if statements with for loops, to make a snippet that is almost interesting:

In [57]:

```
for day in days_of_the_week:
statement = "Today is " + day
print statement
if day == "Sunday":
print " Sleep in"
elif day == "Saturday":
print " Do chores"
else:
print " Go to work"
```

**bool()** function.

In [58]:

```
bool(1)
```

Out[58]:

True

In [59]:

```
bool(0)
```

Out[59]:

False

In [60]:

```
bool(["This "," is "," a "," list"])
```

Out[60]:

True

The Fibonacci sequence is a sequence in math that starts with 0 and 1, and then each successive entry is the sum of the previous two. Thus, the sequence goes 0,1,1,2,3,5,8,13,21,34,55,89,...

A very common exercise in programming books is to compute the Fibonacci sequence up to some number **n**. First I'll show the code, then I'll discuss what it is doing.

In [61]:

```
n = 10
sequence = [0,1]
for i in range(2,n): # This is going to be a problem if we ever set n <= 2!
sequence.append(sequence[i-1]+sequence[i-2])
print sequence
```

[0, 1, 1, 2, 3, 5, 8, 13, 21, 34]

Let's go through this line by line. First, we define the variable **n**, and set it to the integer 20. **n** is the length of the sequence we're going to form, and should probably have a better variable name. We then create a variable called **sequence**, and initialize it to the list with the integers 0 and 1 in it, the first two elements of the Fibonacci sequence. We have to create these elements "by hand", since the iterative part of the sequence requires two previous elements.

We then have a for loop over the list of integers from 2 (the next element of the list) to **n** (the length of the sequence). After the colon, we see a hash tag "#", and then a **comment** that if we had set **n** to some number less than 2 we would have a problem. Comments in Python start with #, and are good ways to make notes to yourself or to a user of your code explaining why you did what you did. Better than the comment here would be to test to make sure the value of **n** is valid, and to complain if it isn't; we'll try this later.

In the body of the loop, we append to the list an integer equal to the sum of the two previous elements of the list.

After exiting the loop (ending the indentation) we then print out the whole list. That's it!

We might want to use the Fibonacci snippet with different sequence lengths. We could cut an paste the code into another cell, changing the value of **n**, but it's easier and more useful to make a function out of the code. We do this with the **def** statement in Python:

In [62]:

```
def fibonacci(sequence_length):
"Return the Fibonacci sequence of length *sequence_length*"
sequence = [0,1]
if sequence_length < 1:
print "Fibonacci sequence only defined for length 1 or greater"
return
if 0 < sequence_length < 3:
return sequence[:sequence_length]
for i in range(2,sequence_length):
sequence.append(sequence[i-1]+sequence[i-2])
return sequence
```

We can now call **fibonacci()** for different sequence_lengths:

In [63]:

```
fibonacci(2)
```

Out[63]:

[0, 1]

In [64]:

```
fibonacci(12)
```

Out[64]:

[0, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89]

**docstring**, and is a special kind of comment that is often available to people using the function through the python command line:

In [65]:

```
help(fibonacci)
```

If you define a docstring for all of your functions, it makes it easier for other people to use them, since they can get help on the arguments and return values of the function.

Next, note that rather than putting a comment in about what input values lead to errors, we have some testing of these values, followed by a warning if the value is invalid, and some conditional code to handle special cases.

Functions can also call themselves, something that is often called *recursion*. We're going to experiment with recursion by computing the factorial function. The factorial is defined for a positive integer **n** as

First, note that we don't need to write a function at all, since this is a function built into the standard math library. Let's use the help function to find out about it:

In [66]:

```
from math import factorial
help(factorial)
```

This is clearly what we want.

In [67]:

```
factorial(20)
```

Out[67]:

2432902008176640000

However, if we did want to write a function ourselves, we could do recursively by noting that

$$ n! = n(n-1)!$$The program then looks something like:

In [68]:

```
def fact(n):
if n <= 0:
return 1
return n*fact(n-1)
```

In [69]:

```
fact(20)
```

Out[69]:

2432902008176640000

Recursion can be very elegant, and can lead to very simple programs.

Before we end the Python overview, I wanted to touch on two more data structures that are very useful (and thus very common) in Python programs.

A **tuple** is a sequence object like a list or a string. It's constructed by grouping a sequence of objects together with commas, either without brackets, or with parentheses:

In [70]:

```
t = (1,2,'hi',9.0)
t
```

Out[70]:

(1, 2, 'hi', 9.0)

Tuples are like lists, in that you can access the elements using indices:

In [71]:

```
t[1]
```

Out[71]:

2

However, tuples are *immutable*, you can't append to them or change the elements of them:

In [72]:

```
t.append(7)
```

In [73]:

```
t[1]=77
```

In [74]:

```
('Bob',0.0,21.0)
```

Out[74]:

('Bob', 0.0, 21.0)

In [75]:

```
positions = [
('Bob',0.0,21.0),
('Cat',2.5,13.1),
('Dog',33.0,1.2)
]
```

In [76]:

```
def minmax(objects):
minx = 1e20 # These are set to really big numbers
miny = 1e20
for obj in objects:
name,x,y = obj
if x < minx:
minx = x
if y < miny:
miny = y
return minx,miny
x,y = minmax(positions)
print x,y
```

0.0 1.2

Here we did two things with tuples you haven't seen before. First, we unpacked an object into a set of named variables using *tuple assignment*:

```
>>> name,x,y = obj
```

We also returned multiple values (minx,miny), which were then assigned to two other variables (x,y), again by tuple assignment. This makes what would have been complicated code in C++ rather simple.

Tuple assignment is also a convenient way to swap variables:

In [77]:

```
x,y = 1,2
y,x = x,y
x,y
```

Out[77]:

(2, 1)

**Dictionaries** are an object called "mappings" or "associative arrays" in other languages. Whereas a list associates an integer index with a set of objects:

In [78]:

```
mylist = [1,2,9,21]
```

*key*, and the corresponding dictionary entry is the *value*. A dictionary can use (almost) anything as the key. Whereas lists are formed with square brackets [], dictionaries use curly brackets {}:

In [79]:

```
ages = {"Rick": 46, "Bob": 86, "Fred": 21}
print "Rick's age is ",ages["Rick"]
```

Rick's age is 46

There's also a convenient way to create dictionaries without having to quote the keys.

In [80]:

```
dict(Rick=46,Bob=86,Fred=20)
```

Out[80]:

{'Bob': 86, 'Fred': 20, 'Rick': 46}

The **len()** command works on both tuples and dictionaries:

In [81]:

```
len(t)
```

Out[81]:

4

In [82]:

```
len(ages)
```

Out[82]:

3

We can generally understand trends in data by using a plotting program to chart it. Python has a wonderful plotting library called Matplotlib. The IPython notebook interface we are using for these notes has that functionality built in.

As an example, we have looked at two different functions, the Fibonacci function, and the factorial function, both of which grow faster than polynomially. Which one grows the fastest? Let's plot them. First, let's generate the Fibonacci sequence of length 20:

In [83]:

```
fibs = fibonacci(10)
```

Next lets generate the factorials.

In [84]:

```
facts = []
for i in range(10):
facts.append(factorial(i))
```

Now we use the Matplotlib function **plot** to compare the two.

In [85]:

```
figsize(8,6)
plot(facts,label="factorial")
plot(fibs,label="Fibonacci")
xlabel("n")
legend()
```

Out[85]:

<matplotlib.legend.Legend at 0x10d1b4890>

The factorial function grows much faster. In fact, you can't even see the Fibonacci sequence. It's not entirely surprising: a function where we multiply by n each iteration is bound to grow faster than one where we add (roughly) n each iteration.

Let's plot these on a semilog plot so we can see them both a little more clearly:

In [86]:

```
semilogy(facts,label="factorial")
semilogy(fibs,label="Fibonacci")
xlabel("n")
legend()
```

Out[86]:

<matplotlib.legend.Legend at 0x10d2bee90>

There is, of course, much more to the language than I've covered here. I've tried to keep this brief enough so that you can jump in and start using Python to simplify your life and work. My own experience in learning new things is that the information doesn't "stick" unless you try and use it for something in real life.

You will no doubt need to learn more as you go. I've listed several other good references, including the Python Tutorial and Learn Python the Hard Way. Additionally, now is a good time to start familiarizing yourself with the Python Documentation, and, in particular, the Python Language Reference.

Tim Peters, one of the earliest and most prolific Python contributors, wrote the "Zen of Python", which can be accessed via the "import this" command:

In [87]:

```
import this
```

No matter how experienced a programmer you are, these are words to meditate on.

Numpy contains core routines for doing fast vector, matrix, and linear algebra-type operations in Python. Scipy contains additional routines for optimization, special functions, and so on. Both contain modules written in C and Fortran so that they're as fast as possible. Together, they give Python roughly the same capability that the Matlab program offers. (In fact, if you're an experienced Matlab user, there a guide to Numpy for Matlab users just for you.)

Fundamental to both Numpy and Scipy is the ability to work with vectors and matrices. You can create vectors from lists using the **array** command:

In [88]:

```
array([1,2,3,4,5,6])
```

Out[88]:

array([1, 2, 3, 4, 5, 6])

**array** that gives the numeric type. There are a number of types listed here that your matrix can be. Some of these are aliased to single character codes. The most common ones are 'd' (double precision floating point number), 'D' (double precision complex number), and 'i' (int32). Thus,

In [89]:

```
array([1,2,3,4,5,6],'d')
```

Out[89]:

array([ 1., 2., 3., 4., 5., 6.])

In [90]:

```
array([1,2,3,4,5,6],'D')
```

Out[90]:

array([ 1.+0.j, 2.+0.j, 3.+0.j, 4.+0.j, 5.+0.j, 6.+0.j])

In [91]:

```
array([1,2,3,4,5,6],'i')
```

Out[91]:

array([1, 2, 3, 4, 5, 6], dtype=int32)

To build matrices, you can either use the array command with lists of lists:

In [92]:

```
array([[0,1],[1,0]],'d')
```

Out[92]:

array([[ 0., 1.], [ 1., 0.]])

**zeros** command:

In [93]:

```
zeros((3,3),'d')
```

Out[93]:

array([[ 0., 0., 0.], [ 0., 0., 0.], [ 0., 0., 0.]])

In [94]:

```
zeros(3,'d')
```

Out[94]:

array([ 0., 0., 0.])

In [95]:

```
zeros((1,3),'d')
```

Out[95]:

array([[ 0., 0., 0.]])

or column vectors:

In [96]:

```
zeros((3,1),'d')
```

Out[96]:

array([[ 0.], [ 0.], [ 0.]])

There's also an **identity** command that behaves as you'd expect:

In [97]:

```
identity(4,'d')
```

Out[97]:

array([[ 1., 0., 0., 0.], [ 0., 1., 0., 0.], [ 0., 0., 1., 0.], [ 0., 0., 0., 1.]])

as well as a **ones** command.

The **linspace** command makes a linear array of points from a starting to an ending value.

In [98]:

```
linspace(0,1)
```

Out[98]:

array([ 0. , 0.02040816, 0.04081633, 0.06122449, 0.08163265, 0.10204082, 0.12244898, 0.14285714, 0.16326531, 0.18367347, 0.20408163, 0.2244898 , 0.24489796, 0.26530612, 0.28571429, 0.30612245, 0.32653061, 0.34693878, 0.36734694, 0.3877551 , 0.40816327, 0.42857143, 0.44897959, 0.46938776, 0.48979592, 0.51020408, 0.53061224, 0.55102041, 0.57142857, 0.59183673, 0.6122449 , 0.63265306, 0.65306122, 0.67346939, 0.69387755, 0.71428571, 0.73469388, 0.75510204, 0.7755102 , 0.79591837, 0.81632653, 0.83673469, 0.85714286, 0.87755102, 0.89795918, 0.91836735, 0.93877551, 0.95918367, 0.97959184, 1. ])

In [99]:

```
linspace(0,1,11)
```

Out[99]:

array([ 0. , 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1. ])

**linspace** is an easy way to make coordinates for plotting. Functions in the numpy library (all of which are imported into IPython notebook) can act on an entire vector (or even a matrix) of points at once. Thus,

In [100]:

```
x = linspace(0,2*pi)
sin(x)
```

Out[100]:

array([ 0.00000000e+00, 1.27877162e-01, 2.53654584e-01, 3.75267005e-01, 4.90717552e-01, 5.98110530e-01, 6.95682551e-01, 7.81831482e-01, 8.55142763e-01, 9.14412623e-01, 9.58667853e-01, 9.87181783e-01, 9.99486216e-01, 9.95379113e-01, 9.74927912e-01, 9.38468422e-01, 8.86599306e-01, 8.20172255e-01, 7.40277997e-01, 6.48228395e-01, 5.45534901e-01, 4.33883739e-01, 3.15108218e-01, 1.91158629e-01, 6.40702200e-02, -6.40702200e-02, -1.91158629e-01, -3.15108218e-01, -4.33883739e-01, -5.45534901e-01, -6.48228395e-01, -7.40277997e-01, -8.20172255e-01, -8.86599306e-01, -9.38468422e-01, -9.74927912e-01, -9.95379113e-01, -9.99486216e-01, -9.87181783e-01, -9.58667853e-01, -9.14412623e-01, -8.55142763e-01, -7.81831482e-01, -6.95682551e-01, -5.98110530e-01, -4.90717552e-01, -3.75267005e-01, -2.53654584e-01, -1.27877162e-01, -2.44929360e-16])

In conjunction with **matplotlib**, this is a nice way to plot things:

In [101]:

```
plot(x,sin(x))
```

Out[101]:

[<matplotlib.lines.Line2D at 0x10d49f510>]

Matrix objects act sensibly when multiplied by scalars:

In [102]:

```
0.125*identity(3,'d')
```

Out[102]:

array([[ 0.125, 0. , 0. ], [ 0. , 0.125, 0. ], [ 0. , 0. , 0.125]])

as well as when you add two matrices together. (However, the matrices have to be the same shape.)

In [103]:

```
identity(2,'d') + array([[1,1],[1,2]])
```

Out[103]:

array([[ 2., 1.], [ 1., 3.]])

In [104]:

```
identity(2)*ones((2,2))
```

Out[104]:

array([[ 1., 0.], [ 0., 1.]])

To get matrix multiplication, you need the **dot** command:

In [105]:

```
dot(identity(2),ones((2,2)))
```

Out[105]:

array([[ 1., 1.], [ 1., 1.]])

**dot** can also do dot products (duh!):

In [106]:

```
v = array([3,4],'d')
sqrt(dot(v,v))
```

Out[106]:

5.0

as well as matrix-vector products.

**determinant**, **inverse**, and **transpose** functions that act as you would suppose. Transpose can be abbreviated with ".T" at the end of a matrix object:

In [107]:

```
m = array([[1,2],[3,4]])
m.T
```

Out[107]:

array([[1, 3], [2, 4]])

**diag()** function that takes a list or a vector and puts it along the diagonal of a square matrix.

In [108]:

```
diag([1,2,3,4,5])
```

Out[108]:

array([[1, 0, 0, 0, 0], [0, 2, 0, 0, 0], [0, 0, 3, 0, 0], [0, 0, 0, 4, 0], [0, 0, 0, 0, 5]])

We'll find this useful later on.

You can solve systems of linear equations using the **solve** command:

In [109]:

```
A = array([[1,1,1],[0,2,5],[2,5,-1]])
b = array([6,-4,27])
solve(A,b)
```

Out[109]:

array([ 5., 3., -2.])

There are a number of routines to compute eigenvalues and eigenvectors

**eigvals**returns the eigenvalues of a matrix**eigvalsh**returns the eigenvalues of a Hermitian matrix**eig**returns the eigenvalues and eigenvectors of a matrix**eigh**returns the eigenvalues and eigenvectors of a Hermitian matrix.

In [110]:

```
A = array([[13,-4],[-4,7]],'d')
eigvalsh(A)
```

Out[110]:

array([ 5., 15.])

In [111]:

```
eigh(A)
```

Out[111]:

(array([ 5., 15.]), array([[-0.4472136 , -0.89442719], [-0.89442719, 0.4472136 ]]))

Now that we have these tools in our toolbox, we can start to do some cool stuff with it. Many of the equations we want to solve in Physics involve differential equations. We want to be able to compute the derivative of functions:

$$ y' = \frac{y(x+h)-y(x)}{h} $$by *discretizing* the function $y(x)$ on an evenly spaced set of points $x_0, x_1, \dots, x_n$, yielding $y_0, y_1, \dots, y_n$. Using the discretization, we can approximate the derivative by

We can write a derivative function in Python via

In [208]:

```
def nderiv(y,x):
"Finite difference derivative of the function f"
n = len(y)
d = zeros(n,'d') # assume double
# Use centered differences for the interior points, one-sided differences for the ends
for i in range(1,n-1):
d[i] = (y[i+1]-y[i-1])/(x[i+1]-x[i-1])
d[0] = (y[1]-y[0])/(x[1]-x[0])
d[n-1] = (y[n-1]-y[n-2])/(x[n-1]-x[n-2])
return d
```

Let's see whether this works for our sin example from above:

In [209]:

```
x = linspace(0,2*pi)
dsin = nderiv(sin(x),x)
plot(x,dsin,label='numerical')
plot(x,cos(x),label='analytical')
title("Comparison of numerical and analytical derivatives of sin(x)")
legend()
```

Out[209]:

<matplotlib.legend.Legend at 0x110b39c10>

Pretty close!

Now that we've convinced ourselves that finite differences aren't a terrible approximation, let's see if we can use this to solve the one-dimensional harmonic oscillator.

We want to solve the time-independent Schrodinger equation

$$ -\frac{\hbar^2}{2m}\frac{\partial^2\psi(x)}{\partial x^2} + V(x)\psi(x) = E\psi(x)$$for $\psi(x)$ when $V(x)=\frac{1}{2}m\omega^2x^2$ is the harmonic oscillator potential. We're going to use the standard trick to transform the differential equation into a matrix equation by multiplying both sides by $\psi^*(x)$ and integrating over $x$. This yields

$$ -\frac{\hbar}{2m}\int\psi(x)\frac{\partial^2}{\partial x^2}\psi(x)dx + \int\psi(x)V(x)\psi(x)dx = E$$We will again use the finite difference approximation. The finite difference formula for the second derivative is

$$ y'' = \frac{y_{i+1}-2y_i+y_{i-1}}{x_{i+1}-x_{i-1}} $$We can think of the first term in the Schrodinger equation as the overlap of the wave function $\psi(x)$ with the second derivative of the wave function $\frac{\partial^2}{\partial x^2}\psi(x)$. Given the above expression for the second derivative, we can see if we take the overlap of the states $y_1,\dots,y_n$ with the second derivative, we will only have three points where the overlap is nonzero, at $y_{i-1}$, $y_i$, and $y_{i+1}$. In matrix form, this leads to the tridiagonal Laplacian matrix, which has -2's along the diagonals, and 1's along the diagonals above and below the main diagonal.

The second term turns leads to a diagonal matrix with $V(x_i)$ on the diagonal elements. Putting all of these pieces together, we get:

In [114]:

```
def Laplacian(x):
h = x[1]-x[0] # assume uniformly spaced points
n = len(x)
M = -2*identity(n,'d')
for i in range(1,n):
M[i,i-1] = M[i-1,i] = 1
return M/h**2
```

In [115]:

```
x = linspace(-3,3)
m = 1.0
ohm = 1.0
T = (-0.5/m)*Laplacian(x)
V = 0.5*(ohm**2)*(x**2)
H = T + diag(V)
E,U = eigh(H)
h = x[1]-x[0]
# Plot the Harmonic potential
plot(x,V,color='k')
for i in range(4):
# For each of the first few solutions, plot the energy level:
axhline(y=E[i],color='k',ls=":")
# as well as the eigenfunction, displaced by the energy level so they don't
# all pile up on each other:
plot(x,-U[:,i]/sqrt(h)+E[i])
title("Eigenfunctions of the Quantum Harmonic Oscillator")
xlabel("Displacement (bohr)")
ylabel("Energy (hartree)")
```

Out[115]:

<matplotlib.text.Text at 0x10d7b9650>

We've made a couple of hacks here to get the orbitals the way we want them. First, I inserted a -1 factor before the wave functions, to fix the phase of the lowest state. The phase (sign) of a quantum wave function doesn't hold any information, only the square of the wave function does, so this doesn't really change anything.

But the eigenfunctions as we generate them aren't properly normalized. The reason is that finite difference isn't a real basis in the quantum mechanical sense. It's a basis of Dirac δ functions at each point; we interpret the space betwen the points as being "filled" by the wave function, but the finite difference basis only has the solution being at the points themselves. We can fix this by dividing the eigenfunctions of our finite difference Hamiltonian by the square root of the spacing, and this gives properly normalized functions.

The solutions to the Harmonic Oscillator are supposed to be Hermite polynomials. The Wikipedia page has the HO states given by

$$\psi_n(x) = \frac{1}{\sqrt{2^n n!}} \left(\frac{m\omega}{\pi\hbar}\right)^{1/4} \exp\left(-\frac{m\omega x^2}{2\hbar}\right) H_n\left(\sqrt{\frac{m\omega}{\hbar}}x\right)$$Let's see whether they look like those. There are some special functions in the Numpy library, and some more in Scipy. Hermite Polynomials are in Numpy:

In [116]:

```
from numpy.polynomial.hermite import Hermite
def ho_evec(x,n,m,ohm):
vec = [0]*9
vec[n] = 1
Hn = Hermite(vec)
return (1/sqrt(2**n*factorial(n)))*pow(m*ohm/pi,0.25)*exp(-0.5*m*ohm*x**2)*Hn(x*sqrt(m*ohm))
```

Let's compare the first function to our solution.

In [117]:

```
plot(x,ho_evec(x,0,1,1),label="Analytic")
plot(x,-U[:,0]/sqrt(h),label="Numeric")
xlabel('x (bohr)')
ylabel(r'$\psi(x)$')
title("Comparison of numeric and analytic solutions to the Harmonic Oscillator")
legend()
```

Out[117]:

<matplotlib.legend.Legend at 0x10da6b950>

The agreement is almost exact.

We can use the **subplot** command to put multiple comparisons in different panes on a single plot:

In [118]:

```
phase_correction = [-1,1,1,-1,-1,1]
for i in range(6):
subplot(2,3,i+1)
plot(x,ho_evec(x,i,1,1),label="Analytic")
plot(x,phase_correction[i]*U[:,i]/sqrt(h),label="Numeric")
```