def measure_equivalent_width(wavelength, flux, display_xmin, display_xmax, ret_fit=False):
"""Measure the equivalent width for an emission or absorption line.
A matplotlib window will pop up, allowing the user to click on
either side of the line feature where it ends. A continuum is then
fit using a linear fit.
"""
coords = []
def onclick(event):
"""Click twice on a matplotlib figure and record the x coordinates."""
x = event.xdata
y = event.ydata
coords.append(x)
ax.plot(x, y, "x")
fig.canvas.draw()
fig.canvas.flush_events()
if len(coords) == 2:
fig.canvas.mpl_disconnect(cid)
plt.close()
return coords
# If the wavelength array looks like it's been passed in descending order,
# then reverse it to make it in ascending order. This also assumes the
# flux and continuum flux were also in the same order, so reverses them
# as well
if wavelength[1] < wavelength[0]:
wavelength = wavelength[::-1]
flux = flux[::-1]
# There are sanity checks to make sure all array values are in ascending
# order before we go on any further
is_increasing = np.all(np.diff(wavelength) > 0)
if not is_increasing:
raise ValueError("the values for the wavelength bins provided are not increasing")
# Plot the spectrum, then allow the user to click on the edges of the line
# to label where it starts and stops
fig, ax = plt.subplots(figsize=(12, 5))
wavelength, flux = get_xy_subset(wavelength, flux, display_xmin, display_xmax)
ax.loglog(wavelength, flux, linewidth=2, label="Spectrum")
# ax.set_xlim(display_xmin, display_xmax)
# ax.set_ylim(get_y_lims_for_x_lims(wavelength, flux, display_xmin, display_xmax, scale=2.0))
ax.set_xlabel("Wavelength")
ax.set_ylabel("Flux")
ax = add_line_ids(ax, common_lines())
ax.set_title("Mark the blue and then red end of the line")
cid = fig.canvas.mpl_connect("button_press_event", onclick)
plt.show()
# Extract a portion of the spectrum to both the left and right of the line,
# by 500 Angstroms, this is used to create a linear fit to estimate the
# continuum
if len(coords) == 0:
print("you didn't click on anything!")
exit(EXIT_FAIL)
if coords[0] > coords[1]:
tmp = coords[0]
coords[0] = coords[1]
coords[1] = tmp
wavelength = np.array(wavelength)
flux = np.array(flux)
i1 = get_array_index(wavelength, coords[0] - 500)
j1 = get_array_index(wavelength, coords[0])
i2 = get_array_index(wavelength, coords[1])
j2 = get_array_index(wavelength, coords[1] + 500)
a = np.concatenate([wavelength[i1:j1], wavelength[i2:j2]])
b = np.concatenate([flux[i1:j1], flux[i2:j2]])
fit = np.poly1d(np.polyfit(a, b, 1))
# Plot the spectrum and the fit to see how well it's doing, although if it's
# a shit fit I do nothing about it and just let the code run its course :-)
# fig, ax = plt.subplots(figsize=(12, 5))
# wavelength, flux = get_xy_subset(wavelength, flux, display_xmin, display_xmax)
# ax.loglog(wavelength, flux, linewidth=2, label="Spectrum")
# # ax.plot(a, b, linewidth=2, label="Extracted bit")
# wavelength, flux = get_xy_subset(a, fit(a), display_xmin, display_xmax)
# ax.plot(wavelength, flux, label="Linear fit")
# # ax.set_xlim(display_xmin, display_xmax)
# # ax.set_ylim(get_y_lims_for_x_lims(wavelength, flux, display_xmin, display_xmax, scale=2.0))
# ax.legend()
# ax.set_xlabel("Wavelength")
# ax.set_ylabel("Flux")
# plt.show()
# Restrict the wavelength range to be around the line we are interested in
# The way we do it is probably really slow, but in the case where the
# spectrum is very finely gridded, we have a bit of trouble with the
# original method
# i_line = get_array_index(wavelength, coords[0])
# j_line = get_array_index(wavelength, coords[1])
i_line = 1
j_line = 2
for i_line, ww in enumerate(wavelength):
if ww > coords[0]:
break
for j_line, ww in enumerate(wavelength):
if ww > coords[1]:
break
i_line -= 1
j_line -= 1
wavelength_ew = wavelength[i_line:j_line]
flux_ew = flux[i_line:j_line]
flux_continuum_ew = fit(wavelength_ew)
# Now we can calculate the equivalent width, remember the formula for this is,
# W = \int_{\lambda_1}^{\lambda_2} \frac{F_{c} - F_{\lambda}}{F_{c}} d\lambda
# W = Sum(Fc - Flamda / Fc * dlambda)
# todo: see if this can be sped up, but n_bins is usually "small" so
# probably this doesn't matter
w = 0
n_bins = len(wavelength_ew)
for i in range(1, n_bins):
d_wavelength = wavelength_ew[i] - wavelength_ew[i - 1]
w += (flux_continuum_ew[i] - flux_ew[i]) / flux_continuum_ew[i] * d_wavelength
# Take the absolute value, as can be negative depending on if the line is in
# emission or absorption
w = np.abs(w)
if ret_fit:
return w, fit
else:
return w
def reprocessing(spectrum, xmin=None, xmax=None, scale="loglog", label_edges=True, alpha=0.75, display=False):
"""Create a plot to show the amount of reprocessing in the model.
Parameters
----------
Returns
-------
fig: plt.Figure
matplotlib Figure object.
ax: plt.Axes
matplotlib Axes object.
"""
if "spec_tau" not in spectrum.available:
raise ValueError("There is no spec_tau spectrum so cannot create this plot")
if "spec" not in spectrum.available and "log_spec" not in spectrum.available:
raise ValueError("There is no observer spectrum so cannot create this plot")
fig, ax = plt.subplots(figsize=(12, 7))
ax2 = ax.twinx()
ax = set_axes_scales(ax, scale)
ax2 = set_axes_scales(ax2, scale)
# Plot the optical depth
spectrum.set("spec_tau")
for n, inclination in enumerate(spectrum.inclinations):
y = spectrum[inclination]
if np.count_nonzero == 0:
continue
ax.plot(spectrum["Freq."], y, label=str(inclination) + r"$^{\circ}$", alpha=alpha)
ax.legend(loc="upper left")
ax.set_xlim(xmin, xmax)
ax.set_xlabel("Rest-frame Frequency [Hz]")
ax.set_ylabel("Continuum Optical Depth")
if label_edges:
ax = add_line_ids(ax, photoionization_edges(spectrum=spectrum), "none")
ax.set_zorder(ax2.get_zorder() + 1)
ax.patch.set_visible(False)
# Plot the emitted and created spectrum
spectrum.set("spec")
for thing in ["Created", "Emitted"]:
x, y = get_xy_subset(spectrum["Freq."], spectrum[thing], xmin, xmax)
if spectrum.units == SpectrumUnits.f_lm:
y *= spectrum["Lambda"]
else:
y *= spectrum["Freq."]
ax2.plot(x, y, label=thing, alpha=alpha)
ax2.legend(loc="upper right")
ax2 = set_axes_labels(ax2, spectrum)
fig = finish_figure(fig)
fig.savefig(f"{spectrum.fp}/{spectrum.root}_reprocessing.png")
if display:
plt.show()
return fig, ax
def multiple_models(output_name,
spectra,
spectrum_type,
things_to_plot,
xmin=None,
xmax=None,
multiply_by_spatial_units=False,
alpha=0.7,
scale="loglog",
label_lines=True,
log_spec=False,
smooth=None,
distance=None,
display=False):
"""Plot multiple spectra, from multiple models, given in the list of
spectra provided.
Spectrum file paths are passed and then each spectrum is loaded in as a
Spectrum object. Each spectrum must have the same units and are also assumed
to be at the same distance.
In this function, it is possible to compare Emitted or Created spectra. It
is agnostic to the type of spectrum file being plotted, unlike
spectrum_observer.
todo: label absorption edges if spec_tau is selected
Parameters
----------
output_name: str
The name of the output .png file.
spectra: str
The file paths of the spectra to plot.
spectrum_type: str
The type of spectrum to plot, i.e. spec or spec_tot.
things_to_plot: str or list of str or tuple of str
The things which will be plotted, i.e. '45' or ['Created', '45', '60']
xmin: float [optional]
The lower x boundary of the plot
xmax: float [optional]
The upper x boundary for the plot
multiply_by_spatial_units: bool [optional]
Plot in flux units, instead of flux density.
alpha: float [optional]
The transparency of the plotted spectra.
scale: str [optional]
The scaling of the axes.
label_lines: bool [optional]
Label common emission and absorption features, will not work with
spec_tau.
log_spec: bool [optional]
Use either the linear or logarithmically spaced spectra.
smooth: int [optional]
The amount of smoothing to apply to the spectra.
distance: float [optional]
The distance to scale the spectra to in parsecs.
display: bool [optional]
Display the figure after plotting, or don't.
Returns
-------
fig: plt.Figure
matplotlib Figure object.
ax: plt.Axes
matplotlib Axes object.
"""
if type(spectra) is str:
spectra = list(spectra)
spectra_to_plot = []
for spectrum in spectra:
if type(spectrum) is not Spectrum:
root, fp = get_root_name(spectrum)
spectra_to_plot.append(Spectrum(root, fp, log_spec, smooth, distance, spectrum_type))
else:
spectra_to_plot.append(spectrum)
units = list(dict.fromkeys([spectrum.units for spectrum in spectra_to_plot]))
if len(units) > 1:
msg = ""
for spectrum in spectra_to_plot:
msg += f"{spectrum.units} : {spectrum.fp}{spectrum.root}.{spectrum.current}\n"
raise ValueError(f"Some of the spectra have different units, unable to plot:\n{msg}")
# Now, this is some convoluted code to get the inclination angles
# get only the unique, sorted, values if inclination == "all"
if things_to_plot == "all":
things_to_plot = ()
for spectrum in spectra_to_plot: # have to do it like this, as spectrum.inclinations is a tuple
things_to_plot += spectrum.inclinations
if len(things_to_plot) == 0:
raise ValueError("\"all\" does not seem to have worked, try specifying what to plot instead")
things_to_plot = tuple(sorted(list(dict.fromkeys(things_to_plot)))) # Gets sorted, unique values from tuple
else:
things_to_plot = things_to_plot.split(",")
things_to_plot = tuple(things_to_plot)
n_to_plot = len(things_to_plot)
n_rows, n_cols = subplot_dims(n_to_plot)
if n_rows > 1:
figsize = (12, 12)
else:
figsize = (12, 5)
fig, ax = plt.subplots(n_rows, n_cols, figsize=figsize, squeeze=False)
fig, ax = remove_extra_axes(fig, ax, n_to_plot, n_rows * n_cols)
ax = ax.flatten()
for n, thing in enumerate(things_to_plot):
for spectrum in spectra_to_plot:
try:
y = spectrum[thing]
except KeyError:
continue # We will skip key errors, as models may have different inclinations
if np.count_nonzero(y) == 0: # skip arrays which are all zeros
continue
if multiply_by_spatial_units:
if spectrum.units == SpectrumUnits.f_lm:
y *= spectrum["Lambda"]
else:
y *= spectrum["Freq."]
if spectrum.units == SpectrumUnits.f_lm:
x = spectrum["Lambda"]
else:
x = spectrum["Freq."]
x, y = get_xy_subset(x, y, xmin, xmax)
label = spectrum.fp.replace("_", r"\_") + spectrum.root.replace("_", r"\_")
ax[n].plot(x, y, label=label, alpha=alpha)
ax[n] = set_axes_scales(ax[n], scale)
ax[n] = set_axes_labels(ax[n], spectra_to_plot[0], multiply_by_spatial_units=multiply_by_spatial_units)
if thing.isdigit():
ax[n].set_title(f"{thing}" + r"$^{\circ}$")
else:
ax[n].set_title(f"{thing}")
if label_lines:
ax[n] = add_line_ids(ax[n], common_lines(spectrum=spectra_to_plot[0]), "none")
ax[0].legend(fontsize=10).set_zorder(0)
fig = finish_figure(fig)
fig.savefig(f"{output_name}.png")
if display:
plt.show()
return fig, ax
def optical_depth(spectrum,
inclinations="all",
xmin=None,
xmax=None,
scale="loglog",
label_edges=True,
frequency_space=True,
display=False):
"""Plot the continuum optical depth spectrum.
Create a plot of the continuum optical depth against either frequency or
wavelength. Frequency space is the default and preferred. This function
returns the Figure and Axes object.
Parameters
----------
spectrum: pypython.Spectrum
The spectrum object.
inclinations: str or list or tuple [optional]
A list of inclination angles to plot, but all is also an acceptable
choice if all inclinations are to be plotted.
xmin: float [optional]
The lower x boundary for the figure
xmax: float [optional]
The upper x boundary for the figure
scale: str [optional]
The scale of the axes for the plot.
label_edges: bool [optional]
Label common absorption edges of interest onto the figure
frequency_space: bool [optional]
Create the figure in frequency space instead of wavelength space
display: bool [optional]
Display the final plot if True.
Returns
-------
fig: plt.Figure
matplotlib Figure object.
ax: plt.Axes
matplotlib Axes object.
"""
fig, ax = plt.subplots(1, 1, figsize=(12, 7))
current = spectrum.current
spectrum.set("spec_tau")
# Determine if we're plotting in frequency or wavelength space and then
# determine the inclinations we want to plot
if frequency_space:
xlabel = "Freq."
units = SpectrumUnits.f_nu
else:
xlabel = "Lambda"
units = SpectrumUnits.f_lm
if not xmin:
xmin = np.min(spectrum[xlabel])
if not xmax:
xmax = np.max(spectrum[xlabel])
inclinations = _get_inclinations(spectrum, inclinations)
# Now loop over the inclinations and plot each one, skipping ones which are
# completely empty
for inclination in inclinations:
if inclination != "all":
if inclination not in spectrum.inclinations: # Skip inclinations which don't exist
continue
x, y = get_xy_subset(spectrum[xlabel], spectrum[inclination], xmin, xmax)
if np.count_nonzero(y) == 0: # skip arrays which are all zeros
continue
ax.plot(x, y, label=f"{inclination}" + r"$^{\circ}$")
ax = set_axes_scales(ax, scale)
ax.set_ylabel(r"Continuum Optical Depth")
if frequency_space:
ax.set_xlabel(r"Rest-frame Frequency [Hz]")
else:
ax.set_xlabel(r"Rest-frame Wavelength [$\AA$]")
ax.legend(loc="upper left")
if label_edges:
ax = add_line_ids(ax, photoionization_edges(units=units), linestyle="none", offset=0)
spectrum.set(current)
fig = finish_figure(fig)
fig.savefig(f"{spectrum.fp}/{spectrum.root}_spec_optical_depth.png")
if display:
plt.show()
return fig, ax
def _plot_subplot(ax, spectrum, things_to_plot, xmin, xmax, alpha, scale, use_flux):
"""Plot some things to a provided matplotlib ax object.
This function is used to do a lot of the plotting heavy lifting in this
sub-module. It's still fairly flexible to be used outside of the main
plotting functions, however. You are just required to pass an axes to
plot onto.
Parameters
----------
ax: plt.Axes
A matplotlib axes object to plot onto.
spectrum: pypython.Spectrum
The spectrum object to plot. The current spectrum wishing to be set
must be correct, otherwise the wrong thing may be plotted.
things_to_plot: str or list or tuple of str
A collection of names of things to plot to iterate over.
xmin: float
The lower x boundary of the plot.
xmax: float
The upper x boundary of the plot.
alpha: float
The transparency of the spectra plotted.
scale: str
The scaling of the plot axes.
use_flux: bool
Plot the spectrum as a flux or nu Lnu instead of flux density or
luminosity.
Returns
-------
ax: pyplot.Axes
The modified matplotlib Axes object.
"""
if type(things_to_plot) is str:
things_to_plot = things_to_plot,
for thing in things_to_plot:
y = spectrum[thing]
# If plotting in frequency space, of if the units then the flux needs
# to be converted in nu F nu
if use_flux:
if spectrum.units == SpectrumUnits.f_lm:
y *= spectrum["Lambda"]
else:
y *= spectrum["Freq."]
if spectrum.units == SpectrumUnits.f_lm:
x = spectrum["Lambda"]
else:
x = spectrum["Freq."]
x, y = get_xy_subset(x, y, xmin, xmax)
ax.plot(x, y, label=thing, alpha=alpha)
ax.legend(loc="lower left")
ax = set_axes_scales(ax, scale)
ax = set_axes_labels(ax, spectrum, multiply_by_spatial_units=use_flux)
return ax