Python for Scientists, 2nd Edition: Code Samples¶

a = 3
b = 2.7
c = 12 + 5j
s = "Hello World!"
L = [a, b, c, s, 77.77]

#%%writefile fib.py
# File: fib.py Fibonacci numbers

"""Module fib: Fibonacci numbers
   Contains one function fib(n).
"""

def fib(n):
    """ Returns n'th Fibonacci number. """
    a, b = 0, 1
    for i in range(n):
        a, b = b, a+b
    return a

##########################################################

if __name__ == "__main__":
    for i in range(1001):
        print "fib(", i, ") = ", fib(i)

#%%writefile gcd.py
#  File gcd.py  Implementing the GCD Euclidean algorithm.

""" Module gcd: contains two implementations of the Euclid
    GCD algorithm, gcdr and gcd.
""" 

def gcdr(a, b):
    """ Euclidean algorithm, recursive vers., returns GCD. """
    if b == 0:
        return a
    else:
        return gcdr(b, a%b)

def gcd(a, b):
    """ Euclidean algorithm, non-recursive vers., returns GCD. """
    while b:
        a, b = b, a%b
    return a

##########################################################
if __name__ == "__main__":
    import fib

    for i in range(963):
        print i, '  ', gcdr(fib.fib(i), fib.fib(i+1))

a = 4; b = 5.5; c = 1.5+2j; d = 'a'
e = 6.0*a - b*b + \
        c**(a+b+c)
f = 6.0*a - b*b + c**(
         a+b+c)
a, b, c, d, e, f

p =3.14
p
q = p
p = 'pi'
p
q

[1, 4.0, 'a']
u = [1, 4.0, 'a']
v = [3.14, 2.78, u, 42]
v
len(v)             
len?                # or help(len) 
v*2
v + u
v.append('foo')
v

u = [0, 1, 2, 3, 4, 5, 6, 7]
su = u[2 : 4]
su
su[0] = 17
su
u
u[2 : 4] = [3.14,' a']
u

a = 4
b = a
b = 'foo'
a
b
u = [0, 1, 4, 9, 16]
v = u
v[2] = 'foo'
v
u

u = [0, 1, 4, 9, 16]
v = u[ : ]
v[2] = 'foo'
v
u

empty={} 
parms = {'alpha':1.3, 'beta':2.74} 
#parms = dict(alpha=1.3,beta=2.74) 
parms['gamma']=0.999
parms

y = 107
for x in range(2, y):
    if y%x == 0:
        print y, " has a factor ",  x 
        break
else:
        print y, "is prime."

def add_one(x):
    """ Takes x and returns x + 1. """
    x = x + 1
    return x

x=23
#add_one?                   # or help(add_one)
add_one(0.456)
x

def add_x_and_y(x, y):
    """ Add x and y  and return them and their sum. """
    z = x + y
    return x, y, z

a, b, c = add_x_and_y(1,0.456)
a, b, c

L = [0, 1, 2]
print id(L)
def add_with_side_effects(M):
    """ Increment first element of list. """
    M[0] += 1

add_with_side_effects(L)
print L
id(L)

L = [0, 1, 2]
id(L)
def add_without_side_effects(M):
    """ Increment first element of list. """
    MC = M[ : ]
    MC[0] +=1 
    return MC

L = add_without_side_effects(L)
L
id(L)

def average(*args):
    """ Return mean of a non-empty tuple of numbers. """
    print args
    sum = 0.0
    for x in args:
        sum += x
    return sum/len(args)

print average(1,2,3,4)
print average(1,2,3,4,5)

#%%writefile frac.py
# File: frac.py

import gcd

class Frac:
    """ Fractional class.  A Frac is a pair of integers num, den
        (with den!=0) whose GCD is 1.
    """

    def __init__(self, n, d):
        """ Construct a Frac from integers n and d.
            Needs error message if d = 0!
        """
        hcf=gcd.gcd(n, d)
        self.num, self.den = n/hcf, d/hcf

    def __str__(self):
        """ Generate a string representation of a Frac. """
        return "%d/%d" % (self.num,self.den)

    def __mul__(self, another):
        """ Multiply two Fracs to produce a Frac. """
        return Frac(self.num*another.num, self.den*another.den)

    def __add__(self,another):
        """ Add two Fracs to produce a Frac. """
        return Frac(self.num*another.den + self.den*another.num,
                    self.den*another.den)

    def to_real(self):
        """ Return floating point value of Frac. """
        return float(self.num)/float(self.den)

