44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib
|
|
from matplotlib.animation import FuncAnimation
|
|
import ast
|
|
|
|
wave_arr = []
|
|
fig = plt.figure()
|
|
ax = plt.gca()
|
|
img = ax.imshow([[]])
|
|
|
|
def plot():
|
|
with open("test.txt") as f:
|
|
lines = f.readlines();
|
|
size = int(lines[0])
|
|
for line in lines[1:]:
|
|
arr = line.strip().split("\t")
|
|
arr = np.asarray(list(map(lambda x: ((a := complex(*ast.literal_eval(x)))*a.conjugate()).real, arr)))
|
|
|
|
# print(sum(arr))
|
|
arr = arr.reshape(size,size)
|
|
wave_arr.append(arr.T)
|
|
# print(arr)
|
|
|
|
# plt.imshow(arr, cmap="hot", interpolation="nearest")
|
|
# plt.show()
|
|
|
|
|
|
def animation(i):
|
|
norm = matplotlib.cm.colors.Normalize(vmin=0, vmax=np.max(wave_arr[i]))
|
|
img.set_norm(norm)
|
|
img.set_data(wave_arr[i])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
plot()
|
|
|
|
norm = matplotlib.cm.colors.Normalize(vmin=0, vmax=np.max(wave_arr[0]))
|
|
img = ax.imshow(wave_arr[0], extent=[0,1,0,1], cmap=plt.get_cmap("viridis"), norm=norm)
|
|
anim = FuncAnimation(fig, animation, interval=1, frames=np.arange(0,len(wave_arr)), repeat=True, blit=0)
|
|
# plt.show()
|
|
|
|
anim.save("./animation.mp4", writer="ffmpeg", bitrate=10000, fps=15)
|