heat1Dq_solution

import pylab as plt
from scipy.special import erfc

# Subfuctions: 
def oneDdiff (fin,dz,kappa,dt):
    # Performs 1 diffusion timestep on array fin of size nz and discr dz
    # using timestep dt and diffusion coeff kappa
    # fixed essential b.c. used.
    
    # Copy old solution into new one:
    fout = np.array(fin)
    
    # 2 ways to calculate internal d2f/dx2:
    fout[1:-1] = fout[1:-1] + kappa*dt*(fin[2:] - 2*fin[1:-1] + fin[0:-2]) / dz**2
    # df[1:-1]=kappa*dt*np.diff(fin,n=2)/dz**2
    
    # Zero basal heatflux boundary condition:
    fout[-1]   = fout[-1] + 2*kappa*dt*(fin[-2]-fin[-1])/dz**2

    return fout

def halfspacecooling (Tm, z, kappa, t):
    # Calculates Halfspace cooling solution:
    fout = np.array(Tm-Tm*erfc(z/(2*np.sqrt(kappa*t))))
    return fout

# Main code: 
# Initialisation:
# Time variables:
dt       = 0.15                # timestep in Myrs
tmax     = 100
nt       = int(tmax/dt)+1      # number of tsteps to reach tmax Myrs
secinmyr = 1e6*365*24*3600     # amount of seconds in 1 Myr
dt       = dt*secinmyr         # unit conversion to SI: time in sec
time     = np.zeros(nt)
t        = 0                   # set initial time to zero
nplot    = 20                  # plotting interval: plot every nplot timesteps

# Mesh setup:
h        = 1.5e5                 # height of box: 3x10^5 m = 300 km
dz       = 1e4                 # discretization step in meters
nz       = h/dz+1         
dz       = h/(nz-1)            # Adjust reqested dz to fit in equidistant grid space
z        = np.linspace(0,h,nz) # array for the finite difference mesh

# Heat diffusion variables:
kappa    = 1e-6                # thermal diffusivity
Tm       = 1350                # mantle temperature in degC
Ttop     = 0                   # surface T
Told     = Tm*np.ones(nz)      # initial T=Tm everywhere ...
Told[0]  = Ttop;               # ... except at surface, where T=0

# Timestepping
for it in range(1,nt):
    #update time
    t=t+dt 
    time[it]=t

    # numerical solution
    Tnew = oneDdiff(Told, dz, kappa, dt)
                   
    # analytical solution
    Tana = halfspacecooling (Tm, z, kappa, t)

    if (it%nplot==0):
        # plot solution:
        plt.clf()
        plt.figure(1)
        plt.plot (Tnew,-z,'b')
        plt.plot (Tana,-z,'r--')
        plt.xlabel('T [^oC]')
        plt.ylabel('z [m]')
        tmyrs=round(t*10/secinmyr)/10
        plt.title(' T after '+str(tmyrs)+' Myrs')
        plt.pause(0.0005)

    # prepare for next time step:
    Told = Tnew