Python cookbook
Michael Graupner
Created: November 8th, 2011
Last Modified: November 8th, 2011
Modified and extended from James Battat's cookbook
Python is a very flexible programming language which I use
for data analysis, visualizations, and simulations. It's a
high level language which allows for a steep learning curve
and fast solutions. Moreover, python is open-source and
therefore freely available for everybody.
This is a list of python functions and short scripts which
I use often. Note that they require the following pyton
packages to be installed : matplotlib, scipy, numpy and
brian (http://www.briansimulator.org/).
Pick 100 random numbers with a gaussian distribution
Fit a polynomial to data
Fit an arbitrary function to data
Two different y axes
Eliminate axis tick marks
Sophisticated figure layout 'ready to publish'
Multi-panel figure with adjusted layout
Multi-panel figure with varying cell sizes
Arrange plots (.svg files) into composite figure
Multi-panel figure with all plotting know-how
Numerically integrate ODEs
Next header
Pick 100 random numbers with a gaussian distribution
nums = scipy.randn(100)
#Plot the data and a histogram of the data with 10 bins
pylab.plot(nums, 'ro')
pylab.savefig('randnplot.png')
pylab.cla() # clear the axes
pylab.hist(nums, 10)
pylab.savefig('randnhist.png')

Fit a polynomial to data
xdata = scipy.linspace(0, 9, num=10)
ydata = 0.5+xdata*2.0
ydata += scipy.rand(10)
polycoeffs = scipy.polyfit(xdata, ydata, 1)
print polycoeffs
# [ 2.00710807, 1.09204496]
yfit = scipy.polyval(polycoeffs, xdata)
pylab.plot(xdata, ydata, 'k.')
pylab.plot(xdata, yfit, 'r-')
Fit an arbitrary function to data
import scipy, pylab
import scipy.optimize
# make x data
num = 100
x = scipy.linspace(-10, 10, num=num)
distancePerLag = x[1]-x[0]
# make two gaussians, with different means
offset = 2.0
y1 = scipy.exp(-x**2/8.0)
y2 = scipy.exp(-(x-offset)**2/1.0)
# compute the cross-correlation between y1 and y2
ycorr = scipy.correlate(y1, y2, mode='full')
xcorr = scipy.linspace(0, len(ycorr)-1, num=len(ycorr))
# define a gaussian fitting function where
# p[0] = amplitude
# p[1] = mean
# p[2] = sigma
fitfunc = lambda p, x: p[0]*scipy.exp(-(x-p[1])**2/(2.0*p[2]**2))
errfunc = lambda p, x, y: fitfunc(p,x)-y
# guess some fit parameters
p0 = scipy.c_[max(ycorr), scipy.where(ycorr==max(ycorr))[0], 5]
# fit a gaussian to the correlation function
p1, success = scipy.optimize.leastsq(errfunc, p0.copy()[0],
args=(xcorr,ycorr))
# compute the best fit function from the best fit parameters
corrfit = fitfunc(p1, xcorr)
# get the mean of the cross-correlation
xcorrMean = p1[1]
# convert index to lag steps
# the first point has index=0 but the largest (negative) lag
# there is a simple mapping between index and lag
nLags = xcorrMean-(len(y1)-1)
# convert nLags to a physical quantity
# note the minus sign to ensure that the
# offset is positive for y2 is shifted to the right of y1
# a negative offset means that y2 is shifted to the left of y1
# I don't know what the standard notation is (if there is one)
offsetComputed = -nLags*distancePerLag
# see how well you have done by comparing the actual
# to the computed offset
print 'xcorrMean, nLags = ', \
xcorrMean, ', ', nLags
print 'actualOffset, computedOffset = ', offset,', ', offsetComputed
# visualize the data
# plot the initial functions
pylab.subplot(211)
pylab.plot(x, y1, 'ro')
pylab.plot(x, y2, 'bo')
# plot the correlation and fit to the correlation
pylab.subplot(212)
pylab.plot(xcorr, ycorr, 'k.')
pylab.plot(xcorr, corrfit, 'r-')
pylab.plot([xcorrMean, xcorrMean], [0, max(ycorr)], 'g-')
pylab.show()
Two different y axes
"""
Demonstrate how to do two plots on the same axes with different left
right scales.
The trick is to use *2 different axes*. Turn the axes rectangular
frame off on the 2nd axes to keep it from obscuring the first.
Manually set the tick locs and labels as desired. You can use
separate matplotlib.ticker formatters and locators as desired since
the two axes are independent.
This is acheived in the following example by calling pylab's twinx()
function, which performs this work. See the source of twinx() in
pylab.py for an example of how to do it for different x scales. (Hint:
use the xaxis instance and call tick_bottom and tick_top in place of
tick_left and tick_right.)
"""
from pylab import *
ax1 = subplot(111)
t = arange(0.01, 10.0, 0.01)
s1 = exp(t)
plot(t, s1, 'b-')
xlabel('time (s)')
ylabel('exp')
# turn off the 2nd axes rectangle with frameon kwarg
ax2 = twinx()
s2 = sin(2*pi*t)
plot(t, s2, 'r.')
ylabel('sin')
ax2.yaxis.tick_right()
show()
Eliminate Tick Marks from Plots
import pylab
# first axes
y = pylab.arange(10)
ax1=pylab.subplot(211)
pylab.plot(y)
ax1.xaxis.set_major_locator(pylab.NullLocator())
pylab.title('This x-axis has no ticks')
# second axes
ax2=pylab.subplot(212)
pylab.plot(y)
pylab.title('This x-axis has ticks')
Sophisticated figure layout 'ready to publish'
from matplotlib import rcParams
import numpy as np
import matplotlib.pyplot as plt
#set plot attributes
fig_width = 7 # width in inches
fig_height = 4.2 # height in inches
fig_size = [fig_width,fig_height]
params = {'axes.labelsize': 14,
'axes.titlesize': 14,
'font.size': 10,
'xtick.labelsize': 12,
'ytick.labelsize': 12,
'figure.figsize': fig_size,
'savefig.dpi' : 600,
'axes.linewidth' : 1.3,
'ytick.major.size' : 6, # major tick size in points
'xtick.major.size' : 6 # major tick size in points
#'xtick.major.size' : 2,
#'ytick.major.size' : 2,
}
rcParams.update(params)
# set sans-serif font to Arial
rcParams['font.sans-serif'] = 'Arial'
def sigmoid(x):
return 1./(1+np.exp(-(x-5)))+1
# data to be displayed
np.random.seed(1234)
t = np.arange(0.1, 9.2, 0.15)
y = sigmoid(t) + 0.2*np.random.randn(len(t))
residuals = y - sigmoid(t)
t_fitted = np.linspace(0, 10, 100)
###########################
#adjust plots position
###########################
fig = plt.figure()
ax1 = plt.axes((0.14, 0.12, 0.65, 0.8))
plt.plot(t, y, 'b.', ms=5., clip_on=False)
plt.plot(t_fitted, sigmoid(t_fitted), 'r-', lw=3)
plt.text(5, 1.0, r"$L = \frac{1}{1+\exp(-V+5)}+10$",
fontsize=14,
transform=ax1.transData, clip_on=False,
va='top', ha='left')
#set axis limits
ax1.set_xlim((0, t.max()))
ax1.set_ylim((0.5, y.max()))
#hide right and top axes
ax1.spines['top'].set_visible(False)
ax1.spines['right'].set_visible(False)
ax1.spines['bottom'].set_position(('outward', 10))
ax1.spines['left'].set_position(('outward', 10))
ax1.yaxis.set_ticks_position('left')
ax1.xaxis.set_ticks_position('bottom')
#set labels
plt.xlabel(r'voltage (V, $\mu$V)')
plt.ylabel('luminescence (L)')
#make inset
ax_inset = plt.axes((0.25, 0.74, 0.2, 0.2), frameon=False)
plt.hist(residuals, fc='0.8',ec='w', lw=2)
plt.xticks([-0.5, 0, 0.5],[-0.5, 0, 0.5], size=6)
plt.xlim((-0.5, 0.5))
plt.yticks([5, 10], size=6)
plt.xlabel("residuals", size=8)
ax_inset.xaxis.set_ticks_position("none")
ax_inset.yaxis.set_ticks_position("left")
# change ticks for inset
for line in ax_inset.get_yticklines():
line.set_markersize(4)
line.set_markeredgewidth(0.5)
ax_inset.yaxis.grid(lw=1, color='w', ls='-')
plt.text(0, 0.9, "frequency", transform=ax_inset.transAxes,
va='center', ha='right', size=8)
#export to svg and png
plt.savefig('sigmoid_fit.png')
plt.savefig('sigmoid_fit.svg')
Multi-panel figure with adjusted layout
"""
A gridspec instance provides array-like (2d or 1d) indexing that returns
the SubplotSpec instance. For, SubplotSpec that spans multiple cells,
use slice.
When a GridSpec is explicitly used, you can adjust the layout parameters
of subplots that are created from the gridspec.
This is similar to subplots_adjust, but it only affects the subplots that
are created from the given GridSpec.
"""
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
def make_ticklabels_invisible(fig):
for i, ax in enumerate(fig.axes):
ax.text(0.5, 0.5, "ax%d" % (i+1), va="center", ha="center")
for tl in ax.get_xticklabels() + ax.get_yticklabels():
tl.set_visible(False)
# demo 3 : gridspec with subplotpars set.
f = plt.figure()
plt.suptitle("GirdSpec w/ different subplotpars")
gs1 = GridSpec(3, 3)
# right spacing creates space for gs2
gs1.update(left=0.05, right=0.48, wspace=0.05)
ax1 = plt.subplot(gs1[:-1, :])
ax2 = plt.subplot(gs1[-1, :-1])
ax3 = plt.subplot(gs1[-1, -1])
gs2 = GridSpec(3, 3)
# left spacing creates space for gs1
gs2.update(left=0.55, right=0.98, hspace=0.05)
ax4 = plt.subplot(gs2[:, :-1])
ax5 = plt.subplot(gs2[:-1, -1])
ax6 = plt.subplot(gs2[-1, -1])
make_ticklabels_invisible(plt.gcf())
#plt.show()
#export to svg and png
plt.savefig('gridspec.png')
plt.savefig('gridspec.svg')
Multi-panel figure with varying cell sizes and automatic labeling
"""
By default, GridSpec creates cells of equal sizes. You can adjust relative
heights and widths of rows and columns. Note that absolute values are
meaningless, only their relative ratios matter.
The labelling code is taken from Dan Goodman.
More infos on GridSpec can be found here.
"""
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
def make_ticklabels_invisible(fig):
for i, ax in enumerate(fig.axes):
ax.text(0.5, 0.5, "ax%d" % (i+1), va="center", ha="center")
for tl in ax.get_xticklabels() + ax.get_yticklabels():
tl.set_visible(False)
def label_panel(ax, letter, *,offset_left=0.6, offset_up=0.1, prefix='', postfix=')', **font_kwds):
kwds = dict(fontsize=16)
kwds.update(font_kwds)
# this mad looking bit of code says that we should put the code offset a certain distance in
# inches (using the fig.dpi_scale_trans transformation) from the top left of the frame
# (which is (0, 1) in ax.transAxes transformation space)
fig = ax.figure
trans = ax.transAxes + transforms.ScaledTranslation(-offset_left, offset_up, fig.dpi_scale_trans)
ax.text(0, 1, prefix+letter+postfix, transform=trans, **kwds)
def label_panels(axes, letters=None, **kwds):
if letters is None:
letters = axes.keys()
for letter in letters:
ax = axes[letter]
label_panel(ax, letter, **kwds)
f = plt.figure()
gs = gridspec.GridSpec(2, 2,
width_ratios=[1,2],
height_ratios=[4,1]
)
ax1 = plt.subplot(gs[0])
ax2 = plt.subplot(gs[1])
ax3 = plt.subplot(gs[2])
ax4 = plt.subplot(gs[3])
make_ticklabels_invisible(f)
#plt.show()
label_panel(ax1,'A')
label_panel(ax2,'B')
label_panel(ax3,'C')
label_panel(ax4,'D')
plt.savefig('gridspec3.png')
plt.savefig('gridspec3.svg')
Arrange plots (.svg files) into composite figure
"""
The procedure uses the small Python package svgutils. It is written completely
in Python and uses only standard libraries. You may download it from github.
More infos can be found here.
"""
import svgutils.transform as sg
import sys
#create new SVG figure
fig = sg.SVGFigure("20.5cm", "6.5cm")
# load matpotlib-generated figures
fig1 = sg.fromfile('sigmoid_fit.svg')
fig2 = sg.fromfile('anscombe.svg')
# get the plot objects
plot1 = fig1.getroot()
plot2 = fig2.getroot()
plot2.moveto(370, 0, scale=0.65)
# add text labels
txt1 = sg.TextElement(25,20, "A", size=12, weight="bold")
txt2 = sg.TextElement(420,20, "B", size=12, weight="bold")
# append plots and labels to figure
fig.append([plot1, plot2])
fig.append([txt1, txt2])
# save generated SVG and PNG files
fig.save("fig_final.svg")
fig.save("fig_final.png")
Multi-panel figure with all plotting know-how
"""
All the know-how I have accumulated over time to generate multi-panel figures with various layout changes.
"""
from scipy import *
from numpy import *
from numpy.fft import *
from math import *
from pylab import *
import scipy.signal
import scipy.stats
import os
import pdb
import random
import sys
import time
import pickle
from matplotlib import rcParams
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
# data to plot
x = arange(-10,10,0.1)
siny = sin(x)
cosy = cos(x)
# set plot attributes
fig_width = 12 # width in inches
fig_height = 8 # height in inches
fig_size = [fig_width,fig_height]
params = {'axes.labelsize': 14,
'axes.titlesize': 13,
'font.size': 11,
'xtick.labelsize': 11,
'ytick.labelsize': 11,
'figure.figsize': fig_size,
'savefig.dpi' : 600,
'axes.linewidth' : 1.3,
'ytick.major.size' : 4, # major tick size in points
'xtick.major.size' : 4 # major tick size in points
#'edgecolor' : None
#'xtick.major.size' : 2,
#'ytick.major.size' : 2,
}
rcParams.update(params)
# set sans-serif font to Arial
rcParams['font.sans-serif'] = 'Arial'
# create figure instance
fig = plt.figure()
# define sub-panel grid and possibly width and height ratios
gs = gridspec.GridSpec(2, 2,
width_ratios=[1,1.2],
height_ratios=[1,1]
)
# define vertical and horizontal spacing between panels
gs.update(wspace=0.3,hspace=0.4)
# possibly change outer margins of the figure
#plt.subplots_adjust(left=0.14, right=0.92, top=0.92, bottom=0.18)
# sub-panel enumerations
plt.figtext(0.06, 0.92, 'A',clip_on=False,color='black', weight='bold',size=22)
plt.figtext(0.47, 0.92, 'B',clip_on=False,color='black', weight='bold',size=22)
plt.figtext(0.06, 0.47, 'C',clip_on=False,color='black', weight='bold',size=22)
plt.figtext(0.47, 0.47, 'D',clip_on=False,color='black', weight='bold',size=22)
# first sub-plot #######################################################
ax0 = plt.subplot(gs[0])
# title
ax0.set_title('sub-plot 1')
# diplay of data
ax0.axhline(y=0,ls='--',color='0.5',lw=2)
ax0.axvline(x=0,ls='--',color='0.5',lw=2)
ax0.plot(x,siny,label='sin')
ax0.plot(x,cosy,label='cos')
# removes upper and right axes
# and moves left and bottom axes away
ax0.spines['top'].set_visible(False)
ax0.spines['right'].set_visible(False)
ax0.spines['bottom'].set_position(('outward', 10))
ax0.spines['left'].set_position(('outward', 10))
ax0.yaxis.set_ticks_position('left')
ax0.xaxis.set_ticks_position('bottom')
# legends and labels
plt.legend(loc=1,frameon=False)
plt.xlabel('time (sec)')
plt.ylabel('oscillations')
# second sub-plot #######################################################
# creates sub-panels in second sub-plot
gssub = gridspec.GridSpecFromSubplotSpec(2, 2, subplot_spec=gs[1],hspace=0.4)
# sub-panel 1 #############################################
ax10 = plt.subplot(gssub[0])
# diplay of data
ax10.axhline(y=0,ls='--',color='0.5',lw=2)
ax10.axvline(x=0,ls='--',color='0.5',lw=2)
ax10.plot(x,siny,label='sin')
# removes upper, lower and right axes
# and moves left axis away
ax10.spines['top'].set_visible(False)
ax10.spines['right'].set_visible(False)
ax10.spines['bottom'].set_visible(False)
ax10.spines['left'].set_position(('outward', 10))
ax10.yaxis.set_ticks_position('left')
ax10.axes.get_xaxis().set_visible(False)
# sub-panel 2 #############################################
ax11 = plt.subplot(gssub[1])
# diplay of data
ax11.axhline(y=0,ls='--',color='0.5',lw=2)
ax11.axvline(x=0,ls='--',color='0.5',lw=2)
ax11.plot(x,cosy,label='sin')
# removes all axes
ax11.spines['top'].set_visible(False)
ax11.spines['right'].set_visible(False)
ax11.spines['bottom'].set_visible(False)
ax11.spines['left'].set_visible(False)
ax11.axes.get_xaxis().set_visible(False)
ax11.axes.get_yaxis().set_visible(False)
text11 = ax11.annotate(r'$y = cos(x)$', xy=(1,1.1), annotation_clip=False,
xytext=None, textcoords='data',fontsize=15,
arrowprops=None
)
# sub-panel 3 #############################################
ax12 = plt.subplot(gssub[2])
# diplay of data
ax12.axhline(y=0,ls='--',color='0.5',lw=2)
ax12.axvline(x=0,ls='--',color='0.5',lw=2)
ax12.plot(x,cosy,label='sin')
# removes upper and right axes
# and moves left and bottom axes away
ax12.spines['top'].set_visible(False)
ax12.spines['right'].set_visible(False)
ax12.spines['bottom'].set_position(('outward', 10))
ax12.spines['left'].set_position(('outward', 10))
ax12.yaxis.set_ticks_position('left')
ax12.xaxis.set_ticks_position('bottom')
plt.xlabel('time (sec)',position=(1.1,-0.2))
plt.ylabel('oscillations',position=(-0.2,1.1))
# sub-panel 4 #############################################
ax13 = plt.subplot(gssub[3])
# diplay of data
ax13.axhline(y=0,ls='--',color='0.5',lw=2)
ax13.axvline(x=0,ls='--',color='0.5',lw=2)
ax13.plot(x,siny,label='sin')
# removes upper, left and right axes
# and moves bottom axis away
ax13.spines['top'].set_visible(False)
ax13.spines['right'].set_visible(False)
ax13.spines['left'].set_visible(False)
ax13.spines['bottom'].set_position(('outward', 10))
ax13.xaxis.set_ticks_position('bottom')
ax13.axes.get_yaxis().set_visible(False)
# third sub-plot #######################################################
ax2 = plt.subplot(gs[2])
# diplay of data
ax2.axhline(y=0,ls='--',color='0.5',lw=2)
ax2.axvline(x=0,ls='--',color='0.5',lw=2)
ax2.plot(x,siny,label='sin')
ax2.plot(x,cosy,label='cos')
# removes upper and right axes
# and moves left and bottom axes away
ax2.spines['top'].set_visible(False)
ax2.spines['right'].set_visible(False)
ax2.spines['bottom'].set_position(('outward', 10))
ax2.spines['left'].set_position(('outward', 10))
ax2.yaxis.set_ticks_position('left')
ax2.xaxis.set_ticks_position('bottom')
# legends and labels
plt.legend(loc=3,frameon=False)
plt.xlabel('time (sec)')
plt.ylabel('oscillations')
# change tick spacing
majorLocator_x = MultipleLocator(10)
ax2.xaxis.set_major_locator(majorLocator_x)
bracket = ax2.annotate('', xy=(-4.8, 0.98), xycoords='data', annotation_clip=False,
xytext=(0, 0.98), textcoords='data',
arrowprops=dict(arrowstyle="-",
connectionstyle='bar,fraction=0.22', linewidth=2,
ec='k',
shrinkA=10, shrinkB=10,
)
)
text = ax2.annotate(r'$\ast\ast$', xy=(-3.2,1.16), annotation_clip=False,
xytext=None, textcoords='data',fontsize=15,
arrowprops=None
)
# fourth sub-plot #######################################################
ax3 = plt.subplot(gs[3])
# diplay of data
ax3.axhline(y=0,ls='--',color='0.5',lw=2)
ax3.axvline(x=0,ls='--',color='0.5',lw=2)
ax3.plot(x,siny,label='sin')
ax3.plot(x,cosy,label='cos')
# removes upper and right axes
# and moves left and bottom axes away
ax3.spines['top'].set_visible(False)
ax3.spines['right'].set_visible(False)
ax3.spines['bottom'].set_position(('outward', 10))
ax3.spines['left'].set_position(('outward', 10))
ax3.yaxis.set_ticks_position('left')
ax3.xaxis.set_ticks_position('bottom')
# legends and labels
plt.xlabel('time (sec)')
plt.ylabel('oscillations')
# add arrows (or lines) to plot
ax3.annotate('', xy=(-7.5, 0.6), xytext=(-9, 0.6), annotation_clip=False,
arrowprops=dict(arrowstyle="->",facecolor='black',lw=2),
)
ax3.annotate('', xy=(-4, -0.7), xytext=(-5.5, -0.7), annotation_clip=False,
arrowprops=dict(arrowstyle="->",facecolor='black',lw=2),
)
# legend position defined by coordinates
plt.legend(loc=(0.2,0.5),frameon=False)
# change legend text size
leg = plt.gca().get_legend()
ltext = leg.get_texts()
plt.setp(ltext, fontsize=11)
## save figure ############################################################
fname = 'fig_example_plot'
savefig(fname+'.png')
savefig(fname+'.pdf')
Numerically integrate ODEs
"""
Use 'odeint' to ingreate odinary differential equations.
"""
from scipy.integrate import odeint
from pylab import * # for plotting commands
def deriv0(y,t): # right hand side of equation (2)
taur = 10.
taud = 30.
Cpre = 1.
npre = Cpre*(1./taud - 1./taur)*( (taur/taud)**(1./(1.-taur/taud)) - (taur/taud)**(1./(taud/taur-1.)) )**(-1.)
return array([-y[0]/taur + y[1]*npre,-y[1]/taud])
once = True
time0 = linspace(0.0,20.,2001)
yinit0 = array([0.0,1.]) # initial value
y0 = odeint(deriv0,yinit0,time0[:-1])
def deriv1(y,t):
taur = 10.
taud = 30.
Cpre = 1.
npre = Cpre*(1./taud - 1./taur)*( (taur/taud)**(1./(1.-taur/taud)) - (taur/taud)**(1./(taud/taur-1.)) )**(-1.)
#
taurp = 2.
taudp = 12.
Cpost = 2.5
npost = Cpost*(1./taudp - 1./taurp)*( (taurp/taudp)**(1./(1.-taurp/taudp)) - (taurp/taudp)**(1./(taudp/taurp-1.)) )**(-1.)
#
return array([-y[0]/taur + y[1]*npre,-y[1]/taud,-y[2]/taurp + y[3]*npost,-y[3]/taudp])
eta = 1.5
time1 = linspace(20.,100.,8001)
yinit1 = array([y0[-1,0],y0[-1,1],0.0,1.+eta*y0[-1,0]])
y1 = odeint(deriv1,yinit1,time1)
time = hstack((time0[:-1],time1))
y = vstack((y0,y1[:,:2]))
ytot = copy(y[:,0])
ytot[-len(y1[:,2]):] += y1[:,2]
figure()
plot(time,y[:,0],label='A')
plot(time,y[:,1],label='B')
plot(time1,y1[:,2],label='E')
plot(time1,y1[:,3],label='F')
plot(time,ytot,c='k',lw=2,label='A+E')
xlabel('time (ms)')
ylabel('calcium')
legend(frameon=False)
fname = 'odeint_example'
savefig(fname+'.png')
savefig(fname+'.pdf')