Python bevezető fizika tanárszakosoknak

Csabai István, ELTE Komplex Rendszerek Fizikája Tanszék

Based on: (http://nbviewer.ipython.org/gist/rpmuller/5920182)

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

Miért Python?

A Python egy dinamikusan fejlődő kutatók és nem kutatók körében egyre népszerűbb programozási nyelv. Népszerűsége többek közt annak köszönhető, hogy nagyon összetett feladatokat is könnyen meg lehet vele oldani, és viszonylag gyorsan bele lehet tanulni, egyszerűen telepíthető, számos nagyon kényelmes keretrendszer áll rendelkezésre. Szokás mondani, hogy ez egy olyan nyelv, amit pár nap alatt megtanulhatsz, és egész életed során használhatod.

A Python Tutorial (angolul) egy kiváló oldal, ahol általában a Python tanulását el lehet kezdeni, rengeteg példa található ott. Ez a jegyzet a nemrég megjelent IPython Project keretrendszerét használja fel, amely lehetővé teszi egy Notebook, azaz jegyzőkönyv formájában a szövegesen, képletekkel vagy ábrák segítségével adott információk és a Python kód integrációját.

További hasznos linkek (angolul):

Mit kell telepíteni

A gyakorlaton leginkább egy felhő (cloud) szolgáltatást a SageMathCloud-ot használjuk. Ennek hatalmas előnye, hogy semmit nem kell gépünkre telepíteni, regisztráció után, operációs rendszertől függetlenül, web-böngészőből használható. Ráadásul nem csak a Python programozást lehet rajta keresztül kipróbálni, hanem a Sage szimbolikus matematikai keretrendszert, a LaTeX dokument-szerkesztőt vagy akár egy virtuális terminálon a Linux rendszert is.

Álljon azért itt röviden (egyelőre angolul) egy tájékoztató, hogy hogyan lehet saját gépre telepíteni a rendszert:

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.

További hasznos Python bevezetők (angolul)

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.

Az iPython Notebook használata

A Notebook-ok cellákból állnak. Az In[] cellákba gépelhetünk szöveget (menü:"Markdown"), vagy utasításokat (menü:"Code"), míg az Out[] cellákban jelenik meg a kimenet. Az utasítások végrehajtását a meüben lévő lejátszás gombbal, vagy a Shift+ENTER (utána ugrás a következő cellára) vagy a Ctrl+ENTER (marad a cellában) leütésével indíthatjuk. A régebbi cellák átírhatóak, újra futtathatóak. A kimenet nem csak számokból állhat, hanem grafikonok és ábrák is lehetnek.

A Python mint számológép

A Python úgy nevezett interpreteres nyelv, azaz az utasítások a több más nyelvre (pl. C, C++, Fortran, Java) jellemző továbbí fordítási lépés nélkül azonnal értelmezésre kerül, Ennek megfelelően, legegyszerűbb esetben használhatjuk számológépnek is:

In [1]:
2+2
Out[1]:
4
In [2]:
(50-5*6)/4
Out[2]:
5

Figyelni kell azonban, mert a Python máshogy kezeli az egész számokat (integer v. int) és a törteket reprezentáló úgynevezett lebegőpontos (float v. double) számokat:

In [3]:
7/3
Out[3]:
2

A C és Fortran nyelvekhet hasonlóan az egész-osztást az egész számokon végrehajtva kerekített egész számot kapunk végeredményként. Ha azt szeretnénk, hogy a Python valós (azaz lebegőpontos) számnak tekintse a beadott értéket, akkor azt a tizedes pont (!! Figyelem: itt is, mint ahogy szinte minden számítógépes környezetben, az angolszász jelölést használják, azaz nem tizedes vesszőt, hanem tizedes pontot!!) kitételével jelezhetjük.

In [4]:
7/3.
Out[4]:
2.3333333333333335

Másik lehetőség a szám típusának explicit, azaz direkt megadása:

In [2]:
7/float(3)
Out[2]:
2.3333333333333335

A Python nyelv alapkészletében csak nagyon egyszerű műveletek állnak rendelkezésre, így például már olyan egyszerű műveleteket, mint a gyökvonás se tartalmazza. Ha nem ismert műveletet írunk be, hibaüzenetet kapunk.

In [3]:
sqrt(81)
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-3-e1bf934cd6c7> in <module>()
----> 1 sqrt(81)

NameError: name 'sqrt' is not defined

Számtalan kiegészítő program-könyvtár, azaz library hívható be az import parancs segítségével. Pl. a math könyvtár sqrt függvénye így hívható be:

In [5]:
 from math import sqrt
In [6]:
sqrt(81)
Out[6]:
9.0

Alternatív megoldás, hogy nem csak egy függvényt, hanem a teljes könyvtárat beolvastatjuk:

In [2]:
import math
math.sqrt(81)
Out[2]:
9.0

Változók használata

Mint ahogy a matematikai egyenletekben az ismeretleneket, a programnyelvekben is elnevezhetjük az úgy nevezett változókat. Az egyenlőség (=) jel segítségével adhatunk nekik értéket:

In [7]:
x=3

A változóknak adhatunk hosszabb, több betűből álló nevet is. Kerüljük az ékezetes és speciális karakterek használatát!

In [3]:
width = 20
length = 30
area = length*width
area
Out[3]:
600

Ha olyan változó értékét szeretnénk használni, amely még nem lett definiálva, hibát kapunk:

In [9]:
volume
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-9-0c7fc58f9268> in <module>()
----> 1 volume

NameError: name 'volume' is not defined

Más származtatott változókat is használhatunk az új változók defininíciójához:

In [4]:
depth = 10
volume = area*depth
volume
Out[4]:
6000

Szinte bármilyen változónevet adhatunk, de vannak szabályok. Ne használjunk ékezetes karaktereket! A változónak betűvel vagy "_" jellel kell kezdődnie. A ("_")-n kívül más speciális karaktert ne használjunk!. Bizonyos szavak már uatsítások céljára foglaltak:

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

Ha ilyeneket akarunk definiálni, szintaktikai hibát kapunk:

In []:
return = 0

A kis és NAGY betűk, illetve KoMbináciÓik külön kezelődnek:

In [8]:
Alma=11
alma=27
AlMa=9
ALma=3
alMA=15
In [9]:
Alma
Out[9]:
11
In [10]:
alma
Out[10]:
27
In [12]:
ALma+alMA
Out[12]:
18

String-ek

A String-ek karakterek (betűk számok, írásjelek, illetve egyéb egzotikus jelek mint pl @ & %) egymás után fűzött láncai, melyeket szimpla idézőjelekkel

In [10]:
'Hello, World!'
Out[10]:
'Hello, World!'

vagy dupla idézőjelekkel definiálhatunk

In [11]:
"Hello, World!"
Out[11]:
'Hello, World!'

Mivel mind a két fajta idézőjel is karakter ezért ha ügyelnünk kell olyan stringek definiálásánál amely valamelyiket tartalmazhatja.

In [12]:
"He's a Rebel"
Out[12]:
"He's a Rebel"
In [13]:
'She asked, "How are you today?"'
Out[13]:
'She asked, "How are you today?"'

Mint ahogy a már megismert adatípusoknál is megszokhattuk (int-ek és float-ok) egy változónak adhatunk string títusú értéket.

In [14]:
greeting = "Hello, World!"

A print parancs segítségével megjeleníthetünk string-eket:

In [15]:
print greeting
Hello, World!

Ezen kívül más típusú változók megjelenítésére is használhatjuk:

In [16]:
print "The area is ",area
The area is  600

A fenti példában az area nevű változó értéke 600, melyet a print parancs automatikusan string tipúsra konvertál megjelenítés előtt.

A "+" operátor segítségével két string típusú változót egymás után fűzhetünk:

In [17]:
statement = "Hello," + "World!"
print statement
Hello,World!

Figyelem a szóköz is karakter, ha a kiírandó stringeket nem szorosan egymás akarjuk kiírni hanem szóközzel akarjuk elválasztani őket akkor a megelelő helyre szóközöket kell tenni!

In [1]:
statement = "Hello," + "World!"
print statement
statement = "Hello, " + "World!"
print statement
statement = "Hello,"+ " " + "World!"
print statement
Hello,World!
Hello, World!
Hello, World!

A fenti példák közzül az utolsóban 1 sorban 3 stringet fűztünk össze egymás után. Hasonló módon több string-et is egymás után tudunk fűzni:

In [19]:
print "This " + "is " + "a " + "longer " + "statement."
This is a longer statement.

Ha sok stringet akarunk össze fűzni akkor arra vannak elegánsabb megoldások is de a fenti megoldás néhány string össze fűzésére elegendő.

List

Programozás közben sokszor előfordul hogy hasonló típusú objektumokat, szeretnénk együtt kezelni. A python ezt a problémát a list típusú változók segítségével kezeli. Például szeretnénk a hét napjait (mondjuk egy naptár alkalmazás keretein belül) együtt kezelni, Ekkor a [,] szögletes zárójelek segítségével definálhatjuk a következő változót:

In [2]:
days_of_the_week = ["Sunday","Monday","Tuesday","Wednesday","Thursday","Friday","Saturday"]

azaz a days_of_the_week egy olyan list melynek string típusú elemei vannak és minden elem megfelel egy napnak.

A megfelelő elemhez való hozzáférést az elem index-én keresztül történik, szintén a [,] szögletes zárójelek segítségével:

In [21]:
days_of_the_week[2]
Out[21]:
'Tuesday'

A python list-ek, első eleme a 0-as index-hez tartozik. (Ez hasonló a C nyelvbeni konvencióhoz de például ellenkezik a Fortran nyelv konvenciójával). Tehát a fenti páldában a 0-as elem a "Sunday", az 1-es a "Monday" és így tovább. Pythonban egy lista végén szereplő elemeit negatív index-ekkel érhetjük el. Például a -1-es elem a fenti listában:

In [22]:
days_of_the_week[-1]
Out[22]:
'Saturday'

Egy listához az ** .append() ** parancs segítségével fűzhetünk további elemeket:

In [23]:
languages = ["Fortran","C","C++"]
languages.append("Python")
print languages
['Fortran', 'C', 'C++', 'Python']

A range() parancs segítségével számok egyszerű sorozatait hozhatjuk létre:

In [24]:
range(10)
Out[24]:
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

Figyeljük meg hogy a fenti lista 0-val kezdődik és csak 9-ig tart, azaz 0-tól a 10-nél kisebb számok sorozata. Ha a range parancsot két paraméterrel hívjuk meg range(start,stop) akkor a lista elemei nem 0-val hanem a start értékkel kezdődnek:

In [2]:
range(2,8)
Out[2]:
[2, 3, 4, 5, 6, 7]

Egy harmadik paraméter segítségével az egymást követő tagok lépés különbségét változtathatjuk. range(start,stop,step) Az alábbi példában a harmadik paraméter miatt a listában 0-tól 20-ig csak minden második elem jelenik meg:

In [4]:
evens = range(0,20,2)
evens
Out[4]:
[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]
In [27]:
evens[3]
Out[27]:
6

Egy list-be nem csak ugyan olyan típusú objektumokat tehetünk. Például:

In [28]:
["Today",7,99.3,""]
Out[28]:
['Today', 7, 99.3, '']

Enek ellenére azért mindíg célszerű olyan objektumokat egy listába foglalni amik valamilyen logikus kapcsolatban állnak egymással. Programozás technikailag célszerűbb különböző típusú objektumokat nem list-ekbe hanem inkább tuple-k ba foglani, ezekre a későbbiekben fogunk látni konkrét példát.

Egy lista hosszát a len() parancs adja meg:

In [29]:
help(len)
Help on built-in function len in module __builtin__:

len(...)
    len(object) -> integer
    
    Return the number of items of a sequence or mapping.


In [30]:
len(evens)
Out[30]:
10

Iterációk és tömbök

A számítógépek legfontosabb tulajdonságai közt szerepel az, hogy nagyon gyorsak és "fáradhatatlanok". Olyan feladatok megoldásában a leghatékonyabbak, amikor a feladatot kevés munkával meg lehet fogalmazni ("az alkotó pihen") de végrehajtása nagyon sok ismétlést, iterációt igényel ("a gép forog"). Az iteráció (angol: iterate) azt jelenti, hogy például egy lista elemein egyesével végigmegy a program, és műveleteket végez el rajtuk. A Python-ban, az egyik erre használható utasítás a for parancs (magyarul kb. a ...-ra, azaz pl. a hét minden napjára, a lista minden elemére):

In [3]:
for day in days_of_the_week:
    print day
Sunday
Monday
Tuesday
Wednesday
Thursday
Friday
Saturday

Ez a kódrészlet a days_of_the_week listán megy végig és a meglátogatott elemet hozzárendeli a day változóhoz, amit ciklusváltozónak is neveznek. Ezekután mindent végrehajt amit a beljebb tabulált (angolul: indented) parancsblokkban írtunk (most csak egy print utasítás), amihez felhasználhatja a ciklusváltozót is. Miután vége a beljebb tabulált régiónak, kilép a ciklusból.

Szinte minden programnyelv használ hasonló ciklusokat. A C, C++, Java nyelvekben kapcsos zárójeleket {} használnak a ciklusok elkülönítésére, kb. így:

for (i=1..10) {print i}

A FORTRAN nyelvben az END szócska kiírása jelzi a ciklus végét. A pythonban a kerrőspont (":"), majd vagy tabulátor-ral, vagy szóközökkel beljebb írt sorok szolgálnak erre. FIGYELEM: a kettőt egy programban ne keverjük, vagy TAB vagy szóközök! Mivel a szövegszerkesztőben a képernyőn pár szóköz ugyanúgy néz ki mint egy TAB, ez a megoldás gyakran megtréfálhatja a kezdő programozót. Különösen gond lehet, ha más programokból másolunk át részleteket.

Annak semmi jelentősége nincs, hogy a példában a day nevet adtunk az iterációban szereplő ciklusváltozónak. A program semmit se tud az emberi időszámításról, például, hogy a hétben napok vannak és nem kiscicák:

In [3]:
for macska in days_of_the_week:
    print macska
Sunday
Monday
Tuesday
Wednesday
Thursday
Friday
Saturday

A ciklus utasításblokkja állhat több utasításból is:

In [32]:
for day in days_of_the_week:
    statement = "Today is " + day
    print statement
Today is Sunday
Today is Monday
Today is Tuesday
Today is Wednesday
Today is Thursday
Today is Friday
Today is Saturday

A range() parancs remekül használható ha a for ciklusban adott számú műveletet szeretnénk elvégezni:

In [33]:
for i in range(20):
    print "The square of ",i," is ",i*i
The square of  0  is  0
The square of  1  is  1
The square of  2  is  4
The square of  3  is  9
The square of  4  is  16
The square of  5  is  25
The square of  6  is  36
The square of  7  is  49
The square of  8  is  64
The square of  9  is  81
The square of  10  is  100
The square of  11  is  121
The square of  12  is  144
The square of  13  is  169
The square of  14  is  196
The square of  15  is  225
The square of  16  is  256
The square of  17  is  289
The square of  18  is  324
The square of  19  is  361

Szeletek, tartományok kijelölése listákból

A karakterláncok (string) és listák közös tulajdonsága, hogy mindkettőt lehet elemek sorozatának tekinteni. Ugyanúgy ahogy a listák elemein, a string-ek elemein is végig lehet iterálni:

In [4]:
for letter in "Sunday":
    print letter
S
u
n
d
a
y

(Próbáld ki, hogy tudja-e a gép, hogy a szavak betűkből és nem kisnyulakból állnak össze :-)

Talán ennél is hasznosabb a szeletelés (slicing) művelet, amivel adott tartományok vághatóak ki. A legegyszerűbb szelet egyetlen elem kiválasztása az index-e alapján:

In [35]:
days_of_the_week[0]
Out[35]:
'Sunday'

Ugye emlékszünk, az indexelés 0-val (és nem 1-gyel kezdődik)! Az első két elem így választható ki:

In [36]:
days_of_the_week[0:2]
Out[36]:
['Sunday', 'Monday']

Figyelem: a tartomány végét jelző indexhez tartozó elem nem kerül a szeletbe. Tehát itt a lista 2-es indexű eleme, azaz a lista 3. eleme ('Tuesday') nem része. Megj.: hasonló logika alapján viselkedett a range() utasítás is, lásd fenn.

Ha 0-tól kezdődik a szelet egyszerűsítve így is írhatjuk:

In [37]:
days_of_the_week[:2]
Out[37]:
['Sunday', 'Monday']

Hasonlóan a lista végének indexét is elhagyhatjuk:

In [7]:
days_of_the_week[4:]
Out[7]:
['Thursday', 'Friday', 'Saturday']

A lista végéről negatív indexekkel kérhetünk le értékeket:

In [6]:
days_of_the_week[-2:]
Out[6]:
['Friday', 'Saturday']

A kiválasztott szeletek újabb változókba tárolhatóak el:

In [39]:
workdays = days_of_the_week[1:6]
print workdays
['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday']

Ugyanez természetesen stringekkel is megtehető:

In [40]:
day = "Sunday"
abbreviation = day[:3]
print abbreviation
Sun

A range() utasításhoz hasonlóan, egy újabb ':' után a tartomány indexelésének lépésköze is megadható

In [8]:
numbers = range(0,40)
evens = numbers[2:40:2]
evens
Out[8]:
[2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38]

Ilyenkor is elhagyható a tartomány elejét és/vagy végét jelző index, ha az első elemtől vagy az utolsó elemig történik a kiválsaztás:

In [11]:
numbers = range(0,40)
evens = numbers[2::2]
evens
Out[11]:
[2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38]

Logikai Bool-változók és feltételes végrehajtás

Ahhoz hogy a program reagálni tudjon bemenetekre, vagy egy már kiszámolt részeredménytől függhessen további működése, szükség van feltételek vizsgálatára. Gondoljunk például a másodfokú egyenlet megoldóképletére! Attól függően, hogy a diszkrimináns pozitív vagy negatív léteznek ill. nem léteznek valós gyökök, így értelemszerűen egy ilyen programnak "döntenie" kell. A döntés, egy feltétel vizsgálatán alapul (pl. poz. vagy neg. diszkrimináns), ami lehet igaz, vagy hamis, angolul True ill. False. Az igaz és hamis értéket tárolni képes változó típusa: boolean (Bool-változó). (Emlékezzünk vissza az int ill. float típusokra).

A feltétek vizsgálatára az angol 'ha' szócskából alkotott if utasítás szolgál:

In [42]:
if day == "Sunday":
    print "Sleep in"
else:
    print "Go to work"
Sleep in

(Kérdés: miért "Sleep in" (alhatunk) választ kaptunk? Mi a "day" változó értéke? Hol állítottuk be?)

Nézzük végig részletesen a fenti kódot! Az első lépés a feltétel vizsgálata, ezt önmagában is elvégezhetjük:

In [43]:
day == "Sunday"
Out[43]:
True

A for ciklushoz hasonlóan a feltételes utasításoknál is a ':' és a tabulálás jelzi a feltételes kód blokkot. Ha a feltétel igaz ('True'), akkor a blokk végrehajtódik, ellenkező esetben, nem.

Az else (='egyébként') utasítás után lévő blokk a hamis ('False') feltétel esetén hajtódik végre. Ez a rész nem kötelező, ha hiányzik, hamis feltétel esetén egyszerűen továbblép a program.

Más típusú változókon is végezhető vizsgálat:

In [44]:
1 == 2
Out[44]:
False

Itt jegyezzük meg, hogy a feltétek vizsgálatnál dupla egyenlőség jel szerepel ==, nem pedig az értékadáshoz használt szimpla:

In [13]:
1 = 2
  File "<ipython-input-13-b922a9ca1612>", line 1
    1 = 2
SyntaxError: can't assign to literal
In [45]:
50 == 2*25
Out[45]:
True
In [46]:
3 < 3.14159
Out[46]:
True
In [47]:
1 == 1.0
Out[47]:
True
In [48]:
1 != 0
Out[48]:
True
In [49]:
1 <= 2
Out[49]:
True
In [50]:
1 >= 1
Out[50]:
True

'ÉS' (and) ill. 'VAGY' (or) operátorokkal több logikai feltétel kombinálható:

In [18]:
2>1 and 5>4 
Out[18]:
True
In [19]:
1>2 or 2>1
Out[19]:
True

Több összehasonlítás kombinálható így is:

In [54]:
hours = 5
0 < hours < 24
Out[54]:
True

A fentiekben szerepelt egy érdekes vizsgálat: 1 == 1.0 ami ugyan a várt eredményt adta de az egynlőség két oldalán "különböző" szám áll, nevezetesen típusuk nem egyezik meg egyikük egész (integer) a másik valós (float) típusú. Értékük megegyezett, ezért helyes eredményt kaptunk. Van egy speciális operátor az is (kb. "az-e mint"), ami a típust is vizsgálja:

In [51]:
1 is 1.0
Out[51]:
False

Összetett objektumok, például listák is összehasonlíthatóak. Akár a '<' ill '>' operátorok is definiálhatóak (Nézz utána!).

In [52]:
[1,2,3] == [1,2,4]
Out[52]:
False

Ha sokfelé elágazó vizsgálatot szeretnénk végezni az ("else if") kombinációból származtatható elif használható:

In [55]:
if day == "Sunday":
    print "Sleep in"
elif day == "Saturday":
    print "Do chores"
else:
    print "Go to work"
Sleep in

Akkor válik mindez még érdekesebbé, ha az eddig tanult iterációt és feltétel vizsgálatot kombináljuk:

In [56]:
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"
Today is Sunday
   Sleep in
Today is Monday
   Go to work
Today is Tuesday
   Go to work
Today is Wednesday
   Go to work
Today is Thursday
   Go to work
Today is Friday
   Go to work
Today is Saturday
   Do chores

Egy példa program: Fibbonacci sorozat

A Fibonacci sequence sorozat első két eleme 0 és 1, majd a következő elemet mindig az előző kettő összegéből számoljuk ki: 0,1,1,2,3,5,8,13,21,34,55,89,...

Ha nagyobb n értékekre is ki akarjuk számolni a sorozatot, ez kiváló feladat lehet egy fáradhatatlan és gyors számítógépnek!

In [5]:
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]

