def plot_2d(w1, w3, signal, path, invert_w1=False, scale=None, axlim=(None, None)): signal2 = -1 * signal.copy() # plot in 1000 of wn w1 = w1.copy() / 1e3 w3 = w3.copy() / 1e3 if invert_w1: w1 = -w1 if scale is None: # calculate scale, return it in meta scale = np.max(np.abs(signal2)) signal2 /= scale # fiddle parameters to make the plots look good linthresh = 0.01 linscale = 0.1 nlevels = 100 norm = SymLogNorm(linthresh=linthresh, linscale=linscale, vmin=-1, vmax=1) logloc = SymmetricalLogLocator(linthresh=linthresh, base=2) levels = logloc.tick_values(-1, 1) fig = figure() ax = fig.add_subplot(111, aspect='equal', adjustable='box-forced') qset = ax.contourf(w1, w3, signal2, nlevels, norm=norm, cmap='RdBu_r') c = ax.contour(w1, w3, signal2, levels=levels, colors='k', alpha=0.4) ax.set_xlabel(r'$\omega_\tau$ ($\times 10^3\ \mathrm{cm}^{-1}$)') ax.set_ylabel(r"$\omega_t$ ($\times 10^3\ \mathrm{cm}^{-1}$)") ax.text(0.99, 0.01, r'$\mathrm{{scale:}} {!s}$'.format(latex_float(scale)), transform=ax.transAxes, horizontalalignment='right', verticalalignment='bottom') loc = matplotlib.ticker.MaxNLocator(11) fmt = matplotlib.ticker.ScalarFormatter(useOffset=False, useMathText=True) cb = fig.colorbar(qset, ax=ax, ticks=loc, format=fmt) cb.add_lines(c) ax.set_xlim(*axlim) ax.set_ylim(*axlim) #ax.grid(True) ax.add_artist( Line2D((0, 1), (0, 1), linewidth=2, color='k', alpha=0.5, transform=ax.transAxes)) #ax.relim() #ax.autoscale_view() fig.savefig(str(path))
def plot_2d(w1, w3, signal, path, invert_w1=False, scale=None, axlim=None): signal2 = -1 * signal rcParams.update(params) # plot in 1000 of wn w1 = w1.copy() / 1e3 w3 = w3.copy() / 1e3 if invert_w1: w1 = -w1 if scale is None: # calculate scale, return it in meta scale = np.max(np.abs(signal2)) signal2 /= scale # fiddle parameters to make the plots look good linthresh = 0.01 linscale = 0.1 nlevels = 100 norm = SymLogNorm(linthresh=linthresh, linscale=linscale, vmin=-1, vmax=1) logloc = SymmetricalLogLocator(linthresh=linthresh, base=2) levels = logloc.tick_values(-1, 1) fig = figure() ax = fig.add_subplot(111, aspect='equal', adjustable='box-forced') qset = ax.contourf(w1, w3, signal2, nlevels, norm=norm, cmap='RdBu_r') ax.contour(w1, w3, signal2, levels=levels, colors='k', alpha=0.4) cb = fig.colorbar(qset, ax=ax) if axlim: ypts = xpts = np.array(sorted(axlim)) / 1e3 print('Using limits ', xpts, ypts) ax.set_xlim(xpts) ax.set_ylim(ypts) else: ypts = xpts = sorted([np.min(w3), np.max(w3)]) ax.set_xlim(xpts) ax.set_ylim(ypts) ax.add_artist( Line2D(xpts, ypts, linewidth=2, color='k', alpha=0.5, transform=ax.transData)) mainpath = Path(path) fullpath = mainpath.with_suffix('.full.png') cbpath = mainpath.with_suffix('.cb.png') boundsinfo = mainpath.with_suffix('.info') ax.xaxis.set_visible(False) ax.yaxis.set_visible(False) cb.ax.xaxis.set_visible(False) cb.ax.yaxis.set_visible(False) dpi_trans = fig.dpi_scale_trans.inverted() main_extent = ax.get_window_extent().transformed(dpi_trans) cb_extent = cb.ax.get_window_extent().transformed(dpi_trans) fig.savefig(str(mainpath), bbox_inches=main_extent) fig.savefig(str(cbpath), bbox_inches=cb_extent) fig.savefig(str(fullpath), bbox_inches='tight') with boundsinfo.open('w') as f: print('Main axis bounds:', file=f) print(repr(ax.viewLim.get_points()), file=f) print(file=f) print('Axis limits:', file=f) print(repr(ax.get_xlim()), file=f) print(repr(ax.get_ylim()), file=f) print(file=f) print('Signal scale: ', file=f) print(scale, file=f) print(file=f) print('Colorbar bounds:', file=f) print([cb.boundaries[x] for x in [0, -1]], file=f)