import numpy as np
import matplotlib.pyplot as plt
import os
import sys
import shutil
import natsort
import imageio
from .ProcData import ProcData
import copy
import re
class Plotter:
ani_pngs_dir = "plotter_lib_pngs/"
def __init__(self):
self.ani_period = float(1./20.)
self.filename = None
self.out = None
self.oname = None
self.var = None
self.mval = None
self.function = None
self.file_column_names = None
self.title = None
self.file_data = None
self.plt_max_val = None
self.plt_min_val = None
self.min_y = None
self.max_y = None
self.min_max_var_vals = None
self.plane_type = None
self.slice_position = None
self.ifTranspose = False
def __check_arg_dim_equiv(self, args):
if args.var != None and args.mval != None:
if len(args.var) != len(args.mval):
print("The count of var assumed to be equal to the count of mval")
sys.exit(-1)
def __get_variable_names(self):
var_names = []
for d in self.file_data:
names = d.variable_names
var_names.append(np.array(names, dtype=object))
if(len(var_names) > 1):
for i in range(1, len(var_names)):
if np.equal(var_names[0], var_names[i]).any() != True:
print("All files must have the same variable names")
sys.exit(-1)
elif len(var_names) == 0:
print("Undefined variable names")
sys.exit(-1)
self.file_column_names = var_names[0]
def __set_fig_names(self):
if self.oname == None:
self.fig_names = copy.deepcopy(self.var)
self.fig_end = ".png"
else:
self.fig_names = copy.deepcopy(self.oname)
self.fig_end = ""
special_characters = '/'
for i in range(len(self.fig_names)):
for symbol in special_characters:
position = 0
while position < len(self.fig_names[i]):
if self.fig_names[i][position] == symbol:
self.fig_names[i] = self.fig_names[i][:position] + '\\' + self.fig_names[i][position+1:]
position += 2
else:
position += 1
def set(self, args, **kwargs):
self.__check_arg_dim_equiv(args)
self.filename = args.filename
self.ndim = args.ndim
self.out = args.out
self.oname = args.oname
self.var = args.var
self.mval = args.mval
self.function = args.func
self.title = args.title
self.min_y = args.min_y
self.max_y = args.max_y
self.plane_type = args.plane_type
self.slice_position = args.slice_position
self.ifTranspose = args.ifTranspose
if args.levels != None:
self.levels = args.levels
else:
self.levels = 25
self.if_manual_plot = kwargs.get('if_manual_plot', False)
self.if_save_result = kwargs.get('if_save_result', True)
pData = []
for fname in self.filename:
p = ProcData(fname)
p.get_variable_names()
pData.append(p)
self.file_data = copy.deepcopy(pData)
self.__get_variable_names()
if args.func != self.dump:
if args.func == self.plot or args.func == self.joint_plot or args.func == self.ani_plot or args.func == self.multiple_plot:
self.ndim = 1
elif args.func == self.plot_contour or args.func == self.ani_plot_contour:
self.ndim = 2
elif args.func == self.avg_plot or args.func == self.slice:
self.ndim = 3
if self.var == None and len(self.file_column_names) != 0:
self.var = [self.file_column_names[i] for i in range(self.ndim, len(self.file_column_names))]
for p in self.file_data:
p.process_file(self.ndim, self.var, self.mval)
if args.func == self.joint_plot:
self.fig_count = len(self.file_data)
else:
self.fig_count = len(self.var)
self.__set_fig_names()
def __plot(self):
os.system("mkdir -p " + self.out)
x_name = self.file_column_names[0]
# if self.oname == None:
# fig_names = self.var
# fig_end = ".png"
# else:
# fig_names = self.oname
# fig_end = ""
for i in range(self.fig_count):
y_name = self.var[i]
fig = plt.figure()
if self.ifTranspose:
plt.plot(self.file_data[0].data[y_name], self.file_data[0].data[x_name], linewidth=4)
plt.xlabel(y_name, fontsize=10, fontweight='bold')
plt.ylabel(x_name, fontsize=10, fontweight='bold')
else:
plt.plot(self.file_data[0].data[x_name], self.file_data[0].data[y_name], linewidth=4)
plt.xlabel(x_name, fontsize=10, fontweight='bold')
plt.ylabel(y_name, fontsize=10, fontweight='bold')
if self.if_manual_plot: plt.show()
else: plt.close(fig)
if self.if_save_result: fig.savefig(self.out + self.fig_names[i] + self.fig_end)
def __joint_plot(self):
os.system("mkdir -p " + self.out)
x_name = self.file_column_names[0]
# if self.oname == None:
# fig_names = self.var
# fig_end = ".png"
# else:
# fig_names = self.oname
# fig_end = ""
for i in range(self.fig_count):
fig = plt.figure()
for y_name in self.var:
if self.ifTranspose:
plt.plot(self.file_data[0].data[y_name], self.file_data[0].data[x_name], linewidth=4, label = x_name)
else:
plt.plot(self.file_data[0].data[x_name], self.file_data[0].data[y_name], linewidth=4, label = y_name)
if self.ifTranspose:
plt.xlabel(y_name, fontsize=10, fontweight='bold')
else:
plt.xlabel(x_name, fontsize=10, fontweight='bold')
plt.legend()
if self.if_manual_plot: plt.show()
else: plt.close(fig)
if self.if_save_result: fig.savefig(self.out + self.fig_names[i] + self.fig_end)
def __multiple_plot(self):
os.system("mkdir -p " + self.out)
x_name = self.file_column_names[0]
# if self.oname == None:
# fig_names = self.var
# fig_end = ".png"
# else:
# fig_names = self.oname
# fig_end = ""
for i in range(self.fig_count):
y_name = self.var[i]
fig = plt.figure()
for read_data in self.file_data:
if self.ifTranspose:
plt.plot(read_data.data[y_name], read_data.data[x_name], linewidth=4)
else:
plt.plot(read_data.data[x_name], read_data.data[y_name], linewidth=4)
plt.legend(self.filename)
if self.ifTranspose:
plt.xlabel(x_name, fontsize=10, fontweight='bold')
plt.ylabel(y_name, fontsize=10, fontweight='bold')
else:
plt.xlabel(y_name, fontsize=10, fontweight='bold')
plt.ylabel(x_name, fontsize=10, fontweight='bold')
if self.if_manual_plot: plt.show()
else: plt.close(fig)
if self.if_save_result: fig.savefig(self.out + self.fig_names[i] + self.fig_end)
def __get_min_max_ax(self):
min_max_var_vals = {var : [] for var in self.var}
for var in self.var:
max_val = np.nanmax(self.file_data[0].data[var])
min_val = np.nanmin(self.file_data[0].data[var])
for data in self.file_data:
maval = np.nanmax(data.data[var])
mival = np.nanmin(data.data[var])
if maval > max_val:
max_val = maval
if mival < min_val:
min_val = mival
min_max_var_vals[var] = np.array([min_val, max_val])
max_vals = np.array([min_max_var_vals[var][1] for var in self.var])
min_vals = np.array([min_max_var_vals[var][0] for var in self.var])
max_val = np.max(max_vals)
min_val = np.min(min_vals)
self.min_max_var_vals = np.array([min_val, max_val])
def __ani_plot(self):
if self.if_save_result:
png_names = []
os.system("mkdir -p " + self.out)
os.system("mkdir -p " + self.ani_pngs_dir)
names = natsort.natsorted(self.filename,reverse=False)
x_name = self.file_column_names[0]
data_i = 0
if self.max_y == None:
max_val = self.min_max_var_vals[1]
else:
max_val = self.max_y[0]
if self.min_y == None:
min_val = self.min_max_var_vals[0]
else:
min_val = self.min_y[0]
duration = self.ani_period * len(names)
for datafile in names:
fig = plt.figure()
if self.ifTranspose:
plt.ylim([min_val, max_val])
else:
plt.xlim([min_val, max_val])
for y_name in self.var:
if self.ifTranspose:
plt.plot(self.file_data[data_i].data[y_name], self.file_data[data_i].data[x_name], linewidth=4)
else:
plt.plot(self.file_data[data_i].data[x_name], self.file_data[data_i].data[y_name], linewidth=4)
plt.legend(self.var)
plt.xlabel(x_name, fontsize=10, fontweight='bold')
figname = os.path.basename(datafile)
plt.close(fig)
fig.savefig(self.ani_pngs_dir + figname.split('.')[0] + '.png')
name = self.ani_pngs_dir + figname.split('.')[0] + '.png'
png_names.append(name)
data_i = data_i + 1
images = []
for file_name in png_names:
images.append(imageio.v2.imread(file_name))
# imageio.mimsave(self.oname[0], images, fps = 5, loop = 0) , duration = 0.04
imageio.mimsave(self.oname[0], images, duration = 0.25, loop = 0)
shutil.rmtree(self.ani_pngs_dir)
def __plot_contour(self):
os.system("mkdir -p " + self.out)
x_name = self.file_column_names[0]
y_name = self.file_column_names[1]
# if self.oname == None:
# fig_names = self.var
# fig_end = ".png"
# else:
# fig_names = self.oname
# fig_end = ""
if self.max_y == None:
max_val = {var:self.min_max_var_vals[var][1] for var in self.var}
else:
max_val = {var:self.max_y[i] for var, i in zip(self.var, list(range(len(self.var))))}
if self.min_y == None:
min_val = {var:self.min_max_var_vals[var][0] for var in self.var}
else:
min_val = {var:self.min_y[i] for var, i in zip(self.var, list(range(len(self.var))))}
for i in range(self.fig_count):
fig,ax=plt.subplots(1,1)
if self.title == None:
title = self.var[i]
else:
title = self.title
X = self.file_data[0].data[x_name]
Y = self.file_data[0].data[y_name]
Z = self.file_data[0].data[self.var[i]]
vmin = min_val[self.var[i]]; vmax = max_val[self.var[i]];
levels = np.linspace(vmin, vmax, self.levels)
if self.ifTranspose:
if self.function == self.plot_diff:
cp = ax.contourf(Y, X, Z.T)
else:
cp = ax.contourf(Y, X, Z.T, vmin=vmin, vmax=vmax, levels=levels)
ax.set_ylabel(x_name)
ax.set_xlabel(y_name)
else:
if self.function == self.plot_diff:
cp = ax.contourf(X, Y, Z)
else:
cp = ax.contourf(X, Y, Z, vmin=vmin, vmax=vmax, levels=levels)
ax.set_xlabel(x_name)
ax.set_ylabel(y_name)
fig.colorbar(cp) # Add a colorbar to a plot
ax.set_title(title)
if self.if_manual_plot: plt.show()
else: plt.close(fig)
if self.if_save_result: fig.savefig(self.out + self.fig_names[i] + self.fig_end)
def __get_min_max_bar(self):
self.filename = natsort.natsorted(self.filename,reverse=False)
self.min_max_var_vals = {var : [] for var in self.var}
for var in self.var:
max_val = np.nanmax(self.file_data[0].data[var])
min_val = np.nanmin(self.file_data[0].data[var])
for data in self.file_data:
maval = np.nanmax(data.data[var])
mival = np.nanmin(data.data[var])
if maval > max_val:
max_val = maval
if mival < min_val:
min_val = mival
self.min_max_var_vals[var] = np.array([min_val, max_val])
def __ani_plot_contour(self):
if self.if_save_result:
# png_names = {var:[] for var in self.var}
os.system("mkdir -p " + self.out)
os.system("mkdir -p " + self.ani_pngs_dir)
x_name = self.file_column_names[0]
y_name = self.file_column_names[1]
if self.oname == None:
fig_names = self.var
fig_end = ".gif"
else:
fig_names = self.oname
fig_end = ""
X = self.file_data[0].data[x_name]
Y = self.file_data[0].data[y_name]
if self.max_y == None:
max_val = {var:self.min_max_var_vals[var][1] for var in self.var}
else:
max_val = {var:self.max_y[i] for var, i in zip(self.var, list(range(len(self.var))))}
if self.min_y == None:
min_val = {var:self.min_max_var_vals[var][0] for var in self.var}
else:
min_val = {var:self.min_y[i] for var, i in zip(self.var, list(range(len(self.var))))}
duration = self.ani_period * len(self.filename)
i = 0
for var in self.var:
if self.title == None:
title = var
else:
title = self.title
counter = 0
png_names = []
for data in self.file_data:
levels = np.linspace(min_val[var], max_val[var], 30)
fig,ax=plt.subplots(1,1)
ax.set_title(title)
Z = data.data[var]
if self.ifTranspose:
cp = ax.contourf(Y, X, Z.T, vmin=min_val[var], vmax=max_val[var], levels=levels)
ax.set_ylabel(x_name)
ax.set_xlabel(y_name)
else:
cp = ax.contourf(X, Y, Z, vmin=min_val[var], vmax=max_val[var], levels=levels)
ax.set_xlabel(x_name)
ax.set_ylabel(y_name)
fig.colorbar(cp) # Add a colorbar to a plot
plt.close(fig)
figname = var + str(counter)
fig.savefig(self.ani_pngs_dir + figname + '.png')
name = self.ani_pngs_dir + figname + '.png'
png_names.append(name)
counter += 1
images = []
for file_name in png_names:
images.append(imageio.v2.imread(file_name))
# imageio.mimsave(self.out + fig_names[i] + fig_end, images, fps = 5, loop=0)
imageio.mimsave(self.out + fig_names[i] + fig_end, images, duration = 0.25, loop=0)
# imageio.mimsave(self.oname[0], images, duration = 0.25, loop = 0)
i += 1
shutil.rmtree(self.ani_pngs_dir)
def __avg(self, data, var_name):
cx = data['cx']
cy = data['cy']
cz = data['cz']
flat_matrix_data = np.zeros((cx * cy * cz))
flat_data = data[var_name].flatten()
for k in range(cz):
for j in range(cy):
for i in range(cx):
flat_matrix_data[k * cy * cx + j * cx + i] = flat_data[k * cy * cx + j * cx + i]
matrix_data = np.reshape(flat_matrix_data, (cx, cy, cz), order='F')
avg_data = np.average(matrix_data, axis=(1, 0))
return avg_data
def __avg_plot(self):
os.system("mkdir -p " + self.out)
fig = plt.figure()
x_name = self.file_column_names[2]
x = self.file_data[0].data[x_name]
for var in self.var:
avg_data = self.__avg(self.file_data[0].data, var)
if self.ifTranspose:
plt.plot(avg_data, x, linewidth=4)
else:
plt.plot(x, avg_data, linewidth=4)
plt.legend(self.var)
if self.ifTranspose:
plt.ylabel(x_name, fontsize=10, fontweight='bold')
else:
plt.xlabel(x_name, fontsize=10, fontweight='bold')
if self.if_manual_plot: plt.show()
else: plt.close(fig)
if self.if_save_result: fig.savefig(self.out + self.oname[0])
def __dump(self):
for variable_name in self.file_column_names:
print(variable_name, end='; ')
print('\n')
def __plot_diff(self):
for name in self.file_column_names:
if self.file_data[0].data[name].shape != self.file_data[1].data[name].shape:
print("Data dimensions do not match")
sys.exit(-1)
diff = {}
for name in self.var:
diff[name] = self.file_data[0].data[name] - self.file_data[1].data[name]
dim_variables = list(set(self.file_column_names) - set(self.var))
for name in dim_variables:
diff[name] = self.file_data[0].data[name]
diffProcData = ProcData()
diffProcData.data = diff
self.file_data = [diffProcData]
basename0 = os.path.basename(self.filename[0])
basename1 = os.path.basename(self.filename[1])
self.title = str(basename0) + ' - ' + str(basename1)
if self.ndim == 1:
self.__get_min_max_ax()
elif self.ndim == 2:
self.__get_min_max_bar()
if self.ndim == 1:
self.__plot()
elif self.ndim == 2:
self.__plot_contour()
def __find_closest(self, arr, val):
idx = np.abs(arr - val).argmin()
return idx
def __slice(self):
os.system("mkdir -p " + self.out)
if self.plane_type == 'xy':
x_name = self.file_column_names[0]
y_name = self.file_column_names[1]
z_name = self.file_column_names[2]
arr = self.file_data[0].data[z_name]
pos = self.__find_closest(arr, self.slice_position)
elif self.plane_type == 'yz':
x_name = self.file_column_names[1]
y_name = self.file_column_names[2]
z_name = self.file_column_names[0]
arr = self.file_data[0].data[z_name]
pos = self.__find_closest(arr, self.slice_position)
elif self.plane_type == 'xz':
x_name = self.file_column_names[0]
y_name = self.file_column_names[2]
z_name = self.file_column_names[1]
arr = self.file_data[0].data[z_name]
pos = self.__find_closest(arr, self.slice_position)
for i in range(self.fig_count):
fig,ax=plt.subplots(1,1)
if self.title == None:
title = self.var[i]
else:
title = self.title
if self.plane_type == 'xy':
Z = self.file_data[0].data[self.var[i]][pos, :, :]
elif self.plane_type == 'xz':
Z = self.file_data[0].data[self.var[i]][:, pos, :]
elif self.plane_type == 'yz':
Z = self.file_data[0].data[self.var[i]][:, :, pos]
X = self.file_data[0].data[x_name]
Y = self.file_data[0].data[y_name]
if self.ifTranspose:
cp = ax.contourf(Y, X, Z.T)
ax.set_ylabel(x_name)
ax.set_xlabel(y_name)
else:
cp = ax.contourf(X, Y, Z)
ax.set_xlabel(x_name)
ax.set_ylabel(y_name)
fig.colorbar(cp) # Add a colorbar to a plot
ax.set_title(title)
if self.if_manual_plot: plt.show()
else: plt.close(fig)
if self.if_save_result: fig.savefig(self.out + self.fig_names[i] + self.fig_end)
def dump(self):
self.__dump()
def plot(self):
self.__get_min_max_ax()
self.__plot()
def joint_plot(self):
self.__get_min_max_ax()
self.__joint_plot()
def ani_plot(self):
self.__get_min_max_ax()
self.__ani_plot()
def avg_plot(self):
self.__get_min_max_ax()
self.__avg_plot()
def plot_contour(self):
self.__get_min_max_bar()
self.__plot_contour()
def ani_plot_contour(self):
self.__get_min_max_bar()
self.__ani_plot_contour()
def plot_diff(self):
self.__plot_diff()
def get_data(self):
return_data = [copy.deepcopy(data.data) for data in self.file_data]
return return_data
def multiple_plot(self):
self.__get_min_max_ax()
self.__multiple_plot()
def slice(self):
self.__get_min_max_bar()
self.__slice()