add interaction_pairs_spin_balanced to variational.util (#344)

python/ffsim/variational/util.py

@@ -126,3 +126,24 @@ def orbital_rotation_from_t1_amplitudes(t1: np.ndarray) -> np.ndarray:
     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