if __name__ == "__main__":
    a=Frac(3,7)
    b=Frac(24,56)
    print "a.num= ", a.num, ", b.den= ", b.den
    print a
    print b
    print "floating point value of a is ", a.to_real()
    print "product= ",a*b,", sum= ",a + b

def sieve_v1(n):
     """
        Use Sieve of Eratosthenes to compute list of primes <= n.
        Version 1
     """
     primes = range(2, n+1)
     for p in primes:
         if p*p > n:
             break
         product = 2*p
         while product <= n:
             if product in primes:
                 primes.remove(product)
             product += p
     return len(primes), primes

def sieve_v2(n):
    """ 
        Sieve of Eratosthenes to compute list of primes <= n.
        Version 2.
    """
    sieve = [True]*(n+1)
    for i in xrange(3, n+1, 2):
        if i*i > n:
            break
        if sieve[i]:
            sieve[i*i : : 2*i] = [False]*((n - i*i) // (2*i) + 1)
    answer = [2] + [i for i in xrange(3, n+1, 2) if sieve[i]]
    return len(answer), answer

import numpy as np
xl = np.linspace(0, 0.5, 5, endpoint=False)
xm = np.linspace(0.5, 0.6, 10,endpoint=False)
xr = np.linspace(0.6, 1.0, 5)
xs = np.hstack((xl, xm, xr))
xs

def m1(x):
    return np.exp(2*x)
def m2(x):
    return 1.0
def m3(x):
    return np.exp(1.0-x)
  
x = np.linspace(-10, 10, 21) 
conditions=[ x >= 0,  x >= 1, x < 0 ]
functions=[ m2, m3, m1 ]
m = np.piecewise(x, conditions, functions)
m

x=np.array([[0, 1, 2, 3], [10, 11, 12, 13], [20, 21, 22, 23]])  
r=np.array([2, 3, 4, 5])
c=np.array([[5], [6], [7]])
print "2*x = \n", 2*x
print "r*x = \n", r*x 
print "x*r = \n", x*r
print "c*x = \n", c*x

xv=np.linspace(-1, 1, 5)
yv=np.linspace(0, 1, 3)
[xa, ya] = np.meshgrid(xv, yv)
print "xa = \n", xa
print "ya = \n", ya
print "xa*ya = \n", xa*ya

[xo, yo]=np.ogrid[-1: 1: 5j, 0: 1: 3j] 
print "xo = \n", xo
print "yo = \n", yo
print "xo.shape = ", xo.shape, " yo.shape = ", yo.shape
print "xo*yo = \n",  xo*yo

import numpy as np
quarter = np.array([1, 2, 3, 4],dtype=int)
results=np.array([37.4, 47.3, 73.4, 99])
outfile=open("q4.txt","w")
outfile.write("The results for the first four quarters\n\n")
for q,r in zip(quarter, results):
    outfile.write("For quarter %d the result is %5.1f\n" % (q,r))
outfile.close()

infile=open("q4.txt", "r")
lquarter = []
lresult = []
temp = infile.readline()
temp = infile.readline()
for line in infile:
    words = line.split()
    lquarter.append(int(words[2]))
    lresult.append(float(words[6]))
infile.close()
import numpy as np
aquarter = np.array(lquarter, dtype=int)
aresult = np.array(lresult)
aquarter, aresult

len = 21
x = np.linspace(0, 2*np.pi, len)
c = np.cos(x)
s = np.sin(x)
t = np.tan(x)
arr = np.empty((4, len), dtype=float)
arr[0, : ] = x
arr[1, : ] = c
arr[2, : ] = s
arr[3, : ] = t
np.savetxt('x.txt', x)
np.savetxt('xcst.txt', (x, c, s, t))
np.savetxt('xarr.txt', arr)

np.savez('test.npz',x = x, c = c, s = s, t = t)

temp=np.load('test.npz')
temp.files
xc = temp['x']
cc = temp['c']
sc = temp['s']
tc = temp['t']

import numpy as np

roots = [0, 1, 1, 2]
coeffs = np.poly(roots)
print "coeffs = \n", coeffs
print "np.roots(coeffs) = \n", np.roots(coeffs)
x = np.linspace(0,0.5*np.pi,7)
y = np.sin(x)
c = np.polyfit(x, y, 3)
print "c = \n", c
y1 = np.polyval(c, x)
print "y = \n", y, "\n y1 = \n", y1, "\n y1-y = \n", y1-y

import numpy as np
import numpy.linalg as npl

a=np.array([[4, 2, 0], [9, 3, 7], [1, 2, 1]])
print "a = ", a
print "det(a) = ", npl.det(a)
b=npl.inv(a)
print "b = ", b, "\n b*a = ", np.dot(b,a)

import numpy as np
import numpy.linalg as npl

a = np.array([[-2, -4, 2], [-2, 1, 2], [4, 2, 5]])
evals, evecs = npl.eig(a)
eval1 = evals[0]
evec1 = evecs[:,0]
eval1, evec1

npl.norm(evec1), np.dot(evec1, evec1)

np.dot(a, evec1) - eval1*evec1

import numpy as np
import numpy.linalg as npl
a=np.array([[3,2,1],[5,5,5],[1,4,6]])
b=np.array([[5,1],[5,0],[-3,-7.0/2]])
x=npl.solve(a,b)
x

np.dot(a,x)-b

import numpy as np
import scipy.optimize as sco
def fun(x):
    return np.cosh(x)/np.sinh(x)-x

roots = sco.fsolve(fun, 1.0)
root = roots[0]
print "root is %15.12f and value is %e" % (root, fun(root))

#%%writefile foo.py 
import numpy as np
import matplotlib.pyplot as plt

#  Use the next command in Ipython terminal mode
#plt.ion()
#  Use the next command in Ipython non-inline notebook mode
#matplotlib
#  Use the next command in Ipython inline notebook mode
%matplotlib notebook

x = np.linspace(-np.pi, np.pi, 101)
y = np.sin(x) + np.sin(3*x)/3.0
 
plt.plot(x, y)

plt.xlabel('x')
plt.ylabel('y')
plt.title('A simple plot')

#plt.show()
plt.savefig('foo.pdf')

import numpy as np
import matplotlib.pyplot as plt
%matplotlib notebook

x = np.linspace(-np.pi, np.pi, 101)
y = np.sin(x) + np.sin(3*x)/3.0

fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(x, y)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_title('A simple plot')
plt.plot(x, y)

fig.savefig('foo.pdf')

import numpy as np
import matplotlib.pyplot as plt
%matplotlib notebook

x=np.linspace(-np.pi,np.pi,101)
f=np.ones_like(x)
f[x<0] = -1
y1 = (4/np.pi)*(np.sin(x) + np.sin(3*x)/3.0)
y2 = y1 + (4/np.pi)*(np.sin(5*x)/5.0 + np.sin(7*x)/7.0)
y3 = y2 + (4/np.pi)*(np.sin(9*x)/9.0 + np.sin(11*x)/11.0)

plt.ion()
fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(x,f,'b-',lw=3,label='f(x)')
ax.plot(x,y1,'c--',lw=2,label='two terms')
ax.plot(x,y2,'r-.',lw=2,label='four terms')
ax.plot(x, y3,'b:',lw=2,label='six terms')
ax.legend(loc='best')
ax.set_xlabel('x',style='italic')
ax.set_ylabel('partial sums',style='italic')
fig.suptitle('Partial sums for Fourier series of f(x)',
            size=16,weight='bold')                                       

import numpy as np
import matplotlib.pyplot as plt 
%matplotlib notebook

theta = np.linspace(0,2*np.pi,201) 
r1 = np.abs(np.cos(5.0*theta) - 1.5*np.sin(3.0*theta)) 
r2 = theta/np.pi
r3 = 2.25*np.ones_like(theta)

fig = plt.figure()
ax = fig.add_subplot(111,projection='polar')
ax.plot(theta, r1,label='trig')
ax.plot(5*theta, r2,label='spiral')
ax.plot(theta, r3,label='circle')
ax.legend(loc='best')                  

import numpy as np
import numpy.random as npr 
x = np.linspace(0,4,21) 
y = np.exp(-x) 
xe = 0.08*npr.randn(len(x)) 
ye = 0.1*npr.randn(len(y))

import matplotlib.pyplot as plt
%matplotlib notebook

fig = plt.figure()
ax = fig.add_subplot(111)
ax.errorbar(x, y, fmt='bo', lw=2, xerr=xe, yerr=ye, ecolor='r', elinewidth=1)

import numpy as np 
x = np.linspace(0,2,101) 
y = (x-1)**3+1

import matplotlib.pyplot as plt
%matplotlib notebook

fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(x,y)
ax.annotate('point of inflection at x=1', xy=(1,1), xytext=(0.8,0.5), 
            arrowprops=dict(facecolor='black',width=1, shrink=0.05))

import numpy as np
import matplotlib.pyplot as plt
%matplotlib notebook

fig = plt.figure()
ax = fig.add_subplot(111)
[X,Y] = np.mgrid[-2.5:2.5:51j, -3:3:61j]
Z=X**2-Y**2
curves = ax.contour(X, Y, Z, 12,colors='k')
ax.clabel(curves)
fig.suptitle(r'The level contours of $z=x^2-y^2$', fontsize=20)

import numpy as np
import matplotlib.pyplot as plt
%matplotlib notebook

fig = plt.figure()
ax = fig.add_subplot(111)
[X,Y] = np.mgrid[-2.5:2.5:51j, -3:3:61j]
Z=X**2-Y**2
im = ax.contourf(X, Y, Z, 12)
fig.colorbar(im, orientation='vertical')

ax.set_title(r'${\rm The\ level\ contours\ of\:} z=x^2-y^2$', fontsize=20)

import numpy as np
import matplotlib.pyplot as plt
%matplotlib notebook

from matplotlib import rc
rc('font',family='serif')
rc('text',usetex = True)

fig = plt.figure()
ax = fig.add_subplot(111)
[X,Y] = np.mgrid[-2.5:2.5:51j, -3:3:61j]
Z=X**2-Y**2
im = ax.imshow(Z)
fig.colorbar(im, orientation='vertical')

ax.set_title(r'The level contours of $z=x^2-y^2$', fontsize=20)

import numpy as np
import matplotlib.pyplot as plt
%matplotlib notebook 

x = np.linspace(0, 2*np.pi, 301)
y = np.cos(x)
z = np.sin(x)

plt.ion()
fig1 = plt.figure()
ax1 = fig1.add_subplot(111) 
ax1.plot(x,y) # First figure 
fig2 = plt.figure()
ax2 = fig2.add_subplot(111) 
ax2.plot(x,z) # Second figure

import numpy as np
import matplotlib.pyplot as plt
%matplotlib notebook

x = np.linspace(0, 5, 101)
y1 = 1.0/(x + 1.0)
y2 = np.exp(-x)
y3 = np.exp(-0.1*x**2)
y4 = np.exp(-5*x**2)
    
fig = plt.figure()
ax1 = fig.add_subplot(2, 2, 1)
ax1.plot(x, y1)
ax1.set_xlabel('x')       
ax1.set_ylabel('y1')
ax2 = fig.add_subplot(222)
ax2.plot(x, y2)
ax2.set_xlabel('x')
ax2.set_ylabel('y2')
ax3 = fig.add_subplot(223)
ax3.plot(x, y3)
ax3.set_xlabel('x')
ax3.set_ylabel('y3')
ax4 = fig.add_subplot(224)
ax4.plot(x, y4)
ax4.set_xlabel('x')
ax4.set_ylabel('y4')
fig.suptitle('Various decay functions')

import numpy as np 
from time import time

# Set the parameters
max_iter = 256            # maximum number of iterations
nx, ny = 1024, 1024       # x- and y-image resolutions
x_lo, x_hi = -2.0, 1.0    # x bounds in complex plane
y_lo, y_hi = -1.5, 1.5    # y bounds in complex plane

start_time = time()

# Construct the two dimensional arrays
ix, iy = np.mgrid[0:nx, 0:ny] 
x, y = np.mgrid[x_lo:x_hi:1j*nx, y_lo:y_hi:1j*ny]
c = x + 1j*y
esc_parm = np.zeros((ny, nx, 3),dtype='uint8') # holds pixel rgb data

# Flattened arrays
nxny = nx*ny
ix_f = np.reshape(ix,nxny)
iy_f = np.reshape(iy,nxny)
c_f = np.reshape(c,nxny)
z_f=c_f.copy()                          # the iterated variable

for iter in xrange(max_iter):           # do the iterations
    if not len(z_f):                    # all points have escaped
        break
    # rgb values for this choice of iter
    n = iter + 1
    r, g, b = n % 4 * 64, n % 8 * 32, n % 16 * 16
    
    # Mandelbrot evolution
    z_f *= z_f
    z_f += c_f
    escape=np.abs(z_f) > 2.0          # points which are escaping
    # Set the rgb pixel value for the escaping points
    esc_parm[iy_f[escape], ix_f[escape], :] = r, g, b
    escape = -escape                   # points not escaping
    # Remove batch of newly escaped points from flattened arrays
    ix_f = ix_f[escape]
    iy_f = iy_f[escape]
    c_f = c_f[escape]
    z_f = z_f[escape]

print "Time taken = ", time() - start_time

from PIL import Image

picture = Image.fromarray(esc_parm)
picture.show()
picture.save("mandelbrot.jpg")

import numpy as np
import matplotlib.pyplot as plt 
%matplotlib notebook
from matplotlib.widgets import Slider, Button, RadioButtons

def solwave(x, t, c):
    """ Solitary wave solution of the K deV equation.""" 
    return c/(2*np.cosh(np.sqrt(c)*(x-c*t)/2)**2)

fig, ax = plt.subplots()
plt.subplots_adjust(left=0.15, bottom=0.30)
plt.xlabel("x")
plt.ylabel("u")
x = np.linspace(-5.0, 20.0, 1001)
t0 = 5.0
c0 = 1.0
line, = plt.plot(x, solwave(x, t0, c0), lw=2, color='blue')
plt.axis([-5, 20, 0, 2])

axcolor = 'lightgoldenrodyellow'
axtime = plt.axes([0.20, 0.15, 0.65, 0.03], axisbg=axcolor)
axvely = plt.axes([0.20, 0.1, 0.65, 0.03], axisbg=axcolor)

stime = Slider(axtime, 'Time', 0.0, 20.0, valinit=t0)
svely = Slider(axvely, 'Vely', 0.1, 3.0, valinit=c0)

def update(val):
    time = stime.val
    vely = svely.val
    line.set_ydata(solwave(x, time, vely))
    fig.canvas.draw_idle()

svely.on_changed(update)
stime.on_changed(update)
    
resetax = plt.axes([0.75, 0.025, 0.1, 0.04])
button = Button(resetax, 'Reset', color=axcolor, 
                hovercolor='0.975')

def reset(event):
    svely.reset()
    stime.reset()
    
button.on_clicked(reset)

plt.show()

import numpy as np
import matplotlib.pyplot as plt 
%matplotlib inline
from matplotlib import animation
from JSAnimation import IPython_display

def solwave(t, x, c=1):
    """ Solitary wave solution of the K deV equation.""" 
    return c/(2*np.cosh(np.sqrt(c)*(x-c*t)/2)**2)

# Initialization
fig = plt.figure()
ax = plt.axes(xlim=(-5, 20), ylim=(0, 0.6))
line, = ax.plot([], [], lw=2)

t=np.linspace(-10,25,91)
x = np.linspace(-5, 20.0, 101)

def init():
    line.set_data([], [])
    return line,

def animate(i):
    y = solwave(t[i], x)
    line.set_data(x, y)
    return line,

animation.FuncAnimation(fig, animate, init_func=init, 
                            frames=90, interval=30, blit=True)

import numpy as np

# Fix speed
c = 1.0

def solwave(t, x):
    """ Solitary wave solution of the K deV equation.""" 
    return c/(2*np.cosh(np.sqrt(c)*(x-c*t)/2)**2)
 
import matplotlib.pyplot as plt 

def plot_solwave(t, x):
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(x, solwave(t, x))
    ax.set_ylim(0, 0.6*c)
    ax.text(-4, 0.55*c, "t = " + str(t))
    ax.text(-4, 0.50*c, "c = " + str(c))

x = np.linspace(-5, 20.0, 101)
t = np.linspace(-10, 25, 701)

for i in range(len(t)):
    file_name='_temp%05d.png' % i
    plot_solwave(t[i], x)
    plt.savefig(file_name)
    plt.clf()

import os
os.system("rm _movie.mpg")
os.system("/opt/local/bin/ffmpeg -r 25 " +
          " -i _temp%05d.png  -b:v 1800  _movie.mpg")
os.system("rm _temp*.png")

import numpy as np
 
theta = np.linspace(0, 2*np.pi, 401)
a = 0.3    # specific but arbitrary choice of the parameters
m = 11
n = 9
x = (1 + a*np.cos(n*theta))*np.cos(m*theta)
y = (1 + a*np.cos(n*theta))*np.sin(m*theta)
z = a*np.sin(n*theta)

import matplotlib.pyplot as plt
%matplotlib notebook
from mpl_toolkits.mplot3d import Axes3D

plt.ion() 
fig = plt.figure()
ax = Axes3D(fig)
ax.plot(x, y, z, 'g', linewidth=4)
ax.set_zlim3d(-1.0, 1.0)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
ax.set_title('A spiral as a parametric curve', weight='bold', size=16)
#ax.elev, ax.azim = 60, -120
fig.savefig('torus.pdf')

%%writefile torus.py 

import numpy as np

theta = np.linspace(0, 2*np.pi, 401)
a = 0.3 # specific but arbitrary choice of the parameters
m = 11
n = 9
x = (1 + a*np.cos(n*theta))*np.cos(m*theta)
y = (1 + a*np.cos(n*theta))*np.sin(m*theta)
z = a*np.sin(n*theta)

from mayavi import mlab 
mlab.plot3d(x, y, z, np.sin(n*theta),
          tube_radius=0.025, colormap='spectral')
mlab.axes(line_width=2, nb_labels=5)
mlab.title('A spiral wrapped around a torus',size=1.2)

import numpy as np

xx, yy = np.mgrid[-2:2:81j, -3:3:91j]
zz = np.exp(-2*xx**2 - yy**2)*np.cos(2*xx)*np.cos(3*yy) 

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
%matplotlib notebook

plt.ion()

fig = plt.figure()
ax = Axes3D(fig)
ax.plot_surface(xx, yy, zz, rstride=4, cstride=3, color='c',alpha=0.9)
ax.contour(xx, yy, zz, zdir='x',offset=-3.0, colors='black')
ax.contour(xx, yy, zz, zdir='y', offset=4.0, colors='blue')
ax.contour(xx, yy, zz, zdir='z', offset=-2.0)
ax.set_xlim3d(-3.0, 2.0)
ax.set_ylim3d(-3.0, 4.0)
ax.set_zlim3d(-2.0, 1.0)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')

fig.savefig('surf1.pdf')

%%writefile surf2.py
import numpy as np
 
xx, yy = np.mgrid[-2:2:81j, -3:3:91j]
zz = np.exp(-2*xx**2 - yy**2)*np.cos(2*xx)*np.cos(3*yy)

from mayavi import mlab
fig = mlab.figure()
s = mlab.surf(xx, yy, zz, representation='surface')
ax = mlab.axes(line_width=2, nb_labels=5)
#mlab.title('Simple surface plot',size=0.4)

import numpy as np

[u, v] = np.mgrid[-2:2:51j, -2:2:61j] 
x, y, z = u*(1 - u**2/3 + v**2), v*(1 - v**2/3 + u**2), u**2 - v**2

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
plt.ion()

fig = plt.figure()
ax = Axes3D(fig)
ax.plot_surface(x.T, y.T, z.T, rstride=2, cstride=2, color='r',
              alpha=0.2, linewidth=0.5)
ax.elev, ax.azim = 50, -80
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
#ax.set_title(’A parametric surface plot’, # weight=’bold’,size=18)
fig.savefig('surf3.pdf')

%%writefile surf4.py

import numpy as np
[u, v] = np.mgrid[-2:2:51j, -2:2:61j]
x, y, z = u*(1 - u**2/3 + v**2), v*(1 - v**2/3 + u**2), u**2 - v**2

from mayavi import mlab
fig = mlab.figure()
s = mlab.mesh(x, y, z, representation='surface',
             line_width=0.5, opacity=0.5)
ax = mlab.axes(line_width=2, nb_labels=5)

%%writefile Julia.py

import numpy as np

# Set up initial grid
x, y = np.ogrid[-1.5:0.5:1000j, -1.0:1.0:1000j]
z= x + 1j*y
julia = np.zeros(z.shape)
c = -0.7 - 0.4j

# Build the Julia set
for it in range(1,101): 
    z = z**2 + c
    escape = z*z.conj()>4 
    julia += (1/float(it))*escape
    
from mayavi import mlab 
mlab.figure(size = (800,600)) 
mlab.surf(julia, colormap='gist_ncar', warp_scale='auto', vmax=1.5)
mlab.view(15,30,500,[-0.5,-0.5,2.0])
mlab.show()

import sympy as sy
sy.init_printing()
x, y = sy.symbols("x y")
D = (x + y)*sy.exp(x)*sy.cos(y)
D

f, g = sy.symbols("f g", cls=sy.Function)
f(x), f(x, y), g(x)
i,  j = sy.symbols("i j", integer=True)
u, v = sy.symbols("u v", real=True)
i.is_integer, j*j,  (j*j).is_integer, u.is_real

D_s = sy.sympify('(x+y)*exp(x)*cos(y)')
cosdiff = sy.sympify("cos(x)*cos(y) + sin(x)*sin(y)")
D_s, cosdiff

import numpy as np
func = sy.lambdify((x, y), cosdiff, "numpy") 
xn = np.linspace(0, np.pi, 13)
func(xn, 0.0) # should return np.cos(xn)

M = sy.Matrix([[1,x],[y,1]])
V = sy.Matrix([[u],[v]])
M, V, M*V     

M[0,1]= u
M

from sympy.solvers import dsolve
dsolve?

ode1 = f(x).diff(x, 2) + 4*f(x)
sol1 = dsolve(ode1, f(x))
ode1, sol1     

C1, C2 = sy.symbols("C1, C2")
fun = sol1.rhs
fund = fun.diff(x)
fun.subs(x, 0), fund.subs(x, 0)

psol = sol1.subs([(C2, 2), (C1, 0)])
psol

 ode2 = sy.sin(f(x)) + (x*sy.cos(f(x)) + f(x))*f(x).diff(x)
 ode2, dsolve(ode2, f(x))

ode3 = x*f(x).diff(x) + f(x) - sy.log(x)*f(x)**2
ode3, dsolve(ode3)

 ode5 = f(x).diff(x, 2) + (f(x).diff(x))**2/f(x) + f(x).diff(x)/x
 ode5, dsolve(ode5)

 ode6 = x*(f(x).diff(x, 2)) + 2*(f(x).diff(x)) + x*f(x)
 ode6, dsolve(ode6)

ode7 = [f(x).diff(x) - 2*f(x) - g(x), g(x).diff(x) -f(x) - 2*g(x)]
ode7, dsolve(ode7)

ode8 = f(x).diff(x)  - (x + f(x))**2
ode8, dsolve(ode8)

ode9 = g(x).diff(x) -1 - (g(x))**2
dsolve(ode9)

%matplotlib notebook
import sympy.plotting as syp

s1 = syp.plot(sy.sin(x), x, x-x**3/6, x-x**3/6 + x**5/120, (x, -4, 4),
         title='sin(x) and its first three Taylor approximants')

xc = sy.cos(u) + sy.cos(7*u)/2 + sy.sin(17*u)/3
yc = sy.sin(u) + sy.sin(7*u)/2 + sy.cos(17*u)/3
fig = syp.plot_parametric(xc, yc, (u, 0, 2*sy.pi))
fig.save("sympy2.pdf")

x, y = sy.symbols("x y", real=True)
z = x + sy.I*y
w = sy.cos(z**2).expand(complex=True)
wa = sy.Abs(w).expand(complex=True)
syp.plot_implicit(wa**2- 1)

fig = syp.plot3d_parametric_surface((3 + sy.sin(u) + sy.cos(v))*sy.cos(2*u),
                                    (3 + sy.sin(u) + sy.cos(v))*sy.sin(2*u),
                                    2*sy.cos(u) + sy.sin(v),
                                    (u, 0, 2*sy.pi), (v, 0, 2*sy.pi))
fig.save("foo.pdf")

 import numpy as np
 import matplotlib.pyplot as plt 
 from scipy.integrate import odeint
 %matplotlib notebook

 def rhs(Y, t, omega):
      y, ydot = Y
      return ydot, -omega**2*y

 t_arr = np.linspace(0, 2*np.pi, 101)
 y_init = [1, 0]
 omega = 2.0
 y_arr=odeint(rhs, y_init, t_arr, args=(omega,))
 y, ydot = y_arr[:, 0], y_arr[:, 1]

 fig = plt.figure()
 ax1 = fig.add_subplot(121)
 ax1.plot(t_arr, y, t_arr, ydot)
 ax1.set_xlabel('t')
 ax1.set_ylabel('y and ydot')
 ax2 = fig.add_subplot(122)
 ax2.plot(y, ydot)
 ax2.set_xlabel('y')
 ax2.set_ylabel('ydot')
 plt.suptitle("Solution curve when omega = %5g" % omega)
 fig.tight_layout()
 fig.subplots_adjust(top=0.90)

%%writefile traj.py

import numpy as np
import matplotlib.pyplot as plt 
from scipy.integrate import odeint

def rhs(Y, t, omega):
    y,ydot=Y
    return ydot, -omega**2*y

t_arr = np.linspace(0, 2*np.pi, 101)
y_init = [1, 0]
omega = 2.0

fig = plt.figure()
y, ydot = np.mgrid[-3:3:21j, -6:6:21j]
u, v = rhs(np.array([y, ydot]), 0.0, omega)
mag = np.hypot(u, v)
mag[mag==0] = 1.0
plt.quiver(y, ydot, u/mag, v/mag, color='red')

# Enable drawing of arbitrary number of trajectories
print "\n\n\nUse mouse to select each starting point"
print "Timeout after 30 seconds"
choice = [(0, 0)]
while len(choice) > 0:
    y01 = np.array([choice[0][0], choice[0][1]])
    y = odeint(rhs, y01, t_arr, args=(omega,))
    plt.plot(y[:, 0], y[:, 1], lw=2)
    choice = plt.ginput()
print "Timed out!"

import numpy as np
from scipy.integrate import odeint

def rhs(y, t, mu):
    return [ y[1], mu*(1-y[0]**2)*y[1] - y[0]]

def jac(y, t, mu):
    return [ [0, 1], [-2*mu*y[0]*y[1]-1, mu*(1-y[0]**2)] ]

mu = 1.0
t_final = 15.0 if mu < 10 else 4.0*mu
n_points = 1001 if mu < 10 else 1001*mu
t = np.linspace(0,t_final,n_points)
y0 = np.array([2.0, 0.0])
y, info = odeint(rhs, y0, t, args=(mu,), Dfun=jac, full_output=True)

print " mu = %g, number of Jacobian calls is %d" % \
        (mu, info['nje'][-1])

import matplotlib.pyplot as plt 
%matplotlib notebook
plt.plot(y[:,0], y[:,1])

import numpy as np
from scipy.integrate import odeint

def rhs(u, t, beta, rho, sigma):
    x, y, z = u
    return [sigma*(y-x), rho*x-y-x*z, x*y-beta*z]

sigma = 10.0
beta = 8.0/3.0
rho1 = 29.0
rho2 = 28.8

u01 = [1.0, 1.0, 1.0]
u02 = [1.0, 1.0, 1.0]

t = np.linspace(0.0, 50.0, 10001)
u1 = odeint(rhs, u01, t, args=(beta,rho1,sigma))
u2 = odeint(rhs, u02, t, args=(beta,rho2,sigma))

x1, y1, z1 = u1[:, 0], u1[:, 1], u1[:, 2]
x2, y2, z2 = u2[:, 0], u2[:, 1], u2[:, 2]

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
%matplotlib notebook

plt.ion()
fig = plt.figure()
ax=Axes3D(fig)
ax.plot(x1, y1, z1, 'b-')
ax.plot(x2, y2, z2, 'r:')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
ax.set_title('Lorenz equations with rho = %g, %g' % (rho1,rho2))

import numpy as np
import scikits.bvp1lg.colnew as colnew

degrees = [2]
boundary_points = np.array([0,np.pi])
tol = 1.0e-8*np.ones_like(boundary_points)

def fsub(x, Z):
    """The equations"""
    u, du = Z 
    return np.array([-u])

def gsub(Z):
    """The boundary conditions"""
    u, du = Z
    return np.array([u[0], du[1]-1.0])

solution = colnew.solve(boundary_points, degrees, fsub, gsub,
                        is_linear=True, tolerances=tol,
                        vectorized=True, maximum_mesh_size=300)
x = solution.mesh
u_exact = -np.sin(x)

import matplotlib.pyplot as plt
%matplotlib notebook
plt.ion()

fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(x,solution(x)[:,0],'b.')
ax.plot(x,u_exact,'g-')

import numpy as np
import scikits.bvp1lg.colnew as colnew

degrees = [2, 1, 1]
boundary_points = np.array([0.0, 0.0, 1.0, 1.0])
tol = 1.0e-5*np.ones_like(boundary_points)

def fsub(x, Z):
    """The equations"""
    u, du, v, w = Z
    ddu = -u*v
    dv = np.zeros_like(x)
    dw = u*u
    return np.array([ddu, dv, dw])

def gsub(Z):
    """The boundary conditions"""
    u, du, v, w = Z
    return np.array([u[0], w[1], u[2], w[3]-1.0])


guess_lambda = 100.0
def guess(x):
    u = x*(1.0/3.0 - x)*(2.0/3.0 - x)*(1.0 - x)
    du = 2.0*(1.0 - 2.0*x)*(1.0 - 9.0*x + 9.0*x*x)/9.0
    v = guess_lambda*np.ones_like(x)
    w = u*u
    Z_guess = np.array([u, du, v, w])
    f_guess = fsub(x, Z_guess)
    return Z_guess, f_guess


solution=colnew.solve(boundary_points, degrees,fsub,gsub,
                      is_linear=False, initial_guess=guess,
                       tolerances=tol, vectorized=True,
                       maximum_mesh_size=300)
x= solution.mesh
u_exact = np.sqrt(2)*np.sin(3.0*np.pi*x)


# plot solution

import matplotlib.pyplot as plt
%matplotlib notebook

plt.ion()
plt.plot(x, solution(x)[:, 0], 'b.', x, u_exact, 'g-')
print "Third eigenvalue is %16.10e ." % solution(x)[0, 2]