scientific-programming-exer.../exam/ex24/main.py

65 lines
1.2 KiB
Python
Raw Permalink Normal View History

2019-02-19 13:13:46 +00:00
import sys
sys.path.append("../../")
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from util.io import readvalue
def positive_int(s):
i = int(s)
if(not i > 0):
raise ValueError("{} <= 0".format(i))
return i
def prepare_A(n):
A = -2 * np.identity(n)
A += np.roll(np.identity(n), 1, axis=0)
A += np.roll(np.identity(n), -1, axis=0)
return A
N = readvalue("N > ", positive_int)
M = readvalue("M > ", positive_int)
kappa = 0.05
A = prepare_A(N)
eigenvalues, eigenvectors = np.linalg.eigh(A)
#print(*zip(*enumerate(eigenvalues)))
plt.scatter(*zip(*enumerate(eigenvalues)))
plt.savefig("output/eigenvalues.png")
plt.clf()
for n in range(N):
plt.scatter(*zip(*enumerate(eigenvectors[n])))
plt.savefig("output/eigenvector_{}.png".format(n))
plt.clf()
v = np.zeros(N)
v[0] = 1
fig, ax = plt.subplots()
max_eigv, = ax.plot(*zip(*enumerate(eigenvectors[0])), "bo")
scatter, = ax.plot(*zip(*enumerate(v)), "ro")
def on_tick(data):
global v
v = (np.identity(N) + kappa*A).dot(v)
scatter.set_data(*zip(*enumerate(v)))
return scatter
def run():
for i in range(M):
yield i
ani = animation.FuncAnimation(fig, on_tick, run, interval=1)
ani.save("output/animation.gif", dpi=80, writer='imagemagick')