-
Notifications
You must be signed in to change notification settings - Fork 0
/
plot_map.py
109 lines (86 loc) · 3.93 KB
/
plot_map.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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
import imageio ## to convert to mp4, please also run "pip install imageio-ffmpeg"
import matplotlib.pyplot as plt
import numpy as np
import os
import datetime
class Plot2DArray(object):
def __init__(self, filename_prefix="", output_dir="imgfiles"):
super().__init__()
# use the current time as filename if not specified
if filename_prefix:
self.filename_prefix = filename_prefix
else:
self.filename_prefix = "simulation_{}".format(datetime.datetime.now().strftime('%m_%d_%H_%M'))
self.output_dir=output_dir
self.max_digit = 4
self.plotted_img_paths = []
def plot_map(self, map, t, cmap="magma", figure_size=(9, 9)):
"""
Param
- map: np.array
an 2d numpy array to plot
- t: float
current timestep t
- cmap:
the color set for meshcolor.
you can choose the one you like at https://matplotlib.org/stable/tutorials/colors/colormaps.html
"""
title = "t = {:.3f}".format(t)
output_path = os.path.join(os.getcwd(), self.output_dir, self.filename_prefix)
filename = "{}_{:.3f}.png".format(self.filename_prefix, t)
plt.figure(figsize=figure_size, dpi=80)
plt.title(title)
plt.imshow(map, vmin=0, vmax=32, cmap=cmap)
plt.colorbar()
self.plotted_img_paths.append(self._save_fig(output_path, filename, t))
plt.close()
def _save_fig(self, output_path, fn, t):
if not os.path.exists(output_path):
os.makedirs(output_path)
file_path = os.path.join(output_path, fn)
plt.savefig(file_path)
print("|t={}| figrue saved to {}".format((str(t)+' '*self.max_digit)[:self.max_digit], file_path))
return file_path
def save_gif(self, fps=30, img_dir=""):
filename = "{}.gif".format(self.filename_prefix)
file_path = os.path.join(os.getcwd(), self.output_dir, filename)
# img paths
all_img_paths = self.plotted_img_paths
if img_dir:
filename_prefix = os.path.split(img_dir)[-1]
all_t = sorted([float(os.path.splitext(f)[0].split('_')[-1]) for f in os.listdir(img_dir) if os.path.isfile(os.path.join(img_dir, f))])
all_img_paths = [os.path.join(img_dir, "{}_{:.3f}.png".format(filename_prefix, t)) for t in all_t]
images = [imageio.imread(img_path) for img_path in all_img_paths]
imageio.mimsave(file_path, images, duration=1/fps)
print("gif saved to {}".format(file_path))
def save_mp4(self, fps=30, img_dir=""):
filename = "{}.mp4".format(self.filename_prefix)
file_path = os.path.join(os.getcwd(), self.output_dir, filename)
# img paths
all_img_paths = self.plotted_img_paths
all_img_paths = self.plotted_img_paths
if img_dir:
filename_prefix = os.path.split(img_dir)[-1]
all_t = sorted([float(os.path.splitext(f)[0].split('_')[-1]) for f in os.listdir(img_dir) if os.path.isfile(os.path.join(img_dir, f))])
all_img_paths = [os.path.join(img_dir, "{}_{:.3f}.png".format(filename_prefix, t)) for t in all_t]
writer = imageio.get_writer(file_path, fps=20)
for img_path in all_img_paths:
writer.append_data(imageio.imread(img_path))
writer.close()
print("mp4 saved to {}".format(file_path))
if __name__ == "__main__":
filename_prefix = "expset(a)_eta_0.2_theta_0.56_Gamma_0.019"
img_dir = os.path.join(os.getcwd(), 'imgfiles', filename_prefix)
plotter = Plot2DArray(filename_prefix=filename_prefix)
plotter.save_gif(img_dir=img_dir)
plotter.save_mp4(img_dir=img_dir)
'''
# usage example
t = 60
lots_of_data = np.random.randint(256, size=(t, 128, 128))
plotter = Plot2DArray()
for i in range(t):
plotter.plot_map(lots_of_data[i], i)
plotter.save_gif()
plotter.save_mp4()
'''