Nézzük végig lépésről lépésre! Először n értékét, azaz a kiszámolandó sorozat hosszát állítjuk be 10-re. A sorozatot majdan tároló listát sequence-nek neveztük el, és inicializáltuk az eslő két értékkel. A "kézi munka" után következhet a gép automatikus munkája, az iteráció.

Az iterációt 2-vel kezdjük (ez ugye a 0-s indexelés miatt a 3. elem lesz, hisz az első kettőt már mi megadtuk) és n-ig, a megadott lista méretig számolunk. A "#" 'megjegyzést' (angol: comment) jelöl, ezt a gép nem 'olvassa', csak a programozónak szól emlékeztetőül, vagy más programozók megértését sgíti. Noha gyakran elmarad, a kódok rendszeres dokumentálása megjegyzésekkel, később (amikor már a kód részleteit a feledés homály borítja) nagyon kifizetődő lehet.

A ciklus törzsében az addig kiszámolt lista végére hozzátűzzük (append) az előző két tag összegét. A ciklus vége után kiíratjuk az eredményt.

Függvények

Ha más hosszúságú sorozatot szeretnénk átmásolhatjuk a fenti kódot egy új cellába és átírhatjuk az n=10-et pl. n=100-ra. Van azonban egy hatékonyabb módszer, új függvény-ként definiálhatjuk a def utasítás segítségével:

