# Copyright © 2024-2025 HQS Quantum Simulations GmbH. All Rights Reserved.
# LICENSE PLACEHOLDER
"""NMR spectra postprocessing routines."""
from __future__ import annotations
import numpy as np
from scipy import interpolate as interpol
from typing import Union, overload, Any, TYPE_CHECKING
import numpy.typing as npt
if TYPE_CHECKING:
from hqs_nmr_parameters import NMRParameters
from hqs_nmr_parameters import nmr_hamiltonian
from hqs_nmr import datatypes, utils, conversion, hamiltonian_tools
from hqs_nmr.calculate import calculate_greens_function, calculate_eom_greens_function
[docs]
def prepare_shifting(
greens_function: datatypes.NMRGreensFunction1D,
num_omegas_spline: int = 2000,
offset: float = 1e-12,
) -> tuple[np.ndarray, list[interpol.CubicSpline], list[interpol.CubicSpline]]:
r"""Prepare splines and rediscretized frequency points for shifting the shifts.
Args:
greens_function: Green's function calculated using hqs_nmr.
num_omegas_spline: Number of frequency points added to the rediscretized.
frequencies. Defaults to 2000.
offset: Small factor to avoid division by zero. Defaults to 1e-12.
Returns:
rediscretized_omegas_ppm: Original omegas with added linear discretization in ppm.
splines_real: List with splines fitting the real part entries of the
Green's function.
splines_complex: List with splines fitting the imaginary part entries
of the Green's function.
"""
greens_function_inv = np.zeros_like(greens_function.spin_contributions)
for i in range(greens_function.spin_contributions.shape[0]):
greens_function_inv[i, :] = 1 / (greens_function.spin_contributions[i, :] + offset)
rediscretized_omegas_ppm = rediscretize_frequency_range(
greens_function.omegas_ppm, num_omegas_spline
)
splines_real, splines_imag = fit_spline_to_vector_function(
greens_function.omegas_ppm, greens_function_inv
)
return rediscretized_omegas_ppm, splines_real, splines_imag
[docs]
def shift_greens_function(
shifts: np.ndarray,
greens_function: datatypes.NMRGreensFunction1D,
rediscretized_omegas_ppm: np.ndarray,
splines_real: list[interpol.CubicSpline],
splines_imag: list[interpol.CubicSpline],
number_relevant_spins: int,
) -> datatypes.NMRSpectrum1D:
r"""Shift Green's function using spline fit.
Args:
shifts: The shifts of the magnetic shifts.
greens_function: Green's function that is supposed to be shifted.
rediscretized_omegas_ppm: Rediscretized frequency points in ppm.
splines_real: Splines for the real part of the Green's function.
splines_imag: Splines for the imaginary part of the Green's function.
number_relevant_spins: Number of times the homo-isotope appears in the molecule.
Returns:
Normalized and shifted spectrum.
"""
num_omegas_spline = rediscretized_omegas_ppm.size
nspins_total = len(splines_real)
shifted_greens_function = np.zeros((nspins_total, num_omegas_spline), dtype=complex)
for i in range(nspins_total):
shift = shifts[i]
shifted_greens_function[i, :] = 1 / (
splines_real[i](rediscretized_omegas_ppm - shift)
+ 1j * splines_imag[i](rediscretized_omegas_ppm - shift)
)
shifted_spectral_function = shifted_greens_function.imag
utils.normalize_correlator(
rediscretized_omegas_ppm, shifted_spectral_function, number_relevant_spins
)
return datatypes.NMRSpectrum1D(
omegas_ppm=rediscretized_omegas_ppm,
spin_contributions=shifted_spectral_function,
fwhm_ppm=greens_function.fwhm_ppm,
)
[docs]
def rediscretize_frequency_range(
omegas_ppm: np.ndarray, num_omegas_spline: int = 2000
) -> np.ndarray:
r"""Rediscretize the frequency range for the spline fits.
Add an additional linear discretization to the current distribution.
Args:
omegas_ppm: Frequency discretization without the linear discretization.
num_omegas_spline: Minimal number of points in the linear discretization.
Returns.:
Vector with new discretization.
"""
# Number of linear discretization points that should be at least achieved
omegas_spline = np.zeros(omegas_ppm.size + num_omegas_spline)
delta_omega_spline = (omegas_ppm[-1] - omegas_ppm[0]) / num_omegas_spline
# merge lineare discretization with the adapted mesh
num_omegas_spline = 0
for n in range(omegas_ppm.size - 1):
omegas_spline[num_omegas_spline] = omegas_ppm[n]
num_omegas_spline += 1
delta_omega_n = omegas_ppm[n + 1] - omegas_ppm[n]
if delta_omega_n > delta_omega_spline:
num_new_omegas = int(delta_omega_n / delta_omega_spline)
delta_omega_n = (omegas_ppm[n + 1] - omegas_ppm[n]) / (num_new_omegas + 1)
for _ in range(num_new_omegas):
omegas_spline[num_omegas_spline] = (
omegas_spline[num_omegas_spline - 1] + delta_omega_n
)
num_omegas_spline += 1
omegas_spline[num_omegas_spline] = omegas_ppm[-1]
return omegas_spline[: num_omegas_spline + 1]
@overload
def fit_spline_to_vector_function(
parameters: np.ndarray, vector_function: npt.NDArray[np.floating]
) -> list[interpol.CubicSpline]: ...
@overload
def fit_spline_to_vector_function(
parameters: np.ndarray, vector_function: npt.NDArray[np.complexfloating]
) -> tuple[list[interpol.CubicSpline], list[interpol.CubicSpline]]: ...
[docs]
def fit_spline_to_vector_function(
parameters: np.ndarray, vector_function: np.ndarray
) -> Union[
list[interpol.CubicSpline],
tuple[list[interpol.CubicSpline], list[interpol.CubicSpline]],
]:
r"""Fit spline to real or complex vector function.
Args:
parameters: Parameters at which the vector function was evaluated.
vector_function: Values of the vector function.
Returns:
If the vector function is real, it returns one list with a spline associated with each
vector function entry, otherwise it returns one list for the real and one for the complex
part.
"""
vector_dim = vector_function.shape[0]
spline_real_part = []
spline_imag_part = []
for x in range(vector_dim):
spline_real_x = interpol.CubicSpline(
parameters, np.real(vector_function[x, :]), bc_type="natural"
)
spline_real_part.append(spline_real_x)
if vector_function.dtype == complex:
for x in range(vector_dim):
spline_imag_x = interpol.CubicSpline(
parameters, np.imag(vector_function[x, :]), bc_type="natural"
)
spline_imag_part.append(spline_imag_x)
else:
return spline_real_part
return spline_real_part, spline_imag_part
[docs]
def create_new_molecule(
molecule_parameters: NMRParameters,
shifts_of_Jz_ppm: np.ndarray,
) -> NMRParameters:
r"""Create a new molecular NMR parameter representation, with adapted shifts.
Args:
molecule_parameters: Molecular NMR parameters that should be adapted.
shifts_of_Jz_ppm: Shifts of the magnetic shifts in ppm.
Returns:
Molecular NMR parameters with the adapted shifts.
"""
molecule_parameters_adapted = molecule_parameters.model_copy()
# Note: original shifts are also defined in ppm.
molecule_parameters_adapted.shifts += shifts_of_Jz_ppm
return molecule_parameters_adapted
[docs]
def get_fitted_greens_function(
greens_function: datatypes.NMRGreensFunction1D,
num_omegas_spline: int = 2000,
offset: float = 1e-12,
) -> tuple[np.ndarray, np.ndarray]:
"""Calculate a spline fit of a Green's function after rediscretizing the frequencies.
Args.:
greens_function: Green's function to fit.
num_omegas_spline: Number of frequency points added to the rediscretized frequencies.
Defaults to 2000.
offset: Small factor to avoid division by zero when fitting the Green's function.
Defaults to 1e-12.
Returns:
Rediscretized frequency range and fitted Green's function.
"""
(
rediscretized_omegas_ppm,
splines_real_greens_function,
splines_imag_greens_function,
) = prepare_shifting(greens_function, num_omegas_spline, offset)
num_spins = greens_function.spin_contributions.shape[0]
fitted_greens_function = np.zeros((num_spins, rediscretized_omegas_ppm.size), complex)
for i in range(num_spins):
fitted_greens_function[i, :] = 1 / (
splines_real_greens_function[i](rediscretized_omegas_ppm)
+ 1j * splines_imag_greens_function[i](rediscretized_omegas_ppm)
)
return rediscretized_omegas_ppm, fitted_greens_function
[docs]
def j_coupling_adaptation_parameters(
result_greens_function: datatypes.NMRResultGreensFunction1D,
result_eom_greens_function: datatypes.NMRResultGreensFunction1D,
result_eom_correction_greens_function: datatypes.NMRResultGreensFunction1D,
) -> dict[str, Any]:
"""Determine all parameters needed to adapt a Green's function with respect to some J-coupling.
Note that all Green's functions used here should be non normalized.
Args.:
result_greens_function: Green's function that should be adapted.
result_eom_greens_function: EOM Green's function associated with the standard NMR Green's
function.
result_eom_correction_greens_function: EOM Green's function to correct the standard NMR
Green's function.
Returns:
A dictionary with all necessary parameters.
"""
# Fit the Green's functions
rediscretized_omegas_ppm, fitted_greens_function = get_fitted_greens_function(
result_greens_function.greens_function
)
_, fitted_eom_greens_function = get_fitted_greens_function(
result_eom_greens_function.greens_function
)
_, fitted_eom_correction_greens_function = get_fitted_greens_function(
result_eom_correction_greens_function.greens_function
)
# Frequencies will be needed in rad_per_s and ppm.
rediscretized_omegas_rad_per_s = conversion.ppm_to_rad_per_s(
rediscretized_omegas_ppm, result_greens_function.calculation_parameters
)
# Get broadening used in Green's function
fwhm_rad_per_s = conversion.ppm_to_rad_per_s(
result_greens_function.greens_function.fwhm_ppm,
result_greens_function.calculation_parameters,
)
broadening_rad_per_s = fwhm_rad_per_s / 2
# Get Sz terms in rad per s.
hamiltonian = nmr_hamiltonian(
result_greens_function.molecule_parameters,
result_greens_function.calculation_parameters.field_T,
reference_isotope=result_greens_function.calculation_parameters.reference_isotope,
gyromagnetic_ratios=result_greens_function.calculation_parameters.gyromagnetic_ratios,
)
sz_terms_rad_per_s, _, _ = hamiltonian_tools.extract_coupling(hamiltonian)
return {
"rediscretized_omegas_rad_per_s": rediscretized_omegas_rad_per_s,
"rediscretized_omegas_ppm": rediscretized_omegas_ppm,
"broadening_rad_per_s": broadening_rad_per_s,
"sz_terms_rad_per_s": sz_terms_rad_per_s,
"greens_function": fitted_greens_function,
"eom_greens_function": fitted_eom_greens_function,
"eom_correction_greens_function": fitted_eom_correction_greens_function,
}
[docs]
def prepare_j_coupling_adaptation(
molecule_parameters: datatypes.NMRParameters,
calculation_parameters: datatypes.NMRCalculationParameters,
adaptation_indices: Union[tuple[int, int], list[tuple[int, int]]],
) -> tuple[datatypes.NMRResultGreensFunction1D, dict[str, Any]]:
"""Calculate the Green's function and all parameters needed to perform a J-coupling adaptation.
Args.:
molecule_parameters: The NMR parameters of the molecule, which should be adapted.
calculation_parameters: The NMR calculation parameters as defined for a standard NMR
Green's function calculation.
adaptation_indices: Tuple or list of tuple of all J-coupling which are supposed to be
adapted at once.
Returns:
The result of a standard Green's function calculation using the input parameters and
a dictionary with all necessary parameters to perform a J-coupling adaptation.
"""
# We do not normalize directly to be able to do the J-coupling adaptation later.
normalize_spectrum_input = calculation_parameters.normalize_spectrum
calculation_parameters.normalize_spectrum = False
# The standard Green's function obtained typically from DFT input data.
result_dft_greens_function = calculate_greens_function(
molecule_parameters, calculation_parameters
)
# The EOM Green's function F as it comes out directly from the derivation.
result_eom_greens_function = calculate_eom_greens_function(result_dft_greens_function)
# The correction to the EOM Green's function F.
result_eom_correction_greens_function = calculate_eom_greens_function(
result_dft_greens_function, adaptation_indices
)
# Get parameters needed to perform the J-coupling adaptation
parameters_j_coupling_adaptation = j_coupling_adaptation_parameters(
result_dft_greens_function,
result_eom_greens_function,
result_eom_correction_greens_function,
)
# After calculating all parameters needed for the J-coupling adaptation we can normalize
# the calculated Green's function.
result_dft_greens_function.greens_function.spin_contributions = utils.normalize_correlator(
result_dft_greens_function.greens_function.omegas_ppm,
result_dft_greens_function.greens_function.spin_contributions,
)
# Change the calculation_parameters back to what they were in the beginning.
calculation_parameters.normalize_spectrum = normalize_spectrum_input
return result_dft_greens_function, parameters_j_coupling_adaptation
[docs]
def adapt_spectrum_by_j_coupling(
delta_j_coupling: float,
rediscretized_omegas_rad_per_s: np.ndarray,
rediscretized_omegas_ppm: np.ndarray,
broadening_rad_per_s: Union[float, np.ndarray],
sz_terms_rad_per_s: np.ndarray,
greens_function: np.ndarray,
eom_greens_function: np.ndarray,
eom_correction_greens_function: np.ndarray,
) -> np.ndarray:
"""Adapt a spectrum with respect to a shift in a J-coupling.
Which J-coupling is shifted depends on the EOM correction Green's function.
Args:
delta_j_coupling: Shift of the J-coupling.
rediscretized_omegas_rad_per_s: Rediscretized frequency range in rad per s.
rediscretized_omegas_ppm: Rediscretized frequency range in ppm.
broadening_rad_per_s: Artificial broadening (half of the fwhm) in rad per s. If a vector
is handed over a spin-dependent broadening is assumed.
sz_terms_rad_per_s: Array with the coefficients of the I^z terms in the Hamiltonian.
greens_function: The spin contributions of the Green's function to be adapted.
eom_greens_function: The spin contributions of the EOM Green's function associated with
the standard Green's function.
eom_correction_greens_function: The spin contributions of the EOM correction Green's
function.
Returns:
The normalized Green's function spin contributions with the adapted J-coupling taken
into account.
"""
num_spins = sz_terms_rad_per_s.size
if isinstance(broadening_rad_per_s, float):
broadening_rad_per_s = broadening_rad_per_s * np.ones(num_spins)
eom_greens_function_spin_contributions = np.zeros_like(greens_function)
for k in range(num_spins):
eom_greens_function_spin_contributions[k, :] = 0.5 / (
rediscretized_omegas_rad_per_s
+ 1j * broadening_rad_per_s[k]
+ sz_terms_rad_per_s[k]
- (eom_greens_function[k, :] + delta_j_coupling * eom_correction_greens_function[k, :])
/ greens_function[k, :]
)
utils.normalize_correlator(rediscretized_omegas_ppm, eom_greens_function_spin_contributions)
return eom_greens_function_spin_contributions