Skip to content

Commit 749fde5

Browse files
author
corrooli
committed
Cleaned up TODOs, refactor AnimationPlotter
1 parent dbace8f commit 749fde5

14 files changed

+309
-412
lines changed

‎common/animation_plotter.py

+258
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,258 @@
1+
from common.parameters import SimulationParameters
2+
3+
import numpy as np
4+
from matplotlib.animation import FuncAnimation, FFMpegWriter
5+
import matplotlib.pyplot as plt
6+
from datetime import datetime
7+
8+
9+
class AnimationPlotter():
10+
'''
11+
Plotting class with FFMPEG video support.
12+
Caution: The plotter is static and can be a bit of a chore to display everything correctly.
13+
You have to edit the code for each setup you want to display.
14+
'''
15+
16+
@staticmethod
17+
def plot_3D(pressure_field: np.ndarray, sim_param: SimulationParameters, title: str = '', interval=0, video_output: bool = False, file_name: str = '', source_zyx: tuple = None, direction: np.character = [None, 'x', 'y', 'z'][1]):
18+
'''
19+
Plots 3D domain in real-time with video output.
20+
21+
Parameters
22+
----------
23+
pressure_field : ndarray
24+
Pressure field data.
25+
sim_param : SimulationParameters
26+
Instance of simulation parameter class.
27+
title : str
28+
Title of the plot.
29+
interval : int
30+
Delay between frames in milliseconds.
31+
video_output: bool
32+
Displays the video on screen.
33+
file_name : str
34+
File name of video to write on disk.
35+
zyx : tuple
36+
Z, Y, X of source location.
37+
direction : char
38+
Direction. Either None, x, y or z.
39+
40+
Returns
41+
-------
42+
tuple
43+
Animation and FuncAnimation instances.
44+
'''
45+
46+
plt.close() # close any existing plots
47+
48+
if source_zyx is not None:
49+
(z, y, x) = source_zyx
50+
fig, (X, Y, Z) = plt.subplots(nrows=3, ncols=1, figsize=(10, 10))
51+
52+
animation = None
53+
if direction is not None:
54+
p_t = list()
55+
if direction == 'x':
56+
for frame in pressure_field:
57+
p_t.append(frame[z, y, :])
58+
if direction == 'y':
59+
for frame in pressure_field:
60+
p_t.append(frame[z, :, x])
61+
if direction == 'z':
62+
for frame in pressure_field:
63+
p_t.append(frame[:, y, x])
64+
65+
animation = AnimationPlotter.plot_1D(
66+
p_t, sim_param, interval=0, video_output=False, file_name='')
67+
68+
p_x = [pressure_field[i][:, :, x] for i in range(len(pressure_field))]
69+
p_y = [pressure_field[i][:, y, :] for i in range(len(pressure_field))]
70+
p_z = [pressure_field[i][z, :, :] for i in range(len(pressure_field))]
71+
72+
fig.suptitle(title, fontsize=14, fontweight='bold')
73+
text = fig.text(0.1, 0.9, '', # X, Y; 1-top or right
74+
verticalalignment='center', horizontalalignment='center',
75+
color='green', fontsize=15)
76+
77+
k = np.max([np.min(np.abs([pressure_field])),
78+
np.max(np.abs([pressure_field]))])
79+
k = 0.5*k
80+
ma = k
81+
mi = -k
82+
83+
colormap = ['Greys', 'seismic', 'coolwarm', 'twilight'][1]
84+
im_x = X.imshow(np.zeros_like(
85+
p_x[0]), vmin=mi, vmax=ma, aspect='equal', cmap=colormap)
86+
im_y = Y.imshow(np.zeros_like(
87+
p_y[0]), vmin=mi, vmax=ma, aspect='equal', cmap=colormap)
88+
im_z = Z.imshow(np.zeros_like(
89+
p_z[0]), vmin=mi, vmax=ma, aspect='equal', cmap=colormap)
90+
91+
# Color Bar
92+
fig.subplots_adjust(right=0.85)
93+
cbar_ax = fig.add_axes([0.88, 0.15, 0.04, 0.7])
94+
fig.colorbar(im_x, cax=cbar_ax)
95+
96+
def init_func():
97+
X.set_title('YZ-Plane')
98+
Y.set_title('ZX-Plane')
99+
Z.set_title('XY-Plane')
100+
101+
def update_plot(time_step):
102+
time = sim_param.delta_t * time_step
103+
text.set_text("Time: %.2f sec" % time)
104+
im_x.set_data(p_x[time_step])
105+
im_y.set_data(p_y[time_step])
106+
im_z.set_data(p_z[time_step])
107+
return [im_x, im_y, im_z]
108+
109+
# keep the reference
110+
func_animation = FuncAnimation(
111+
fig,
112+
update_plot,
113+
frames=range(sim_param.number_of_samples),
114+
init_func=init_func,
115+
interval=interval, # Delay between frames in milliseconds
116+
blit=False)
117+
if video_output:
118+
AnimationPlotter.write_video(func_animation, file_name)
119+
return [animation, func_animation]
120+
else:
121+
pass
122+
123+
@staticmethod
124+
def plot_2D(pressure_field: np.ndarray, sim_param: SimulationParameters, interval: int = 0, video_output: bool = False, file_name: str = ''):
125+
'''
126+
Plots 3D domain in real-time with video output.
127+
128+
Parameters
129+
----------
130+
pressure_field : ndarray
131+
Pressure field data.
132+
sim_param : SimulationParameters
133+
Instance of simulation parameter class.
134+
interval : int
135+
Delay between frames in milliseconds.
136+
video_output: bool
137+
Displays the video on screen.
138+
file_name : str
139+
File name of video to write on disk.
140+
141+
Returns
142+
-------
143+
FuncAnimation
144+
FuncAnimation instance.
145+
'''
146+
147+
fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(5, 5))
148+
149+
fig.suptitle("Time: %.2f sec" % 0)
150+
151+
mi = np.min(-np.abs([pressure_field]))
152+
ma = np.max(np.abs([pressure_field]))
153+
154+
im = ax.imshow(np.zeros_like(pressure_field[0]), vmin=mi, vmax=ma)
155+
156+
# Color Bar
157+
fig.subplots_adjust(right=0.85)
158+
cbar_ax = fig.add_axes([0.88, 0.15, 0.04, 0.7])
159+
fig.colorbar(im, cax=cbar_ax)
160+
161+
def init_func():
162+
'''
163+
No implementation.
164+
'''
165+
pass
166+
167+
def update_plot(time_step):
168+
time = sim_param.dt * time_step
169+
fig.suptitle("Time: %.2f sec" % time)
170+
im.set_data(pressure_field[time_step])
171+
return [im]
172+
173+
# keep the reference
174+
animation = FuncAnimation(fig,
175+
update_plot,
176+
frames=sim_param.time_steps,
177+
init_func=init_func,
178+
interval=interval, # Delay between frames in milliseconds
179+
blit=False)
180+
if video_output:
181+
AnimationPlotter.write_video(animation, file_name)
182+
return animation
183+
184+
@staticmethod
185+
def plot_1D(p_field_t: np.ndarray, sim_param: SimulationParameters, interval: int = 0, video_output: bool = False, file_name: str = ''):
186+
'''
187+
Plots 3D domain in real-time with video output.
188+
189+
Parameters
190+
----------
191+
pressure_field : ndarray
192+
Pressure field data.
193+
sim_param : SimulationParameters
194+
Instance of simulation parameter class.
195+
interval : int
196+
Delay between frames in milliseconds.
197+
video_output: bool
198+
Displays the video on screen.
199+
file_name : str
200+
File name of video to write on disk.
201+
202+
Returns
203+
-------
204+
FuncAnimation
205+
FuncAnimation instance.
206+
'''
207+
fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(5, 5))
208+
209+
fig.suptitle("Time: %.2f sec" % 0)
210+
211+
k = np.max([np.min(np.abs([p_field_t])), np.max(np.abs([p_field_t]))])
212+
k = 0.5*k
213+
ma = k
214+
mi = -k
215+
216+
ln, = ax.plot(0, 0)
217+
218+
def init_func():
219+
ax.set_ylim(mi, ma)
220+
ax.set_xlim(0, len(p_field_t[0]))
221+
ln.set_xdata(np.arange(len(p_field_t[0])))
222+
223+
def update_plot(time_step):
224+
time = sim_param.delta_t * time_step
225+
fig.suptitle("Time: %.2f sec" % time)
226+
ln.set_ydata(p_field_t[time_step])
227+
return [ln]
228+
229+
animation = FuncAnimation(fig,
230+
update_plot,
231+
frames=range(
232+
sim_param.number_of_samples),
233+
init_func=init_func,
234+
interval=interval, # Delay between frames in milliseconds
235+
blit=False)
236+
if video_output:
237+
AnimationPlotter.write_video(animation, file_name)
238+
return animation
239+
240+
@staticmethod
241+
def write_video(animation: FuncAnimation, file_name: str):
242+
'''
243+
Writes video to disk.
244+
245+
Parameters
246+
----------
247+
animation : FuncAnimation
248+
FuncAnimation instance, contains the animation.
249+
file_name : str
250+
Name of the file to be written on disk.
251+
'''
252+
253+
writervideo = FFMpegWriter(fps=60)
254+
fileloc = "videos/"
255+
filename = file_name + '_' + datetime.now().strftime("%d-%m-%Y_%H-%M-%S") + ".mp4"
256+
animation.save(fileloc+filename,
257+
dpi=300,
258+
writer=writervideo)

