qrunch.tools.acceleration.interface

Interface definitions for hardware-accelerated tensor operations.

Classes

AcceleratedFunctions

Interface for hardware-accelerated tensor operations.

class AcceleratedFunctions

Bases: Protocol

Interface for hardware-accelerated tensor operations.

__init__(*args, **kwargs)

static calculate_boson_orbital_optimization_hessian(h_core: ndarray[tuple[Any, ...], dtype[float64]], gamma: ndarray[tuple[Any, ...], dtype[float64]], gamma_2d: ndarray[tuple[Any, ...], dtype[float64]], integrals: ndarray[tuple[Any, ...], dtype[float64]]) → ndarray[tuple[Any, ...], dtype[float64]]

Compute the reduced Hessian for boson orbital optimization.

Computes the full 4-index Hessian tensor from the single-body Hamiltonian, the 1-RDM (gamma), the 2-RDM (gamma_2d), and the 4-index ERI integrals, then applies 8-fold symmetrization and extracts the upper-triangular block.

The returned matrix has shape (P, P) where P = N*(N-1)/2 and N is the number of orbitals.

This function automatically dispatches to the best available hardware backend: - CUDA (CuPy): Used on Linux with NVIDIA GPUs. - CPU (NumPy): Used as a fallback on other systems or if GPU is unavailable.

Parameters:

h_core (ndarray[tuple[Any, ...], dtype[float64]]) – Single-body Hamiltonian matrix of shape (N, N). Symmetric.
gamma (ndarray[tuple[Any, ...], dtype[float64]]) – 1-RDM of shape (N, N). Symmetric.
gamma_2d (ndarray[tuple[Any, ...], dtype[float64]]) – 2-RDM of shape (N, N). Symmetric.
integrals (ndarray[tuple[Any, ...], dtype[float64]]) – 4-index ERI tensor of shape (N, N, N, N) with 8-fold symmetry.

Return type:

ndarray[tuple[Any, …], dtype[float64]]

static calculate_rotated_mixed_spin_eri_integral_tensor(rotation_matrix_12: ndarray[tuple[Any, ...], dtype[float64]], rotation_matrix_34: ndarray[tuple[Any, ...], dtype[float64]], eri_integral: ndarray[tuple[Any, ...], dtype[float64]]) → ndarray[tuple[Any, ...], dtype[float64]]

Compute the rotated two-body electron repulsion integral (ERI) tensor.

This function transforms a 4-index ERI tensor from one orbital basis to another by applying the given rotation matrix to each of the four indices. Mathematically, this performs the contraction:

\[\begin{split}\\tilde{V}_{ijkl} = \\sum_{abcd} U_{ai}^{(1)} U_{bj}^{(1)} U_{ck}^{(2)} U_{dl}^{(2)} V_{abcd}\end{split}\]

where $V$ is the input integral tensor, $U$ is the rotation matrix, and $\tilde{V}$ is the rotated tensor. This corresponds to the einsum string "ai,bj,ck,dl,abcd->ijkl".

This function automatically dispatches to the best available hardware backend: - Apple Silicon (MLX): Used on macOS with Apple Silicon. - CUDA (CuPy): Used on Linux with NVIDIA GPUs. - CPU (NumPy): Used as a fallback on other systems or if GPU is unavailable.

Parameters:

rotation_matrix_12 (ndarray[tuple[Any, ...], dtype[float64]]) – The rotation matrix of shape (N, N) to rotate dimension 1 and 2.
rotation_matrix_34 (ndarray[tuple[Any, ...], dtype[float64]]) – The rotation matrix of shape (N, N) to rotate dimension 3 and 4.
eri_integral (ndarray[tuple[Any, ...], dtype[float64]]) – The 4-index ERI tensor of shape (N, N, N, N).

Return type:

ndarray[tuple[Any, …], dtype[float64]]

static calculate_rotated_pure_spin_eri_integral_tensor(rotation_matrix: ndarray[tuple[Any, ...], dtype[float64]], eri_integral: ndarray[tuple[Any, ...], dtype[float64]]) → ndarray[tuple[Any, ...], dtype[float64]]

Compute the rotated two-body electron repulsion integral (ERI) tensor.

This function transforms a 4-index ERI tensor from one orbital basis to another by applying the given rotation matrix to each of the four indices. Mathematically, this performs the contraction:

\[\begin{split}\\tilde{V}_{ijkl} = \\sum_{abcd} U_{ai} U_{bj} U_{ck} U_{dl} V_{abcd}\end{split}\]

where $V$ is the input integral tensor, $U$ is the rotation matrix, and $\tilde{V}$ is the rotated tensor. This corresponds to the einsum string "ai,bj,ck,dl,abcd->ijkl".

This function automatically dispatches to the best available hardware backend: - Apple Silicon (MLX): Used on macOS with Apple Silicon. - CUDA (CuPy): Used on Linux with NVIDIA GPUs. - CPU (NumPy): Used as a fallback on other systems or if GPU is unavailable.

Parameters:

rotation_matrix (ndarray[tuple[Any, ...], dtype[float64]]) – The rotation matrix of shape (N, N).
eri_integral (ndarray[tuple[Any, ...], dtype[float64]]) – The 4-index ERI tensor of shape (N, N, N, N).

Return type:

ndarray[tuple[Any, …], dtype[float64]]