Introduction
heat.py
"""
file: heat.py
Based on steady-state.py.
This defines not only the steady state model but also the PyGame code necessary to
visualize it using Surfarray.
"""
import sys
import math
import time
import Numeric
import pygame
from pygame.locals import *
from pygame import surfarray
def drawfield( screen, scale_surface, pixels ):
surfarray.blit_array( scale_surface, pixels )
temp = pygame.transform.scale(scale_surface, screen.get_size())
screen.blit(temp, (0,0))
return
def createPalette( ):
r,g,b = -1,0,256
palette = []
for i in xrange(0,256):
r += 1
b -= 1
palette.append( (r,g,b) )
return palette
def main( N=100, EPSILON = 0.01 ):
# Stores Old(previous timestep) Values
u = Numeric.zeros( (N,N),Numeric.Float32 )
# New(current/next timestep) Values
w = Numeric.zeros( (N,N),Numeric.Float32 )
pixels = Numeric.zeros( (N,N), Numeric.Int32 )
# Set boundary values and compute mean boundary value
mean = 0.0
# This "magic" factor is just a scaling to get the full range of our (255.0/N)
# color palette.
for i in xrange(N):
# 3 sides @ 100.0 degrees
u[i][0] = 100.0
pixels[i][0] = int(100 * (255.0/N))
u[i][N-1] = 100.0
pixels[i][N-1] = int(100 * (255.0/N))
u[0][i] = 100.0
pixels[0][i] = int(100 * (255.0/N))
# 1 side @ 0 degrees
for i in xrange(N):
u[N-1][i] = 0.0
pixels[N-1][i] = 0.0
mean += Numeric.sum(u[:][0])
mean += Numeric.sum(u[:][N-1])
mean += Numeric.sum(u[0][:])
mean += Numeric.sum(u[N-1][:])
mean /= (4.0 * N)
# Initialize interior values:
print mean
for i in xrange(1,N-1):
for j in xrange(1,N-1):
u[i][j] = mean
pixels[i][j] = int(mean * (255.0/N))
WINSIZE = 640,480
ARRAYSIZE = N,N
# Initialize the Pygame Engine!
pygame.init()
screen = pygame.display.set_mode(WINSIZE,0,8)
scale_screen = pygame.surface.Surface( ARRAYSIZE,0,8 )
pygame.display.set_caption('Heat')
black = 20,20, 40
palette = createPalette()
screen.fill(black)
scale_screen.fill(black)
screen.set_palette( palette )
scale_screen.set_palette( palette )
# Compute Steady-State solution:
done = False # Is true when we reach steady state
userquit = False # is true only when the user is done watching
iterations = 0
while not userquit:
iterations += 1
delta = 0.0
if not done:
for i in xrange(1,N-1):
for j in xrange(1,N-1):
w[i][j] = u[i-1][j] + u[i+1][j] + u[i][j-1] + u[i][j+1]
w[i][j] /= 4.0
d = math.fabs( w[i][j] - u[i][j] )
if( d > delta ):
delta = d
if( delta <= EPSILON ):
done = True
# Copy new interior state to old:
for i in xrange(1,N-1):
for j in xrange(1,N-1):
u[i][j] = w[i][j]
pixels[i][j] = int(w[i][j] * (255.0/N))
# Draw
drawfield( screen, scale_screen, pixels )
pygame.display.update()
# Handle Events
events = pygame.event.get( )
for e in events:
if( e.type == QUIT ):
userquit = True
break
elif (e.type == KEYDOWN):
if( e.key == K_ESCAPE ):
userquit = True
break
if( e.key == K_f ):
pygame.display.toggle_fullscreen()
# Print Solution:
print u
return iterations
if __name__=="__main__":
print "Starting Steady State Example:"
if( len(sys.argv) >= 2 ):
N = int( sys.argv[1] )
else:
N = 100
start = time.time()
iterations = main(N)
end = time.time()
print "Finished Steady State Example in",end - start,"and",iterations,"iterations."
syntax highlighted by Code2HTML, v. 0.9.1
Download
Download this example: heat.py
