""" Illustration of the heat equation Solve the heat equation using finite differences and Forward Euler. Based on: https://commons.wikimedia.org/wiki/File:Heat_eqn.gif """ from __future__ import division, print_function import numpy as np from mayavi import mlab import subprocess path_to_convert = "C:\Program Files\ImageMagick-6.9.3\convert.exe" def step_function(N, scale, X, Y, shape="crescent"):  """Function that is 1 on a set and 0 outside of it""" shapes = ["crescent", "cylinder", "hexagon", "superquadric", "smiley"] if shape not in shapes: shape = "crescent" if shape == "cylinder": Z = np.ones_like(X) Z[X**2 + Y**2 < 0.5] = 0 Z[X**2 + Y**2 > 2] = 0 if shape == "superquadric": Z = np.ones_like(X) Z[np.abs(X)**0.5 + np.abs(Y)**0.5 > 1.5] = 0 if shape == "hexagon": Z = np.ones_like(X) hexa = 2*np.abs(X) + np.abs(X - Y*np.sqrt(3)) +\ np.abs(X + Y*np.sqrt(3)) Z[hexa > 6] = 0 if shape == "crescent": c = 2 d = -1 e = 1 f = 0.5 k = 1.2 shift = 10 Z = (c**2 - (X/e - d)**2 - (Y/f)**2)**2 + k*(c + d - X/e)**3 - shift Z = 1 - np.maximum(np.sign(Z), 0) if shape == "smiley": Z = np.ones_like(X) fac = 1.2 x_eye = 0.5 y_eye = 0.4 bicorn = fac**2*(Y + 0.3)**2*(1 - fac**2*X**2) -\ (fac**2*X**2 - 2*fac*(Y + 0.3) - 1)**2 left_eye = (X + x_eye)**2/0.1 + (Y - y_eye)**2/0.4 - 1 right_eye = (X - x_eye)**2/0.1 + (Y - y_eye)**2/0.4 - 1 Z[X**2 + Y**2 > 2] = 0 Z[bicorn > 0] = 0 Z[left_eye < 0] = 0 Z[right_eye < 0] = 0 Z = scale * Z return Z def data_gen(num): # Solve the heat equation with zero boundary conditions for cont in range(ntime_anim): Z[1:N-1, 1:N-1] = Z[1:N-1, 1:N-1] + dt*(Z[2:N, 1:N-1] + Z[0:N-2, 1:N-1] + Z[1:N-1, 0:N-2] + Z[1:N-1, 2:N] - 4*Z[1:N-1, 1:N-1])/dx**2 surf = mlab.surf(X, Y, Z, colormap='autumn', warp_scale=1) # Change the visualization parameters. surf.actor.property.interpolation = 'phong' surf.actor.property.specular = 0.3 surf.actor.property.specular_power = 20 surf.module_manager.scalar_lut_manager.reverse_lut = True surf.module_manager.scalar_lut_manager.data_range = np.array([ 0., scale]) return surf N = 500 # Grid points L = 2.5 # Box size X, Y = np.mgrid[-L:L:N*1j, -L:L:N*1j] scale = 2 Z = step_function(N, scale, X, Y, shape="crescent") CFL = 0.125 dx = X[1, 0] - X[0, 0] dy = dx dt = CFL*dx**2 end_time = 0.05 time = np.arange(0, end_time, dt) nframes = 50 ntime = time.shape[0] ntime_anim = int(ntime/nframes) #%% Plot frames fname = "heat" bgcolor = (1, 1, 1) fig = mlab.figure(size=(1200, 1000), bgcolor=bgcolor) fig.scene.camera.azimuth(180) mlab.get_engine() engine = mlab.get_engine() scene = engine.scenes[0] for cont in range(nframes): mlab.clf() surf = data_gen(cont) scene.scene.camera.position = [-8, -8, 7] scene.scene.camera.clipping_range = [7, 22] scene.scene.camera.focal_point = [0, 0, 1] print(cont) mlab.savefig("{}_{n:02d}.png".format(fname, n=cont)) #%% Generate video args = [path_to_convert, "-delay", "10", "-loop" , "0", fname + "_*.png", fname + ".gif"] subprocess.call(args, shell=True) subprocess.call(["del", "/Q", fname + "*.png"], shell=True) print("Done!")