Notes and Examples for Python

For Linux the packages numpy, scipy, Gnuplot, and matplotlib should be in usr/lib/Python2.x/site-packages.

1-d Arrays, Matrices, Numerical Integration, Numerical Solution of ODEs, Curve Fitting, Fit to line, Reading and Writing Array files, Finding zeros of functions, Graphing with Gnuplot, Fast Fourier Transform, Waveforms: Square, Sawtooth, Time Delay, Noise, Create Postscript Graph, Simple Plots with matplotlib, Plot Functions and Data, Interactive Plots with matplotlib, Plotting with log or linear axes, Subplots, 2 Y axes, Inset Graph

Instructions for installing a python-based system for scientific/mathematic use, here

Good references: http://www.scipy.org/Cookbook, for example, to learn plotting: http://www.scipy.org/Cookbook/Matplotlib, for gnuplot, see demo.py and test.py in the gnuplot-py directory

To learn about packages, subpackages, and routines, run python from a terminal, then follow these example commands to learn about the scipy package:

To make programs executable for Linux 1st line of program is:

#! /usr/bin/env python {and change permissions to include execute}

May be useful to run as:

python filename.py --verbose-helpful

Programs should load packages as follows:

from scipy import *
from numpy import *   # This not necessary since scipy loads numpy
from numpy.random import *   # If need random number generator
import Gnuplot              # If use gnuplot for graphs
from matplotlib import *    # If use matplotlib for graphs
from pylab import * 


Note, if using Gnuplot.py, then gnuplot must also be installed. http://www.gnuplot.info