In [6]:
def fibonacci(sequence_length):
    "Return the Fibonacci sequence of length *sequence_length*" # ez csak a 'help'-hez kell
    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

Most már meghívhatjuk a fibonacci() függvényt különböző hosszakra:

In [15]:
fibonacci(2)
Out[15]:
[0, 1]
In [20]:
fibonacci(18)
Out[20]:
[0, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 377, 610, 987, 1597]

Elemezzük a fenti kódot! A már megszokott módon a kettőspont és behúzás (TAB) határozza meg a függvény definícióhoz tartozó kód blokkot. A 2. sorban idézőjelek közt szerepel a "docstring", ami a függvény működését magyarázza el röviden, és később a help paranccsal hívható elő:

In [64]:
help(fibonacci)
Help on function fibonacci in module __main__:

fibonacci(sequence_length)
    Return the Fibonacci sequence of length *sequence_length*


Ez később nagyon hasznos lehet, érdemes megszokni, hogy dokumentáljuk munkánkat!

A függvényt megpróbáltuk "bolond-biztossá" is tenni azáltal, hogy leellenőrizzük, hogy a megadott sorozat hossz érvényes pozitív szám-e.

Rekurzív függvények

A függvények saját magukat is meg tudják hívni, amit rekurzió-nak nevezünk. A rekurziót a faktoriális függvényen demonstráljuk. n faktoriálisa definíció szerint:

