"""
Re-plot the three plots for all data directories given on the command line.
"""
import numpy as np
import os
from typing import Callable
import matplotlib.pyplot as plt
from matplotlib.ticker import AutoMinorLocator
from matplotlib.cm import ScalarMappable
from matplotlib.colors import Normalize, ListedColormap
import matplotlib.colors
import logging
log = logging.getLogger(__name__)
from .calc import calc_modulus_spectrum
from .prsdata import Data, PrsData, REFLECTION, TRANSMISSION, ABSORPTION
from . import calc
from .functions import nm_to_eV, eV_to_nm
from .offset import merge_offset_transients, get_offset_data
# Fallback color for all undefined keys
DEFAULT_COLOR = "#444"
from devctrl.data.spectrum.plot import _plot_spectrum_on_ax as plot_spectrum_on_ax
from devctrl.data.spectrum.plot import plot_spectrum as _plot_spectrum
# def plot_spectrum_on_ax(spectrum_data: np.ndarray, ax, key, columns=None, xkey="E", ylabel: str|bool=True, scale_factor: None|float|int=None, error_shade=False, mode=REFLECTION, max_ylim: None|tuple[float,float]=None, **plot_kw):
# """
# Main plotting function for PrsData
# :param spectrum_data: Numpy array with <columns> as returned by PrsData.get_spectrum_data()
# :param ax: Matplotlib axes
# :param key: Which key of <columns> to plot
# :param columns: List of column keys
# :param xkey: Which key of <columns> to plot on the x-axis (wl or E)
# :param ylabel: Either a label (str) or True to use the PrsData.get_column_tex_label_with_unit_and_scale() or False to not use a label
# :param scale_factor: Which scaling to apply to the ylabel.
# :param error_shade: Whether to include the errors by drawing a semi-transparent shade around the plot line
# :param plot_kw: Keyword arguments to pass to the axes.plot() function
# :param mode: Whether to use reflection or transmission labels
# """
# if columns is None:
# columns = PrsData.default_spectrum_columns
# if key in ["theta", "theta-corr", "imag", "L"] or key.startswith("dI"):
# ax.axhline(0, color="black", linewidth=plt.rcParams["axes.linewidth"])
# x_idx = columns.index(xkey)
# data_idx = columns.index(key)
# ydata = spectrum_data[:, data_idx]
#
# # # apply scaling
# # if scale_factor is not None:
# # log.debug(f"Applying scale: {scale_factor} to key '{key}'")
# # ydata = ydata.copy()/scale_factor
#
# # ydata = remove_outliers(ydata, threshold=0.5)
# # ydata = spectrum_data[:, data_idx]
# ylim_before = ax.get_ylim()
#
# if error_shade:
# if f"s{key}" in columns:
# sdata_idx = columns.index(f"s{key}")
# sydata = spectrum_data[:, sdata_idx]
# # if scale_factor is not None:
# # sydata = sydata.copy()/scale_factor
# linecolor = plot_kw["color"] if "color" in plot_kw else PrsData.get_column_color(key, default=DEFAULT_COLOR, mode=mode)
# color = "#00000030"
# if not linecolor.startswith("#"): # if a named color was given, try getting the #rrggbb
# try:
# linecolor = matplotlib.colors.get_named_colors_mapping()[linecolor]
# if not type(linecolor) == str:
# log.debug(f"Non-string color: {linecolor}")
# linecolor = color
# raise ValueError()
# except:
# pass
#
# if linecolor.startswith("#"):
# if len(linecolor) in [7,9]: # #rrggbb(aa)
# color = linecolor[:7] + "33" # make semi-transparent
# elif len(linecolor) in [4,5]: # #rgb(a):
# color = linecolor[:4] + "3" # make semi-transparent
#
# ax.fill_between(spectrum_data[:, x_idx], ydata-sydata, ydata+sydata, color=color)
# else:
# log.warning(f"Can not draw error shade, since 's{key}' is not in columns")
#
#
# line, = ax.plot(spectrum_data[:, x_idx], ydata, **(dict(color=PrsData.get_column_color(key, default=DEFAULT_COLOR, mode=mode)) | plot_kw))
# if error_shade:
# # set y limits
# # if the error_shade is at some point much larger than the line, we dont want to show it (likely an outlier)
# # therefore we ignore the error shade ylimits
# ymin = np.nanmin(ydata)
# ymax = np.nanmax(ydata)
# # extra margins: 5% each side
# margin = 0.05 * (ymax-ymin)
# # if the limits were previously larger, dont re-set them (might be form another plot line)
# # defaults are (0.0, 1.0), ignore those
# default_lim = ylim_before[0] == 0.0 and ylim_before[1] == 1.0
# ymin = min(ymin-margin, ylim_before[0]) if not default_lim else ymin-margin
# ymax = max(ymax+margin, ylim_before[1]) if not default_lim else ymax+margin
# if not (np.isnan(ymin) or np.isinf(ymin) or np.isnan(ymax) or np.isinf(ymax)):
# ax.set_ylim(ymin, ymax)
# else:
# log.debug(f"Not setting limits since at least one is invalid: ymin={ymin}, ymax={ymax}")
# if max_ylim: # hard limit to given values
# ylims = ax.get_ylim()
# ymin = max(ylims[0], max_ylim[0])
# ymax = min(ylims[1], max_ylim[1])
# ax.set_ylim(ymin, ymax)
#
# # label
# if type(ylabel) == bool and ylabel: # if True, use auto ylabel
# ylabel = PrsData.get_column_tex_label_with_unit_and_scale(key, scale=scale_factor, mode=mode)
# if ylabel: # if string
# ax.set_ylabel(ylabel)
# return line
[docs]
def plot_spectrum(data: PrsData|str|np.ndarray, columns=None, title: str|None=None, what=["dI-X_I"], fig=None, ax=None, subplot_kw={}, use_scales: dict|bool=True, mode=None, transforms=None, xlabel=True, add_second_xax=True, **plot_spectrum_kw):
"""
Plot recorded data
:param data: Path to the data directory or numpy array with columns PrsData.default_spectrum_columns
:param title: Title for the plot. If None, try to use the metadata[name] of the data if the type is PrsData. The default is None.
:return: fig Matplotlib figure object.
"""
if columns is None:
columns = PrsData.default_spectrum_columns
if type(data) == str:
prsdata = PrsData(data)
mode = mode or prsdata.mode
_data, columns = prsdata.get_spectrum_data(columns=columns, transforms=transforms)
elif isinstance(data, PrsData):
_data, columns = data.get_spectrum_data(columns=columns, transforms=transforms)
mode = mode or data.mode
else:
_data = data
if mode is None: mode = REFLECTION
if isinstance(what, str):
what = [what]
n_plots = len(what)
if fig is None or ax is None:
fig, ax = plt.subplots(n_plots, 1, sharex=True, squeeze=False, **subplot_kw)
ax = ax[:, 0]
elif ax is not None:
ax = np.asanyarray(ax)
if "E" in columns:
xkey = "E"
_xlabel = PrsData.get_column_tex_label_with_unit("E")
elif "wl" in columns:
xkey = "wl"
_xlabel = PrsData.get_column_tex_label_with_unit("wl")
else:
raise ValueError(f"Data columns must contain either 'wl' or 'E' or both.")
if xlabel == True:
xlabel = _xlabel
if xlabel:
ax[-1].set_xlabel(xlabel)
i = 0
for key in what:
scale = None
if use_scales == True:
if isinstance(data, Data):
scale = data.get_column_scale(key)
else:
scale = Data.get_column_scale(key)
elif isinstance(use_scales, dict):
if key in use_scales:
scale = use_scales[key]
else:
log.warning(f"Invalid value for use_scales: '{use_scales}'")
if scale is not None:
log.debug(f"Using scale {scale} for key {key}")
ax[i].yaxis.offsetText.set_visible(False) # Hide the default offset text
# ax.yaxis.set_major_formatter(
# plt.FuncFormatter(lambda val, pos: f'{val/scale:g}') # Divide by 10^4 for display
# )
_scale = scale
ax[i].yaxis.set_major_formatter(plt.FuncFormatter(lambda val, pos: '{:g}'.format(val / _scale)))
plot_spectrum_on_ax(_data, columns, ax[i], key=key, xkey=xkey, mode=mode, scale_factor=scale, **plot_spectrum_kw)
i += 1
# add second axes here to avoid divide by zero errors, which occur when the axis is added while no data has been plotted yet
xax_nms = []
if "wl" in columns and "E" in columns and add_second_xax:
for j in range(0, ax.shape[0]):
xax_nms.append(add_wl_axis(ax[j], label=j == 0))
if title is None and isinstance(data, PrsData):
if "name" in data.metadata:
title = data.metadata["name"]
if title:
fig.suptitle(title)
if fig.get_layout_engine() != "constrained":
fig.tight_layout()
try:
from mqutil.plot.label import remove_overlapping_xticklabels
fig.draw_without_rendering()
for xax in xax_nms: remove_overlapping_xticklabels(xax)
except ImportError as e:
log.debug(f"mqutil not available -> can not check for overlapping tick labels (ImportError: {e})")
return fig, ax
#
# OFFSETS DURING SPECTRAL MEASUREMENTS
#
[docs]
def plot_offset_data(data: PrsData, what=["lock-in-R"], with_errors=True, relative_time=True, subplots_kw=None, **plot_kw):
"""
:return: np.ndarray of [index, [timestamp, *what]], fig, ax
"""
ewhat = []
edata = None
if with_errors:
ewhat = ["s"+w for w in what]
_, edata = get_offset_data(data, what=ewhat)
offset_keys, odata = get_offset_data(data, what=what)
print(offset_keys)
print(what, odata)
print(ewhat, edata)
fig, axs = plt.subplots(len(what), squeeze=False, **(subplots_kw or {"sharex": True}))
axs = axs[:,0]
if relative_time:
odata[:,0] -= odata[0,0]
for i, c in enumerate(what):
if relative_time:
axs[i].set_xlabel("Time [s]")
else:
axs[i].set_xlabel("Timestamp")
if with_errors:
axs[i].errorbar(odata[:,0], odata[:,i+1], edata[:,i+1], **plot_kw)
else:
axs[i].plot(odata[:,0], odata[:,i+1], **plot_kw)
axs[i].set_ylabel(data.get_column_tex_label_with_unit(c.removesuffix("_raw")))
# make sure there is some space between line and axis so that the error lines are visible
if with_errors:
trange = odata[-1,0] - odata[0,0]
space = 0.05
for ax in axs:
ax.set_xlim(odata[0,0] - trange*space, odata[-1,0] + trange*space)
fig.tight_layout()
return fig, axs
# OFFSETS TRANSIENTS
[docs]
def plot_offset_transients(data:PrsData, what=["lock-in-R_raw"], convert_idx2time=True, subplots_kw=None, **plot_kw):
offset_keys, odata = merge_offset_transients(data, what)
fig, axs = plt.subplots(len(what), squeeze=False, **(subplots_kw or {}))
axs = axs[:,0]
srat = 0
if convert_idx2time:
try:
srat = data.metadata["lock-in-amp"]["sample_rate_Hz"]
except KeyError as e:
raise KeyError(f"Can not convert indices to time because the sample rate can not be determined.") from e
for i, c in enumerate(what):
xdata = np.array(range(len(odata[i])), dtype=float)
if convert_idx2time:
xdata /= srat
axs[i].set_xlabel("Time [s]")
else:
axs[i].set_xlabel("Index")
axs[i].plot(xdata, odata[i], **plot_kw)
axs[i].set_ylabel(data.get_column_tex_label_with_unit(c.removesuffix("_raw")))
fig.tight_layout()
return fig, axs
#
# SPECTRA
#
[docs]
def plot_data(data: PrsData, basename: str|None=None, out_dir: str|None=None, title=None, **plot_spectrum_kw):
"""
Convenience function that generates and saves useful plots
:param data: data object
:param basename: base filename. Defaults to the directory name of the PrsData
:param out_dir: output directory. Defaults to PrsData directory
"""
if basename is None:
basename: str = os.path.basename(data.dirpath)
if out_dir is None:
out_dir: str = data.dirpath
# supress divide by zero warnings, which happen when matplotlib transforms axes
import warnings
with warnings.catch_warnings():
warnings.simplefilter("ignore")
# check if has necessary data
# if all(map(lambda col: col in has_cols, ["lock-in-R_raw", "lock-in-aux_raw", "lock-in-theta_raw"])):
if data.has_AC() and data.has_DC():
if data.has_wavelengths():
# only dR/R
cols1 = ["dI_I", "dI-X_I", "dI-Y_I", "theta-corr", "I"] # also get I so that points with invalid I get removed
cols1 += ["s" + col for col in cols1] + ["E", "wl"]
sdata1, cols1 = data.get_spectrum_data(columns=cols1, transforms=[calc.remove_spectrum_outliers, calc.remove_invalid_points_from_spectrum])
fig1, _ = plot_spectrum(sdata1, title=title, columns=cols1, what=["dI-X_I"],
mode=data.mode, **(dict(max_ylim=(-100, 100)) | plot_spectrum_kw))
fig1_path = os.path.join(out_dir, basename + ".pdf")
fig1.savefig(fig1_path)
# plot with interesting things
fig2, _ = plot_spectrum(sdata1, title=title, columns=cols1, what=["dI_I", "dI-X_I", "dI-Y_I", "theta-corr"],
subplot_kw={"figsize": (6, 9)}, mode=data.mode, **plot_spectrum_kw)
fig2_path = os.path.join(out_dir, basename + "-X-Y" + ".pdf")
fig2.savefig(fig2_path)
# save plot with all params
# use magnitude dI/I for easier comparison with lock-in-R
cols2 = ["dI_I", "dI", "lock-in-R", "theta-corr", "lock-in-theta", "I"]
cols2 += ["s" + col for col in cols2] + ["E", "wl"]
sdata2, cols2 = data.get_spectrum_data(columns=cols2, transforms=[calc.remove_spectrum_outliers, calc.remove_noisy_points_from_spectrum])
fig3, _ = plot_spectrum(sdata2, title=title, columns=cols2, what=["dI_I", "dI", "lock-in-R", "theta-corr", "lock-in-theta", "I"],
subplot_kw={"figsize": (6, 12)}, mode=data.mode, **plot_spectrum_kw)
fig3_path = os.path.join(out_dir, basename + "-all" + ".pdf")
fig3.savefig(fig3_path)
# save plot with all and errors
fig4, _ = plot_spectrum(sdata2, title=title, columns=cols2, what=["sdI_I", "sdI", "slock-in-R", "stheta-corr", "slock-in-theta", "sI"],
subplot_kw={"figsize": (6, 12)}, mode=data.mode, **(plot_spectrum_kw | dict(error_shade=False)))
fig4_path = os.path.join(out_dir, basename + "-all-errors" + ".pdf")
fig4.savefig(fig4_path)
# modulus. Errors can currently not be calculated
_, _ = plot_modulus_spectrum(data, save_fig=True, **(plot_spectrum_kw | dict(error_shade=False, title=title)))
# XY
_, _ = plot_XY(data, title=title, save_fig=True)
# offset plot
if len(data._get_time_sorted_offset_keys()) >= 2:
fig, ax = plot_offset_data(data, what=["lock-in-R", "lock-in-theta", "lock-in-aux"], relative_time=True, with_errors=True, subplots_kw=dict(figsize=(5.2, 8)))
if title: fig.suptitle(title)
fig.tight_layout()
fig.savefig(os.path.join(out_dir, basename + "-offset-all.pdf"))
return
# return fig1, fig2, fig3, fig4
# Reference
elif data.has_DC():
cols = ["E", "wl", "lock-in-aux", "slock-in-aux"]
sdata, cols = data.get_spectrum_data(columns=cols)
fig1, _ = plot_spectrum(sdata, title=title, columns=cols, what=["lock-in-aux"], **plot_spectrum_kw)
fig1_path = os.path.join(out_dir, basename + ".pdf")
fig1.savefig(fig1_path)
return fig1, None, None, None
else:
log.warning(f"Not plotting anything since data object might be lacking necessary data.")
[docs]
def plot_rta(ref_data: PrsData, tra_data: PrsData, abs_data: PrsData, fig=None, axs=None, title=True, single_ax=False, subplots_kw={}, label=None, save_fig=True, transforms=calc.sane_transforms, **plot_spectrum_kw):
"""
Convenience function that plots reflectance, transmittance and absorbance
"""
n_rows = 1 if single_ax else 3
if not (fig is not None and axs is not None):
figsize = (6, 5) if single_ax else (6, 8)
fig, axs = plt.subplots(n_rows, 1, **(dict(figsize=figsize, sharex=True) | subplots_kw))
wavelengths = abs_data.wavelengths
colors = ["red", "green", "blue"]
cols = ["wl", "E", "dI-X_I", "sdI-X_I", "I", "sI", "dI-X", "sdI-X"]
for i, data in enumerate([ref_data, tra_data, abs_data]):
use_label = label
if use_label is None:
use_label = data.mode
ax = axs if single_ax else axs[i]
sdata, cols = data.get_spectrum_data(wavelengths=wavelengths, columns=cols, transforms=transforms)
plot_spectrum_on_ax(sdata, columns=cols, ax=ax, key="dI-X_I", label=use_label, color=colors[i], mode=data.mode, **(dict(max_ylim=(-100, 100)) | plot_spectrum_kw))
if single_ax:
axs.legend()
add_wl_axis(axs)
axs.set_xlabel(PrsData.get_column_tex_label_with_unit("E"))
else:
add_wl_axis(axs[0], label=True)
add_wl_axis(axs[1], label=False)
add_wl_axis(axs[2], label=False)
axs[-1].set_xlabel(PrsData.get_column_tex_label_with_unit("E"))
if title == True:
title = ref_data.name.removesuffix("-" + REFLECTION)
if title:
fig.suptitle(title.replace("_", " "))
fig.tight_layout()
if save_fig == True:
save_fig = os.path.join(os.path.dirname(ref_data.dirpath), os.path.basename(ref_data.dirpath).removesuffix("-"+REFLECTION) + ".pdf")
if save_fig:
log.debug(f"Saving RTA plot to '{save_fig}'")
fig.savefig(save_fig)
return fig, axs
[docs]
def plot_modulus_spectrum(data: PrsData, min_n_wavelengths=100, which_real="dI_I", save_fig=False, **plot_spectrum_kw):
"""
Plot the modulus spectrum of data and save it in the data directory (by default).
If data has less than `min_n_wavelengths`, the modulus spectrum is not plotted.
This is because the Kramers-Kronig transformation requires a certain energy range to be useful.
:param min_n_wavelengths: Minimum number of wavelengths the data must have.
:param which_real: Which column to use as real spectrum. Can be "dI_I", "dI-X_I" or "dI-Y_I"
:return: (fig, ax) or (None, None)
"""
if len(data.wavelengths) < min_n_wavelengths:
log.info(f"Not plotting modulus spectrum because the number of wavelengths is lower than the minimum: {len(data.wavelengths)} < {min_n_wavelengths}")
return None, None
if save_fig == True:
save_fig = os.path.join(data.dirpath, data.dirname + "_mod.pdf")
sdata, cols = calc_modulus_spectrum(data, which_real=which_real)
fig1, axs = plot_spectrum(sdata, columns=cols, what=["L"], mode=data.mode, **plot_spectrum_kw)
if save_fig:
log.debug(f"Saving modulus spectrum as '{save_fig}'")
fig1.savefig(save_fig)
return fig1, axs
[docs]
def plot_XY(data: PrsData, min_n_wavelengths=50, save_fig=False, get_spectrum_kw=None, fig=None, ax=None, title=True, colorbar=True, subplots_kw=None, blue_above=True):
"""
Plot dI-Y/I against dI-Y/I. X-Y plots can provide information on where a spectrum component comes from, since components with different time constants will have different angles in the plot.
mqutil is required plot the data points in the color of the respective wavelength.
:param min_n_wavelengths: Minimum number of wavelengths the data must have
:param save_fig: If True, save in data dir as '<data dir name>_XY-plane.pdf'. If str, save_fig is the plot file path.
:param get_spectrum_kw: Keyword arguments for get_spectrum_data (for example transforms)
:param title: If True, use data name as plot title. If str, use title as title.
:param colorbar: If True, draw a colorbar next to the plot.
:param blue_above: If True, plot order of points from red to blue, otherwise from blue to red.
:return: (fig, ax) or (None, None)
"""
if len(data.wavelengths) < min_n_wavelengths:
log.info(f"Not plotting XY spectrum because the number of wavelengths is lower than the minimum: {len(data.wavelengths)} < {min_n_wavelengths}")
return None, None
default_figsize = (5.2, 4.8) if colorbar else (4.8, 4.8)
if fig is None or ax is None:
fig, ax = plt.subplots(**(subplots_kw or dict(figsize=default_figsize)))
ax.axhline(0, zorder=0.5)
ax.axvline(0, zorder=0.5)
# columns = ["dI-X", "dI-Y", "dI-X_I", "dI-Y_I", "theta-corr"]
if get_spectrum_kw and "columns" in get_spectrum_kw:
columns = get_spectrum_kw["columns"]
del get_spectrum_kw["columns"]
else:
columns = ["dI-X_I", "dI-Y_I", "theta-corr"]
if data.has_DC(): columns.append("I") # required for transforms that remove invalid points
columns += [f"s{col}" for col in columns]
columns += ["wl", "E"]
sdata, columns = data.get_spectrum_data(columns=columns, **(get_spectrum_kw or {}))
xcol = "dI-X_I"
ycol = "dI-Y_I"
X = sdata[:,columns.index(xcol)]
Y = sdata[:,columns.index(ycol)]
sX = sdata[:,columns.index("s"+xcol)]
sY = sdata[:,columns.index("s"+ycol)]
wls = sdata[:,columns.index("wl")]
# The colormap must be linear, so we generate a list of wavelengths in case the spectrum does not have constant wl spacing
_color_wls = np.linspace(min(wls), max(wls), 400)
# init Colormap, Norm and Mappable for the colorbar
try:
from mqutil.color.util import rgb_int_to_hex, wavelength_nm_to_rgb_int
colors = [rgb_int_to_hex(*map(int, wavelength_nm_to_rgb_int(wl))) for wl in _color_wls]
cmap = ListedColormap(colors)
except ImportError as e:
log.debug(f"mqutil not available -> can map wavelengths to colors (ImportError: {e})")
cmap = plt.cm.get_cmap("plasma")
norm = Normalize(vmin=min(_color_wls), vmax=max(_color_wls))
sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])
# plot points: colored background with black point
# using for loop to plot points in desired order
r = range(len(wls))
if blue_above: r = reversed(r)
for i in r:
pts = ax.scatter(X[i], Y[i], c=wls[i], cmap=cmap, norm=norm, s=12) # color=color) #, c=wl, cmap="plasma")
pts = ax.scatter(X[i], Y[i], color="black", marker=".", s=8) #, c=wl, cmap="plasma")
ax.set_xlabel(data.get_column_tex_label_with_unit(xcol))
ax.set_ylabel(data.get_column_tex_label_with_unit(ycol))
if colorbar:
cbar = plt.colorbar(sm, ax=ax)
cbar.set_label(data.get_column_tex_label_with_unit("wl"))
# set equal limits everywhere and use 1:1 ratio
lim = np.max(np.abs(list(ax.get_ylim()) + list(ax.get_xlim())))
ax.set_xlim(-lim, lim)
ax.set_ylim(-lim, lim)
ax.set_aspect('equal', "box", anchor="C")
if title == True:
title = data.name
if title:
fig.suptitle(title)
fig.tight_layout()
if save_fig == True:
save_fig = os.path.join(data.dirpath, data.dirname + "_XY-plane.pdf")
if save_fig:
log.debug(f"Saving XY spectrum as '{save_fig}'")
fig.savefig(save_fig)
return fig, ax
[docs]
def replot_in_place(data_paths: list[str]|str=None, base_dir=None, **plot_spectrum_kw):
"""
Load the data from the directory and replot the default plots.
You can either specify the data paths directly or provide a base_dir to search for data
:param data_paths: (List of) directories to load data from
:param base_dir: Base directory in which all subdirectories containing data will be loaded
:param plot_spectrum_kw: Keyword arguments for the plot_spectrum function
:return: List of processed data dirs with a status message
"""
if data_paths is None and base_dir is None:
raise ValueError(f"Either data_paths or base_dir must be specified")
paths = []
if data_paths is not None:
if type(data_paths) == str:
paths.append(data_paths)
else:
paths += data_paths
if base_dir is not None:
paths += [os.path.join(base_dir, d) for d in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, d))]
result = []
for data_path in paths:
try:
data = PrsData(data_path=data_path)
plot_data(data, basename=os.path.basename(data_path), out_dir=data_path, **plot_spectrum_kw)
result.append((data_path, "Success"))
except Exception as e:
result.append((data_path, e))
continue
return result
[docs]
def plot_many(
data_dirs: list[str],
out_dir=None,
fig_suptitle="",
what=["dI-X_I", "theta-corr"],
n_cols=1,
n_rows=None,
direction: str="top-bottom",
get_name: Callable[[PrsData], str]=None,
sharex=True,
sharey=True,
):
"""
Plot multiple PrsData next to each other.
:param data_dirs: Directories containing data to be plotted.
:param out_dir: Directory into which the plots are saved
:param fig_suptitle: Figure suptitle
:param what: For which quantities to create plots
:param n_cols: Number of columns to use in the plot
:param n_rows: Number of roms to use in the plot
:param direction: The data will be plotted top to bottom, left to right
:param get_name: A function that takes a PrsData object returns the name. If None, then data.metadat["name"] is used
:param sharey
:param sharex
:return:
"""
if out_dir is None:
log.debug("No output directory given ,using current working directory")
out_dir = os.getcwd()
if not os.path.isdir(out_dir):
log.info(f"Output directory '{out_dir}' does not exist, creating it")
os.makedirs(out_dir)
datas = [] # (name, spectrum data)
for i, d in enumerate(data_dirs):
log.debug(f"Getting data from {d}")
pd = PrsData(data_path=d)
if get_name:
name = get_name(pd)
else:
name = pd.metadata["name"]
datas.append((name, pd.get_spectrum_data()[0]))
# if wls is None: wls = sdata[:,0]
n_data = len(datas)
if not n_rows:
n_rows = n_data // n_cols + n_data % n_cols
if not n_cols:
n_cols = n_data // n_rows + n_data % n_rows
for qty in what:
fig, axs = plt.subplots(n_rows, n_cols, figsize=(3 * n_cols, 2.5 * n_rows), sharex=sharex, sharey=sharey)
if direction == "top-bottom":
axs = axs.T
axs = axs.flatten()
xkey = "E"
for i, (name, data) in enumerate(datas):
if direction == "top-bottom":
i_row = i // n_rows
i_col = i % n_rows
else:
i_row = i // n_cols
i_col = i % n_cols
ax = axs[i]
# ylabel only on left column if sharey
ylabel = True
if sharey and i_col > 0:
ylabel = False
plot_spectrum_on_ax(data, ax, qty, xkey=xkey, ylabel=ylabel)
# xlabel only bottom row if sharex
if not sharex or i_row == n_rows - 1:
xlabel = PrsData.get_column_tex_label_with_unit(xkey)
ax.set_xlabel(xlabel)
if name:
ax.set_title(name)
if fig_suptitle:
plt.suptitle(fig_suptitle)
if fig.get_layout_engine() != "constrained":
fig.tight_layout()
fig.savefig(os.path.join(out_dir, f"result_{qty}.pdf"))
[docs]
def add_wl_axis(ax, label=True, draw_ticklabels=None):
"""
Add a top axis that shows wavelength in nm. Primary axis must be energy in eV
:param ax: Matplotlib axes
:param label: True to use the PrsData label for 'wl', False for no label or str for custom label
:param draw_xticklabels: True/False or None to draw them if a label is drawn
"""
xax_nm = ax.secondary_xaxis("top", functions=(eV_to_nm, nm_to_eV))
if draw_ticklabels is None:
draw_ticklabels = label
if not draw_ticklabels:
xax_nm.set_xticklabels([])
# xax_nm.xaxis.set_minor_locator(AutoMinorLocator(2))
if label == True:
label = PrsData.get_column_tex_label_with_unit("wl")
if label:
xax_nm.set_xlabel(label)
# remove the top ticks from the E axis
ax.xaxis.set_tick_params(which="both", top=False)
return xax_nm
if __name__ == "__main__":
import sys
paths = sys.argv[1:]
for p in paths:
data = PrsData(data_path=p)
data.remove_calculated_values()
print(data.metadata["name"], os.path.basename(p))
name = os.path.basename(p)
outdir = p
plot_data(data, name, outdir)
