# Untitled

unknown
python
3 years ago
2.7 kB
4
Indexable
Never
```import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator
import scipy.sparse as ssp

def matrix(N):
dx = 1/(N+1)
matrix = (1/dx**2)*(np.diag(np.ones(N-1), -1) + np.diag(-2*np.ones(N)) + np.diag(np.ones(N-1), 1))
return matrix

def eulerstep(Tdx, uold, dt):
unew = uold + dt*np.dot(Tdx,uold)
return unew

def trapztep(Tdx, uold, dt):
unew = np.linalg.solve(np.eye(len(Tdx)) - np.multiply(Tdx, dt/2), uold + dt/2*np.dot(Tdx,uold))
return unew

def step(choice, Tdx, uold, dt):
if choice == 0:
return eulerstep(Tdx, uold, dt)
if choice == 1:
return trapztep(Tdx, uold, dt)

def plot1(m, T, X):
fig = plt.figure()
ax = plt.axes(projection = '3d')
wireframe = ax.plot_surface(T,X, np.transpose(m), cmap = cm.inferno)
ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter('{x:.02f}')
fig.colorbar(wireframe, shrink = 0.5, aspect = 5)
ax.set_xlabel('T')
ax.set_ylabel('X')
ax.set_zlabel('u(t,x)')
ax.set_title('Solution to Diffusion Equation Custom \$u(0,x) = g(x)\$')
plt.show()

def task1(N , M, g, t0, tf, choice):
# Next line should compute a 2-dimensional vector u(m∆t,n∆x)
# u(t,x) has u(t,0) = u(t,1) = 0
# u(0,x) = g(x)
dt = (tf-t0)/M
xx = np.linspace(0,1,N+2)
tt = np.linspace(t0,tf,M+1)
T,X = np.meshgrid(tt, xx)
uold = [g(i, i) for i in range(1, N+1)]
m = []
m.append(np.array([0, *uold, 0]))
for _ in range(1, np.size(tt)):
uold = step(choice, Tdx = matrix(N), uold = uold, dt = dt )
m.append(np.array([0, *uold, 0]))
m = np.array(m)
return m, T, X

if __name__ == '__main__':
c = 1e-6
N = 100
tstart = 0
tend= 0.6
t = tend - tstart
M_CFL =  2*t*(N+1)**2 # CFL-condition dt/dx**2 <= 1/2 \iff t/M * (N+1)^2
chc = 1
if chc == 0:
Mnbr = int(np.ceil(M_CFL))
if chc == 1:
Mnbr = 10*N
# m, T, X = task1(N = 100, M = Mnbr, g = lambda x,t: c*(x)**3 + c/((t)**3), t0= tstart, tf=tend, choice = chc )
# m, T, X = task1(N = N, M = Mnbr, g = lambda x,t: c*(x)**3 + c/((t)**3), t0= tstart, tf=tend, choice = chc )
# m, T, X = task1(N = N, M = Mnbr-10, g = lambda x,t: c*(x)**3 + c/((t)**3), t0= tstart, tf=tend, choice = chc) #BREAK IT!
m, T, X = task1(N = N, M = Mnbr, g = lambda x,t: np.exp(x /t), t0= tstart, tf = tend, choice = chc )
# m, T, X = task1(N = N, M = Mnbr-16, g = lambda x,t: np.exp(x /t), t0= tstart, tf = tend, choice = chc) #BREAK IT AND GET PANTS!
plot1(m = m, T = T, X = X)

```