\[ n! = n(n-1)(n-2)\cdots 1 \]

(Persze létezik beépített faktoriális függvény is, amin egyből a help utasítást is gyakorolhatjuk):

In [24]:
from math import factorial
help(factorial)
Help on built-in function factorial in module math:

factorial(...)
    factorial(x) -> Integral
    
    Find x!. Raise a ValueError if x is negative or non-integral.


In [66]:
factorial(20)
Out[66]:
2432902008176640000

Ha magunk akarjuk megírni a függvényt ahhoz használjuk fel, hogy:

\[ n! = n(n-1)!\]

Ezek után a program:

In [25]:
def fact(n):
    "Calculate the factorial (n!) of a positive integer"
    if n <= 0:
        return 1
    return n*fact(n-1)
In [26]:
fact(20)
Out[26]:
2432902008176640000L

A rekurzió segítségével nagyon tömör, nagyon elegáns (néha emberi ésszel már nehezen átlátható) programokat lehet írni.

EDDIG --------------------------------

Two More Data Structures: Tuples and Dictionaries

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 [69]:
t = (1,2,'hi',9.0)
t
Out[69]:
(1, 2, 'hi', 9.0)

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

In [70]:
t[1]
Out[70]:
2

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

In [71]:
t.append(7)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-71-50c7062b1d5f> in <module>()
----> 1 t.append(7)

