# A numerical solution of 2D Gaussian advection equation by Upwind Scheme.
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as pl
import numpy as np
t0 = 0.0
dt = 0.045
t_final = 1
x0 = -2.0
dx = 0.04
x_max = 4.0
y0 = -2.0
dy = 0.04
y_max = 4.0
T = np.arange(t0, t_final+dt, dt)
X = np.arange(x0, x_max+dx, dx)
Y = np.arange(y0, y_max+dy, dy)
U = [[0.0 for _ in Y] for _ in X]
fig = pl.figure()
ax = fig.add_subplot(111, projection='3d')
nframe = 0
for t in T:
if t == t0:
for i in range(len(X)):
for j in range(len(Y)):
U[i][j] = np.exp(-X[i]**2-Y[j]**2)
else:
U_prime = [[0.0 for _ in Y] for _ in X]
for i in range(len(X)):
for j in range(len(Y)):
if i == 0 or j == 0 or i == len(X)-1 or j == len(Y)-1:
U_prime[i][j] = np.exp(-(X[i]-(1-np.cos(t)))**2-(Y[j]-np.sin(t))**2)
else:
ux = np.sin(t)
uy = np.cos(t)
if ux > 0 and uy > 0:
U_prime[i][j] = U[i][j] - ux*dt/dx*(U[i][j] - U[i-1][j]) - uy*dt/dy*(U[i][j] - U[i][j-1])
elif ux > 0 and uy <= 0:
U_prime[i][j] = U[i][j] - ux*dt/dx*(U[i][j] - U[i-1][j]) - uy*dt/dy*(U[i][j+1] - U[i][j])
elif ux <= 0 and uy > 0:
U_prime[i][j] = U[i][j] - ux*dt/dx*(U[i+1][j] - U[i][j]) - uy*dt/dy*(U[i][j] - U[i][j-1])
else:
U_prime[i][j] = U[i][j] - ux*dt/dx*(U[i+1][j] - U[i][j]) - uy*dt/dy*(U[i][j+1] - U[i][j])
U = U_prime
ax.clear()
gridX, gridY = np.meshgrid(X, Y)
ax.plot_surface(gridX, gridY, U, cmap = 'summer', lw = 0.0)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('phi')
ax.set_xlim([x0, x_max])
ax.set_ylim([y0, y_max])
ax.set_zlim([0, 1])
fig.savefig('upwind2d/frame_'+"%03d" % nframe+'.png')
nframe += 1
# To generate GIF: convert frame_* GaussianUpwind2D.gif