"""Modules for the preconditioning of the Hamiltonian prior to Schrieffer-Wolff transformation."""
# Copyright © 2021-2025 HQS Quantum Simulations GmbH. All Rights Reserved.
import sys
import textwrap
import numpy as np
from typing import Tuple, cast
from opt_einsum import contract
from lattice_functions.transformer import init_transformer
from hqs_spin_mapper.linear_algebraic_schrieffer_wolff import (
get_fermion_spinful_expression,
extract_block_offdiagonal_hamiltonian,
)
try:
from hqs_spin_mapper.spin_mapper_protocols import (
Supports_SW_Transformation,
Supports_Spin_Optimization,
)
except ImportError as e:
if "mkl" in str(e):
print(
textwrap.dedent(
"""\
===== ERROR: Missing MKL =====
The Intel Math Kernel Library (MKL) could not be found. Please ensure
that the MKL is installed properly.
When using HQStage, run
hqstage envs install-mkl
to install the MKL. In other cases, please consult the manual."""
),
file=sys.stderr,
)
sys.exit(1)
else:
raise
__all__ = [
"preconditioning",
"precondition_interaction",
"diagonalize_bath",
"estimate_required_memory",
]
MATRIX_CONTRACTION = "ij, iq, jp -> qp"
TENSOR_CONTRACTION = "ijkl, iq, jp, kr, ls -> qprs"
ROUNDING_ACCURACY = 10
[docs]
def preconditioning(
transformable_system: Supports_SW_Transformation,
apply_pre_diagonalization: bool = False,
apply_cross_geometry: bool = False,
) -> None:
r"""Return a system suitable for our specific Schrieffer-Wolff transformation method.
The preconditioning includes the transformation to an optimized basis, the removal of
intractable coupling constants, and the pre-diagonalization of the fermionic bath. The results
are applied to the ``transformable_system``.
Args:
transformable_system (Supports_SW_Transformation): Container object for the system data
apply_pre_diagonalization (bool): Flag for the pre-diagonalization of the bath.
apply_cross_geometry (bool): Flag for transformation to cross geometry.
"""
# Rotate Quadratic Hamiltonians into parity-optimized basis
# In the spin-asymmetric case take as input for the quadratic Hamiltonian a tuple of matrices.
hamiltonians = transformable_system.hamiltonians
transformation = transformable_system.rotation_matrix
hamiltonians.H0_uu = transformation.T.conj() @ hamiltonians.H0_uu @ transformation
hamiltonians.H0_dd = transformation.T.conj() @ hamiltonians.H0_dd @ transformation
hamiltonians.H0_ud = transformation.T.conj() @ hamiltonians.H0_ud @ transformation
# If the input for HU (n-n interaction) is a matrix, transform it into a rank-4 tensor.
dim = transformable_system.system_size
if len(hamiltonians.HU.shape) == 2:
HU4 = np.zeros(shape=(dim, dim, dim, dim), dtype=hamiltonians.HU.dtype)
indices = np.nonzero(hamiltonians.HU)
for ip in zip(indices[0], indices[1]):
# TODO: Hubbard- / n.n.-interaction handling correct
HU4[ip[0], ip[1], ip[1], ip[0]] += hamiltonians.HU[ip[0], ip[1]]
hamiltonians.HU = HU4
# If the input for Jz is a matrix, add it onto a rank-4 tensor for HU.
if isinstance(hamiltonians.Jz, np.ndarray):
if len(hamiltonians.Jz.shape) == 2 and len(np.nonzero(hamiltonians.Jz)[0]) > 1:
indices = np.nonzero(hamiltonians.Jz)
for ip in zip(indices[0], indices[1]):
hamiltonians.HU[ip[0], ip[1], ip[1], ip[0]] += hamiltonians.Jz[ip[0], ip[1]]
hamiltonians.Jz = None
# If the input for Jp is a matrix, add it onto a rank-4 tensor for HU.
if isinstance(hamiltonians.Jp, np.ndarray):
if len(hamiltonians.Jp.shape) == 2 and len(np.nonzero(hamiltonians.Jp)[0]) > 1:
indices = np.nonzero(hamiltonians.Jp)
for ip in zip(indices[0], indices[1]):
# Add entries for both S^+ S^- and S^- S^+
hamiltonians.HU[ip[0], ip[1], ip[0], ip[1]] += hamiltonians.Jp[ip[0], ip[1]]
hamiltonians.HU[ip[1], ip[0], ip[1], ip[0]] += hamiltonians.Jp[ip[0], ip[1]]
hamiltonians.Jp = None
# If the input for Jc is a matrix, add it onto a rank-4 tensor for HU.
if isinstance(hamiltonians.Jc, np.ndarray):
if len(hamiltonians.Jc.shape) == 2 and len(np.nonzero(hamiltonians.Jc)[0]) > 1:
indices = np.nonzero(hamiltonians.Jc)
for ip in zip(indices[0], indices[1]):
hamiltonians.HU[ip[0], ip[1], ip[0], ip[1]] += hamiltonians.Jc[ip[0], ip[1]]
hamiltonians.HU[ip[1], ip[0], ip[1], ip[0]] += hamiltonians.Jc[ip[0], ip[1]]
hamiltonians.Jc = None
# Rotate the rank-4 tensor into the spin-optimized basis
hamiltonians.HU = cast(
np.ndarray,
contract(
TENSOR_CONTRACTION,
hamiltonians.HU,
transformation,
transformation,
transformation,
transformation,
),
)
transformable_system.hamiltonians = hamiltonians
if apply_pre_diagonalization:
# Separate the bath-exclusive interaction (and apply a Hartree-Fock decoupling)
precondition_interaction(transformable_system)
# If the system contains finite number of bath sites, rotate into the eigen basis of the bath.
if len(transformable_system.spin_indices) != transformable_system.system_size:
diagonalize_bath(transformable_system, apply_cross_geometry)
[docs]
def precondition_interaction(transformable_system: Supports_SW_Transformation) -> None:
r"""Separate the interaction into terms suitable for the Schrieffer-Wolff transformation
and remainder.
Args:
transformable_system (Supports_SW_Transformation): Container object for the system data
"""
dim = transformable_system.system_size
spin_indices = transformable_system.spin_indices
hamiltonians = transformable_system.hamiltonians
HR = np.zeros(shape=(dim, dim, dim, dim), dtype=hamiltonians.HU.dtype)
# Check every finite entry of the tensor HU, and move all terms that are non-density-like
# on the spin-like sites into the remainder HR or the spin-spin interaction tensors.
indices = np.nonzero(hamiltonians.HU)
for ids in zip(indices[0], indices[1], indices[2], indices[3]):
# Move all interaction that completely remain in the bath to HR.
if all(list(map(lambda i: i not in spin_indices, ids))):
HR[ids[0], ids[1], ids[2], ids[3]] += hamiltonians.HU[ids[0], ids[1], ids[2], ids[3]]
hamiltonians.HU[ids[0], ids[1], ids[2], ids[3]] = 0.0
# We employ a Hartree-Fock decoupling of HR if conditions are met (i.e. RDMs have been provided)
if isinstance(transformable_system, Supports_Spin_Optimization):
# Check whether the 1RDM is spin-resolved
if isinstance(transformable_system.rdm1, tuple):
rdm1_uu = transformable_system.rdm1[0]
rdm1_dd = transformable_system.rdm1[1]
if len(transformable_system.rdm1) == 3:
rdm1_ud = transformable_system.rdm1[2]
else:
rdm1_uu = 0.5 * cast(np.ndarray, transformable_system.rdm1)
rdm1_dd = 0.5 * cast(np.ndarray, transformable_system.rdm1)
rdm1_ud = None
# Retain energy from double counting removal <A><B>
energy_shift = 0.0
indices = np.nonzero(HR)
for ids in zip(indices[0], indices[1], indices[2], indices[3]):
if abs(HR[ids[0], ids[1], ids[2], ids[3]]) > transformable_system.prefactor_cutoff:
hamiltonians.H0_uu[ids[0], ids[3]] += (
rdm1_dd[ids[1], ids[2]] * HR[ids[0], ids[1], ids[2], ids[3]]
)
hamiltonians.H0_dd[ids[1], ids[2]] += (
rdm1_uu[ids[0], ids[3]] * HR[ids[0], ids[1], ids[2], ids[3]]
)
energy_shift -= (
rdm1_uu[ids[0], ids[3]]
* rdm1_dd[ids[1], ids[2]]
* HR[ids[0], ids[1], ids[2], ids[3]]
)
if rdm1_ud is not None:
hamiltonians.H0_ud[ids[1], ids[3]] -= (
rdm1_ud[ids[0], ids[2]] * HR[ids[0], ids[1], ids[2], ids[3]]
)
hamiltonians.H0_ud[ids[0], ids[2]] -= (
rdm1_ud[ids[1], ids[3]] * HR[ids[0], ids[1], ids[2], ids[3]]
)
energy_shift += (
rdm1_ud[ids[0], ids[2]]
* rdm1_ud[ids[1], ids[3]]
* HR[ids[0], ids[1], ids[2], ids[3]]
)
HR[ids[0], ids[1], ids[2], ids[3]] = 0.0
# transformable_system.energy_shift += energy_shift
print(energy_shift)
hamiltonians.HR = HR
transformable_system.hamiltonians = hamiltonians
[docs]
def diagonalize_bath(
transformable_system: Supports_SW_Transformation, apply_cross_geometry: bool = False
) -> None:
r"""Diagonalize the quadratic bath Hamiltonian as preparation of the Schrieffer-Wolff
transformation.
Args:
transformable_system (Supports_SW_Transformation): Container object for the system data
apply_cross_geometry (bool): Flag for transformation to cross geometry.
"""
dim = transformable_system.system_size
spin_indices = transformable_system.spin_indices
hamiltonians = transformable_system.hamiltonians
# Use the star- and cross- transformer to diagonalize the bath subspace whilst
# keeping the spin-like subspace fixed, resulting in modified hopping terms.
# All coupling constants connecting spin- and bath subspaces are later added to
# the hybridization that is removed via the Schrieffer-Wolff transformation.
transformer_up = init_transformer(hamiltonians.H0_uu, sorted(spin_indices))
_ = transformer_up.to_stargeometry()
x0 = int((dim // 2) - (len(spin_indices) // 2))
if apply_cross_geometry:
transformer_up.to_crossgeometry(x0)
hamiltonians.H0_uu = transformer_up.get_hamiltonian()
transformer_dn = init_transformer(hamiltonians.H0_dd, sorted(spin_indices))
_ = transformer_dn.to_stargeometry()
if apply_cross_geometry:
transformer_dn.to_crossgeometry(x0)
hamiltonians.H0_dd = transformer_dn.get_hamiltonian()
hamiltonians.H0_ud = cast(
np.ndarray,
contract(
MATRIX_CONTRACTION,
hamiltonians.H0_ud,
transformer_up.get_transformer(),
transformer_dn.get_transformer(),
),
).round(ROUNDING_ACCURACY)
hamiltonians.HU = cast(
np.ndarray,
contract(
TENSOR_CONTRACTION,
hamiltonians.HU,
transformer_up.get_transformer(),
transformer_dn.get_transformer(),
transformer_dn.get_transformer(),
transformer_up.get_transformer(),
),
).round(ROUNDING_ACCURACY)
hamiltonians.HR = cast(
np.ndarray,
contract(
TENSOR_CONTRACTION,
hamiltonians.HR,
transformer_up.get_transformer(),
transformer_dn.get_transformer(),
transformer_dn.get_transformer(),
transformer_up.get_transformer(),
),
).round(ROUNDING_ACCURACY)
if isinstance(transformable_system, Supports_Spin_Optimization):
if isinstance(transformable_system.rdm1, np.ndarray):
transformable_system.rdm1 = cast(
np.ndarray,
contract(
MATRIX_CONTRACTION,
transformable_system.rdm1,
transformer_up.get_transformer(),
transformer_up.get_transformer(),
),
).round(ROUNDING_ACCURACY)
if isinstance(transformable_system.rdm1, tuple):
rdm1_uu = cast(
np.ndarray,
contract(
MATRIX_CONTRACTION,
transformable_system.rdm1[0],
transformer_up.get_transformer(),
transformer_up.get_transformer(),
),
).round(ROUNDING_ACCURACY)
rdm1_dd = cast(
np.ndarray,
contract(
MATRIX_CONTRACTION,
transformable_system.rdm1[1],
transformer_dn.get_transformer(),
transformer_dn.get_transformer(),
),
).round(ROUNDING_ACCURACY)
rdm1_ud = cast(
np.ndarray,
contract(
MATRIX_CONTRACTION,
transformable_system.rdm1[2],
transformer_up.get_transformer(),
transformer_dn.get_transformer(),
),
).round(ROUNDING_ACCURACY)
transformable_system.rdm1 = (rdm1_uu, rdm1_dd, rdm1_ud)
transformable_system.rdm2 = cast(
np.ndarray,
contract(
TENSOR_CONTRACTION,
transformable_system.rdm2,
transformer_up.get_transformer(),
transformer_dn.get_transformer(),
transformer_dn.get_transformer(),
transformer_up.get_transformer(),
),
).round(ROUNDING_ACCURACY)
# Transform all interaction tensors into the spin-bath basis.
if apply_cross_geometry:
transformable_system.spin_indices = {(x0 + n) for n in range(len(spin_indices))}
else:
transformable_system.spin_indices = {(n) for n in range(len(spin_indices))}
# Store the transformation matrix for use later.
# TODO: Introduce spin direction dependent rotation matrices
transformable_system.rotation_matrix = (
transformable_system.rotation_matrix @ transformer_up.get_transformer()
)
transformable_system.hamiltonians = hamiltonians
[docs]
def estimate_required_memory(transformable_system: Supports_SW_Transformation) -> None:
r"""Gives an estimate for the required memory necessary for the Schrieffer-Wolff transformation
at the current value of ``prefactor_cutoff``.
Args:
transformable_system (Supports_SW_Transformation): The container object for the Spin Mapper
"""
H0 = get_fermion_spinful_expression(
transformable_system.hamiltonians.H0_uu,
transformable_system.hamiltonians.H0_dd,
transformable_system.hamiltonians.H0_ud,
transformable_system.hamiltonians.HU.reshape(pow(transformable_system.system_size, 4)),
transformable_system.spin_indices,
)
V = extract_block_offdiagonal_hamiltonian(
H0, transformable_system.spin_indices, transformable_system.prefactor_cutoff
)
mem = np.round((len(H0) + np.power(len(V), 1.8)) / (1024**2), 2)
print(
f"Using the current prefactor_cutoff, we estimate a memory requirement of up to {mem} GB"
)