def plot_first_figure(file_list, varname, vmin=None, vmax=None): global verbose if verbose > 1: print('Getting data range') if vmin is None and vmax is None: vmin, vmax = get_var_range_from_sdf_files(file_list, varname) elif vmin is None: vmin = get_var_range_from_sdf_files(file_list, varname)[0] elif vmax is None: vmax = get_var_range_from_sdf_files(file_list, varname)[1] if verbose > 1: print('Found data range ({}, {})'.format(vmin, vmax)) filename = file_list[0] im, fig = plot_figure(filename, varname, vmin, vmax) return im, fig
def plot_figure(filename, varname, vmin=None, vmax=None, scale=False, figsize=None, ax=None): """ Plot the given variable for each file from a simulation """ from matplotlib.ticker import FuncFormatter global verbose, title_info if verbose > 1: print('Plotting {} from file {}'.format(varname, filename)) axis_type = 1 # 0 - off, 1 - left+bottom, 2 - all if ax is not None: fig = ax.get_figure() elif figsize is None: fig, ax = plt.subplots(1, 1, figsize=(8, 8)) else: fig, ax = plt.subplots(1, 1, figsize=figsize) if verbose > 1: print('Adjusting subplot') if axis_type == 0: plt.subplots_adjust(left=0, right=0.9, bottom=0.05, top=0.92, wspace=0, hspace=0) else: plt.subplots_adjust(left=0.1, right=0.9, bottom=0.1, top=0.92, wspace=0, hspace=0) plt.set_cmap(cm.coolwarm) if verbose > 1: print('Reading data') # Draw initial plot f = filename data = sdf.read(f) var = data.__dict__[varname] grid = var.grid num_x_axis_pts, num_y_axis_pts = grid.dims vdata = var.data.reshape(num_x_axis_pts, num_y_axis_pts).T print(vdata) vscale = get_scale(filename, varname, scale) i0 = 0 i1 = 1 if len(grid.dims) == 2: x = grid.data[i0] y = grid.data[i1] else: if (grid.dims[0] == 1): orig = var.grid.name.replace('/', '_') + '_Orig' grid = data.__dict__[orig] i0 = 2 i1 = 1 x = grid.data[i0] y = grid.data[i1] elif (grid.dims[1] == 1): i0 = 0 i1 = 1 x = grid.data[0][:, 0, :] y = grid.data[1][:, 0, :] else: i0 = 0 i1 = 2 x = grid.data[0][:, :, 0] y = grid.data[1][:, :, 0] z = grid.data[2][:, :, 0] x = np.sqrt(x * x + y * y) y = z if verbose > 1: print('Scaling data') vdata = vdata * vscale if vmin is None: vmin, vmax = get_var_range_from_sdf_files((filename,), varname, scale) if verbose > 1: print('Plotting data') im = ax.pcolormesh(x, y, vdata, vmin=vmin, vmax=vmax) if verbose > 1: print('Plotting axes') def fmt(x, pos): return r'${}$'.format(x) xmult, xsym = get_si_prefix(np.max(x) - np.min(x)) ymult, ysym = get_si_prefix(np.max(y) - np.min(y)) if verbose > 1: print('Scale axis by {} ({}, {})'.format(xmult, np.min(x), np.max(x))) ax.xaxis.set_major_formatter(FuncFormatter(lambda x, y: (x * xmult))) ax.yaxis.set_major_formatter(FuncFormatter(lambda x, y: (x * ymult))) ax.set_xlabel(grid.labels[i0] + ' $(' + xsym + grid.units[i0] + ')$') ax.set_ylabel(grid.labels[i1] + ' $(' + ysym + grid.units[i1] + ')$') ax.axis('image') if axis_type == 0: ax.axis('off') elif axis_type == 1: ax.spines['top'].set_visible(False) ax.spines['right'].set_visible(False) ax.get_xaxis().tick_bottom() ax.get_yaxis().tick_left() # Add colorbar mult, sym = get_si_prefix(vmax - vmin) if verbose > 0 and plot_figure.first: plot_figure.first = False print('Scale colorbar by {} ({}, {})'.format(mult, vmin, vmax)) data_label = var.name + ' $(' + sym + var.units + ')$' if title_info: legend_label = '' else: legend_label = data_label divider = make_axes_locatable(ax) cax = divider.append_axes("right", "5%", pad="15%") cbar = fig.colorbar(im, cax=cax, ax=ax, label=legend_label, format=FuncFormatter(lambda x, y: x * mult)) if title_info: title_label = data_label + ', ' else: title_label = '' fig.fs = 'large' fig.suptitle(title_label + 'step={}, time={}'.format(data.Header['step'], data.Header['time']), fontsize=fig.fs, y=0.98) fig.sca(ax) fig.canvas.draw() if verbose > 1: print('Done plotting {} from file {}'.format(varname, filename)) return im, ax
def plot_first_figure(file_list, varname, vmin=None, vmax=None, scale=False): global verbose, title_info if verbose > 1: print('Getting data range') if vmin is None and vmax is None: vmin, vmax = get_var_range_from_sdf_files(file_list, varname, scale) elif vmin is None: vmin = get_var_range_from_sdf_files(file_list, varname, scale)[0] elif vmax is None: vmax = get_var_range_from_sdf_files(file_list, varname, scale)[1] if verbose > 1: print('Found data range ({}, {})'.format(vmin, vmax)) filename = file_list[0] im, ax = plot_figure(filename, varname, vmin, vmax, scale) fig = ax.get_figure() renderer = find_renderer(fig) box_orig = fig.bbox_inches box_raw = fig.get_tightbbox(renderer) pp = box_raw.get_points() figsize = (pp[1, 0] - pp[0, 0], pp[1, 1] - pp[0, 1]) ratio = figsize[0] / figsize[1] if ratio > 1: w = box_orig.width h = w / ratio h = round(10 * h) / 10. figsize = (w, h) else: h = box_orig.height w = h * ratio w = round(10 * w) / 10. figsize = (w, h) im, ax = plot_figure(filename, varname, vmin, vmax, scale, figsize) fig = ax.get_figure() # Get positions of the title's step and time text fields so that they # can be updated when animating data = sdf.read(file_list[-1]) var = data.__dict__[varname] mult, sym = get_si_prefix(vmax - vmin) data_label = var.name + ' $(' + sym + var.units + ')$, ' if title_info: title_label = data_label + ', ' title_label2 = data_label + ',_' else: title_label = '' title_label2 = '' fig._suptitle.set_text(title_label + 'step={}, time={}'.format(data.Header['step'], data.Header['time'])) bbox = fig._suptitle.get_window_extent(renderer) inv = fig.transFigure.inverted() x0 = inv.transform(bbox)[0][0] y = inv.transform(bbox)[0][1] fig.y = y tt = fig.text(x0, y, title_label2, fontsize=fig.fs, ha='left', va='bottom') bbox = tt.get_window_extent(renderer) fig.x1 = inv.transform(bbox)[1][0] fig.y = inv.transform(bbox)[0][1] tt.set_text(title_label) tt = fig.text(fig.x1, y, 'step={},_'.format(data.Header['step']), fontsize=fig.fs, ha='left', va='bottom') bbox = tt.get_window_extent(renderer) fig.x2 = inv.transform(bbox)[1][0] tt.set_visible(False) fig._suptitle.set_visible(False) return im, fig
def composite_field_plot(varname, vmin=None, vmax=None, directory='Data'): global verbose, dpi file_list = get_files(wkdir=directory) file_list = clean_file_list(file_list, varname) file_list.remove(directory + '00000.sdf') file_list.remove(directory + '00002.sdf') file_list.remove(directory + '00003.sdf') file_list.remove(directory + '00004.sdf') file_list.remove(directory + '00005.sdf') file_list.remove(directory + '00006.sdf') file_list.remove(directory + '00007.sdf') file_list.remove(directory + '00008.sdf') file_list.remove(directory + '00009.sdf') file_list.remove(directory + '00010.sdf') file_list.remove(directory + '00011.sdf') file_list.remove(directory + '00012.sdf') file_list.remove(directory + '00013.sdf') file_list.remove(directory + '00014.sdf') file_list.remove(directory + '00015.sdf') file_list.remove(directory + '00016.sdf') file_list.remove(directory + '00017.sdf') file_list.remove(directory + '00018.sdf') file_list.remove(directory + '00019.sdf') file_list.remove(directory + '00020.sdf') file_list.remove(directory + '00021.sdf') file_list.remove(directory + '00022.sdf') file_list.remove(directory + '00023.sdf') file_list.remove(directory + '00024.sdf') file_list.remove(directory + '00025.sdf') file_list.remove(directory + '00027.sdf') file_list.remove(directory + '00028.sdf') file_list.remove(directory + '00029.sdf') file_list.remove(directory + '00030.sdf') file_list.remove(directory + '00031.sdf') file_list.remove(directory + '00032.sdf') file_list.remove(directory + '00033.sdf') file_list.remove(directory + '00034.sdf') file_list.remove(directory + '00035.sdf') file_list.remove(directory + '00037.sdf') file_list.remove(directory + '00038.sdf') file_list.remove(directory + '00039.sdf') file_list.remove(directory + '00040.sdf') file_list.remove(directory + '00041.sdf') file_list.remove(directory + '00042.sdf') file_list.remove(directory + '00043.sdf') file_list.remove(directory + '00044.sdf') file_list.remove(directory + '00045.sdf') file_list.remove(directory + '00047.sdf') file_list.remove(directory + '00048.sdf') file_list.remove(directory + '00049.sdf') file_list.remove(directory + '00050.sdf') file_list.remove(directory + '00051.sdf') file_list.remove(directory + '00052.sdf') file_list.remove(directory + '00053.sdf') file_list.remove(directory + '00054.sdf') file_list.remove(directory + '00055.sdf') if verbose > 0: print('Found {} files to plot'.format(len(file_list))) data = [] for f in file_list: d = sdf.read(f) var = d.__dict__[varname] data.append(var.data) data = np.asarray(data) data = data.T tmin = sdf.read(file_list[0]).Header['time'] tmax = sdf.read(file_list[-1]).Header['time'] grid = var.grid_mid xmin = np.min(grid.data[0]) xmax = np.max(grid.data[0]) shape = data.shape extent = [tmin, tmax, xmax, xmin] xmult, xsym = get_si_prefix(xmax - xmin) # y axis tmult, tsym = get_si_prefix(tmax - tmin) # x axis if vmin is None and vmax is None: vmin, vmax = get_var_range_from_sdf_files(file_list, varname) elif vmin is None: vmin = get_var_range_from_sdf_files(file_list, varname)[0] elif vmax is None: vmax = get_var_range_from_sdf_files(file_list, varname)[1] mult, sym = get_si_prefix(vmax - vmin) fig, ax = plt.subplots() im = ax.imshow(data, extent=extent, aspect=calculate_aspect(shape, extent), interpolation='none', cmap=cm.plasma, vmin=vmin, vmax=vmax) ax.xaxis.set_major_formatter(FuncFormatter(lambda x, y: (x * tmult))) ax.yaxis.set_major_formatter(FuncFormatter(lambda x, y: (x * xmult))) plt.xlabel('t $(' + tsym + 's)$') plt.ylabel(grid.labels[0] + ' $(' + xsym + grid.units[0] + ')$') # data_label = var.name + ' $(' + sym + var.units + ')$' data_label = 'Argon$^{+8}$ Number Density $(' + sym + var.units + ')$' plt.title('Electron Density Evolution') cbar = fig.colorbar(im, label=data_label, format=FuncFormatter(lambda x, y: x * mult)) plt.tight_layout() plt.savefig('electron_comp_thermal_nocoll.png', dpi=600, bbox_inches="tight")