#!/usr/bin/env python

# Code to produce the figure in 
# http://planetmath.org/encyclopedia/ExampleOfSolvingTheHeatEquation.html
#
# To run this, you need Python 2.x, and install the
# numpy and matplotlib packages for that language.

from numpy.core import *
from numpy.lib import linspace
import pylab
import matplotlib
import matplotlib.axes3d

# The sequence 2m+1, m=0,1,... in the infinite series solution
M = arange(1, 80, 2)
M.shape = (-1, 1, 1)

# Discretized points for x and y
x = linspace(0, pi, 150)
y = linspace(0, pi, 150)
X,Y = pylab.meshgrid(x,y)
x.shape = (1,1,-1)
y.shape = (1,-1,1)

# Infinite series solution
#u = (4.0/pi)*sum(sin(M*x)*cosh(M*y)/(M*cosh(M*pi)), 0)
# Expand the exponentials in the cosh function, and cancel factors,
# to avoid numerical problems with huge numbers
u = (4.0/pi)*sum(sin(M*x)/M * \
                 ((exp(M*(y-pi)) + exp(M*(-y-pi)))/(1.0 + exp(-2*M*pi))), 0)

# Set to use TeX for typesetting
golden_mean = (sqrt(5)-1.0)/2.0
pylab.rcParams.update( \
  {'backend': 'ps',
   'ps.usedistiller': 'xpdf',
   'axes.labelsize': 10,
   'text.fontsize': 10,
   'xtick.labelsize': 8,
   'ytick.labelsize': 8,
   'figure.figsize': [ 7.0, golden_mean*7.0 ],
   'text.usetex': True })

# The surface plot
fig_surface = pylab.figure()
ax = matplotlib.axes3d.Axes3D(fig_surface)
ax.plot_surface(X,Y,u)
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.set_zlabel('$u(x,y)$')

# The color temperature plot
fig_color = pylab.figure()
ax = pylab.gca()
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
ax.imshow(u, interpolation='bilinear', origin='lower', 
             cmap=matplotlib.cm.hot, extent=(0,pi,0,pi))

# Show figures
pylab.show()

# Save figures to EPS format
fig_surface.savefig('heat-surface.eps')
fig_color.savefig('heat-color.eps')