# Diffusion on the surface of a torus

An example in Chapter 7 of the scipython book describes the numerical solution of the two-dimensional heat equation for a flat plate with edges held at a fixed temperature.

The code below extends this example by considering the solution of the heat equation on the surface of a torus (i.e. with periodic boundary conditions).

In this example, six evenly-spaced rings around the poloidal direction of the torus are kept at a fixed high temperature and the heat allowed to diffuse out to the intervals between them.

Some care is needed in setting the numerical grid on which the numerical integration of the heat equation is carried out: the cell size length is constant around the poloidal direction, but in the toroidal direction is smaller on the inside of the torus than the outside.

The code requires the files torus.py and palettes.py from the previous blog post defining a class for depicting 3D objects in SVG.

import numpy as np
from torus import Torus
from palettes import rgb_to_html

import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.colors import Normalize
Tcool, Thot = 300, 700
norm = Normalize(vmin=Tcool, vmax=Thot)
cmap = plt.get_cmap('hot')
scalarMap = cm.ScalarMappable(norm=norm, cmap=cmap)

# The major and minor radius of the torus.
c, a = 2.5, 1.5
# Pick a colour scheme from the get_colours dictionary
palette_args = '#ff0000', '#ffff00'
# Image dimensions and scaling factors from torus units to image units.
width, height = 800, 600
scalex, scaley = 120, 120
# Tait-Bryan angles for intrinsic rotation of the figure.
alpha, beta, gamma = 90, 35, 235
# Camera position, C, and projection plane position, E (relative to C).
cx, cy, cz = 0, 6, 0
ex, ey, ez = 0, 3, 0
C, E = np.array((cx,cy,cz)), np.array((ex,ey,ez))

def preamble(fo):
"""The SVG preamble and styles."""

print('<?xml version="1.0" encoding="utf-8"?>\n'

'<svg xmlns="http://www.w3.org/2000/svg"\n' + ' '*5 +
.format(width, height), file=fo)

print("""
<defs>
<style type="text/css"><![CDATA[""", file=fo)

#print('path {stroke-width: 0.2px; stroke: #000;}', file=fo)
print('path {stroke-width: 1px; stroke: #000;}', file=fo)

print("""]]></style>
</defs>""", file=fo)

def get_colour(i, u):
"""Get a colour for quad i from the colourmap."""
rgb = scalarMap.to_rgba(u[i], alpha=False, bytes=True)
return rgb_to_html(rgb)

def draw_torus(itr, u):
"""Draw a torus on iteration itr of the diffusion with array u."""
torus = Torus(c, a, ntheta, nphi)
torus.setup_torus()
torus.rotate_xyz(alpha, beta, gamma)

torus.get_perspective_view(C, E)

# Draw the torus as a SVG image and save.
with open('frames/torus-{:04d}.svg'.format(itr), 'w') as fo:
preamble(fo)
for i in torus.idy:
colour = get_colour(i, u.ravel())
print('<path d="M{},{} L{},{} L{},{} L{},{} Z" fill="{}"/>'.format(
print('</svg>', file=fo)

# Thermal diffusivity of steel, mm2.s-1
D = 4.
# grid size of unravelled torus
w = 2 * np.pi * a
h = 2 * np.pi * c
# Number of grid points around the torus (toroidal direction): phi, y.
ny = nphi = 120
# Number of grid points around the torus cross section (poloidal direction):
# theta, x.
nx = ntheta = int(nphi * a/c)
# grid size: the grid length along the x- (theta-) direction is constant.
dx = w / ntheta
# the grid length along the y- (phi-) direction varies with theta
theta = np.linspace(0, 2.*np.pi, ntheta)
dy = 2 * np.pi * (c + a*np.cos(theta)) / nphi
# We need the minimum y-grid spacing to work out the largest dt we can use.
dymin = 2 * np.pi * (c - a) / nphi

dx2, dymin2, dy2 = dx*dx, dymin*dymin, dy*dy
dt = dx2 * dymin2 / (2 * D * (dx2 + dymin2))

# 2D grid to calculate the diffusion on: read from u0 and write to u on
# each iteration.
u0 = Tcool * np.ones((nx, ny))
u = np.empty((nx, ny))

# We define two different heat input functions here. Pick one.
def set_heat_input_spots(u):
"""Set heat input to two spots of radius r at opposite sides of torus."""
r, cx, cy = 1, w//2, h//2
r2 = r**2
for i in range(nx):
for j in range(ny):
p2 = (i*dx-cx)**2 + (j*dy[nx//2]-cy)**2
if p2 < r2:
u[(i+nx//2)%nx,j] = Thot
u[(i+nx//2)%nx,(j+ny//2)%ny] = Thot

def set_heat_input_rings(u, nrings=6):
"""Set heat input to nrings rings, spaced around the torus."""
spacing = ny // nrings
u[:,::spacing] = Thot

set_heat_input = set_heat_input_rings

def do_timestep(u0, u):
# Propagate with forward-difference in time, central-difference in space
set_heat_input(u0)
idx, idy = np.arange(nx), np.arange(ny)
u = u0 + D * dt * (
(u0[(idx+1)%nx,:] - 2*u0 + u0[(idx-1)%nx,:]) / dx2
+ (u0[:,(idy+1)%ny] - 2*u0 + u0[:,(idy-1)%ny]) / dy2[:,None]
)
u0 = u.copy()
return u0, u

# Number of timesteps
nsteps = 500
fig = plt.figure()
for m in range(nsteps):
u0, u = do_timestep(u0, u)
draw_torus(m+1, u)

Current rating: 5