Script for solving the 1D diffusion equation

We will solve for the cooling of hot dike that was emplaced into cooler country rock. The method we are using is a simple explicit 1D finite differences scheme.

First we load some useful libraries

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from IPython.display import HTML
from matplotlib.patches import Rectangle

Next we define some physics constants

L       = 100           # Length of modeled domain [m]
Tmagma  = 1200          # Temperature of magma [C]
Trock   = 300           # Temperature of country rock [C]
kappa   = 1e-6          # Thermal diffusivity of rock [m2/s]
W       = 5;            # Width of dike [m]
day     = 60*60*24       # # seconds per day
dt      = 1*day         # Timestep [s]

and some numerical constants constants

nx      = 201                                          # Number of gridpoints in x-direction
nt      = 500                                          # Number of timesteps to compute
Xvec,dx = np.linspace(-L/2, L/2, nx,  retstep=True)    # X coordinate vector, constant spacing
beta    = kappa*dt/dx**2

Set the initial conditions

T_init  = np.ones(nx)*Trock;              # everything is cold initially
T_init[np.nonzero(np.abs(Xvec) <= W/2)] = Tmagma   # and hot where the dike is
time    = 0

Make a function that computes temperature

# hide
# only store the latest time step solution to Told
def fdm_solve(Told):
    Tnew=Told.copy()
    for i in range(1,len(Xvec)-1):
        Tnew[i] = beta * (Told[i+1] - 2*Told[i] + Told[i-1]) + Told[i]
    return Tnew

# write a function that solves for temperature!

#def fdm_solve(Told):
#    Tnew=Told.copy()
#    for i in range(1,len(Xvec)-1):
#        Tnew[i] = ???
#    return Tnew

Compute and visualize everything

# First set up the figure, the axis, and the plot element we want to animate
fig = plt.figure(figsize=(10,5))
ax = plt.axes(xlim=(-L/2, L/2), ylim=(0, Tmagma))
ax.add_patch(Rectangle((-W/2, 0), W, 1200, alpha=0.5, color='tab:red'))
line, = ax.plot([], [], lw=1)
timeLabel=ax.text(0.02,0.98,'Time: ',transform=ax.transAxes,va='top')
ax.set_xlabel('X (m)')
ax.set_ylabel('Temperature ($^{\circ}$C)')

# Initialization function: plot the background of each frame
def init():
    line.set_data([], [])
    return line,

# Initialize Tnew
Tnew=T_init
# Animation function which updates figure data.  This is called sequentially
def animate(i):
    timeLabel._text='Time: %.1f day'%(i*dt/day)

    # use global keyword to store the latest solution and update it using fdm_solve function
    global Tnew
    Tnew=fdm_solve(Tnew)
    line.set_data(Xvec, Tnew)

    return line,

# Call the animator.  blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=nt, interval=30, blit=True)

plt.close(anim._fig)

# Call function to display the animation
#HTML(anim.to_html5_video())  # lower resolution
HTML(anim.to_jshtml())  # higher resolution

Once Loop Reflect