Diffusion on the surface of a torus

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 = 'depth gradient'
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 +
       'xmlns:xlink="http://www.w3.org/1999/xlink" width="{}" height="{}" >'
            .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)
    torus.get_quads()

    # 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:
            quad = torus.quads[i] * (scalex, scaley) + (width/2, height/2)
            colour = get_colour(i, u.ravel())
            print('<path d="M{},{} L{},{} L{},{} L{},{} Z" fill="{}"/>'.format(
                *quad.ravel(), colour), file=fo)
        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)