# Copyright © 2025 HQS Quantum Simulations GmbH. All Rights Reserved.
"""A NMR solver that takes into account the full initial density matrix."""
from typing import Any
import numpy as np
from hqs_nmr.datatypes import NMRCalculationParameters
from hqs_nmr import conversion, hamiltonian_tools
from struqture_py.spins import SpinHamiltonianSystem
from hqs_quantum_solver.spins import (
VectorSpace,
Operator,
spin_x,
isotropic_interaction,
magnetic_field_z,
)
[docs]
def complete_nmr_solver(
hamiltonian: SpinHamiltonianSystem,
normalized_gyromagnetic_ratios: np.ndarray,
omegas: np.ndarray,
eta: np.ndarray,
spin_contribution_indices: list[int],
calculation_parameters: NMRCalculationParameters,
**kwargs: dict[str, Any], # noqa: ARG001
) -> np.ndarray:
r"""A complete NMR solver without approximations.
Evaluates a correlator function of the form:
\\sum_n \\frac{\\bra{n} \\gamma_l I^x_l \\ket{m} \\bra{m} P \\rho P^\\dagger \\ket{n}}{
{\\omega + E_n - E_m + i \\eta}
Where P is a pi / 2 pulse and rho a density matrix given by the Boltzmann distribution.
Args:
hamiltonian: Struqture spin Hamiltonian.
normalized_gyromagnetic_ratios: Array of the gyromagnetic factors per site, normalized with
respect to the reference isotope.
omegas: Desired frequencies.
eta: Explicit spin-dependent broadening of the peaks.
spin_contribution_indices: List of indices l of I^+_l for which to evaluate the correlator.
calculation_parameters: Object storing all calculation parameters including solver-specific
settings.
kwargs: Catch all for general interface.
Returns:
np.ndarray[tuple[Any, Any], np.dtype[np.floating]]: The spectrum. First index = spin #,
second index = time index.
"""
num_spins = hamiltonian.number_spins()
dim = 2**num_spins
# Construct and diagonalize Hamiltonian in rotating frame.
Jz, hJ, _ = hamiltonian_tools.extract_coupling(hamiltonian)
vec_space = VectorSpace(sites=num_spins, total_spin_z=(num_spins % 2, 2))
hamiltonian_rotating = Operator(
isotropic_interaction(hJ) + magnetic_field_z(coef=Jz), domain=vec_space
)
eigenvalues, eigenvectors = np.linalg.eigh(hamiltonian_rotating.todense())
# Obtain energy differences for denominator.
energy_differences = np.outer(eigenvalues, np.ones(dim)) - np.outer(np.ones(dim), eigenvalues)
# Construct time evolution operator of hard pi / 2 pulse. Use eigenbasis for efficiency.
sum_spin_x = Operator(spin_x(coef=np.ones(num_spins)), domain=vec_space).todense()
eigenvalues_pulse, eigenvectors_pulse = np.linalg.eigh(sum_spin_x)
pulse = (
eigenvectors_pulse
@ (np.exp(-1j * (np.pi / 2) * eigenvalues_pulse) * eigenvectors_pulse.conjugate()).T
)
# Calculate all vectors: <n| pulse
bra_pulse = eigenvectors.T @ pulse
# Obtain Boltzmann distribution for initial density matrix. The Hamiltonian has to be defined
# in the laboratory frame of reference here.
Jz_laboratory = (
Jz
+ calculation_parameters.field_T
* (calculation_parameters.gyromagnetic_ratios[calculation_parameters.reference_isotope])
)
beta = 1 / (calculation_parameters.temperature_K * conversion.k_boltzmann_J_per_K)
beta_hbar = beta * conversion.h_planck_J_s
hamiltonian_laboratory = Operator(
beta_hbar * (isotropic_interaction(hJ) + magnetic_field_z(coef=Jz_laboratory)),
domain=vec_space,
)
eigenvalues_laboratory, eigenvectors_laboratory = np.linalg.eigh(
hamiltonian_laboratory.todense()
)
boltzmann_weights = np.exp(-(eigenvalues_laboratory - eigenvalues_laboratory[0]))
rho = eigenvectors_laboratory @ (boltzmann_weights * eigenvectors_laboratory.conjugate()).T
Z = np.sum(boltzmann_weights)
# Time evolution of initial density matrix under the pulse.
pulse_rho_pulse = bra_pulse @ rho @ bra_pulse.conjugate().T
# Precalculate the numerator of the final expression for each spin individually.
# Also create a filter to remove small matrix elements.
list_numerator = []
list_filter = []
for spin_index in spin_contribution_indices:
sx = Operator(
normalized_gyromagnetic_ratios[spin_index] * spin_x(site=spin_index),
domain=vec_space,
)
# <n| S^x |m><m| pulse_rho_pulse |n>
numerator = ((eigenvectors.T @ sx @ eigenvectors) * pulse_rho_pulse.T).flatten()
index_filter = np.where(
abs(numerator) / np.max(abs(numerator))
> calculation_parameters.solver_settings.threshold_matrix_elements
)[0]
list_numerator.append(numerator[index_filter])
list_filter.append(index_filter)
# Evaluate the correlator.
correlator_function = np.zeros((len(spin_contribution_indices), omegas.size), dtype=complex)
for w, omega in enumerate(omegas):
for i, _ in enumerate(spin_contribution_indices):
denominator = (omega + 1j * eta[i] + energy_differences).flatten()
correlator_function[i, w] = np.sum(list_numerator[i] / denominator[list_filter[i]])
## we calculate <\gamma_l S^x_l \rho(t)>
correlator_function *= calculation_parameters.gyromagnetic_ratios[
calculation_parameters.reference_isotope
]
## Return result.
if calculation_parameters.solver_settings.calc_greens_function:
# The factor 1j is necessary for this result to be compatible with other functions
# concerning Green's function
return (-1j * np.pi) * correlator_function / Z
else:
return correlator_function.real / Z