To enter values interactively use raw_input(text). (Doesn't work on emacs!! Use Eric)

x = float(raw_input('Enter x value: '))   #converts text input into real.

1-d Arrays

x = arange(0,10)   # To make integer array vector of x values. 
x = linspace(0,10,101)    # To make more general array. 
y = x**2
x=zeros(100,float)    # creates x array filled with 0

Example: Matrices

from scipy import *
from scipy.linalg import *
from numpy.random import randn

#Example of using 2-dim arrays as matrices

#Nr=6;Nc=4         # Number of rows and columns
Nr=4;Nc=Nr         # Square matrix
A=zeros([Nr,Nc])      # Makes a Nr by Nc matrix of zeros
#A=ones([Nr,Nc])

print A

for i in arange(0,Nr):
    for j in arange(0,Nc):
        A[i,j]=i+j

print A

print ' Now make array of random numbers and do some linear algebra functions'

A = randn(Nr,Nc)
B = trace(A)     #Really only for square matrix
C = det(A)       #determinant
D = inv(A)       #inverse
print A
print 'Trace = ',B,'   determinant = ', C
print D
print 'Now should get identity matrix from A x Inv(A)'
print dot(A,D)  #This does matrix multiplication as opposed to element

Example: To include Time Delay

import time
time.sleep(2.0)   #Causes 2.0 second delay.

Example: Numerical integration

from scipy import * 
from scipy.integrate import quad
"""
Example of integration of a function using quad
Do the exponential function.
"""
result = quad(lambda x: exp(-x),0,5)   #Uses anonymous function, lambda
disp(['Numerical result: ',result])
analytic = 1-exp(-5)
disp(['Analytic result: ',analytic])

Example: Numerical solution of ODEs

from scipy import *
from scipy.integrate import odeint
from pylab import *
"""
Example of solving system of differential equations using odeint.
Do the damped oscillator, b is damping coefficient
"""
def damped_osc(u,t,b):   #defines the system of odes
    x=u[0]; v=u[1]
    return(v,-x-b*v)     #the derivatives of u

t = linspace(0,30,400)
u0 = [1,0]
b = 0.3

u=odeint(damped_osc,u0,t,args=(b,)) #b is in tuple, needs comma

figure(1) 
plot(t,u)                 #plots both columns of u
xlabel('Time'); title('Damped Oscillator')
figure(2)
plot(u[:,0],u[:,1])
title('Phase-space'); xlabel('Position'); ylabel('Velocity')
show()

OR try

from scipy import *
from scipy.integrate import ode
from pylab import *

"""Example using ode integrator, Lorenz equations"""

def foo(t,y,p):                 #system of 1st order ode
    sigma=p[0]; beta=p[1]; rho=p[2]
    dy = zeros([3])
    dy[0] = sigma*(y[1]-y[0])
    dy[1] = y[0]*(rho - y[2])-y[1]
    dy[2] = y[0]*y[1]-beta*y[2]
    return dy

t0 = 0; tEnd = 100.0; dt = 0.01
y0 = [5,-5,10]                  #Initial conditions
Y=[]; T=[]                      #create empty lists
p = [10.0,8.0/3.0,28.0]         #parameters for odes

#Set up integrator 'vode'.  Non-stiff use Adams, stiff use bdf 
r = ode(foo).set_integrator('vode',method='adams')
#r = ode(foo).set_integrator('vode',method='bdf')

#Maybe future version of scipy will have Runge-Kutta methods dopri5 and dop853
#r = ode(foo).set_integrator('dopri5')

r.set_f_params(p).set_initial_value(y0,t0)

while r.successful() and r.t+dt < tEnd:
    r.integrate(r.t+dt)
    Y.append(r.y)        #makes a list of 1d arrays
    T.append(r.t)

Y = array(Y)        #convert from list to 2d array

subplot(2,1,1)
plot(T,Y)
subplot(2,1,2)
plot(Y[:,0],Y[:,2])
show()

Example: Curve-fitting

from scipy import *
from matplotlib import *
from pylab import *
from scipy.optimize import leastsq
"""
Example of curve fitting for
a1*exp(-k1*t) + a2*exp(-k2*t)
"""
def dbexpl(t,p):
    return(p[0]*exp(-p[1]*t) + p[2]*exp(-p[3]*t))
a1,a2 = 1.0, 1.0
k1,k2 = 0.05, 0.2
t=arange(0,100,0.1)
data = dbexpl(t,[a1,k1,a2,k2]) + 0.02*randn(len(t))
def residuals(p,data,t):
    err = data - dbexpl(t,p)
    return err
p0 = [0.5,1,0.5,1] # initial guesses
guessfit = dbexpl(t,p0)
pbest = leastsq(residuals,p0,args=(data,t),full_output=1)
bestparams = pbest[0]
cov_x = pbest[1]
print 'best fit parameters ',bestparams
print cov_x
datafit = dbexpl(t,bestparams)
plot(t,data,'x',t,datafit,'r',t,guessfit)
xlabel('Time')
title('Curve-fitting example')
grid(True)
show()

Example: Fit to straight line

from scipy import *
from scipy import optimize
from pylab import *

"""Least squares fit to straight line
   Also demonstrates useful text placement"""

fitfunc = lambda p, x: p[0]*x + p[1]      #straight line
errfunc = lambda p, x, y: fitfunc(p,x) - y  #error to minimize

x=array([1,2,3,4,5])
y=array([0.8,2.1,2.8,4.1,5.2])

p0 = array([0,0])
p1 = optimize.leastsq(errfunc,p0, args=(x,y),full_output=1)
m,b = p1[0]
mesg = p1[3]
ier = p1[4]
print 'Best fit slope and intercept ',p1[0]
print ier,' ',mesg
fit=fitfunc(p1[0],x)

ax=subplot(1,1,1)
plot(x,y,'x',label='data')
plot(x,fit,'-',label='fit')
legend()
text(0.1,0.9,'slope = '+str(m),transform=ax.transAxes)
text(0.1,0.85,'intercept = '+str(round(b,3)),transform=ax.transAxes)
xlabel(r'$x$')
ylabel(r'$y$')
show()

Example: File IO, Reading and writing arrays

from scipy import *
"""
Example of file io, reading and writing array files. Does csv format.
"""
x=arange(0,10)
y=exp(-0.2*x)*cos(2*pi*x/20.)
data=column_stack((x,y))
print data
filename="test_io.txt"
savetxt(filename,data,delimiter=',',fmt='%.8e')
A = genfromtxt(filename,delimiter=',')
print A
print "done"

Example: Simple Plots with matplotlib To trouble-shoot, use command line $ python filename.py --verbose-helpful

from scipy import *
from pylab import *
"""
Shows how to put a couple of plots on a graph using matplotlib
"""
x = linspace(0,10,50)           #Define x-range: 0 to 10 with 50 points
y = x**2                         # Function x squared
z = 5*x                          # Function 5x

#plot using blue circles with line, and red x's no line
plot(x,y,'bo-',label='$x^2$',markersize=8.0,linewidth=2.0)
plot(x,z,'rx',label='$5x$',markersize=10.0,linewidth=2.0)
legend()   
xlabel('x axis')
ylabel('y axis')
title('Simple Graph')
grid(True)
show()

Example: Plots Functions and Data

from scipy import *
from pylab import *

#Shows how to plot functions and data points

x = linspace(0, 10, 100)        #Define x-range: 0 to 10 with 100 points
y = x**2                        # Function x squared
z = 5*x                         # Function 5x

#plot functions using blue line, and red dashed line
plot(x,y,'b-',label='$x^2$',linewidth=2.0)
plot(x,z,'r--',label='$5x$',linewidth=2.0)

#Include some "data" plotted with black (k) circles
xdat=[0,2,4,6,8,10]
ydat=[1,5,15,38,60,98]
plot(xdat,ydat,'ko',label='data',markersize=8.0)

legend(loc=2)          #puts legend in upper left corner
xlabel('x axis')
ylabel('y axis')
title('Simple Graph')
grid(True)
show()

Example: Interactive Plots with matplotlib

#!/usr/bin/env python
from scipy import *
import matplotlib
matplotlib.use('TkAgg')	  #Try this to work best in linux. Or qt4agg, or gtkagg
from pylab import *

print """Example of interactive plotting with matplotlib's pyplot.
Won't work correctly on Windows using IDLE.  Run from Command Prompt Window"""

x = linspace(0,1,200)       #Defines array of x values from 0 to 1

f = 2
ion()              #Turns interactive mode on.
ans=''
while ans != 'q':
    y = sin(2*pi*f*x)
    plot(x,y,'o-')
    draw()         #Need this to work on linux?

    ans=raw_input('f = '+repr(f)+'. Enter new f, or q: ')
    if ans != 'q':
        f = float(ans)
        clf()       #Clears the figure
ioff()

close()            #Closes the figure

#Use following 2 lines instead of close() if want figure to remain.

#print 'Remember to close figure window.'
#show()

Example: Switching between log and linear axes on plots

#! /usr/bin/env python
from __future__ import division
from scipy import *
import matplotlib
matplotlib.use('TkAgg')     #or maybe qt4agg, or gtkagg
from pylab import *

print "Example of switching between log and linear axes
Won't work correctly on Windows using IDLE.  Run from Command Prompt Window"

RC = 1e-3
f=logspace(1,5)          #good for logarithmic axes

ans=' '
ion()
while ans != 'q':
    y=1/sqrt(1+(2*pi*f*RC)**2)
    plot(f,y,'b+-')
    grid(True)
    draw()
    ans=raw_input('Enter log, lin, RC, q: ')
    if ans == 'log':
        semilogx()
    if ans == 'lin':
        delaxes()
    if ans =='RC':
       RC=float(raw_input('Old RC = '+repr(RC)+'. Enter new RC: '))
       clf()

close()
ioff()

Example: Sub-plots

from scipy import *
from pylab import *

"""Example for subplots"""

x=linspace(0,100,200)
y1=0.9*sin(2*pi*x/25)
y2=0.9*cos(2*pi*x/20)

subplot(1,2,1)
suptitle('Subplot Example Main Title',fontsize='16')
plot(x,y1)
title(r'sin($2\pi x/ \lambda$)')

subplot(1,2,2)
plot(x,y2)
title(r'cos($2\pi x/ \lambda$)')

show()

Example: 2 Y axes

from scipy import *
from pylab import *

"""Example with 2 y axes"""

x=linspace(0,10,200)
y1=x**2

y2=sqrt(x)

subplot(111)
plot(x,y1,'b')
ylabel(r'$x^2$ blue',fontsize=14)

twinx()
plot(x,y2,'r')
ylim([0,4])
ylabel(r'$x^{1/2}$ red',fontsize=14)

show()

Example: Inset Graph

from scipy import *
from pylab import *

"""Example with inset graph"""

x=linspace(0.05,2,2000)
y=sin(1/x)
y2=1/x

plot(x,y,x,y2)
xlabel(r'$x$',fontsize=14)
ylim([-2,5])
title(r'sin$(1/x)$ and $1/x$')

axes([0.6,0.5,0.25,0.25])
plot(x,y)
xlim([0.05,0.1])
title(r'Zoom sin$(1/x)$')

show()

Example: Finding zeros of nonlinear functions using fsolve

from scipy import *
from scipy.optimize import fsolve
from matplotlib import pylab
from pylab import *
"""
Use fsolve to find intersections of exp(x)-1 with cos(x)
Note that need different guesses to find different intersections
"""
def f(x):
    return (exp(x)-1)-cos(x)
result = fsolve(f,-1.5)
print result
x = linspace(-5,1,101)
plot(x,exp(x)-1,x,cos(x))
grid(b=1)
show()

Example: Plotting using Gnuplot

#!/usr/bin/env python
from numpy import *
import Gnuplot
"""
Demonstrates how to use gnuplot for graphs. Plots 2 graphs with different styles.
"""
x = linspace(0,10,31)
y = x**2
y2 = 10*sin(pi*x)
g = Gnuplot.Gnuplot() #!! Won't be able to use 'with' in python 2.6?
d = Gnuplot.Data(x,y,title='squared', with='lp lt 1 pt 6') 
d2=Gnuplot.Data(x,y2,title='sine',with='p pt 7 ps 4')
g('set grid')
g.xlabel('x axis')
g.ylabel('y axis')
g.plot(d,d2)
ans = raw_input('Enter f to create .png file, or Enter to quit ')
if ans == 'f':
    g.hardcopy('filename.png',terminal = 'png')
    #g.hardcopy('filename.ps',terminal='postscript',enhanced=1,color=1)
g.reset()

Example: fft

from scipy import * 
from pylab import * 
"""
Test the fft routine. Add signals, and multiply signals. 
"""
npts = 512            #Use some power of 2
t=linspace(0,1,npts+1)     # Use 2^N + 1 
dt = (t[-1]-t[0])/(len(t)-1)    # Maximum frequency is 1/2dt ?
fmax = 1/(2*dt) 
f1 = 80 
f2 = 90 
#sig = 1 + sin(2*pi*f1*t) + 1 + sin(2*pi*f2*t)  # sum of signals 
sig=(1+sin(2*pi*f1*t))*(1+sin(2*pi*f2*t))      # product of signals

figure(1)
plot(t,sig);xlabel('Time');title('Signal')

ft = fft(sig,n=npts) 
mgft=abs(ft)             #Get magnitude of fft
df = fmax/float(npts/2)
f=linspace(0,fmax,npts/2+1)
print 'fmax = ',fmax,' df = ',df,' ','\n 1st freqs = ',f[0:5]

figure(2)
plot(f,mgft[0:npts/2+1]);title('Fast Fourier Transform Magnitude')
xlabel('frequency')
ylabel('fft magnitude')
show()

Example: Make a square or sawtooth wave using scipy.signal.waveforms

#! /usr/bin/env python
from scipy import *
from scipy.signal.waveforms import sawtooth
from pylab import *

# Make a square wave or sawtooth. Graph it or save data file
# Use sawtooth from scipy, but define our own square wave function

def sqr(t,duty):
    f = 0
    if mod(t,2*pi) < duty*2*pi: f = 1
    return f
square=vectorize(sqr)

t = linspace(0,0.01,501)
f = 700
A = 2.8
duty = 0.5

# Uncomment the desired waveform
y = A*square(2*pi*f*t,duty)
#y = A*sawtooth(2*pi*f*t,width=0.5)

ans = raw_input('Enter g for graph, \n'+
          'd to make data .csv file, \n or press enter ')

if ans== 'g':
    figure(1)
    plot(t,y,'-')
    show()

if ans == 'd':
    filename = raw_input('Enter filename for data file ')
    data=column_stack((t,y))
    savetxt(filename,data,delimiter=',',fmt='%.8e')
Example: Add gaussian noise to signal
#! /usr/bin/env python
from numpy import *
from numpy.random import normal
from pylab import *
"""
Make signal with gaussian noise
"""
npts=200
x=linspace(0,10,npts)
theory=5*sin(2*pi*x/2.0)
noise=normal(0,1,npts)  #mean, std dev, num pts
sig=theory+noise

plot(x,theory,'-',x,sig,'+')
show()