heat2d_solution

# heat2D_solution.py
#
# Subitop modelling course, Edinburgh 2017
# Jeroen van Hunen, March 2017
# purpose: calculates 2D heat diffusion
# method: explicit time integration

import numpy as np
import pylab as plt

# Subfuctions: 
def twoDdiff (fin,dx,dz,kappa,dt):
# Performs 1 diffusion timestep on array fin of size nx X nz 
# and discr dx and dz using timestep dt and diffusion coeff kappa
# Fixed essential b.c. are used.

    # Calculate 2nd derivatives in x- & z-dir.:
    d2fdx2=np.diff(fin,n=2,axis=1)/dx**2
    d2fdz2=np.diff(fin,n=2,axis=0)/dz**2
    
    # Initialize a timestep df/dt vector:
    dfdt=np.zeros(np.shape(fin))
    dfdt[1:-1,1:-1]=kappa*(d2fdx2[1:-1,:]+d2fdz2[:,1:-1])
    
    # Add dt * df/dt-vector to old solution:
    fout=fin+dt*dfdt
    return fout

# Main code: 
# Initialisation:
# Heat diffusion variables:
kappa    = 1e-6                # thermal diffusivity

# Mesh setup:
h        = 1000e3              # height of box: 1000 km
w        = 1.0*h               # box of aspect ratio 1
dx       = 1e4                 # discretization step in meters
dz       = 1e4
nx       = w/dx+1
nz       = h/dz+1         
dx       = w/(nx-1)            # Adjust requested dx & dz to fit in equidistant grid space
dz       = h/(nz-1) 
x        = np.linspace(0,w,nx) # array for the finite difference mesh
z        = np.linspace(0,h,nz)
[xx,zz]  = np.meshgrid(x,z)

# Time variables:
dt       = 0.2*dx**2/kappa     # timestep in Myrs
nt       = 100                  # number of tsteps
secinmyr = 1e6*365*24*3600     # amount of seconds in 1 Myr
t        = 0                   # set initial time to zero
nplot    = 5                   # plotting interval: plot every nplot timesteps

# Initial condition:
Tm       = 1350                # mantle temperature in degC
Ttop     = 0                   # surface T
Told=Tm*np.ones(np.shape(xx))  # initial T=Tm everywhere ...
                               #    ... except for a slab: 
                                   
Told[(xx>=250e3) & (xx<=750e3) & (zz>=600e3) & (zz<=700e3)]=0.

     
# timestepping
for it in range(1,nt):
    # numerical solution
    Tnew = twoDdiff(Told, dx, dz, kappa, dt)
    
    #update time
    t=t+dt                  # in sec
    
    # plot solution:
    if (it%nplot==0):
        tmyrs=int(t/secinmyr)   # time in Myrs
        plt.clf()
        plt.figure(1)
        plt.imshow(Told, 
                   extent=(0,h,0,h),
                   clim=(0,Tm),
                   interpolation='bilinear', 
                   cmap='jet')
        plt.title('T after '+str(tmyrs)+' Myrs')
        plt.pause(0.00005)
  
    # prepare for next time step:
    Told = Tnew