Coverage for pySDC/projects/AllenCahn_Bayreuth/visualize_temp.py: 0%

1import json 

2import glob 

3import numpy as np 

4import matplotlib 

5 

6matplotlib.use('TkAgg') 

7import matplotlib.pyplot as plt 

8from mpl_toolkits.axes_grid1 import AxesGrid 

9from argparse import ArgumentParser 

10 

11import imageio 

12 

13 

14def plot_data(path='./data', name='', output='.'): 

15 """ 

16 Visualization using numpy arrays (written via MPI I/O) and json description 

17 

18 Produces one png file per time-step, combine as movie via e.g. 

19 > ffmpeg -i data/name_%08d.png name.mp4 

20 

21 Args: 

22 path (str): path to data files 

23 name (str): name of the simulation (expects data to be in data path) 

24 output (str): path to output 

25 """ 

26 

27 json_files = sorted(glob.glob(f'{path}/{name}_*.json')) 

28 data_files = sorted(glob.glob(f'{path}/{name}_*.dat')) 

29 

30 for json_file, data_file in zip(json_files, data_files): 

31 with open(json_file, 'r') as fp: 

32 obj = json.load(fp) 

33 

34 index = json_file.split('_')[-1].split('.')[0] 

35 print(f'Working on step {index}...') 

36 

37 array = np.fromfile(data_file, dtype=obj['datatype']) 

38 array = array.reshape(obj['shape'], order='C') 

39 

40 fig = plt.figure() 

41 

42 grid = AxesGrid( 

43 fig, 111, nrows_ncols=(1, 2), axes_pad=0.15, cbar_mode='single', cbar_location='right', cbar_pad=0.15 

44 ) 

45 

46 im = grid[0].imshow(array[..., 0], vmin=0, vmax=1) 

47 im = grid[1].imshow(array[..., 1], vmin=0, vmax=1) 

48 

49 grid[0].set_title(f"Field - Time: {obj['time']:6.4f}") 

50 grid[1].set_title(f"Temperature - Time: {obj['time']:6.4f}") 

51 

52 grid[1].yaxis.set_visible(False) 

53 

54 grid.cbar_axes[0].colorbar(im) 

55 

56 plt.savefig(f'{output}/{name}_{index}.png', bbox_inches='tight') 

57 plt.close() 

58 

59 

60def make_movie(path='./data', name='', output='.'): 

61 """ 

62 Visualization using numpy arrays (written via MPI I/O) and json description 

63 

64 Produces one png file per time-step, combine as movie via e.g. 

65 > ffmpeg -i data/name_%08d.png name.mp4 

66 

67 Args: 

68 path (str): path to data files 

69 name (str): name of the simulation (expects data to be in data path) 

70 output (str): path to output 

71 """ 

72 

73 json_files = sorted(glob.glob(f'{path}/{name}_*.json')) 

74 data_files = sorted(glob.glob(f'{path}/{name}_*.dat')) 

75 

76 img_list = [] 

77 for json_file, data_file in zip(json_files, data_files): 

78 with open(json_file, 'r') as fp: 

79 obj = json.load(fp) 

80 

81 index = json_file.split('_')[1].split('.')[0] 

82 print(f'Working on step {index}...') 

83 

84 array = np.fromfile(data_file, dtype=obj['datatype']) 

85 array = array.reshape(obj['shape'], order='C') 

86 

87 fig, ax = plt.subplots(1, 2) 

88 

89 ax[0].imshow(array[..., 0], vmin=0, vmax=1) 

90 ax[1].imshow(array[..., 1], vmin=0, vmax=1) 

91 

92 # ax.set_colorbar() 

93 ax[0].set_title(f"Field - Time: {obj['time']:6.4f}") 

94 ax[1].set_title(f"Temperature - Time: {obj['time']:6.4f}") 

95 

96 fig.tight_layout() 

97 

98 # draw the canvas, cache the renderer 

99 fig.canvas.draw() 

100 image = np.frombuffer(fig.canvas.tostring_rgb(), dtype='uint8') 

101 img_list.append(image.reshape(fig.canvas.get_width_height()[::-1] + (3,))) 

102 plt.close() 

103 

104 # c += 1 

105 # if c == 3: 

106 # break 

107 

108 fname = f'{output}/{name}.mp4' 

109 imageio.mimsave(fname, img_list, fps=8) 

110 

111 

112def make_movie_from_files(path='./data', name='', output='.'): 

113 """ 

114 Visualization using numpy arrays (written via MPI I/O) and json description 

115 

116 Produces one png file per time-step, combine as movie via e.g. 

117 > ffmpeg -i data/name_%08d.png name.mp4 

118 

119 Args: 

120 path (str): path to data files 

121 name (str): name of the simulation (expects data to be in data path) 

122 output (str): path to output 

123 """ 

124 

125 img_files = sorted(glob.glob(f'{path}/{name}_*.png')) 

126 print(f'{path}{name}') 

127 

128 images = [] 

129 for fimg in img_files: 

130 img = imageio.imread(fimg) 

131 print(fimg, img.shape) 

132 images.append(imageio.imread(fimg)) 

133 fname = f'{output}/{name}.mp4' 

134 imageio.mimsave(fname, images, fps=8) 

135 

136 

137if __name__ == "__main__": 

138 parser = ArgumentParser() 

139 parser.add_argument("-p", "--path", help='Path to data files', type=str, default='./data') 

140 parser.add_argument("-n", "--name", help='Name of the simulation', type=str) 

141 parser.add_argument("-o", "--output", help='Path for output file', type=str, default='.') 

142 args = parser.parse_args() 

143 

144 # name = 'AC-test-tempforce' 

145 name = 'AC-bench-tempforce' 

146 

147 plot_data(path=args.path, name=args.name, output=args.output) 

148 # make_movie(path=args.path, name=args.name, output=args.output) 

149 make_movie_from_files(path=args.path, name=args.name, output=args.output)