# Copyright © 2025 HQS Quantum Simulations GmbH. All Rights Reserved.
"""Visualization module."""
from __future__ import annotations
from typing import TYPE_CHECKING, Literal, Optional, Union
import matplotlib.pyplot as plt
import numpy as np
from hqs_nmr.datatypes import NMRExperimentalSpectrum1D, NMRGreensFunction1D, NMRSpectrum1D
if TYPE_CHECKING:
from matplotlib.axes import Axes
[docs]
def maximum_excluding_region(
x: np.ndarray, y: np.ndarray, exclude: Optional[tuple[float, float]] = None
) -> float:
"""Determine the maximum of y excluding a specific region of x values.
Args:
x: Values of x.
y: Values of y.
exclude: Region of x values to exclude. Find the overall maximum of y if set to None.
Returns:
The appropriate maximum value of y.
"""
if exclude is not None:
maximum1 = np.max(y[x < min(exclude)], initial=0.0)
maximum2 = np.max(y[x > max(exclude)], initial=0.0)
maximum = max(maximum1, maximum2)
else:
maximum = np.max(y)
return maximum
[docs]
def plot_spectra(
simulated: Union[NMRSpectrum1D, NMRGreensFunction1D],
experimental: NMRExperimentalSpectrum1D,
left: float = 10.0,
right: float = 0.0,
exclude: Optional[tuple[float, float]] = None,
ax: Optional[Axes] = None,
scale_experimental: Literal[
"max-intensity", "max-intensity-window", "factor"
] = "max-intensity",
scale_experimental_factor: float = 1.0,
) -> None:
"""Plot simulated and experimental spectra in one graph.
The simulated spectrum is in orange and upright. The experimental spectrum is in grey and
upside down.
Args:
simulated: Object storing the simulated spectrum.
experimental: Object storing the experimental spectrum.
left: Left margin of the plot in ppm.
right: Right margin of the plot in ppm.
exclude: Region to exclude for normalizing the spectrum (for the solvent peak).
ax: Draw spectra in the specified Matplotlib axes.
scale_experimental: Defines the scaling mode for the experimental spectrum:
- `"max-intensity"` rescales the experimental spectrum so that the largest intensity
value is equal to one.
- `"max-intensity-window"` rescales the experimental spectrum so that the largest
intensity value within the current viewport (defined by `left` and `right`) is equal to
one.
- `"factor"` rescales the experimental spectrum with respect to the given scaling factor.
scale_experimental_factor: Scaling factor for experimental spectrum that is only applied if
the scaling mode (`scale_experimental`) is chosen to be `"factor"`.
"""
if ax is None:
cf = plt.figure(figsize=(9.0, 6.0))
ax = cf.add_axes((0, 0, 1, 1))
ax.tick_params(axis="both", labelsize=16)
ax.tick_params(axis="x", width=2)
for axis in "top", "bottom", "left", "right":
ax.spines[axis].set_linewidth(2)
ax.invert_xaxis()
ax.set_xlim(left, right)
ax.set_ylim(-1.1, 1.1)
ax.set_xlabel(r"$\delta$ (ppm)")
ax.get_yaxis().set_visible(False)
sim_x = np.array(simulated.omegas_ppm)
if isinstance(simulated, NMRGreensFunction1D):
sim_y = -1 * np.sum(simulated.spin_contributions.imag, axis=0)
else:
sim_y = np.sum(simulated.spin_contributions, axis=0)
sim_y = sim_y / maximum_excluding_region(sim_x, sim_y, exclude)
exp_x = experimental.omegas_ppm
exp_y = experimental.intensity
if scale_experimental == "max-intensity-window":
window = (exp_x >= right) & (exp_x <= left)
exp_x = exp_x[window]
exp_y = exp_y[window]
exp_y = -exp_y / maximum_excluding_region(exp_x, exp_y, exclude)
elif scale_experimental == "max-intensity":
exp_y = -exp_y / maximum_excluding_region(exp_x, exp_y, exclude)
else:
exp_y = -exp_y * scale_experimental_factor
ax.plot(sim_x, sim_y, color="#d97815", linewidth=2, label="simulated")
ax.plot(exp_x, exp_y, color="#4d4d4d", linewidth=2, label="experiment")
ax.legend(loc="best", fontsize=16, frameon=False)