diff --git a/python/ffsim/contract/diag_coulomb.py b/python/ffsim/contract/diag_coulomb.py index 1f779bc69..9e5854f72 100644 --- a/python/ffsim/contract/diag_coulomb.py +++ b/python/ffsim/contract/diag_coulomb.py @@ -233,8 +233,6 @@ def diag_coulomb_linop( If passing a pair, you can use ``None`` for one of the values in the pair to indicate that no operation should be applied to that spin sector. - mat_alpha_beta: A matrix of coefficients to use for interactions between - orbitals with differing spin. z_representation: Whether the input matrices are in the "Z" representation. Returns: diff --git a/python/ffsim/variational/ucj_spin_balanced.py b/python/ffsim/variational/ucj_spin_balanced.py index 69d2807c4..11db14097 100644 --- a/python/ffsim/variational/ucj_spin_balanced.py +++ b/python/ffsim/variational/ucj_spin_balanced.py @@ -17,11 +17,11 @@ from typing import cast import numpy as np -import scipy.linalg from ffsim import gates, linalg from ffsim.variational.util import ( orbital_rotation_from_parameters, + orbital_rotation_from_t1_amplitudes, orbital_rotation_to_parameters, validate_interaction_pairs, ) @@ -451,10 +451,7 @@ def from_t_amplitudes( final_orbital_rotation = None if t1 is not None: - final_orbital_rotation_generator = np.zeros((norb, norb), dtype=complex) - final_orbital_rotation_generator[:nocc, nocc:] = t1 - final_orbital_rotation_generator[nocc:, :nocc] = -t1.T - final_orbital_rotation = scipy.linalg.expm(final_orbital_rotation_generator) + final_orbital_rotation = orbital_rotation_from_t1_amplitudes(t1) # Zero out diagonal coulomb matrix entries if requested if pairs_aa is not None: diff --git a/python/ffsim/variational/ucj_spin_unbalanced.py b/python/ffsim/variational/ucj_spin_unbalanced.py index 2279e5c68..51609206f 100644 --- a/python/ffsim/variational/ucj_spin_unbalanced.py +++ b/python/ffsim/variational/ucj_spin_unbalanced.py @@ -17,11 +17,11 @@ from typing import cast import numpy as np -import scipy.linalg from ffsim import gates, linalg from ffsim.variational.util import ( orbital_rotation_from_parameters, + orbital_rotation_from_t1_amplitudes, orbital_rotation_to_parameters, validate_interaction_pairs, ) @@ -579,20 +579,8 @@ def from_t_amplitudes( final_orbital_rotation = None if t1 is not None: t1a, t1b = t1 - - final_orbital_rotation_generator_a = np.zeros((norb, norb), dtype=complex) - final_orbital_rotation_generator_a[:nocc_a, nocc_a:] = t1a - final_orbital_rotation_generator_a[nocc_a:, :nocc_a] = -t1a.T - final_orbital_rotation_a = scipy.linalg.expm( - final_orbital_rotation_generator_a - ) - - final_orbital_rotation_generator_b = np.zeros((norb, norb), dtype=complex) - final_orbital_rotation_generator_b[:nocc_b, nocc_b:] = t1b - final_orbital_rotation_generator_b[nocc_b:, :nocc_b] = -t1b.T - final_orbital_rotation_b = scipy.linalg.expm( - final_orbital_rotation_generator_b - ) + final_orbital_rotation_a = orbital_rotation_from_t1_amplitudes(t1a) + final_orbital_rotation_b = orbital_rotation_from_t1_amplitudes(t1b) final_orbital_rotation = np.stack( [final_orbital_rotation_a, final_orbital_rotation_b] ) diff --git a/python/ffsim/variational/ucj_spinless.py b/python/ffsim/variational/ucj_spinless.py index e274348a3..ea1a23fd0 100644 --- a/python/ffsim/variational/ucj_spinless.py +++ b/python/ffsim/variational/ucj_spinless.py @@ -17,11 +17,11 @@ from typing import cast import numpy as np -import scipy.linalg from ffsim import gates, linalg from ffsim.variational.util import ( orbital_rotation_from_parameters, + orbital_rotation_from_t1_amplitudes, orbital_rotation_to_parameters, validate_interaction_pairs, ) @@ -385,10 +385,7 @@ def from_t_amplitudes( final_orbital_rotation = None if t1 is not None: - final_orbital_rotation_generator = np.zeros((norb, norb), dtype=complex) - final_orbital_rotation_generator[:nocc, nocc:] = t1 - final_orbital_rotation_generator[nocc:, :nocc] = -t1.T - final_orbital_rotation = scipy.linalg.expm(final_orbital_rotation_generator) + final_orbital_rotation = orbital_rotation_from_t1_amplitudes(t1) # Zero out diagonal coulomb matrix entries if requested if interaction_pairs is not None: diff --git a/python/ffsim/variational/util.py b/python/ffsim/variational/util.py index 7f4ff27ca..78ecd2ee3 100644 --- a/python/ffsim/variational/util.py +++ b/python/ffsim/variational/util.py @@ -107,3 +107,43 @@ def orbital_rotation_from_parameters( generator[rows, cols] += vals generator[cols, rows] -= vals return scipy.linalg.expm(generator) + + +def orbital_rotation_from_t1_amplitudes(t1: np.ndarray) -> np.ndarray: + """Construct an orbital rotation from t1 amplitudes. + + The orbital rotation is constructed as exp(t1 - t1†). + + Args: + t1: The t1 amplitudes. + + Returns: + The orbital rotation. + """ + nocc, nvrt = t1.shape + norb = nocc + nvrt + generator = np.zeros((norb, norb), dtype=t1.dtype) + generator[:nocc, nocc:] = t1 + generator[nocc:, :nocc] = -t1.T.conj() + return scipy.linalg.expm(generator) + + +def interaction_pairs_spin_balanced( + connectivity: str, norb: int +) -> tuple[list[tuple[int, int]] | None, list[tuple[int, int]] | None]: + """Returns alpha-alpha and alpha-beta diagonal Coulomb interaction pairs.""" + if connectivity == "all-to-all": + pairs_aa = None + pairs_ab = None + elif connectivity == "square": + pairs_aa = [(p, p + 1) for p in range(norb - 1)] + pairs_ab = [(p, p) for p in range(norb)] + elif connectivity == "hex": + pairs_aa = [(p, p + 1) for p in range(norb - 1)] + pairs_ab = [(p, p) for p in range(norb) if p % 2 == 0] + elif connectivity == "heavy-hex": + pairs_aa = [(p, p + 1) for p in range(norb - 1)] + pairs_ab = [(p, p) for p in range(norb) if p % 4 == 0] + else: + raise ValueError(f"Invalid connectivity: {connectivity}") + return pairs_aa, pairs_ab