AttributeError: 'tuple' object has no attribute 'append'
In [72]:
t[1]=77
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-72-03cc8ba9c07d> in <module>()
----> 1 t[1]=77

TypeError: 'tuple' object does not support item assignment

Tuples are useful anytime you want to group different pieces of data together in an object, but don't want to create a full-fledged class (see below) for them. For example, let's say you want the Cartesian coordinates of some objects in your program. Tuples are a good way to do this:

In [73]:
('Bob',0.0,21.0)
Out[73]:
('Bob', 0.0, 21.0)

Again, it's not a necessary distinction, but one way to distinguish tuples and lists is that tuples are a collection of different things, here a name, and x and y coordinates, whereas a list is a collection of similar things, like if we wanted a list of those coordinates:

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

Tuples can be used when functions return more than one value. Say we wanted to compute the smallest x- and y-coordinates of the above list of objects. We could write:

In [75]:
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 [76]:
x,y = 1,2
y,x = x,y
x,y
Out[76]:
(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 [77]:
mylist = [1,2,9,21]

The index in a dictionary is called the 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 [78]:
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 [79]:
dict(Rick=46,Bob=86,Fred=20)
Out[79]:
{'Bob': 86, 'Fred': 20, 'Rick': 46}

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

In [80]:
len(t)
Out[80]:
4
In [81]:
len(ages)
Out[81]:
3

Plotting with Matplotlib

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 [82]:
fibs = fibonacci(10)

Next lets generate the factorials.

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

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

In [84]:
figsize(8,6)
plot(facts,label="factorial")
plot(fibs,label="Fibonacci")
xlabel("n")
legend()
Out[84]:
<matplotlib.legend.Legend at 0x10eeb2e50>

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 [85]:
semilogy(facts,label="factorial")
semilogy(fibs,label="Fibonacci")
xlabel("n")
legend()
Out[85]:
<matplotlib.legend.Legend at 0x10eedd7d0>

There are many more things you can do with Matplotlib. We'll be looking at some of them in the sections to come. In the meantime, if you want an idea of the different things you can do, look at the Matplotlib Gallery. Rob Johansson's IPython notebook Introduction to Matplotlib is also particularly good.

Conclusion of the Python Overview

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 [86]:
import this
The Zen of Python, by Tim Peters

Beautiful is better than ugly.
Explicit is better than implicit.
Simple is better than complex.
Complex is better than complicated.
Flat is better than nested.
Sparse is better than dense.
Readability counts.
Special cases aren't special enough to break the rules.
Although practicality beats purity.
Errors should never pass silently.
Unless explicitly silenced.
In the face of ambiguity, refuse the temptation to guess.
There should be one-- and preferably only one --obvious way to do it.
Although that way may not be obvious at first unless you're Dutch.
Now is better than never.
Although never is often better than *right* now.
If the implementation is hard to explain, it's a bad idea.
If the implementation is easy to explain, it may be a good idea.
Namespaces are one honking great idea -- let's do more of those!

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

II. Numpy and Scipy

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.)