‎common/finite_differences.py

+18-18
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import numpy as np
22

3+
34
class FiniteDifferences():
45
'''
56
Helper class for getting FD coefficients and generating the laplacian matrix.
@@ -8,30 +9,29 @@ class FiniteDifferences():
89
[2] "Finite Difference Coefficients Calculator", Taylor, Cameron R. (2016). URL: https://web.media.mit.edu/~crtaylor/calculator.html
910
'''
1011

11-
FD_COEFFICIENTS = { # [1]
12-
1 : { 2 : np.array([-1/2, 0, 1/2]),
13-
4 : np.array([1/12, -2/3, 0, 2/3, -1/12]),
14-
6 : np.array([-1/60, 3/20, -3/4, 0, 3/4, -3/20, 1/60]),
15-
8 : np.array([1/280, -4/105, 1/5, -4/5, 0, 4/5, -1/5, 4/105, -1/280 ])
12+
FD_COEFFICIENTS = { # [1]
13+
1: {2: np.array([-1/2, 0, 1/2]),
14+
4: np.array([1/12, -2/3, 0, 2/3, -1/12]),
15+
6: np.array([-1/60, 3/20, -3/4, 0, 3/4, -3/20, 1/60]),
16+
8: np.array([1/280, -4/105, 1/5, -4/5, 0, 4/5, -1/5, 4/105, -1/280])
1617
},
17-
2 : { 2 : np.array([1, -2, 1]),
18-
4 : np.array([-1/12, 4/3, -5/2, 4/3, -1/12]),
19-
6 : np.array([1/90, -3/20, 3/2, -49/18, 3/2, -3/20, 1/90]),
20-
8 : np.array([-1/560, 8/315, -1/5, 8/5, -205/72, 8/5, -1/5, 8/315, -1/560]),
21-
10: np.array([8 ,-125 ,1000 ,-6000 ,42000 ,-73766 ,42000 ,-6000 ,1000 ,-125 ,8])/(25200) # [2]
22-
}
23-
}
18+
2: {2: np.array([1, -2, 1]),
19+
4: np.array([-1/12, 4/3, -5/2, 4/3, -1/12]),
20+
6: np.array([1/90, -3/20, 3/2, -49/18, 3/2, -3/20, 1/90]),
21+
8: np.array([-1/560, 8/315, -1/5, 8/5, -205/72, 8/5, -1/5, 8/315, -1/560]),
22+
# [2]
23+
10: np.array([8, -125, 1000, -6000, 42000, -73766, 42000, -6000, 1000, -125, 8])/(25200)
24+
}
25+
}
2426

2527
def __init__(self):
2628
'''
2729
No implementation.
2830
'''
2931
pass
30-
31-
3232

3333
@staticmethod
34-
def get_fd_coefficients(derivative, accuracy):
34+
def get_fd_coefficients(derivative: int, accuracy: int):
3535
'''
3636
Returns FD coefficients.
3737
@@ -67,7 +67,7 @@ def get_laplacian_matrix(derivative: int, accuracy: int):
6767
np.ndarray
6868
K, the laplacian matrix.
6969
'''
70-
70+
7171
coefs = FiniteDifferences.get_fd_coefficients(derivative, accuracy)
7272
nr_pts = int(accuracy / 2)
7373
k = coefs[0:nr_pts]
@@ -78,6 +78,6 @@ def get_laplacian_matrix(derivative: int, accuracy: int):
7878
dest = line[::-1]
7979
K[z][0:(z+1)] = dest
8080

81-
K = (np.vstack([K,-np.flipud(K)]))
82-
K = (np.hstack([-np.fliplr(K),K]))
81+
K = (np.vstack([K, -np.flipud(K)]))
82+
K = (np.hstack([-np.fliplr(K), K]))
8383
return K

‎common/impulse.py

+3-14
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22

33
import string
44
import numpy as np
5-
import matplotlib.pyplot as plt
65
from scipy.io.wavfile import read
76
from scipy.signal import firwin, freqz, lfilter
87

@@ -182,14 +181,8 @@ def get(self):
182181
ndarray
183182
Impulse over time.
184183
'''
185-
if self.sim_param.visualize_source:
186-
plt.plot(self.impulse)
187-
plt.show()
188-
189184
return self.amplitude * self.impulse
190185

191-
# TODO Maybe choose between ExperimentalUnit and standard Unit, removing the other
192-
193186

194187
class ExperimentalUnit(Impulse):
195188
'''
@@ -225,11 +218,10 @@ def __init__(
225218
self.filter_coeffs = firwin(
226219
filter_order, cutoff_frequency, fs=sim_param.Fs)
227220

228-
uno_numberinos = int(self.sim_param.Fs / (cutoff_frequency * 2))
229-
self.impulse[0: uno_numberinos] = 1
230-
self.impulse[uno_numberinos + 1: 2 * uno_numberinos] = -1
221+
high_signal = int(self.sim_param.Fs / (cutoff_frequency * 2))
222+
self.impulse[0: high_signal] = 1
223+
self.impulse[high_signal + 1: 2 * high_signal] = -1
231224

232-
#self.impulse = lfilter(self.filter_coeffs, [1], self.impulse)
233225

234226
def get(self):
235227
'''
@@ -240,8 +232,5 @@ def get(self):
240232
ndarray
241233
Impulse over time.
242234
'''
243-
if self.sim_param.visualize:
244-
plt.plot(self.impulse)
245-
plt.show()
246235

247236
return self.amplitude * self.impulse

0 commit comments

Comments
 (0)