|
#!/usr/bin/env python |
|
# encoding: utf-8 |
|
""" |
|
*Plot an FITS binary table spectrum (packaged in the ESO data-product standard format)* |
|
|
|
:Author: |
|
David Young |
|
|
|
:Date Created: |
|
March 26, 2021 |
|
|
|
Usage: |
|
plot_fits_binary_table_spectrum [-f] <fits1d>... [--i=<fits2d>] [--o=<outputFilename>] |
|
plot_fits_binary_table_spectrum --option <argument> [<optional_argument>] |
|
|
|
Options: |
|
-h, --help show this help message |
|
-f, --fancy add a colour range to the spectum plot moving from red to blue |
|
--i, --image add panel showing 2D spectral image below plot |
|
--o, --output filename to output the plot to |
|
""" |
|
################# GLOBAL IMPORTS #################### |
|
from __future__ import division |
|
import sys |
|
import os |
|
from fundamentals import tools |
|
from astropy.io import fits |
|
from astropy import units as u |
|
import numpy as np |
|
import matplotlib as mpl |
|
from matplotlib import pyplot as plt |
|
from astropy.visualization import quantity_support |
|
from matplotlib.collections import LineCollection |
|
from matplotlib.colors import ListedColormap, BoundaryNorm |
|
from astropy.nddata import CCDData |
|
from colour import Color |
|
from matplotlib.colors import LinearSegmentedColormap |
|
from ccdproc import ImageFileCollection |
|
|
|
|
|
def main(arguments=None): |
|
""" |
|
*The main function used when ``plot_fits_binary_table_spectrum.py`` is run as a single script from the cl* |
|
""" |
|
|
|
# SETUP THE COMMAND-LINE UTIL SETTINGS |
|
su = tools( |
|
arguments=arguments, |
|
docString=__doc__, |
|
logLevel="WARNING", |
|
options_first=False, |
|
projectName=False |
|
) |
|
arguments, settings, log, dbConn = su.setup() |
|
|
|
# ADDED FOR UNIT SUPPORT |
|
from astropy.visualization import quantity_support |
|
quantity_support() |
|
|
|
# UNPACK REMAINING CL ARGUMENTS USING `EXEC` TO SETUP THE VARIABLE NAMES |
|
# AUTOMATICALLY |
|
a = {} |
|
for arg, val in list(arguments.items()): |
|
if arg[0] == "-": |
|
varname = arg.replace("-", "") + "Flag" |
|
else: |
|
varname = arg.replace("<", "").replace(">", "") |
|
a[varname] = val |
|
if arg == "--dbConn": |
|
dbConn = val |
|
a["dbConn"] = val |
|
log.debug('%s = %s' % (varname, val,)) |
|
|
|
# SET THE VARIOUS MPL SETTINGS |
|
fontColor = None |
|
backgroundColor = None |
|
if a['iFlag']: |
|
axCount = 2 |
|
ratio = [5, 2] |
|
gridspec_kw = {'height_ratios': ratio} |
|
else: |
|
axCount = 1 |
|
gridspec_kw = {} |
|
|
|
if a["fancyFlag"]: |
|
cmap = "gist_rainbow" |
|
# print(mpl.rcParams.keys()) |
|
fontColor = '#93a1a1' |
|
backgroundColor = '#101010' |
|
mpl.rcParams['text.color'] = fontColor |
|
mpl.rcParams['axes.labelcolor'] = fontColor |
|
mpl.rcParams['xtick.color'] = fontColor |
|
mpl.rcParams['ytick.color'] = fontColor |
|
mpl.rcParams['axes.edgecolor'] = fontColor |
|
mpl.rcParams['axes.linewidth'] = .2 |
|
mpl.rcParams['ytick.major.width'] = .2 |
|
mpl.rcParams['xtick.major.width'] = .2 |
|
mpl.rcParams['axes.spines.right'] = False |
|
mpl.rcParams['axes.spines.top'] = False |
|
else: |
|
cmap = False |
|
mpl.rcParams['axes.linewidth'] = .2 |
|
|
|
# FIGURE AND AXES FOR PLOTS |
|
fig, axs = plt.subplots(axCount, 1, sharex=False, |
|
sharey=False, gridspec_kw=gridspec_kw) |
|
if axCount > 1: |
|
ax1 = axs[0] |
|
ax2 = axs[1] |
|
else: |
|
ax1 = axs |
|
ax2 = False |
|
|
|
# PLOT 1D SPECTRA |
|
fitsPaths = plot_1d_spectrum(log=log, fitsPaths=a[ |
|
"fits1d"], ax=ax1, cmap=cmap) |
|
|
|
# ADD OPTIONAL 2D IMAGE PANEL |
|
if ax2: |
|
plot_2d_spectrum(log=log, fitsPath=a[ |
|
"iFlag"], ax=ax2, cmapName=cmap, backgroundColor=backgroundColor) |
|
ax1.spines['bottom'].set_position(('axes', -ratio[1] / ratio[0])) |
|
|
|
# SORT LABELS AT AXES TICKS |
|
if len(fitsPaths) > 1: |
|
ax1.set_ylabel( |
|
u.Unit('erg cm-2 s-1 AA-1').to_string('latex_inline') + " + Const.") |
|
ax1.set_yticklabels([]) |
|
else: |
|
ax1.set_ylabel( |
|
u.Unit('erg cm-2 s-1 AA-1').to_string('latex_inline')) |
|
ax1.set_xlabel(u.AA.to_string('latex_inline')) |
|
ax1.zorder = 20 |
|
ax1.tick_params(axis='x', which='both') |
|
|
|
if a["fancyFlag"]: |
|
ax1.tick_params(color=fontColor) |
|
ax1.set_facecolor(backgroundColor) |
|
fig.patch.set_facecolor(backgroundColor) |
|
ax1.grid(False) |
|
|
|
plt.subplots_adjust(hspace=-.001) |
|
|
|
# SAVE PLOT TO FILE OR SHOW |
|
if a["oFlag"]: |
|
fileName = a["oFlag"] |
|
if a["fancyFlag"]: |
|
plt.savefig(fileName, bbox_inches='tight', |
|
facecolor=backgroundColor) |
|
else: |
|
plt.savefig(fileName, bbox_inches='tight') |
|
plt.clf() # clear figure |
|
else: |
|
plt.show() |
|
|
|
return |
|
|
|
|
|
def plot_1d_spectrum( |
|
log, |
|
fitsPaths, |
|
ax, |
|
cmap=False): |
|
"""*plot the 1D spectrum/spectra* |
|
|
|
**Key Arguments:** |
|
|
|
- `log` -- logger |
|
- `fitsPaths` -- paths to one or more FITS binary tables containing spectra (string or list) |
|
- `ax` -- matplotlib ax to use |
|
- `cmap` -- the colormap to use. Default *False* (monochrome) |
|
|
|
|
|
""" |
|
log.debug('starting the ``plot_1d_spectrum`` function') |
|
|
|
# PACK FITS FILEPATHS INTO A LIST |
|
if isinstance(fitsPaths, str): |
|
fitsPaths = [fitsPaths] |
|
if len(fitsPaths) == 1 and os.path.isdir(fitsPaths[0]): |
|
# GENERATE A LIST OF FILE PATHS |
|
pathToDirectory = fitsPaths[0] |
|
fitsPaths = [] |
|
for d in os.listdir(pathToDirectory): |
|
filepath = os.path.join(pathToDirectory, d) |
|
if os.path.isfile(filepath) and os.path.splitext(filepath)[1] == ".fits": |
|
fitsPaths.append(pathToDirectory + "/" + d) |
|
|
|
# GENERATE THE IMAGECOLLECTION |
|
keys = ['MJD-OBS'] |
|
collection = ImageFileCollection(filenames=fitsPaths, keywords=keys) |
|
# SORT IMAGE COLLECTION |
|
collection.sort(['MJD-OBS']) |
|
|
|
lastX = None |
|
lastY = None |
|
ymin = None |
|
ymax = None |
|
gap = None |
|
shift = 0 |
|
|
|
# FOR EACH FILE (MJD SORTED) |
|
for f in collection.files: |
|
f = fits.open(f) |
|
f.verify('fix') |
|
f[1].verify('fix') |
|
# THE SPECTRAL DATA IS IN THE SECOND HDU |
|
specdata = f[1].data |
|
mjd = f[0].header["MJD-OBS"] |
|
f.close() |
|
|
|
# UNPACK THE FLUX DATA |
|
x = specdata['WAVE'][0] |
|
y = specdata['FLUX'][0] |
|
# YRANGE FOR THIS SINGLE SPECTRUM |
|
yrange = y.max() - y.min() |
|
|
|
# DETERMINE HOW MUCH TO SHIFT SPECTRUM IN Y-AXIS TO AVOID OVERLAPPING |
|
if lastX is None: |
|
lastX = x |
|
lastY = y |
|
gap = yrange * 0.1 |
|
else: |
|
if gap < yrange * 0.1: |
|
gap = yrange * 0.1 |
|
shift = np.amin(lastY - y) |
|
y = y + shift - gap |
|
lastX = x |
|
lastY = y |
|
|
|
# DETERMINE Y-RANGE NEEDED FOR ALL SPECTRA IN SET |
|
if ymin is None or y.min() < ymin: |
|
ymin = y.min() |
|
if ymax is None or y.max() > ymax: |
|
ymax = y.max() |
|
|
|
# ADD MJD LABEL TO SPECTRUM |
|
if len(collection.files) > 1: |
|
mjd = f"{mjd:0.2f}" |
|
ax.text( |
|
x.max() + (x.max() - x.min()) * 0.01, |
|
y[-1] - yrange * 0.1, |
|
mjd, |
|
fontsize=10 |
|
) |
|
|
|
# ADD COLOUR GRADIENT TO SPECTRA OR NOT? |
|
if not cmap: |
|
ax.plot(x, y, color="#dc322f") |
|
else: |
|
# CREATE A SET OF LINE SEGMENTS SO THAT WE CAN COLOR THEM INDIVIDUALLY |
|
# WE ARE CREATING N (LENGTH OF X-AXIS) PACKETS WITH ONE CELL FOR COLOR AND ONE COORDINATE |
|
# SET OF THE DATA POINT |
|
points = np.array([x, y]).T.reshape(-1, 1, 2) |
|
segments = np.concatenate([points[:-1], points[1:]], axis=1) |
|
# CREATE A CONTINUOUS NORM TO MAP FROM DATA POINTS TO COLORS - FIND MIN |
|
# AND MAX OF THE DATA-ARRAY TO CONSTRAIN COLOR MAP |
|
norm = plt.Normalize(x.min(), x.max()) |
|
lc = LineCollection(segments, cmap=cmap, norm=norm) |
|
# SET THE VALUES USED FOR COLORMAPPING |
|
lc.set_array(x[::-1]) |
|
lc.set_linewidth(1) |
|
line = ax.add_collection(lc) |
|
# fig.colorbar(line, ax=axs) |
|
|
|
# SET AXIS LIMITS BASED ON ALL SPECTRA |
|
yrange = ymax - ymin |
|
ax.set_xlim(x.min(), x.max()) |
|
ax.set_ylim(ymin - yrange * 0.1, ymax + yrange * 0.1) |
|
|
|
log.debug('completed the ``plot_1d_spectrum`` function') |
|
return fitsPaths |
|
|
|
|
|
def plot_2d_spectrum( |
|
log, |
|
fitsPath, |
|
ax, |
|
cmapName=False, |
|
backgroundColor="black"): |
|
"""*plot the 2D spectrum* |
|
|
|
**Key Arguments:** |
|
|
|
- `log` -- logger |
|
- `fitsPath` -- path the the FITS image containing spectrum |
|
- `ax` -- matplotlib ax to use |
|
- `cmapName` -- the colormap to use. Default *False* (monochrome) |
|
- `backgroundColor` -- the background color to use |
|
|
|
""" |
|
log.debug('starting the ``plot_2d_spectrum`` function') |
|
|
|
frame = CCDData.read(fitsPath, hdu=0, unit=u.electron, hdu_uncertainty='ERRS', |
|
hdu_mask='QUAL', hdu_flags='FLAGS', key_uncertainty_type='UTYPE') |
|
data = frame.data |
|
|
|
# FIND THE ROW WITH THE HIGHEST SIGNAL AND CENTRE IMAGE HERE |
|
yWindowSize = 40 |
|
sourceLocation = np.argmax(np.median(frame.data, 0)) |
|
# TRIM DATA IN Y-DIRECTION |
|
fluxRow = data[:, sourceLocation - 2:sourceLocation + 2] |
|
data = data[:, sourceLocation - yWindowSize // |
|
2:sourceLocation + yWindowSize // 2] |
|
# ROTATE THE IMAGE |
|
rotatedImg = np.rot90(data, 1) |
|
std = np.std(fluxRow) |
|
mean = np.mean(fluxRow) |
|
vmax = mean + 9 * std |
|
vmin = mean - 1.5 * std |
|
|
|
if cmapName: |
|
cmap = fancy_cmap([backgroundColor, 'white']) |
|
overlay = np.arange(rotatedImg.shape[ |
|
0] * rotatedImg.shape[1]).reshape(rotatedImg.shape[1], rotatedImg.shape[0]) |
|
overlay = np.rot90(overlay, -1) |
|
colorTint = ax.imshow(overlay, cmap=cmapName, aspect="auto", |
|
interpolation="nearest", alpha=1.0) |
|
else: |
|
cmap = plt.get_cmap('gray') |
|
|
|
im1 = ax.imshow(rotatedImg, vmin=vmin, vmax=vmax, |
|
cmap=cmap, alpha=1.0, aspect="auto", interpolation="spline36") |
|
|
|
ax.set_xlim(0, rotatedImg.shape[1]) |
|
ax.set_ylim(0, rotatedImg.shape[0]) |
|
ax.set_xticks([]) |
|
ax.set_yticks([]) |
|
if cmap: |
|
ax.axis('off') |
|
|
|
log.debug('completed the ``plot_2d_spectrum`` function') |
|
return None |
|
|
|
|
|
def fancy_cmap(ramp_colors): |
|
fancy_cmap = LinearSegmentedColormap.from_list( |
|
'my_list', [Color(c1).rgb for c1 in ramp_colors]) |
|
# Get the colormap colors |
|
my_colors = fancy_cmap(np.arange(fancy_cmap.N)) |
|
my_colors[:, -1] = np.linspace(1, 0, fancy_cmap.N) |
|
fancy_cmap = ListedColormap(my_colors) |
|
return fancy_cmap |
|
|
|
|
|
if __name__ == '__main__': |
|
main() |