Making vectors and matrices

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 [87]:
array([1,2,3,4,5,6])
Out[87]:
array([1, 2, 3, 4, 5, 6])

You can pass in a second argument to 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 [88]:
array([1,2,3,4,5,6],'d')
Out[88]:
array([ 1.,  2.,  3.,  4.,  5.,  6.])
In [89]:
array([1,2,3,4,5,6],'D')
Out[89]:
array([ 1.+0.j,  2.+0.j,  3.+0.j,  4.+0.j,  5.+0.j,  6.+0.j])
In [90]:
array([1,2,3,4,5,6],'i')
Out[90]:
array([1, 2, 3, 4, 5, 6], dtype=int32)

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

In [91]:
array([[0,1],[1,0]],'d')
Out[91]:
array([[ 0.,  1.],
       [ 1.,  0.]])

You can also form empty (zero) matrices of arbitrary shape (including vectors, which Numpy treats as vectors with one row), using the zeros command:

In [92]:
zeros((3,3),'d')
Out[92]:
array([[ 0.,  0.,  0.],
       [ 0.,  0.,  0.],
       [ 0.,  0.,  0.]])

The first argument is a tuple containing the shape of the matrix, and the second is the data type argument, which follows the same conventions as in the array command. Thus, you can make row vectors:

In [93]:
zeros(3,'d')
Out[93]:
array([ 0.,  0.,  0.])
In [94]:
zeros((1,3),'d')
Out[94]:
array([[ 0.,  0.,  0.]])

or column vectors:

In [95]:
zeros((3,1),'d')
Out[95]:
array([[ 0.],
       [ 0.],
       [ 0.]])

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

In [96]:
identity(4,'d')
Out[96]:
array([[ 1.,  0.,  0.,  0.],
       [ 0.,  1.,  0.,  0.],
       [ 0.,  0.,  1.,  0.],
       [ 0.,  0.,  0.,  1.]])

as well as a ones command.

Linspace, matrix functions, and plotting

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

In [97]:
linspace(0,1)
Out[97]:
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.        ])

If you provide a third argument, it takes that as the number of points in the space. If you don't provide the argument, it gives a length 50 linear space.

In [98]:
linspace(0,1,11)
Out[98]:
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 [99]:
x = linspace(0,2*pi)
sin(x)
Out[99]:
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 [100]:
plot(x,sin(x))
Out[100]:
[<matplotlib.lines.Line2D at 0x10f07a150>]

Matrix operations

Matrix objects act sensibly when multiplied by scalars:

In [101]:
0.125*identity(3,'d')
Out[101]:
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 [102]:
identity(2,'d') + array([[1,1],[1,2]])
Out[102]:
array([[ 2.,  1.],
       [ 1.,  3.]])

Something that confuses Matlab users is that the times (*) operator give element-wise multiplication rather than matrix multiplication:

In [103]:
identity(2)*ones((2,2))
Out[103]:
array([[ 1.,  0.],
       [ 0.,  1.]])

To get matrix multiplication, you need the dot command:

In [104]:
dot(identity(2),ones((2,2)))
Out[104]:
array([[ 1.,  1.],
       [ 1.,  1.]])

dot can also do dot products (duh!):

In [105]:
v = array([3,4],'d')
sqrt(dot(v,v))
Out[105]:
5.0

as well as matrix-vector products.

