""" This module contain functions for plotting spectral data and the fits to it.
None of these functions should be called directly by users - these functions are called from
plot methods in spectrum_fitting.
"""
import os
import pathlib
from typing import Tuple, List, Union, TYPE_CHECKING
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
# TYPE_CHECKING is False at runtime but allows Type hints in IDE
if TYPE_CHECKING:
from xrdfit.spectrum_fitting import PeakParams, PeakFit
matplotlib.rc('xtick', labelsize=14)
matplotlib.rc('ytick', labelsize=14)
matplotlib.rc('axes', titlesize=20)
matplotlib.rc('axes', labelsize=20)
matplotlib.rcParams['axes.formatter.useoffset'] = False
[docs]def plot_polar_heat_map(num_cakes: int, rad: Union[np.ndarray, List[int]], z_data: np.ndarray,
first_cake_angle: int, cake_order: str):
"""Plot a polar heat map using matplotlib.
:param num_cakes: The number of segments the polar map is divided into.
:param rad: The radial bin edges.
:param z_data: A num_cakes by rad shaped array of data to plot.
:param first_cake_angle: The angle clockwise from vertical at which to label the first cake.
:param cake_order: The order of cakes in the input data. Valid options are `clockwise` or
`anticlockwise`
"""
degrees_per_cake = 360 / num_cakes
half_cake_angle = degrees_per_cake / 2
azm = np.linspace(0, 2 * np.pi, num_cakes + 1)
r, theta = np.meshgrid(rad, azm)
# theta_offset is anticlockwise regardless of theta direction.
# Add 90 to theta offset since theta zero defaults to east.
# Subtract first cake angle to orient image correctly
# Subtract half of direction * half cake angle to align cake centers in cardinal directions,
# rather than the bin edges.
if cake_order == "clockwise":
direction = -1
else:
direction = 1
plt.subplot(projection="polar", theta_direction=direction,
theta_offset=np.deg2rad(90 - first_cake_angle - (direction * half_cake_angle)))
plt.pcolormesh(theta, r, z_data.T)
plt.plot(azm, r, ls='none')
plt.grid()
# Turn on theta grid lines at the cake edges
plt.thetagrids([theta * 360 / num_cakes for theta in range(num_cakes)], labels=[])
# Turn off radial grid lines
plt.rgrids([])
ax = plt.gca()
# Put the cake numbers in the right places. Rotation is clockwise in accordance with
# theta_direction -1 above.
trans, _, _ = ax.get_xaxis_text1_transform(0)
for label in range(1, num_cakes + 1):
ax.text(
np.deg2rad((label * degrees_per_cake - half_cake_angle)),
-0.1, label, transform=trans, rotation=0, ha="center", va="center")
plt.show()
[docs]def plot_spectrum(data: np.ndarray, cakes_to_plot: List[int], merge_cakes: bool, show_points: bool,
x_range: Union[None, Tuple[float, float]] = None, log_scale=False):
"""Plot a raw spectrum using matplotlib.
:param data: The data to plot, x_data in column 0, y data in columns 1-N where N is the number
of cakes in the dataset.
:param cakes_to_plot: Which cakes (columns of y data) to plot.
:param merge_cakes: If True plot the sum of the selected cakes as a single line. If False plot
all selected cakes individually.
:param show_points: Whether to show data points on the plot.
:param x_range: If supplied, restricts the x-axis of the plot to this range.
:param log_scale: If True, plot y axis on log scale. If False use linear scale.
"""
if show_points:
line_spec = "-x"
else:
line_spec = "-"
if x_range:
x_mask = np.logical_and(x_range[0] < data[:, 0], data[:, 0] < x_range[1])
else:
x_mask = [True] * data.shape[0]
if merge_cakes:
plt.plot(data[x_mask, 0], data[x_mask, 1:], line_spec, linewidth=2)
else:
for cake_num in cakes_to_plot:
plt.plot(data[x_mask, 0], data[x_mask, cake_num], line_spec, linewidth=2,
label=cake_num)
plt.legend()
# Plot formatting
plt.minorticks_on()
plt.xlabel(r'Two Theta ($^\circ$)')
plt.ylabel('Intensity')
if x_range:
plt.xlim(x_range[0], x_range[1])
if log_scale:
plt.yscale("log")
plt.tight_layout()
[docs]def plot_peak_params(peak_params: List["PeakParams"], x_range: Tuple[float, float],
label_angle: float):
"""A visualisation to show the PeakParams. Peak bounds are indicated by a shaded grey area.
Maxima bounds are shown by a dashed green line for the min bound and a dashed red line for
the max bound. This method is called with an active plot environment and plots the peak
params on top.
:param peak_params: The peak params to plot.
:param x_range: If supplied, restricts the x-axis of the plot to this range.
:param label_angle: If supplied, the angle to rotate the maxima labels.
"""
for params in peak_params:
bounds_min = params.peak_bounds[0]
bounds_max = params.peak_bounds[1]
range_center = (bounds_min + bounds_max) / 2
plt.axvline(bounds_min, ls="-", lw=1, color="grey")
plt.axvline(bounds_max, ls="-", lw=1, color="grey")
plt.axvspan(bounds_min, bounds_max, alpha=0.2, color='grey', hatch="/")
for maximum in params.maxima:
min_x = maximum.bounds[0]
max_x = maximum.bounds[1]
center = (min_x + max_x) / 2
plt.axvline(min_x, ls="--", color="green")
plt.axvline(max_x, ls="--", color="red")
if x_range[0] < range_center < x_range[1]:
plt.text(center, plt.ylim()[1], maximum.name, ha="center", va="bottom",
fontsize=matplotlib.rcParams["axes.titlesize"] * 0.8, rotation=label_angle)
plt.xlim(x_range)
[docs]def plot_peak_fit(peak_fit: "PeakFit", time_step: str = None, file_name: str = None,
title: str = None, label_angle: float = None, log_scale=False):
"""Plot the result of a peak fit as well as the raw data.
:param peak_fit: The result of a peak fit
:param time_step: If provided, used to generate the title of the plot.
:param file_name: If provided used as a on disk location to save the plot.
:param title: If provided, can be used to override the auto generated plot title.
:param label_angle: The angle to rotate maxima labels.
:param log_scale: Whether to plot the y axis on a log or linear scale.
"""
data = peak_fit.raw_spectrum
# First plot the raw data
for index, cake_num in enumerate(peak_fit.cake_numbers):
plt.plot(data[:, 0], data[:, index + 1], 'x', ms=10, mew=3, label=f"Cake {cake_num}")
title_size = matplotlib.rcParams["axes.titlesize"]
# Now plot the fit
x_data = np.linspace(np.min(data[:, 0]), np.max(data[:, 0]), 100)
y_fit = peak_fit.result.model.eval(peak_fit.result.params, x=x_data)
plt.plot(x_data, y_fit, 'k--', lw=1, label="Fit")
# Do all the ancillaries to make the plot look good.
plt.minorticks_on()
plt.tight_layout()
plt.xlabel(r'Two Theta ($^\circ$)')
plt.ylabel('Intensity')
plt.legend()
if title is not None:
plt.title(title, va="bottom", fontsize=title_size, pad=title_size * 1.8)
elif time_step is not None:
plt.title(f'Fit at t = {time_step}', va="bottom", fontsize=title_size, pad=title_size * 1.8)
for index, maxima_name in enumerate(peak_fit.maxima_names):
maxima_center = peak_fit.result.params[f"maximum_{index}_center"]
plt.text(maxima_center, plt.ylim()[1] * 1.05, maxima_name, horizontalalignment="center",
fontsize=title_size * 0.8, rotation=label_angle)
if log_scale:
plt.yscale("log")
plt.tight_layout()
if file_name:
file_name = pathlib.Path(file_name)
if not file_name.parent.exists():
os.makedirs(file_name.parent)
plt.savefig(file_name)
else:
plt.show()
[docs]def plot_parameter(data: np.ndarray, fit_parameter: str, show_points: bool,
show_error: bool, scale_by_error: bool = False, log_scale=False):
"""Plot a parameter of a fit against time.
:param data: The data to plot, x data in the first column, y data in the second column and
the y error in the third column.
:param fit_parameter: The name of the parameter being plotted, used to generate the y-axis label
:param show_points: Whether to show data points on the plot.
:param show_error: Whether to show error bars on the plot.
:param scale_by_error: If True auto scale the y-axis to the range of the error bars. If False,
auto scale the y-axis to the range of the data.
:param log_scale: Whether to plot the y-axis on a log or linear scale.
"""
no_covar_mask = data[:, 2] == 0
covar_mask = [not value for value in no_covar_mask]
# Plotting the data
plt.plot(data[:, 0], data[:, 1], "-", mec="red")
# Save the y-range to reapply later if wanted
data_y_range = plt.ylim()
if show_points:
plt.plot(data[covar_mask, 0], data[covar_mask, 1], "x", mec="blue")
plt.plot(data[no_covar_mask, 0], data[no_covar_mask, 1], "^", mec="blue")
# Plotting the error bars
if show_error:
plt.fill_between(data[:, 0], data[:, 1] - data[:, 2], data[:, 1] + data[:, 2], alpha=0.3)
plt.plot(data[:, 0], data[:, 1] - data[:, 2], "--", lw=0.5, color='gray')
plt.plot(data[:, 0], data[:, 1] + data[:, 2], "--", lw=0.5, color='gray')
if not scale_by_error:
plt.ylim(data_y_range)
if log_scale:
plt.yscale("log")
plt.xlabel("Time (s)")
plt.ylabel(fit_parameter.replace('_', ' ').title())
plt.show()