# Copyright © 2022-2025 HQS Quantum Simulations GmbH. All Rights Reserved.
# LICENSE PLACEHOLDER
"""A simple direct solver as a reference implementation."""
from __future__ import annotations
from typing import Any, TYPE_CHECKING
import numpy as np
from scipy.sparse import coo_matrix
from struqture_py.spins import SpinHamiltonianSystem, SpinSystem, PauliProduct
from hqs_nmr.hamiltonian_tools import sh_to_matrix
if TYPE_CHECKING:
from hqs_nmr.datatypes import NMRCalculationParameters
CALC_PARALLEL_DEFAULT = False
try:
from lattice_functions.nmr_tools import (
calculate_spectral_function_for_nmr_spin_dependent,
calculate_greens_function_for_nmr_spin_dependent,
)
CALC_PARALLEL_DEFAULT = True
except ImportError:
pass
[docs]
def calc_correlator_function_serial(
vectors: np.ndarray,
Enm: np.ndarray,
Ims: list[SpinSystem],
Im: np.ndarray,
eta: np.ndarray,
omegas: np.ndarray,
As: np.ndarray,
spin_contribution_indices: list[int],
) -> None:
"""Calculate spectral function using a serial loop.
Args:
vectors: Eigenstates of the hamiltonian.
Enm: Energy differences between the different eigenenergies (E_n - E_m).
Ims: Ladder operators for individual spins (I_l^-).
Im: Expectation values of total spin ladder operator (<n|I_total^+|m>).
eta: Spin-dependent broadening parameters.
omegas: Frequency range.
As: Empty array to save in the spin resolved spectral function.
spin_contribution_indices: List of indices l of I^+_l for which to evaluate the correlator.
"""
nspins_total = As.shape[0]
number_steps = omegas.shape[0]
# precalculate the terms <n|I_x^+|m> to be used in the inner loop
Inm = np.zeros((nspins_total, Im.shape[0], Im.shape[1]), Im.dtype)
for n, x in enumerate(spin_contribution_indices):
Ixy = coo_matrix(Ims[x].sparse_matrix_coo()).real
Inm[x, :, :] = (vectors.transpose() @ Ixy @ vectors) * Im * eta[n]
for n in range(number_steps):
omega = omegas[n]
for i, x in enumerate(spin_contribution_indices):
lorentz = (omega - Enm) ** 2 + eta[i] * eta[i]
B = Inm[x] / lorentz
As[x, n] = np.einsum("ij->", B)
[docs]
def calc_correlator_function_parallel(
vectors: np.ndarray,
Enm: np.ndarray,
Ims: list[SpinSystem],
Im: np.ndarray,
eta: np.ndarray,
omegas: np.ndarray,
As_fortran: np.ndarray,
spin_contribution_indices: list[int],
) -> None:
r"""Calculate spectral function using a serial loop.
Note: This function will only be called, if the lattice_functions are installed.
Args:
vectors: Eigenstates of the hamiltonian.
Enm: Energy differences between the different eigenenergies (E_n - E_m).
Ims: Ladder operators for individual spins (I_l^-).
Im: Expectation values of total spin ladder operator (<n|I_total^+|m>).
eta: Spin-dependent broadening parameters.
omegas: Frequency range.
As_fortran: Empty array to be overwritten with the spin resolved spectral function. Has to
have fortran style.
spin_contribution_indices: List of indices l of I^+_l for which to evaluate the correlator.
"""
Enm = Enm.reshape(Enm.size)
for n, x in enumerate(spin_contribution_indices):
Ix = coo_matrix(Ims[x].sparse_matrix_coo()).real
Inm = (vectors.transpose() @ Ix @ vectors) * Im * eta[n]
Inm = Inm.reshape(Inm.size)
if As_fortran.dtype == complex:
calculate_greens_function_for_nmr_spin_dependent(
x, eta[n], omegas, Enm, Inm, As_fortran
)
else:
calculate_spectral_function_for_nmr_spin_dependent(
x, eta[n], omegas, Enm, Inm, As_fortran
)
[docs]
def reference_nmr_solver(
hamiltonian: SpinHamiltonianSystem,
normalized_gyromagnetic_ratios: np.ndarray,
omegas: np.ndarray,
eta: np.ndarray,
spin_contribution_indices: list[int],
calculation_parameters: NMRCalculationParameters,
calc_parallel: bool = CALC_PARALLEL_DEFAULT,
**kwargs: dict[str, Any], # noqa: ARG001
) -> np.ndarray:
r"""Calculate a NMR spectrum brute force from a SpinHamiltonianSystem.
Evaluates a correlator function of the form
.. math::
\sum_n \frac{\bra{n} I^+ \ket{m} \bra{m} I^- \ket{n}}{\omega + E_n - E_m + i \eta}
where I_^- = \sum_l I^-_l.
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.
calc_parallel: If True, the parallel implementation is used, otherwise the serial one.
kwargs: Catch all for general interface.
Returns:
An array of the spectral function for the specified spin contributions at each
frequency.
Raises:
ValueError: If number of spins does not match the number
of gyromagnetic factors.
"""
number_steps = omegas.size
nspins_total = hamiltonian.number_spins()
if nspins_total != normalized_gyromagnetic_ratios.shape[0]:
raise ValueError(
"Number of spin sites does not match" + "the number of gyromagnetic factors"
)
dim = 2**nspins_total
# Obtain the eigenvalues and eigenvectors of the spin Hamiltonian.
energies, vectors = np.linalg.eigh(sh_to_matrix(hamiltonian).real)
Enm: np.ndarray = np.outer(energies, np.ones(dim)) - np.outer(np.ones(dim), energies)
# We excite all spins, Im = I^m,
Im = SpinSystem(nspins_total)
# site-resolved response
Ims = [SpinSystem(nspins_total) for _ in range(nspins_total)]
for x in range(nspins_total):
sigma_x = PauliProduct().set_pauli(x, "X")
sigma_y = PauliProduct().set_pauli(x, "Y")
Im.add_operator_product(sigma_x, 0.5 * normalized_gyromagnetic_ratios[x])
Im.add_operator_product(sigma_y, -0.5j * normalized_gyromagnetic_ratios[x])
if x in spin_contribution_indices:
Ims[x].add_operator_product(sigma_x, 0.5 * normalized_gyromagnetic_ratios[x])
Ims[x].add_operator_product(sigma_y, -0.5j * normalized_gyromagnetic_ratios[x])
Ixy = coo_matrix(Im.sparse_matrix_coo()).real
Im = vectors.transpose() @ Ixy @ vectors
if not calc_parallel:
As = np.zeros((nspins_total, number_steps), dtype=float)
calc_correlator_function_serial(
vectors, Enm, Ims, Im, eta, omegas, As, spin_contribution_indices
)
else:
dtype = complex if calculation_parameters.solver_settings.calc_greens_function else float
As_fortran = np.zeros((number_steps, nspins_total), dtype=dtype, order="F")
calc_correlator_function_parallel(
vectors, Enm, Ims, Im, eta, omegas, As_fortran, spin_contribution_indices
)
As = np.ascontiguousarray(As_fortran.T)
return As[spin_contribution_indices, :] / dim