def plotSpectralResolution(inFilePath, minWavelength=None, maxWavelength=None, decades=None, *, title=None, outDirPath=None, outFileName=None, outFilePath=None, figSize=(8, 5), interactive=None): # load the wavelength grid inFilePath = ut.absPath(inFilePath) if inFilePath.suffix.lower() == ".stab": table = stab.readStoredTable(inFilePath) if "lambda" not in table: raise ValueError("No wavelength axis in stored table: {}".format(inFilePath)) grid = table["lambda"] elif inFilePath.suffix.lower() == ".dat": if "wavelength" not in sm.getColumnDescriptions(inFilePath)[0].lower(): raise ValueError("First text column is not labeled 'wavelength': {}".format(inFilePath)) grid = sm.loadColumns(inFilePath, "1")[0] elif inFilePath.suffix.lower() == ".fits": axes = sm.getFitsAxes(inFilePath) if len(axes) != 3: raise ValueError("FITS file does not have embedded wavelength axis") grid = axes[2] else: raise ValueError("Filename does not have the .stab, .dat, or .fits extension: {}".format(inFilePath)) # calculate the spectral resolution R = grid[:-1] / (grid[1:] - grid[:-1]) Rmax = R.max() # choose wavelength units from grid wunit = grid.unit # setup the plot plt.figure(figsize=figSize) plt.xlabel(sm.latexForWavelengthWithUnit(wunit), fontsize='large') plt.ylabel(r"$R=\frac{\lambda}{\Delta\lambda}$", fontsize='large') plt.xscale('log') plt.yscale('log') plt.grid(which='major', axis='both', ls=":") plt.xlim(_adjustWavelengthRange(plt.xlim(), wunit, minWavelength, maxWavelength)) if decades is not None: plt.ylim(Rmax* 10 ** (-decades), Rmax * 10 ** 0.2) # plot the spectral resolution if title is None or len(title)==0: title = inFilePath.stem label = "{}\n{} pts from {:g} to {:g} {}".format(title, len(grid), grid[0].to_value(wunit), grid[-1].to_value(wunit), sm.latexForUnit(wunit)) plt.plot(grid[:-1].to_value(wunit), R, label=label) plt.legend() # if not in interactive mode, save the figure; otherwise leave it open if not ut.interactive(interactive): saveFilePath = ut.savePath(inFilePath.stem+".pdf", (".pdf",".png"), outDirPath=outDirPath, outFileName=outFileName, outFilePath=outFilePath) plt.savefig(saveFilePath, bbox_inches='tight', pad_inches=0.25) plt.close() logging.info("Created {}".format(saveFilePath))
def plotDefaultDustTemperatureCuts(simulation, *, outDirPath=None, outFileName=None, outFilePath=None, figSize=None, interactive=None): # find the relevant probe probes = [ probe for probe in simulation.probes() if probe.type() == "DefaultDustTemperatureCutsProbe" ] if len(probes) != 1: return probe = probes[0] # load the temperature cuts and the range of the x and y axes # (there can be one to three cuts depending on symmetries) paths = probe.outFilePaths("dust_T_*.fits") numcuts = len(paths) if not numcuts in (1, 2, 3): return cuts = [path.stem.split("_")[-1] for path in paths] frames = [sm.loadFits(path) for path in paths] grids = [sm.getFitsAxes(path) for path in paths] # determine the maximum temperature value to display Tmax = max([frame.max() for frame in frames]) # setup the figure depending on the number of cuts if figSize is None: figSize = (8 * numcuts, 6) fig, axes = plt.subplots(ncols=numcuts, nrows=1, figsize=figSize) if numcuts == 1: axes = [axes] # plot the cuts and set axis details for each for ax, cut, frame, (xgrid, ygrid) in zip(axes, cuts, frames, grids): extent = (xgrid[0].value, xgrid[-1].value, ygrid[0].value, ygrid[-1].value) im = ax.imshow(frame.value.T, vmin=0, vmax=Tmax.value, cmap='gnuplot', extent=extent, aspect='auto', interpolation='bicubic', origin='lower') ax.set_xlim(xgrid[0].value, xgrid[-1].value) ax.set_ylim(ygrid[0].value, ygrid[-1].value) ax.set_xlabel(cut[0] + sm.latexForUnit(xgrid.unit), fontsize='large') ax.set_ylabel(cut[-1] + sm.latexForUnit(ygrid.unit), fontsize='large') ax.set_ylabel(cut[-1] + sm.latexForUnit(ygrid.unit), fontsize='large') # add a color bar fig.colorbar(im, ax=axes).ax.set_ylabel("T" + sm.latexForUnit(frame.unit), fontsize='large') # if not in interactive mode, save the figure; otherwise leave it open if not ut.interactive(interactive): saveFilePath = ut.savePath(simulation.outFilePath( "{}_dust_T.pdf".format(probe.name())), (".pdf", ".png"), outDirPath=outDirPath, outFileName=outFileName, outFilePath=outFilePath) plt.savefig(saveFilePath, bbox_inches='tight', pad_inches=0.25) plt.close() logging.info("Created {}".format(saveFilePath))
def plotDefaultMediaDensityCuts(simulation, decades=5, *, outDirPath=None, figSize=(18, 6), interactive=None): # find the relevant probe probes = [ probe for probe in simulation.probes() if probe.type() == "DefaultMediaDensityCutsProbe" ] if len(probes) != 1: return probe = probes[0] for medium in ("dust", "elec", "gas"): for cut in ("xy", "xz", "yz"): # load the theoretical and gridded cuts for the requested medium and cut plane tPaths = probe.outFilePaths("{}_t_{}.fits".format(medium, cut)) gPaths = probe.outFilePaths("{}_g_{}.fits".format(medium, cut)) if len(tPaths) == 1 and len(gPaths) == 1: tFrame = sm.loadFits(tPaths[0]) gFrame = sm.loadFits(gPaths[0]) # determine the range of the x and y axes xgrid, ygrid = sm.getFitsAxes(tPaths[0]) # determine the range of density values to display and clip the data arrays vmax = max(tFrame.max(), gFrame.max()) vmin = vmax / 10**decades tFrame[tFrame < vmin] = vmin gFrame[gFrame < vmin] = vmin # setup the figure fig, (ax1, ax2) = plt.subplots(ncols=2, nrows=1, figsize=figSize) # plot the cuts and a color bar (logarithmic normalizer crashes if all values are zero) if vmax > 0: normalizer = matplotlib.colors.LogNorm( vmin.value, vmax.value) else: normalizer = matplotlib.colors.Normalize( vmin.value, vmax.value) extent = (xgrid[0].value, xgrid[-1].value, ygrid[0].value, ygrid[-1].value) im = ax1.imshow(tFrame.value.T, norm=normalizer, cmap='gnuplot', extent=extent, aspect='auto', interpolation='bicubic', origin='lower') fig.colorbar(im, ax=(ax1, ax2)).ax.set_ylabel( "density" + sm.latexForUnit(tFrame.unit), fontsize='large') ax2.imshow(gFrame.value.T, norm=normalizer, cmap='gnuplot', extent=extent, aspect='auto', interpolation='bicubic', origin='lower') # set axis details ax1.set_xlim(xgrid[0].value, xgrid[-1].value) ax1.set_ylim(ygrid[0].value, ygrid[-1].value) ax2.set_xlim(xgrid[0].value, xgrid[-1].value) ax2.set_ylim(ygrid[0].value, ygrid[-1].value) ax1.set_xlabel(cut[0] + sm.latexForUnit(xgrid.unit), fontsize='large') ax1.set_ylabel(cut[-1] + sm.latexForUnit(ygrid.unit), fontsize='large') ax2.set_xlabel(cut[0] + sm.latexForUnit(xgrid.unit), fontsize='large') ax2.set_ylabel(cut[-1] + sm.latexForUnit(ygrid.unit), fontsize='large') # if not in interactive mode, save the figure; otherwise leave it open if not ut.interactive(interactive): defSavePath = simulation.outFilePath("{}_{}_{}.pdf".format( probe.name(), medium, cut)) saveFilePath = ut.savePath(defSavePath, (".pdf", ".png"), outDirPath=outDirPath) plt.savefig(saveFilePath, bbox_inches='tight', pad_inches=0.25) plt.close() logging.info("Created {}".format(saveFilePath))
def plotPolarization(simulation, *, plotLinMap=True, plotDegMap=False, plotDegAvg=False, plotCirMap=False, wavelength=None, binSize=(7, 7), degreeScale=None, decades=5, outDirPath=None, figSize=(8, 6), interactive=None): # loop over all applicable instruments for instrument, filepath in sm.instrumentOutFilePaths( simulation, "stokesQ.fits"): # form the simulation/instrument name insname = "{}_{}".format(instrument.prefix(), instrument.name()) # get the file paths for the frames/data cubes filepathI = instrument.outFilePaths("total.fits")[0] filepathQ = instrument.outFilePaths("stokesQ.fits")[0] filepathU = instrument.outFilePaths("stokesU.fits")[0] filepathV = instrument.outFilePaths("stokesV.fits")[0] # load datacubes with shape (nx, ny, nlambda) Is = sm.loadFits(filepathI) Qs = sm.loadFits(filepathQ) Us = sm.loadFits(filepathU) Vs = sm.loadFits(filepathV) # load the axes grids (assuming all files have the same axes) xgrid, ygrid, wavegrid = sm.getFitsAxes(filepathI) xmin = xgrid[0].value xmax = xgrid[-1].value ymin = ygrid[0].value ymax = ygrid[-1].value extent = (xmin, xmax, ymin, ymax) # determine binning configuration binX = binSize[0] orLenX = Is.shape[0] dropX = orLenX % binX startX = dropX // 2 binY = binSize[1] orLenY = Is.shape[1] dropY = orLenY % binY startY = dropY // 2 # construct arrays with central bin positions in pixel coordinates posX = np.arange(startX - 0.5 + binX / 2.0, orLenX - dropX + startX - 0.5, binX) posY = np.arange(startY - 0.5 + binY / 2.0, orLenY - dropY + startY - 0.5, binY) # determine the appropriate wavelength index or indices if wavelength is None: indices = [0] elif wavelength == 'all': indices = range(Is.shape[2]) else: if not isinstance(wavelength, (list, tuple)): wavelength = [wavelength] indices = instrument.wavelengthIndices(wavelength) # loop over all requested wavelength indices for index in indices: wave = wavegrid[index] wavename = "{:09.4f}".format(wave.to_value(sm.unit("micron"))) wavelatex = r"$\lambda={:.4g}$".format( wave.value) + sm.latexForUnit(wave) # extract the corresponding frame, and transpose to (y,x) style for compatibility with legacy code I = Is[:, :, index].T.value Q = Qs[:, :, index].T.value U = Us[:, :, index].T.value V = Vs[:, :, index].T.value # perform the actual binning binnedI = np.zeros((len(posY), len(posX))) binnedQ = np.zeros((len(posY), len(posX))) binnedU = np.zeros((len(posY), len(posX))) binnedV = np.zeros((len(posY), len(posX))) for x in range(len(posX)): for y in range(len(posY)): binnedI[y, x] = np.sum( I[startY + binY * y:startY + binY * (y + 1), startX + binX * x:startX + binX * (x + 1)]) binnedQ[y, x] = np.sum( Q[startY + binY * y:startY + binY * (y + 1), startX + binX * x:startX + binX * (x + 1)]) binnedU[y, x] = np.sum( U[startY + binY * y:startY + binY * (y + 1), startX + binX * x:startX + binX * (x + 1)]) binnedV[y, x] = np.sum( V[startY + binY * y:startY + binY * (y + 1), startX + binX * x:startX + binX * (x + 1)]) # ----------------------------------------------------------------- # plot a linear polarization map if plotLinMap: fig, ax = plt.subplots(ncols=1, nrows=1, figsize=figSize) # configure the axes ax.set_xlim(xmin, xmax) ax.set_ylim(ymin, ymax) ax.set_xlabel("x" + sm.latexForUnit(xgrid), fontsize='large') ax.set_ylabel("y" + sm.latexForUnit(ygrid), fontsize='large') # determine intensity range for the background image, ignoring pixels with outrageously high flux Ib = I.copy() highmask = Ib > 1e6 * np.nanmedian(np.unique(Ib)) vmax = np.nanmax(Ib[~highmask]) Ib[highmask] = vmax vmin = vmax / 10**decades # plot the background image and the corresponding color bar normalizer = matplotlib.colors.LogNorm(vmin, vmax) cmap = plt.get_cmap('PuRd') cmap.set_under('w') backPlot = ax.imshow(Ib, norm=normalizer, cmap=cmap, extent=extent, aspect='equal', interpolation='bicubic', origin='lower') cbarlabel = sm.latexForSpectralFlux(Is) + sm.latexForUnit( Is) + " @ " + wavelatex plt.colorbar(backPlot, ax=ax).set_label(cbarlabel, fontsize='large') # compute the linear polarization degree degreeLD = np.sqrt(binnedQ**2 + binnedU**2) degreeLD[degreeLD > 0] /= binnedI[degreeLD > 0] # determine a characteristic 'high' degree of polarization in the frame # (this has to be done before degreeLD contains 'np.NaN') charDegree = np.percentile(degreeLD, 99.0) if not 0 < charDegree < 1: charDegree = np.nanmax((np.nanmax(degreeLD), 0.0001)) # remove pixels with minuscule polarization degreeLD[degreeLD < charDegree / 50] = np.NaN # determine the scaling so that the longest arrows do not to overlap with neighboring arrows if degreeScale is None: degreeScale = _roundUp(charDegree) lengthScale = 2.2 * degreeScale * max( float(len(posX)) / figSize[0], float(len(posY)) / figSize[1]) key = "{:.3g}%".format(100 * degreeScale) # compute the polarization angle angle = 0.5 * np.arctan2( binnedU, binnedQ ) # angle from North through East while looking at the sky # create the polarization vector arrays xPolarization = -degreeLD * np.sin( angle ) #For angle = 0: North & x=0, For angle = 90deg: West & x=-1 yPolarization = degreeLD * np.cos( angle ) #For angle = 0: North & y=1, For angle = 90deg: West & y=0 # plot the vector field (scale positions to data coordinates) X, Y = np.meshgrid(xmin + posX * (xmax - xmin) / orLenX, ymin + posY * (ymax - ymin) / orLenY) quiverPlot = ax.quiver(X, Y, xPolarization, yPolarization, pivot='middle', units='inches', angles='xy', scale=lengthScale, scale_units='inches', headwidth=0, headlength=1, headaxislength=1, minlength=0.8, width=0.02) ax.quiverkey(quiverPlot, 0.85, 0.02, degreeScale, key, coordinates='axes', labelpos='E') # if not in interactive mode, save the figure; otherwise leave it open if not ut.interactive(interactive): saveFilePath = ut.savePath( filepath, ".pdf", outDirPath=outDirPath, outFileName="{}_{}_pollinmap.pdf".format( insname, wavename)) plt.savefig(saveFilePath, bbox_inches='tight', pad_inches=0.25) plt.close() logging.info("Created {}".format(saveFilePath)) # ----------------------------------------------------------------- # plot a linear polarization degree map if plotDegMap: fig, ax = plt.subplots(ncols=1, nrows=1, figsize=figSize) # configure the axes ax.set_xlim(xmin, xmax) ax.set_ylim(ymin, ymax) ax.set_xlabel("x" + sm.latexForUnit(xgrid), fontsize='large') ax.set_ylabel("y" + sm.latexForUnit(ygrid), fontsize='large') # calculate polarization degree for each pixel, in percent # set degree to zero for pixels with very low intensity cutmask = I < (np.nanmax( I[I < 1e6 * np.nanmedian(np.unique(I))]) / 10**decades) degreeHD = np.sqrt(Q**2 + U**2) degreeHD[~cutmask] /= I[~cutmask] degreeHD[cutmask] = 0 degreeHD *= 100 # plot the image and the corresponding color bar vmax = degreeScale if degreeScale is not None else np.percentile( degreeHD, 99) normalizer = matplotlib.colors.Normalize(vmin=0, vmax=vmax) backPlot = ax.imshow(degreeHD, norm=normalizer, cmap='plasma', extent=extent, aspect='equal', interpolation='bicubic', origin='lower') plt.colorbar(backPlot, ax=ax).set_label( "Linear polarization degree (%)" + " @ " + wavelatex, fontsize='large') # if not in interactive mode, save the figure; otherwise leave it open if not ut.interactive(interactive): saveFilePath = ut.savePath( filepath, ".pdf", outDirPath=outDirPath, outFileName="{}_{}_poldegmap.pdf".format( insname, wavename)) plt.savefig(saveFilePath, bbox_inches='tight', pad_inches=0.25) plt.close() logging.info("Created {}".format(saveFilePath)) # ----------------------------------------------------------------- # plot the y-axis averaged linear polarization degree if plotDegAvg: # construct the plot fig, ax = plt.subplots(ncols=1, nrows=1, figsize=figSize) degreeHD = np.sqrt( np.average(Q, axis=0)**2 + np.average(U, axis=0)**2) degreeHD /= np.average(I, axis=0) ax.plot(xgrid.value, degreeHD * 100) ax.set_xlim(xmin, xmax) ax.set_ylim(0, degreeScale) ax.set_title("{} {}".format(insname, wavelatex), fontsize='large') ax.set_xlabel("x" + sm.latexForUnit(xgrid), fontsize='large') ax.set_ylabel('Average linear polarization degree (%)', fontsize='large') # if not in interactive mode, save the figure; otherwise leave it open if not ut.interactive(interactive): saveFilePath = ut.savePath( filepathI, ".pdf", outDirPath=outDirPath, outFileName="{}_{}_poldegavg.pdf".format( insname, wavename)) plt.savefig(saveFilePath, bbox_inches='tight', pad_inches=0.25) plt.close() logging.info("Created {}".format(saveFilePath)) # ----------------------------------------------------------------- # plot a circular polarization map if plotCirMap: fig, ax = plt.subplots(ncols=1, nrows=1, figsize=figSize) # configure the axes ax.set_xlim(xmin, xmax) ax.set_ylim(ymin, ymax) ax.set_xlabel("x" + sm.latexForUnit(xgrid), fontsize='large') ax.set_ylabel("y" + sm.latexForUnit(ygrid), fontsize='large') # determine intensity range for the background image, ignoring pixels with outrageously high flux Ib = I.copy() highmask = Ib > 1e6 * np.nanmedian(np.unique(Ib)) vmax = np.nanmax(Ib[~highmask]) Ib[highmask] = vmax vmin = vmax / 10**decades # plot the background image and the corresponding color bar normalizer = matplotlib.colors.LogNorm(vmin, vmax) cmap = plt.get_cmap('PuRd') cmap.set_under('w') backPlot = ax.imshow(Ib, norm=normalizer, cmap=cmap, extent=extent, aspect='equal', interpolation='bicubic', origin='lower') cbarlabel = sm.latexForSpectralFlux(Is) + sm.latexForUnit( Is) + " @ " + wavelatex plt.colorbar(backPlot, ax=ax).set_label(cbarlabel, fontsize='large') # compute the circular polarization degree degreeLD = binnedV.copy() degreeLD[binnedI > 0] /= binnedI[binnedI > 0] # determine the scaling and add legend if degreeScale is None: degreeScale = _roundUp(np.percentile(np.abs(degreeLD), 99)) lengthScale = 0.7 / max(len(posX), len(posY)) _circArrow(ax, 0.84 - lengthScale / 2, 0.01 + lengthScale / 2, lengthScale) key = r'$+{} \%$'.format(100 * degreeScale) ax.text(0.85, 0.01 + lengthScale / 2, key, transform=ax.transAxes, ha='left', va='center') # actual plotting for x in range(len(posX)): for y in range(len(posY)): if np.isfinite(degreeLD[y, x]) and abs( degreeLD[y, x]) > degreeScale / 50: _circArrow( ax, posX[x] / orLenX, posY[y] / orLenY, degreeLD[y, x] / degreeScale * lengthScale) # if not in interactive mode, save the figure; otherwise leave it open if not ut.interactive(interactive): saveFilePath = ut.savePath( filepath, ".pdf", outDirPath=outDirPath, outFileName="{}_{}_polcirmap.pdf".format( insname, wavename)) plt.savefig(saveFilePath, bbox_inches='tight', pad_inches=0.25) plt.close() logging.info("Created {}".format(saveFilePath))
def plotMagneticFieldCuts(simulation, *, binSize=(32, 32), outDirPath=None, figSize=(6, 6), interactive=None): # find the relevant probes probes = [ probe for probe in simulation.probes() \ if probe.type() in ("DefaultMagneticFieldCutsProbe", "PlanarMagneticFieldCutsProbe") ] # iterate over them for probe in probes: for cut in ("xy", "xz", "yz"): # load magnetic field for the this probe and cut paths = probe.outFilePaths("{}.fits".format(cut)) if len(paths) == 1: # load data cube with shape (nx, ny, 3) Bs = sm.loadFits(paths[0]) # load the axes grids xgrid, ygrid, dummygrid = sm.getFitsAxes(paths[0]) xmin = xgrid[0].value xmax = xgrid[-1].value ymin = ygrid[0].value ymax = ygrid[-1].value extent = (xmin, xmax, ymin, ymax) # determine binning configuration binX = binSize[0] orLenX = Bs.shape[0] dropX = orLenX % binX startX = dropX // 2 binY = binSize[1] orLenY = Bs.shape[1] dropY = orLenY % binY startY = dropY // 2 # construct arrays with central bin positions in pixel coordinates posX = np.arange(startX - 0.5 + binX / 2.0, orLenX - dropX + startX - 0.5, binX) posY = np.arange(startY - 0.5 + binY / 2.0, orLenY - dropY + startY - 0.5, binY) # perform the actual binning, while splitting in vector components Bx = np.zeros((len(posX), len(posY))) By = np.zeros((len(posX), len(posY))) Bz = np.zeros((len(posX), len(posY))) for x in range(len(posX)): for y in range(len(posY)): Bx[x, y] = np.mean( Bs[startX + binX * x:startX + binX * (x + 1), startY + binY * y:startY + binY * (y + 1), 0].value) By[x, y] = np.mean( Bs[startX + binX * x:startX + binX * (x + 1), startY + binY * y:startY + binY * (y + 1), 1].value) Bz[x, y] = np.mean( Bs[startX + binX * x:startX + binX * (x + 1), startY + binY * y:startY + binY * (y + 1), 2].value) # start the figure fig, ax = plt.subplots(ncols=1, nrows=1, figsize=figSize) # configure the axes ax.set_xlim(xmin, xmax) ax.set_ylim(ymin, ymax) ax.set_xlabel(cut[0] + sm.latexForUnit(xgrid), fontsize='large') ax.set_ylabel(cut[-1] + sm.latexForUnit(ygrid), fontsize='large') ax.set_aspect('equal') # determine a characteristic 'large' field strength in the cut plane Bmax = np.percentile(np.sqrt(Bx**2 + By**2), 99.0) if Bmax == 0: Bmax = 1 # guard against all zeros # determine the scaling so that the longest arrows do not to overlap with neighboring arrows lengthScale = 2 * Bmax * max( float(len(posX)) / figSize[0], float(len(posY)) / figSize[1]) key = "{:.3g}{}".format(Bmax, sm.latexForUnit(Bs)) # determine the color scheme for the component orthogonal to cut plane Bzmax = np.abs(Bz).max() if Bzmax == 0: Bzmax = 1 # guard against all zeros normalizer = matplotlib.colors.Normalize(-Bzmax, Bzmax) # plot the vector field (scale positions to data coordinates) X, Y = np.meshgrid(xmin + posX * (xmax - xmin) / orLenX, ymin + posY * (ymax - ymin) / orLenY, indexing='ij') quiverPlot = ax.quiver(X, Y, Bx, By, Bz, cmap='jet', norm=normalizer, pivot='middle', units='inches', angles='xy', scale=lengthScale, scale_units='inches', width=0.015, headwidth=2.5, headlength=2, headaxislength=2, minlength=0.8) ax.quiverkey(quiverPlot, 0.8, -0.08, Bmax, key, coordinates='axes', labelpos='E') # if not in interactive mode, save the figure; otherwise leave it open if not ut.interactive(interactive): saveFilePath = ut.savePath(simulation.outFilePath( "{}_B_{}.pdf".format(probe.name(), cut)), (".pdf", ".png"), outDirPath=outDirPath) plt.savefig(saveFilePath, bbox_inches='tight', pad_inches=0.25) plt.close() logging.info("Created {}".format(saveFilePath))