numpy arrays and methods are incredibly helpful. They are usually optimized and much faster than looping in python. Always look for a way to use an existing numpy method for your application.
np.roll()
will allow you to shift and then you just add.
I learned to use convolve()
from comments on How to np.roll() faster?. I haven't checked if this is faster or not, but it may depend on the number of dimensions. See this answer for a 2D relaxation of the Laplace equation (electrostatics, a different problem)
For this kind of relaxation you'll need a bounding box, so the boolean do_me
is False
on the boundary.
I know that for Jacobi relaxation solutions to the Laplace equation, there are two speed-up methods. I don't know if they can be extended to solving the Heat Diffusion equation, but I'm sure something can be done:
Multigrids; solve on a coarse (fast) grid, then interpolate to a fine grid and iterate a little longer
Successive Overrelaxation; "overshoot" the averaging by changing the values by an amount larger than 100% of what you would have with simple averaging. Try values like 110% to 150%, it can go unstable if they are too large.
import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage import convolve
T0 = np.zeros(50, dtype=float)
T0[17:25] = 1.
T1 = T0.copy() # method 1 np.roll()
T2 = T0.copy() # method 2 convolve()
do_me = np.ones_like(T0, dtype=bool)
do_me[[0, -1]] = False # keep the boundaries of your bounding box fixed
a = 0.01
hold_1 = [T0.copy()]
for i in range(10001):
Ttemp = T1 + a*(np.roll(T1, +1) + np.roll(T1, -1) - 2*T1)
T1[do_me] = Ttemp[do_me]
if not i%1000:
hold_1.append(T1.copy())
hold_2 = [T0.copy()]
kernel = np.array([a, (1 - 2.*a), a])
for i in range(10001):
Ttemp = convolve(T2, kernel)
T2[do_me] = Ttemp[do_me]
if not i%1000:
hold_2.append(T2.copy())
if True:
plt.figure()
plt.subplot(1, 2, 1)
for thing in hold_1:
plt.plot(thing)
plt.subplot(1, 2, 2)
for thing in hold_2:
plt.plot(thing)
plt.show()