2D Linear Diffusion

The Heat Equation

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from mpl_toolkits.mplot3d import Axes3D

def twod_diffusion(nt, nx, ny, tmax, xmax, ymax, nu):
    """
    Returns the velocity field and distance for 2D diffusion
    """
    # Increments
    dt = tmax/(nt-1)
    dx = xmax/(nx-1)
    dy = ymax/(ny-1)
    
    # Initialise data structures
    u = np.zeros(((nx,ny,nt)))
    v = np.zeros(((nx,ny,nt)))
    x = np.zeros(nx)
    y = np.zeros(ny)

    # Boundary conditions
    u[0,:,:] = u[nx-1,:,:] = u[:,0,:] = u[:,ny-1,:] = 1
    v[0,:,:] = v[nx-1,:,:] = v[:,0,:] = v[:,ny-1,:] = 1

    # Initial conditions
    u[:,:,:] = v[:,:,:] = 1
    u[int((nx-1)/4):int((nx-1)/2),int((ny-1)/4):int((ny-1)/2),0] = 2
    v[int((nx-1)/4):int((nx-1)/2),int((ny-1)/4):int((ny-1)/2),0] = 2

    # Loop
    for n in range(0,nt-1):
        for i in range(1,nx-1):
            for j in range(1,ny-1):
                u[i,j,n+1] = ( u[i,j,n]+dt*nu*( ( u[i-1,j,n]-2*u[i,j,n]+u[i+1,j,n] ) /dx**2 +
                                          ( u[i,j-1,n]-2*u[i,j,n]+u[i,j+1,n] ) /dy**2 ) )
                v[i,j,n+1] = ( v[i,j,n]+dt*nu*( ( v[i-1,j,n]-2*v[i,j,n]+v[i+1,j,n] ) /dx**2 +
                                          ( v[i,j-1,n]-2*v[i,j,n]+v[i,j+1,n] ) /dy**2 ) )

    # X Loop
    for i in range(0,nx):
        x[i] = i*dx

    # Y Loop
    for j in range(0,ny):
        y[j] = j*dy

    return u, v, x, y

def plot(u,x,y,time,title,label):
    """
    Plots the 2D velocity field
    """
    fig=plt.figure(figsize=(11,7),dpi=100)
    ax=fig.gca(projection='3d')
    ax.set_xlabel('x (m)')
    ax.set_ylabel('y (m)')
    ax.set_zlabel(label)
    ax.set_zlim(1, 2.5)
    X,Y=np.meshgrid(x,y)
    #surf=ax.plot_surface(X,Y,u[:,:,time],rstride=2,cstride=2)
    #surf = ax.plot_surface(X, Y, u[:,:,time], cmap=cm.viridis)
    surf = ax.plot_surface(X, Y, u[:,:,time], rstride=1, cstride=1, cmap=cm.viridis,
        linewidth=0, antialiased=True)
    plt.title(title)
    plt.show()

u,v,x,y = twod_diffusion(151, 51, 51, 0.5, 2.0, 2.0, 0.1)

plot(u,x,y,0,'Figure 1: nu=0.1, nt=151, nx=51, ny=51, t=0sec','u (m/s)')

plot(u,x,y,10,'Figure 2: nu=0.1, nt=151, nx=51, ny=51, t=10sec','u (m/s)')

plot(v,x,y,50,'Figure 3: nu=0.1, nt=151, nx=51, ny=51, t=50sec','v (m/s)')

plot(v,x,y,100,'Figure 3: nu=0.1, nt=151, nx=51, ny=51, t=100sec','v (m/s)')