# Copyright © 2022-2025 HQS Quantum Simulations GmbH. All Rights Reserved.
# LICENSE PLACEHOLDER
"""Solvers module."""
from __future__ import annotations
from typing import Dict, Any, Protocol, runtime_checkable, TYPE_CHECKING
if TYPE_CHECKING:
import numpy as np
from hqs_nmr.datatypes import NMRCalculationParameters
from struqture_py.spins import PauliHamiltonian
NMR_FREQUENCY_SOLVERS: Dict[str, NMRSolver] = {}
COSY_FREQUENCY_SOLVERS: Dict[str, NMRCOSYSolver] = {}
[docs]
@runtime_checkable
class NMRSolver(Protocol):
"""Callable protocol for NMRSolvers in omega."""
def __call__(
hamiltonian: PauliHamiltonian,
normalized_gyromagnetic_ratios: np.ndarray,
omegas: np.ndarray,
eta: np.ndarray,
spin_contribution_indices: list[int],
calculation_parameters: NMRCalculationParameters,
**kwargs: dict[str, Any],
) -> np.ndarray:
"""Defines the signature of a 1D NMR solver.
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.
"""
...
[docs]
@runtime_checkable
class NMRCOSYSolver(Protocol):
"""Callable protocol for NMRCOSYSolvers in omega."""
def __call__(
hamiltonian: PauliHamiltonian,
normalized_gyromagnetic_ratios: np.ndarray,
omegas_1: np.ndarray,
omegas_2: np.ndarray,
eta: np.ndarray,
spin_contribution_indices: list[int],
**kwargs: dict[str, Any],
) -> np.ndarray:
"""Defines the signature of a COSY solver.
Args:
hamiltonian: Struqture spin Hamiltonian.
normalized_gyromagnetic_ratios: Array of the gyromagnetic factors per site, normalized
with respect to the reference isotope.
omegas_1: Frequency points associated with the time evolution after the first pulse.
omegas_2: Frequency points associated with the time evolution after the second pulse.
eta: Explicit spin-dependent broadening of the peaks.
spin_contribution_indices: List of indices l of I^x_l for which to evaluate the
correlator.
kwargs: Catch all for general interface.
"""
...
[docs]
def register_frequency_solver(
name: str, solver: Any, protocol: Any, dictionary: Dict[str, Any]
) -> None:
"""Register a new frequency solver.
Args:
name: The name (key) for the new solver
solver: A callable implementing a solver.
protocol: Protocol that has to be followed.
dictionary: Dictionary to which to add the solver.
Raises:
KeyError: Name exists already.
TypeError: The callable does not implement the solver protocol.
"""
if name in dictionary:
raise KeyError(f"{name} already in registry.")
if not isinstance(solver, protocol):
raise TypeError(f"{solver} does not implement the correct protocol.")
dictionary[name] = solver