/
analyze.py
77 lines (61 loc) · 1.8 KB
/
analyze.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
import h5py
import numpy as np
import pylab as py
import matplotlib.animation as animation
xc = [0,5,0,5]
yc = [0,0,5,5]
def get_data(filename, num_walkers = 4):
data_f = h5py.File("test.h5")
walker1_d = data_f["0"]
T = walker1_d.shape[1]
data = np.zeros((num_walkers,2,T))
for i in range(num_walkers):
walker_d = data_f[str(i)]
walker_a = np.array(walker_d)
data[i] = walker_a
return data
def plot_trajectories(data, Tf = None):
py.xlim([-.5,9.5])
py.ylim([-.5,9.5])
py.axvline(xc[1]-.5, color='black', linestyle="--")
py.axhline(yc[2]-.5, color='black', linestyle="--")
if Tf == None:
Tf = -1
for i in range(num_walkers):
x = data[i,0,:Tf]
y = data[i,1,:Tf]
py.plot(x,y)
py.plot(x[0],y[0],'bo')
py.plot(x[-1],y[-1],'b*')
def video(data,Ti = None, Tf = None):
fig, ax = py.subplots()
py.xlim([-.5,9.5])
py.ylim([-.5,9.5])
py.axvline(xc[1]-.5, color='black', linestyle="--")
py.axhline(yc[2]-.5, color='black', linestyle="--")
if Tf == None:
Tf = data.shape[2]-1
if Ti == None:
Ti = 0
lines = []
for i in range(num_walkers):
x = data[i,0,Ti:Ti+1]
y = data[i,1,Tf:Tf+1]
line, = py.plot(x,y, 'o')
lines.append(line)
def animate(i):
py.title(i)
for j,line in enumerate(lines):
line.set_xdata(data[j,0,i:i+1])
line.set_ydata(data[j,1,i:i+1])
return lines
ani = animation.FuncAnimation(fig, animate, np.arange(Ti,Tf), interval=300,
blit=False)
py.show()
if __name__ == "__main__":
num_walkers = 1
filename = "test.h5"
data = get_data(filename, num_walkers)
plot_trajectories(data,20)
video(data)
# video(data, 90, 96)