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 makeWavelengthMovie(simulation, *, maxPercentile=100, minPercentile=10, decades=None, renormalize=False, outDirPath=None, outFileName=None, outFilePath=None, rate=7): # get the list of instruments and corresponding output file paths instrA, sedPaths = zip(*sm.instrumentOutFilePaths(simulation, "sed.dat")) instrB, cubPaths = zip( *sm.instrumentOutFilePaths(simulation, "total.fits")) instrA = instrA[:3] instrB = instrB[:3] sedPaths = sedPaths[:3] cubPaths = cubPaths[:3] if len(instrA) < 1 or len(instrA) != len(instrB) \ or any([ a.name() != b.name() for a,b in zip(instrA,instrB) ]): return instruments = instrA # get the wavelength grid for the first instrument (assumed to be the same for all instruments) wavelengths = instruments[0].wavelengths() if len(wavelengths) < 3: return nlambda = len(wavelengths) # load the data logging.info("Creating movie for {} ({} wavelengths and {} instruments)..." \ .format(instruments[0].prefix(), nlambda, len(instruments))) sedData = [ sm.loadColumns(sedPath, "total flux")[0] for sedPath in sedPaths ] cubData = [sm.loadFits(cubPath).value for cubPath in cubPaths] # determine the shape (assume that frames in all fits files have the same shape) imgShape = cubData[0].shape[:2] sedShape = (len(cubData) * imgShape[0], max(imgShape[1] // 2, 300)) totalShape = (sedShape[0], imgShape[1] + sedShape[1]) # determine the global surface brightness range fmax = max( [np.percentile(np.unique(cube), maxPercentile) for cube in cubData]) if decades is None: fmin = min([ np.percentile(np.unique(cube), minPercentile) for cube in cubData ]) else: fmin = fmax / 10**decades # determine the global integrated flux range Fmax = max([sed.max() for sed in sedData]) Fmin = Fmax * fmin / fmax # open the movie file defSaveFilePath = sedPaths[0].with_name(instruments[0].prefix() + "_wavemovie.mp4") saveFilePath = ut.savePath(defSaveFilePath, (".mp4", ), outDirPath=outDirPath, outFileName=outFileName, outFilePath=outFilePath) movie = MovieFile(saveFilePath, shape=totalShape, rate=rate) # for each wavelength, construct and add a movie frame for frame in range(nlambda): logging.info(" adding frame " + str(frame + 1) + "/" + str(nlambda) + "...") # determine the surface brightness range for this frame, if needed if renormalize: fmax = max([ np.percentile(np.unique(cube[:, :, frame]), maxPercentile) for cube in cubData ]) if decades is None: fmin = min([ np.percentile(np.unique(cube[:, :, frame]), minPercentile) for cube in cubData ]) else: fmin = fmax / 10**decades # assemble the top panel image = None for cube in cubData: im = RGBImage( np.dstack((cube[:, :, frame], cube[:, :, frame], cube[:, :, frame]))) im.setRange(fmin, fmax) im.applyLog() im.applyColorMap("gnuplot2") if image == None: image = im else: image.addRight(im) # plot the seds in the bottom panel dpi = 100 figure = Figure(dpi=dpi, figsize=(sedShape[0] / dpi, sedShape[1] / dpi), facecolor='w', edgecolor='w') canvas = FigureCanvasAgg(figure) ax = figure.add_subplot(111) colors = ('r', 'g', 'b') for sed, instrument, color in zip(sedData, instruments, colors): ax.loglog(wavelengths.value, sed.value, color=color, label=instrument.name()) ax.axvline(wavelengths[frame].value, color='m') ax.set_ylabel( sm.latexForSpectralFlux(sedData[0]) + sm.latexForUnit(sedData[0])) ax.set_ylim(Fmin.value / 1.1, Fmax.value * 1.1) ax.legend(loc='lower right', title=sm.latexForWavelength(wavelengths) \ + r"$={0:.4g}\,$".format(wavelengths[frame].value)+sm.latexForUnit(wavelengths)) canvas.draw() im = RGBImage(figure) image.addBelow(im) # add the frame to the movie movie.addFrame(image) # close the movie file movie.close() logging.info("Created {}".format(saveFilePath))
def plotSeds(simulation, minWavelength=None, maxWavelength=None, decades=None, *, outDirPath=None, outFileName=None, outFilePath=None, figSize=(8, 6), interactive=None): # private function to get the maximum flux within the wavelength range passed to the plotSeds function def maxFluxInRange(flux, wave): wmin = minWavelength if minWavelength is not None else wave[0] wmax = maxWavelength if maxWavelength is not None else wave[-1] mask = (wave>=wmin) & (wave<=wmax) if np.count_nonzero(mask) > 0: return flux[mask].max() else: return flux.max() # get the (instrument, output file path) tuples instr_paths = sm.instrumentOutFilePaths(simulation, "sed.dat") if len(instr_paths) < 1: return # setup the figure plt.figure(figsize=figSize) # if there is a single output file, and it has components, plot the components if len(instr_paths) == 1 and any(["transparent" in col for col in sm.getColumnDescriptions(instr_paths[0][1])]): instrument, filepath = instr_paths[0] # load the columns (we assume that all components have the same units) wave, tot, tra, dirpri, scapri, dirsec, scasec = sm.loadColumns(filepath, (0, 1, 2, 3, 4, 5, 6)) waveUnit = wave.unit fluxUnit = tot.unit fluxMax = max(maxFluxInRange(tot, wave), maxFluxInRange(tra, wave)) # plot the various components label = "{} {} ".format(instrument.prefix(), instrument.name()) plt.plot(wave.value, tot.value, color='k', ls='solid', label=label + "total") plt.plot(wave.value, tra.value, color='b', ls='dotted', label=label + "transparent") plt.plot(wave.value, (dirpri + scapri).value, color='b', ls='solid', label=label + "primary") plt.plot(wave.value, (dirsec + scasec).value, color='r', ls='solid', label=label + "secondary") # otherwise loop over all SEDs else: colors = ('r', 'g', 'b', 'c', 'm', 'y') colorindex = 0 first = True for instrument, filepath in instr_paths: # load the total flux; first time remember units; thereafter convert units wave, flux = sm.loadColumns(filepath, (0, 1)) if first: waveUnit = wave.unit fluxUnit = flux.unit fluxMax = maxFluxInRange(flux, wave) first = False else: wave <<= waveUnit flux = sm.convertToFlavor(wave, flux, fluxUnit) fluxMax = max(fluxMax, maxFluxInRange(flux, wave)) # plot plt.plot(wave.value, flux.value, color=colors[colorindex], label="{} {} total".format(instrument.prefix(), instrument.name())) # advance color index colorindex = (colorindex + 1) % len(colors) # set axis details and add a legend plt.xscale('log') plt.yscale('log') plt.xlim(_adjustWavelengthRange(plt.xlim(), waveUnit, minWavelength, maxWavelength)) if decades is not None: plt.ylim(fluxMax.value * 10 ** (-decades), fluxMax.value * 10 ** 0.2) plt.xlabel(sm.latexForWavelengthWithUnit(waveUnit), fontsize='large') plt.ylabel(sm.latexForSpectralFluxWithUnit(fluxUnit), fontsize='large') plt.legend(loc='best') # if not in interactive mode, save the figure; otherwise leave it open if not ut.interactive(interactive): # use the first instrument output path; if there are multiple instruments, remove the instrument name defSaveFilePath = instr_paths[0][1] if len(instr_paths) > 1: defSaveFilePath = defSaveFilePath.with_name(instr_paths[0][0].prefix() + "_sed.pdf") saveFilePath = ut.savePath(defSaveFilePath, (".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 plotSources(simulation, minWavelength=None, maxWavelength=None, decades=None, *, outDirPath=None, outFileName=None, outFilePath=None, figSize=(8, 6), interactive=None): # find the required probes probes = simulation.probes() lumiProbes = [ probe for probe in probes if probe.type() == "LuminosityProbe" ] packProbes = [ probe for probe in probes if probe.type() == "LaunchedPacketsProbe" ] if len(lumiProbes) != 1 or len(packProbes) != 1: return lumiProbe = lumiProbes[0] packProbe = packProbes[0] # load the luminosities lumiFilePath = lumiProbe.outFilePaths("luminosities.dat")[0] descriptions = sm.getColumnDescriptions(lumiFilePath) columns = list(range(len(descriptions))) del columns[1] # remove the "specific luminosity column" lumiWave, lumiTot, *lumiFracs = sm.loadColumns(lumiFilePath, columns) # load the number of launched packets packFilePath = packProbe.outFilePaths("launchedpackets.dat")[0] descriptions = sm.getColumnDescriptions(packFilePath) columns = list(range(len(descriptions))) packWave, packTot, *packSplits = sm.loadColumns(packFilePath, columns) packWave <<= lumiWave.unit # setup the figure plt.figure(figsize=figSize) label = "{} ".format(simulation.prefix()) # plot the total lumiMax = lumiTot.max() packMax = packTot.max() plt.plot(lumiWave.value, lumiTot/lumiMax, color='k', ls='solid', label=label + "total") plt.plot(packWave.value, packTot/packMax, color='k', ls='dashed') # loop over all sources colors = ('r', 'g', 'b', 'c', 'm', 'y') colorindex = 0 sourceindex = 1 for lumiFrac, packSplit in zip(lumiFracs, packSplits): plt.plot(lumiWave.value, lumiFrac*lumiTot/lumiMax, color=colors[colorindex], ls='solid', label=label + str(sourceindex)) plt.plot(packWave.value, packSplit/packMax, color=colors[colorindex], ls='dashed') # advance color and source index colorindex = (colorindex + 1) % len(colors) sourceindex += 1 # set axis details and add a legend plt.xscale('log') plt.yscale('log') plt.xlim(_adjustWavelengthRange(plt.xlim(), lumiWave.unit, minWavelength, maxWavelength)) if decades is not None: plt.ylim(10 ** (-decades), 10 ** 0.2) plt.xlabel(sm.latexForWavelengthWithUnit(lumiWave.unit), fontsize='large') plt.ylabel(r"Normalized $L$ and $N_\mathrm{pp}$", fontsize='large') plt.legend(loc='best') # if not in interactive mode, save the figure; otherwise leave it open if not ut.interactive(interactive): saveFilePath = ut.savePath(simulation.outFilePath("sources.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))