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 makeRGBImages(simulation,
wavelengthTuples=None,
*,
fileType="total",
fromPercentile=30,
toPercentile=100,
outDirPath=None):
# get the (instrument, output file path) tuples to be handled
instr_paths = sm.instrumentOutFilePaths(simulation, fileType + ".fits")
if len(instr_paths) < 1:
return
# convert the wavelength tuples argument into a list of names and wavelength tuples
if wavelengthTuples is None:
name_waves = [("", None)]
elif isinstance(wavelengthTuples, (tuple, list)):
name_waves = [("_" + str(index + 1), wavetuple)
for index, wavetuple in enumerate(wavelengthTuples)]
elif isinstance(wavelengthTuples, dict):
name_waves = [("" if not name else "_" + name, wavetuple)
for name, wavetuple in wavelengthTuples.items()]
else:
raise ValueError("Invalid wavelengthTuples type")
# loop over all wavelength tuples
for name, wavetuple in name_waves:
# determine the appropriate pixel range for ALL images
ranges = []
for instrument, filepath in instr_paths:
image = RGBImage(
filepath, frameIndices=instrument.wavelengthIndices(wavetuple))
ranges += list(
image.percentilePixelRange(fromPercentile, toPercentile))
rmin = min(ranges)
rmax = max(ranges)
# create an RGB file for each output file
for instrument, filepath in instr_paths:
image = RGBImage(
filepath, frameIndices=instrument.wavelengthIndices(wavetuple))
image.setRange(rmin, rmax)
image.applyLog()
image.applyCurve()
# determine output file path
saveFilePath = ut.savePath(filepath.with_name(filepath.stem +
name + ".png"),
".png",
outDirPath=outDirPath)
image.saveTo(saveFilePath)
logging.info("Created {}".format(saveFilePath))
def do(infilepath: (str, "filepath pattern of the FSPS files to be converted"),
outfilepath: (str, "filepath of the resulting stored table file")
) -> "Convert FSPS-generated SED family to stored table format":
import glob
import pts.utils as ut
from pts.storedtable.convert_sed import convertFSPSSEDFamily as convertFSPSSEDFamily
infilepath = str(ut.absPath(infilepath))
if len(glob.glob(infilepath)) == 0:
raise ut.UserError(
"No input files found for pattern: '{}'".format(infilepath))
outfilepath = ut.absPath(outfilepath)
if (outfilepath.is_dir()):
outfilepath = ut.savePath(defFilePath="CustomFSPSSEDFamily",
outDirPath=outfilepath,
suffix=".stab")
else:
outfilepath = ut.savePath(defFilePath=None,
outFilePath=outfilepath,
suffix=".stab")
convertFSPSSEDFamily([infilepath], [outfilepath])
def plotGrids(simulation, *, lineWidth=0.1, outDirPath=None, figSize=(8, 8)):
# loop over the (probe, output file path) tuples
for probe, gridFilePath in sm.probeOutFilePaths(simulation, "grid_*.dat"):
# determine the format type from the first nonempty line (3D format has 3 columns, 2D format has 2 columns)
with open(gridFilePath) as gridfile:
for line in gridfile:
form = len(line.split())
if form > 0: break
# determine save file path
saveFilePath = ut.savePath(gridFilePath, ".pdf", outDirPath=outDirPath)
# use 2D or 3D version
if form == 2:
_plotGrid2D(gridFilePath, saveFilePath, lineWidth, figSize)
else:
_plotGrid3D(gridFilePath, saveFilePath, lineWidth, figSize)
def makeConvolvedRGBImages(simulation, contributions, name="", *, fileType="total", decades=3,
fmax=None, fmin=None, outDirPath=None):
# get the (instrument, output file path) tuples to be handled
instr_paths = sm.instrumentOutFilePaths(simulation, fileType+".fits")
if len(instr_paths) > 0:
name = "" if not name else "_" + name
sbunit = sm.unit("MJy/sr")
# construct an image object with integrated color channels for each output file
# and keep track of the largest surface brightness value (in MJy/sr) in each of the images
images = []
fmaxes = []
for instrument, filepath in instr_paths:
logging.info("Convolving for RGB image {}{}".format(filepath.stem, name))
# get the data cube in its intrinsic units
cube = sm.loadFits(filepath)
# initialize an RGB frame
dataRGB = np.zeros( (cube.shape[0], cube.shape[1], 3) ) << sbunit
# add color for each filter
for band,w0,w1,w2 in contributions:
# convolve and convert to per frequency units
data = band.convolve(instrument.wavelengths(), cube, flavor=sbunit, numWavelengths=10)
dataRGB[:,:,0] += w0*data
dataRGB[:,:,1] += w1*data
dataRGB[:,:,2] += w2*data
# construct the image object
images.append(vis.RGBImage(dataRGB.value))
fmaxes.append(dataRGB.max())
# determine the appropriate pixel range for all output images
fmax = max(fmaxes) if fmax is None else fmax << sbunit
fmin = fmax/10**decades if fmin is None else fmin << sbunit
# create an RGB file for each output file
for (instrument, filepath), image in zip(instr_paths,images):
image.setRange(fmin.value, fmax.value)
image.applyLog()
image.applyCurve()
# determine output file path
saveFilePath = ut.savePath(filepath.with_name(filepath.stem+name+".png"), ".png", outDirPath=outDirPath)
image.saveTo(saveFilePath)
logging.info("Created convolved RGB image file {}".format(saveFilePath))
# return the surface brightness range used for these images
return fmin,fmax
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 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))
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 plotBuiltinBands(minWavelength=1e-6 * u.micron,
maxWavelength=1e6 * u.micron,
nameSegments=None,
*,
outDirPath=None,
outFileName=None,
outFilePath=None,
figSize=(20, 6),
interactive=None):
# load all bands that satisfy the specified criteria
bands = [bnd.BroadBand(name) for name in bnd.builtinBandNames()]
bands = [
band for band in bands
if minWavelength <= band.pivotWavelength() <= maxWavelength
]
if nameSegments is not None:
if isinstance(nameSegments, str): nameSegments = [nameSegments]
bands = [
band for band in bands if any([
s.lower() in band.name().lower().split("_")
for s in nameSegments
])
]
# sort the remaining bands on pivot wavelength
bands = sorted(bands, key=bnd.BroadBand.pivotWavelength)
logging.info("Plotting {} built-in bands...".format(len(bands)))
# setup the figure
plt.figure(figsize=figSize)
colors = ('r', 'g', 'b', 'c', 'm', 'y')
# loop over bands
labelpos = 0.25
colorindex = 0
for band in bands:
wavelengths, transmissions = band.transmissionCurve()
wavelengths <<= u.micron # convert to micron
transmissions /= transmissions.max() # normalize to a maximum of 1
plt.plot(wavelengths.value,
transmissions.value,
color=colors[colorindex])
labelpos += 0.05
if labelpos > 0.69: labelpos = 0.25
plt.text(band.pivotWavelength().to_value(wavelengths.unit),
labelpos,
band.name(),
horizontalalignment='center',
fontsize='x-small',
color=colors[colorindex],
backgroundcolor='w')
colorindex = (colorindex + 1) % len(colors)
# set axis details
plt.xscale('log')
plt.grid(True, axis='y')
# add axis labels and a legend
plt.xlabel(r"$\lambda$" + sm.latexForUnit(wavelengths), fontsize='large')
plt.ylabel("Transmission", fontsize='large')
# if not in interactive mode, save the figure; otherwise leave it open
if not ut.interactive(interactive):
saveFilePath = ut.savePath("FigBuiltinBands.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 plotStoredTableCurve(tableFilePath,
horAxis=0,
verAxis=0,
*,
axis0=None,
axis1=None,
axis2=None,
axis3=None,
axis4=None,
outDirPath=None,
outFileName=None,
outFilePath=None,
figSize=(8, 6),
interactive=None):
# load the complete stored table
table = stab.readStoredTable(tableFilePath)
# get info on horizontal axis
horName = table['axisNames'][horAxis]
horUnit = table['axisUnits'][horAxis]
horScale = table['axisScales'][horAxis]
horGrid = table[horName]
# get info on vertical axis
verName = table['quantityNames'][verAxis]
verUnit = table['quantityUnits'][verAxis]
verScale = table['quantityScales'][verAxis]
# get the appropriate slice from the values hypercube
index = []
for axisName, axisScale, axisValue in zip(
table['axisNames'], table['axisScales'],
(axis0, axis1, axis2, axis3, axis4)):
if axisName == horName:
index.append(Ellipsis)
else:
axisGrid = table[axisName]
if axisValue is None:
if axisScale == 'log':
axisValue = np.sqrt(axisGrid[0] * axisGrid[-1])
else:
axisValue = (axisGrid[0] + axisGrid[-1]) / 2
index.append((np.abs(axisGrid - axisValue)).argmin())
verValues = table[verName][tuple(index)]
# create the plot
plt.figure(figsize=figSize)
if horScale == 'log': plt.xscale('log')
if verScale == 'log': plt.yscale('log')
plt.plot(horGrid, verValues)
plt.xlabel(horName + sm.latexForUnit(horUnit))
plt.ylabel(verName + sm.latexForUnit(verUnit))
# if not in interactive mode, save the figure; otherwise leave it open
if not ut.interactive(interactive):
saveFilePath = ut.savePath("FigStoredTable.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 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))