There are 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 [106]:
m = array([[1,2],[3,4]])
m.T
Out[106]:
array([[1, 3],
       [2, 4]])

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

In [107]:
diag([1,2,3,4,5])
Out[107]:
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.

Matrix Solvers

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

In [108]:
A = array([[1,1,1],[0,2,5],[2,5,-1]])
b = array([6,-4,27])
solve(A,b)
Out[108]:
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 [109]:
A = array([[13,-4],[-4,7]],'d')
eigvalsh(A)
Out[109]:
array([  5.,  15.])
In [110]:
eigh(A)
Out[110]:
(array([  5.,  15.]),
 array([[-0.4472136 , -0.89442719],
       [-0.89442719,  0.4472136 ]]))

Example: Finite Differences

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

\[ y_i' \approx \frac{y_{i+1}-y_{i-1}}{x_{i+1}-x_{i-1}} \]

We can write a derivative function in Python via

In [111]:
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])/(x[i+1]-x[i])
    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 [112]:
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[112]:
<matplotlib.legend.Legend at 0x10ef87dd0>

Pretty close!

One-Dimensional Harmonic Oscillator using Finite Difference

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 [113]:
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 [114]:
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[114]:
<matplotlib.text.Text at 0x10efa4110>

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.

Special 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 [115]:
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 [116]:
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[116]:
<matplotlib.legend.Legend at 0x10f2e8510>

The agreement is almost exact.

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

In [117]:
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")

