# Untitled

unknown
python
3 years ago
6.1 kB
2
Indexable
Never
import matplotlib
import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import expm
from scipy import integrate
from scipy.sparse.linalg import spsolve
import scipy.sparse.linalg
from scipy.sparse import linalg as sla
import time
# https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.spsolve.html

def twopBVP(fvec, alpha, beta, L, N):
dx = L/(N+1)
fvec[0] -= alpha/dx**2
fvec[-1] -= beta/dx**2 #(f - b) from lecture notes
Tdx = scipy.sparse.csr_matrix((1/(dx**2))*(np.diag(-2*np.ones(N)) + np.diag(np.ones(N-1),1) +  np.diag(np.ones(N-1),-1) ))
y = spsolve(Tdx, fvec)
return np.array([alpha, *y, beta])

L = 10
alpha = 1
beta = 2
N = 100
x = np.linspace(0, L, N+2)
inx = x[1:-1]
f = lambda x: np.exp(-x)
f = np.vectorize(f)
g = np.exp(-x)

ys = twopBVP(f(inx), g[0], g[-1], L, N )

Nbr = [2**4, 2**6, 2**8, 2**10, 2**12]
err = []
h = []
for i in Nbr:
delx = L/(i+1)
ingrid = np.arange(0, L , delx)
fvec = f(ingrid)
alpha = f(ingrid[0])
beta = f(L)
fvec = np.delete(fvec, 0)
y = twopBVP(fvec, alpha, beta, L, i)
ingrid = np.append(ingrid, L)
err.append(np.linalg.norm(y - f(ingrid) )*np.sqrt(delx))
h.append(1/i)

plt.figure(1)
plt.loglog(h, err, 'm', label='error')
plt.loglog(h, np.multiply(h,h), label='h^2')
plt.title('Error for Custom d^2/dx^2y = f(x,y)')
plt.legend()
plt.grid()
plt.show()

plt.figure(2)
plt.plot(x, ys, 'r--', label='approx.', )
plt.plot(x,g, 'b--', label ='analytical sol.')
plt.title("Approx. vs. analytical sol.")
plt.legend()
plt.grid()
plt.show()

N = 999
alpha = 0
beta = 0
L = 10
x = np.linspace(0, L, N+2)
inx = x[1:-1]

E = 1.9 * 1e11
f = lambda x: ((-25000*x**2 + 25000*L*x)/(E*1e-3 * (3-2*np.cos(np.pi*x/L)**12)))
f = np.vectorize(f)
us = twopBVP(f(inx), alpha, beta, L , N)
print('Deflection at midpoint (x,y):')
np.printoptions(precision = 8) #does not work with correct usage of printoptions.
answr = [inx[500], 1e3*us[500]] #499th step will land exactly on L = 5.00 meters, yet index 500 has min value
print(answr)

plt.figure(1)
plt.plot(x, us, 'r--', label = 'approx.')
plt.ylabel('Deflection (m)')
plt.title('Approximated Solution of The Beam Equation')
plt.legend()
plt.grid()
plt.show()

def eigLaplace(L, N ):
dx = L/(N+1)
c = np.zeros(N)
c[0] = -2
c[1] = 1
A = scipy.linalg.toeplitz(c)
Tdx = (1/(dx**2))*(A)
E = np.zeros((N,N))
E[-1][-2] = 4/3
E[-1][-1] = -1/3
z  = Tdx + (1/dx**2)*E
eigz, eigzvec = scipy.linalg.eig(z)
return [eigz, eigzvec]

def Schrödinger(V, N, x, neigs):
t0 = time.time()
ingrid = x[1:-1]
dx = 1/(N+1)
Tdx = scipy.sparse.diags([1, -2, 1], [-1, 0, 1], shape=(N, N)) / dx**2
T = Tdx - scipy.sparse.csr_matrix(V(ingrid)).multiply(scipy.sparse.identity(N))
# help(sla.eigs)
v, w = sla.eigsh(T, k=neigs, which ='SM')
t1 = time.time()
print(f'time: {t1 - t0}')
z = np.zeros((1, neigs))
S = np.vstack((z, w, z)) #Satisfying Dirichlet-conditions
# print(np.size(S, 0), np.size(S, 1))

return v, S

N = 499
L = 1
eigz, eigzvec = eigLaplace(L,N)
eigenvaluef = lambda k: -1*(np.pi*(0.5 + k))**2

plt.figure(1)
plt.plot(eigzvec[:,np.argsort(abs(eigz))[0]], label = 'k=1')
plt.plot(eigzvec[:,np.argsort(abs(eigz))[1]], label = 'k=2')
plt.plot(eigzvec[:,np.argsort(abs(eigz))[2]], label = 'k=3')
plt.legend()
plt.title("First 3 Eigenvectors")
plt.show()

Z = np.arange(10,500,10)
err = []
h = []
_, (ax1, ax2, ax3) = plt.subplots(3, 1, sharex=True)
for i in Z:

for k in range(1,4):
y = eigLaplace(L, i)
y[0].sort()
err.append(abs(y[0][-1] - eigenvaluef(k)))
h.append(i)

def choice(kvalue):
switch = {
1:  ax1.loglog(h, err, 'm', label = 'k = {k}'),
2:  ax2.loglog(h, err, 'm', label = 'k = {k}'),
3:  ax3.loglog(h, err, 'm', label = 'k = {k}')
}
choice(k)
plt.show()

N = 1000
neigs = 7
V = np.vectorize(V)
xgrid = np.linspace(0,1, N+2)

d, S = Schrödinger(V, N, xgrid, neigs)
# print(f'First {neigs} energy levels: {-1*d}')
# print(f'First {neigs} wave functions: ')
# print(S)

pnorm = 1/N**2
_, (ax1, ax2, ax3) = plt.subplots(3, 1, sharex=True)
ax1.plot(xgrid, V(xgrid), 'k', label = vname)
ax1.legend(loc = 'upper right')
for i in range(neigs):
E = -d[i]
psi = S[:,i]

psi /= pnorm
probDens = np.real(np.conj(psi)*psi)

ax2.plot(xgrid, psi + E)
ax3.plot(xgrid, probDens + E)
psiname = f'$\hat\psi$'
psisquare = f'$|\hat\psi|^2$'
ax1.set(title = f'Smallest {neigs} $\hat\psi$ and $|\hat\psi|^2$ \n in ascending energy magnitude, \n with N = {N} and norm  = ' + r'$N^{-2}$.')
ax1.set(ylabel = '$V$')
ax2.set(ylabel = psiname)
ax3.set(ylabel = psisquare)
ax3.set(xlabel ='$x$')

plt.show()

print("Energy Eigenvalues, V(x) = " + vname)
E = np.sort(-d)
for i,j in enumerate(E):
print("{}: {:.2f}".format(i+1,j))

if __name__ == '__main__':
# task2p2("$800\sin(\pi x)^2$", V = lambda x: 800*np.multiply(np.sin(np.pi*x), np.sin(np.pi*x)) )
# task2p2("$0.5 (x - 0.5)^2$", V = lambda x: 0.5*(x - 0.5)**2)