import numpy as np
from typing import Union, Literal, Tuple, Optional, Sequence
ArrayOrFloat = Union[np.ndarray, float]
[docs]def moving_average(x: Sequence[float], y: Sequence[float], window: int) -> Tuple[np.ndarray, np.ndarray]:
"""
Calculate a moving average of x and y.
:param x: X values
:type x: Sequence[float]
:param y: Y values
:type y: Sequence[float]
:param window: Moving-average window size
:type window: int
:return: ``x_moving_averaged``, ``y_moving_averaged``
:rtype: Tuple[np.ndarray, np.ndarray]
"""
if not window:
return np.array(x), np.array(y)
window = min(len(x) - 1, window)
if window <= 1:
return np.array(x), np.array(y)
ret_x = np.convolve(x, np.ones(window) / window, mode="valid")
ret_y = np.convolve(y, np.ones(window) / window, mode="valid")
return ret_x, ret_y
def _gaussian(x: np.ndarray, A: ArrayOrFloat, x0: ArrayOrFloat, sigma: ArrayOrFloat) -> np.ndarray:
return A * np.exp(-((x - x0) ** 2) / (2 * sigma**2))
def _lorentzian(x: np.ndarray, A: ArrayOrFloat, x0: ArrayOrFloat, sigma: ArrayOrFloat) -> np.ndarray:
return (A / np.pi) * (0.5 * sigma) / ((x - x0) ** 2 + (0.5 * sigma) ** 2)
def _generate_broadening_widths(x_data: np.ndarray, broadening_width: ArrayOrFloat, centers: np.ndarray) -> np.ndarray:
if not isinstance(broadening_width, np.ndarray):
sigmas = x_data * 0 + broadening_width
else:
if len(broadening_width) == len(centers):
sigmas = broadening_width
else:
raise ValueError(f"{len(broadening_width)=} != {len(centers)=} they should be equal")
return sigmas
[docs]def broaden_results(
centers: np.ndarray,
areas: np.ndarray,
broadening_width: Union[float, np.ndarray] = 40,
broadening_type: Literal["gaussian", "lorentzian"] = "gaussian",
x_data: Union[np.ndarray, Tuple[float, float, float]] = (0, 4000, 0.5),
post_process: Optional[Literal["max_to_1"]] = None,
) -> Tuple[np.ndarray, np.ndarray]:
"""convenient function for create xy curve with gaussian or lorentzian peaks. Used to obtain IR or Raman spectra for example.
:param centers: x positions of the centers of the peaks
:type centers: np.ndarray
:param areas: areas of the peaks
:type areas: np.ndarray
:param broadening_width: if is np.ndarray each peak broadening is assigned otherwise apply the same value for all the peaks, defaults to 40
:type broadening_width: Union[float, np.ndarray], optional
:param broadening_type: the line shape of the peak, defaults to "gaussian"
:type broadening_type: Literal['gaussian', 'lorentzian'], optional
:param x_data: it can be a np.ndarray or you can set a tuple with (min,max,step_size), defaults to (0, 4000, 0.5)
:type x_data: Union[np.ndarray, Tuple[float, float, float]], optional
:param post_process: if max_to_1 the resulted spectrum has the max peak height =1, defaults to None
:type post_process: Literal['max_to_1'], optional
:return: the x and y arrays
:rtype: Tuple[np.ndarray, np.ndarray]
"""
Function_Broaden = {
"gaussian": _gaussian,
"lorentzian": _lorentzian,
}
if isinstance(x_data, tuple) and len(x_data) == 3:
x_array = np.arange(*x_data)
else:
x_array = x_data
y_data: np.ndarray = x_array * 0
sigmas = _generate_broadening_widths(x_array, broadening_width, centers)
for freq_i, Abs_i, sigma_i in zip(centers, areas, sigmas):
y_data += Function_Broaden[broadening_type](x_array, A=Abs_i, x0=freq_i, sigma=sigma_i)
if post_process == "max_to_1":
y_data /= np.max(y_data)
return x_array, y_data