Other than phase errors (which I've corrected with a little hack: can you find it?), the agreement is pretty good, although it gets worse the higher in energy we get, in part because we used only 50 points.

The Scipy module has many more special functions:

In [118]:
from scipy.special import airy,jn,eval_chebyt,eval_legendre
subplot(2,2,1)
x = linspace(-1,1)
Ai,Aip,Bi,Bip = airy(x)
plot(x,Ai)
plot(x,Aip)
plot(x,Bi)
plot(x,Bip)
title("Airy functions")

subplot(2,2,2)
x = linspace(0,10)
for i in range(4):
    plot(x,jn(i,x))
title("Bessel functions")

subplot(2,2,3)
x = linspace(-1,1)
for i in range(6):
    plot(x,eval_chebyt(i,x))
title("Chebyshev polynomials of the first kind")

subplot(2,2,4)
x = linspace(-1,1)
for i in range(6):
    plot(x,eval_legendre(i,x))
title("Legendre polynomials")
Out[118]:
<matplotlib.text.Text at 0x10fcf6890>

As well as Jacobi, Laguerre, Hermite polynomials, Hypergeometric functions, and many others. There's a full listing at the Scipy Special Functions Page.

Least squares fitting

Very often we deal with some data that we want to fit to some sort of expected behavior. Say we have the following:

In [119]:
raw_data = """\
3.1905781584582433,0.028208609537968457
4.346895074946466,0.007160804747670053
5.374732334047101,0.0046962988461934805
8.201284796573875,0.0004614473299618756
10.899357601713055,0.00005038370219939726
16.295503211991434,4.377451812785309e-7
21.82012847965739,3.0799922117601088e-9
32.48394004282656,1.524776208284536e-13
43.53319057815846,5.5012073588707224e-18"""

There's a section below on parsing CSV data. We'll steal the parser from that. For an explanation, skip ahead to that section. Otherwise, just assume that this is a way to parse that text into a numpy array that we can plot and do other analyses with.

In [120]:
data = []
for line in raw_data.splitlines():
    words = line.split(',')
    data.append(map(float,words))
data = array(data)
In [121]:
title("Raw Data")
xlabel("Distance")
plot(data[:,0],data[:,1],'bo')
Out[121]:
[<matplotlib.lines.Line2D at 0x10f572390>]

Since we expect the data to have an exponential decay, we can plot it using a semi-log plot.

In [122]:
title("Raw Data")
xlabel("Distance")
semilogy(data[:,0],data[:,1],'bo')
Out[122]:
[<matplotlib.lines.Line2D at 0x10f2da510>]

For a pure exponential decay like this, we can fit the log of the data to a straight line. The above plot suggests this is a good approximation. Given a function \[ y = Ae^{-ax} \] \[ \log(y) = \log(A) - ax\] Thus, if we fit the log of the data versus x, we should get a straight line with slope \(a\), and an intercept that gives the constant \(A\).

There's a numpy function called polyfit that will fit data to a polynomial form. We'll use this to fit to a straight line (a polynomial of order 1)

In [123]:
params = polyfit(data[:,0],log(data[:,1]),1)
a = params[0]
A = exp(params[1])

Let's see whether this curve fits the data.

In [124]:
x = linspace(1,45)
title("Raw Data")
xlabel("Distance")
semilogy(data[:,0],data[:,1],'bo')
semilogy(x,A*exp(a*x),'b-')
Out[124]:
[<matplotlib.lines.Line2D at 0x10f520890>]

If we have more complicated functions, we may not be able to get away with fitting to a simple polynomial. Consider the following data:

In [125]:
gauss_data = """\
-0.9902286902286903,1.4065274110372852e-19
-0.7566104566104566,2.2504438576596563e-18
-0.5117810117810118,1.9459459459459454
-0.31887271887271884,10.621621621621626
-0.250997150997151,15.891891891891893
-0.1463309463309464,23.756756756756754
-0.07267267267267263,28.135135135135133
-0.04426734426734419,29.02702702702703
-0.0015939015939017698,29.675675675675677
0.04689304689304685,29.10810810810811
0.0840994840994842,27.324324324324326
0.1700546700546699,22.216216216216214
0.370878570878571,7.540540540540545
0.5338338338338338,1.621621621621618
0.722014322014322,0.08108108108108068
0.9926849926849926,-0.08108108108108646"""

data = []
for line in gauss_data.splitlines():
    words = line.split(',')
    data.append(map(float,words))
data = array(data)

plot(data[:,0],data[:,1],'bo')
Out[125]:
[<matplotlib.lines.Line2D at 0x10f3f1b10>]

This data looks more Gaussian than exponential. If we wanted to, we could use polyfit for this as well, but let's use the curve_fit function from Scipy, which can fit to arbitrary functions. You can learn more using help(curve_fit).

First define a general Gaussian function to fit to.

In [126]:
def gauss(x,A,a): return A*exp(a*x**2)

Now fit to it using curve_fit:

In [127]:
from scipy.optimize import curve_fit

params,conv = curve_fit(gauss,data[:,0],data[:,1])
x = linspace(-1,1)
plot(data[:,0],data[:,1],'bo')
A,a = params
plot(x,gauss(x,A,a),'b-')
Out[127]:
[<matplotlib.lines.Line2D at 0x110a84410>]

The curve_fit routine we just used is built on top of a very good general minimization capability in Scipy. You can learn more at the scipy documentation pages.

Monte Carlo, random numbers, and computing \(\pi\)

Many methods in scientific computing rely on Monte Carlo integration, where a sequence of (pseudo) random numbers are used to approximate the integral of a function. Python has good random number generators in the standard library. The random() function gives pseudorandom numbers uniformly distributed between 0 and 1:

In [128]:
from random import random
rands = []
for i in range(100):
    rands.append(random())
plot(rands)
Out[128]:
[<matplotlib.lines.Line2D at 0x110ab3610>]

random() uses the Mersenne Twister algorithm, which is a highly regarded pseudorandom number generator. There are also functions to generate random integers, to randomly shuffle a list, and functions to pick random numbers from a particular distribution, like the normal distribution:

In [129]:
from random import gauss
grands = []
for i in range(100):
    grands.append(gauss(0,1))
plot(grands)
Out[129]:
[<matplotlib.lines.Line2D at 0x110b89510>]

It is generally more efficient to generate a list of random numbers all at once, particularly if you're drawing from a non-uniform distribution. Numpy has functions to generate vectors and matrices of particular types of random distributions.

In [130]:
plot(rand(100))
Out[130]:
[<matplotlib.lines.Line2D at 0x110d09fd0>]

One of the first programs I ever wrote was a program to compute \(\pi\) by taking random numbers as x and y coordinates, and counting how many of them were in the unit circle. For example:

In [131]:
npts = 5000
xs = 2*rand(npts)-1
ys = 2*rand(npts)-1
r = xs**2+ys**2
ninside = (r<1).sum()
figsize(6,6) # make the figure square
title("Approximation to pi = %f" % (4*ninside/float(npts)))
plot(xs[r<1],ys[r<1],'b.')
plot(xs[r>1],ys[r>1],'r.')
figsize(8,6) # change the figsize back to 4x3 for the rest of the notebook

The idea behind the program is that the ratio of the area of the unit circle to the square that inscribes it is \(\pi/4\), so by counting the fraction of the random points in the square that are inside the circle, we get increasingly good estimates to \(\pi\).

The above code uses some higher level Numpy tricks to compute the radius of each point in a single line, to count how many radii are below one in a single line, and to filter the x,y points based on their radii. To be honest, I rarely write code like this: I find some of these Numpy tricks a little too cute to remember them, and I'm more likely to use a list comprehension (see below) to filter the points I want, since I can remember that.

As methods of computing \(\pi\) go, this is among the worst. A much better method is to use Leibniz's expansion of arctan(1):

\[\frac{\pi}{4} = \sum_k \frac{(-1)^k}{2*k+1}\]

In [132]:
n = 100
total = 0
for k in range(n):
    total += pow(-1,k)/(2*k+1.0)
print 4*total
3.13159290356

If you're interested a great method, check out Ramanujan's method. This converges so fast you really need arbitrary precision math to display enough decimal places. You can do this with the Python decimal module, if you're interested.

Numerical Integration

Integration can be hard, and sometimes it's easier to work out a definite integral using an approximation. For example, suppose we wanted to figure out the integral:

\[\int_0^\infty\exp(-x)dx=1\]

In [133]:
from numpy import sqrt
def f(x): return exp(-x)
x = linspace(0,10)
plot(x,exp(-x))
Out[133]:
[<matplotlib.lines.Line2D at 0x110d2a250>]

Scipy has a numerical integration routine quad (since sometimes numerical integration is called quadrature), that we can use for this:

In [134]:
from scipy.integrate import quad
quad(f,0,inf)
Out[134]:
(1.0000000000000002, 5.842606742906004e-11)

There are also 2d and 3d numerical integrators in Scipy. See the docs for more information.

Fast Fourier Transform and Signal Processing

Very often we want to use FFT techniques to help obtain the signal from noisy data. Scipy has several different options for this.

In [135]:
from scipy.fftpack import fft,fftfreq

npts = 4000
nplot = npts/10
t = linspace(0,120,npts)
def acc(t): return 10*sin(2*pi*2.0*t) + 5*sin(2*pi*8.0*t) + 2*rand(npts)

signal = acc(t)

FFT = abs(fft(signal))
freqs = fftfreq(npts, t[1]-t[0])

subplot(211)
plot(t[:nplot], signal[:nplot])
subplot(212)
plot(freqs,20*log10(FFT),',')
show()