Skip to content

fix(orbit): use spatial orbital optimization to preserve spin symmetry #223

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hzhangxyz merged 1 commit into from
Mar 27, 2026
Merged

fix(orbit): use spatial orbital optimization to preserve spin symmetry #223

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 51 additions & 31 deletions qmp/algorithms/haar_with_orbit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,66 +66,86 @@ def _optimize_orbit(
) -> tuple[ModelProto, pathlib.Path]:
"""
Perform orbital optimization and return a new model with the optimized basis and its path.

This implementation uses SPATIAL orbital optimization (not spin-orbital) to preserve
spin symmetry in restricted calculations. The same unitary transformation is applied
to both alpha and beta spin-orbitals.
"""
logging.info("Performing orbital optimization for step %d", step)

n_orbit = typing.cast(typing.Any, model).n_qubits // 2
calculator = NaturalOrbitCalculator(n_orbit)

# 1. Calculate RDM
rdm = calculator.calculate_rdm(configs, psi)
# 1. Calculate SPATIAL RDM (this now returns spatial orbital RDM, not spin-orbital)
rdm_spatial = calculator.calculate_rdm(configs, psi)

# 2. Diagonalize RDM to get Natural Orbitals
# eigvals are occupation numbers, U columns are NOs in old basis
eigvals, U = torch.linalg.eigh(rdm)
# 2. Diagonalize SPATIAL RDM to get Natural Orbitals
# eigvals are occupation numbers, U_spatial columns are NOs in old basis
eigvals, U_spatial = torch.linalg.eigh(rdm_spatial)

# 3. Reorder U to be closest to Identity
# 3. Reorder U_spatial to be closest to Identity
# We want to maximize sum(|U_{ii}|^2)
cost_matrix = -(U.abs() ** 2).cpu().numpy()
cost_matrix = -(U_spatial.abs() ** 2).cpu().numpy()
row_ind, col_ind = linear_sum_assignment(cost_matrix)

# row_ind[i] is matched with col_ind[i]
# We want column row_ind[i] of permuted_U to be column col_ind[i] of U
permuted_U = torch.zeros_like(U)
permuted_U[:, row_ind] = U[:, col_ind]
U = permuted_U
# We want column row_ind[i] of permuted_U to be column col_ind[i] of U_spatial
permuted_U = torch.zeros_like(U_spatial)
permuted_U[:, row_ind] = U_spatial[:, col_ind]
U_spatial = permuted_U

# 4. Load and Transform Integrals
# 4. Load and Transform SPATIAL Integrals
(norb, nelec, nspin), e0, h1, h2 = _read_fcidump_tensors(self.src_fcidump)

device = context.device
U_spatial = U_spatial.to(device=device, dtype=torch.complex128)
h1 = h1.to(device=device, dtype=torch.complex128)
h2 = h2.to(device=device, dtype=torch.complex128)
Comment on lines +97 to +103

Copilot AI Mar 27, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Orbital optimization here is now implemented as spin-restricted (shared U_spatial for α/β and spatial-RDM-based rotation), but the FCIDUMP nspin (MS2) / electron parity isn’t checked before applying it. To avoid silently producing incorrect results on open-shell/non-restricted cases, add an explicit validation (or a fallback to the previous spin-orbital/unrestricted transformation) when MS2 != 0 / when α and β occupations differ.

Copilot uses AI. Check for mistakes.

# Transform SPATIAL integrals first (preserves spin symmetry)
logging.info("Transforming spatial orbital integrals...")

# Transform h1: h1' = U^† h1 U
h1_opt_spatial = U_spatial.conj().T @ h1 @ U_spatial

# Transform h2: h2'_{abcd} = sum U^*_{pa} U^*_{qb} h2_{pqrs} U_{rc} U_{sd}
tmp = torch.einsum("sd,pqrs->pqrd", U_spatial, h2)
tmp = torch.einsum("rc,pqrd->pqcd", U_spatial, tmp)
tmp = torch.einsum("qb,pqcd->pbcd", U_spatial.conj(), tmp)
h2_opt_spatial = torch.einsum("pa,pbcd->abcd", U_spatial.conj(), tmp)

# 5. Expand transformed spatial integrals to spin-orbital basis
# For restricted calculations, alpha and beta use the same spatial orbitals
logging.info("Expanding transformed integrals to spin-orbital basis...")
h1_so = torch.zeros((2 * norb, 2 * norb), dtype=torch.complex128, device=device)
h1_so[0::2, 0::2] = h1.to(device)
h1_so[1::2, 1::2] = h1.to(device)
h1_so[0::2, 0::2] = h1_opt_spatial # alpha-alpha
h1_so[1::2, 1::2] = h1_opt_spatial # beta-beta

h2_so = torch.zeros((2 * norb, 2 * norb, 2 * norb, 2 * norb), dtype=torch.complex128, device=device)
h2_so[0::2, 0::2, 0::2, 0::2] = h2.to(device)
h2_so[0::2, 1::2, 1::2, 0::2] = h2.to(device)
h2_so[1::2, 0::2, 0::2, 1::2] = h2.to(device)
h2_so[1::2, 1::2, 1::2, 1::2] = h2.to(device)

# Transform
h1_opt = U.conj().T @ h1_so @ U
h2_so[0::2, 0::2, 0::2, 0::2] = h2_opt_spatial # aaaa
h2_so[0::2, 1::2, 1::2, 0::2] = h2_opt_spatial # abba
h2_so[1::2, 0::2, 0::2, 1::2] = h2_opt_spatial # baab
h2_so[1::2, 1::2, 1::2, 1::2] = h2_opt_spatial # bbbb

tmp = torch.einsum("sd,pqrs->pqrd", U, h2_so)
tmp = torch.einsum("rc,pqrd->pqcd", U, tmp)
tmp = torch.einsum("qb,pqcd->pbcd", U.conj(), tmp)
h2_opt = torch.einsum("pa,pbcd->abcd", U.conj(), tmp)
# Build spin-orbital unitary for storage (block-diagonal: same U for alpha and beta)
U_spin = torch.zeros((2 * norb, 2 * norb), dtype=torch.complex128, device=device)
U_spin[0::2, 0::2] = U_spatial # alpha block
U_spin[1::2, 1::2] = U_spatial # beta block

# 5. Build Hamiltonian Dict and Save
# 6. Build Hamiltonian Dict and Save
ham_dict: dict[tuple[tuple[int, int], ...], complex] = {}
ham_dict[()] = complex(e0)

h1_indices = torch.nonzero(torch.abs(h1_opt) > 1e-12)
h1_indices = torch.nonzero(torch.abs(h1_so) > 1e-12)
for p, q in h1_indices:
p, q = p.item(), q.item()
ham_dict[((p, 1), (q, 0))] = h1_opt[p, q].item()
ham_dict[((p, 1), (q, 0))] = h1_so[p, q].item()

h2_indices = torch.nonzero(torch.abs(h2_opt) > 1e-12)
h2_indices = torch.nonzero(torch.abs(h2_so) > 1e-12)
for p, q, r, s in h2_indices:
p, q, r, s = p.item(), q.item(), r.item(), s.item()
key = ((p, 1), (q, 1), (r, 0), (s, 0))
ham_dict[key] = ham_dict.get(key, 0j) + h2_opt[p, q, r, s].item()
ham_dict[key] = ham_dict.get(key, 0j) + h2_so[p, q, r, s].item()

site, kind, coef = Hamiltonian._prepare(ham_dict)

Expand All @@ -136,7 +156,7 @@ def _optimize_orbit(
"n_spins": nspin,
"ref_energy": 0.0,
"rdm_eigvals": eigvals.cpu(),
"orbit_unitary": U.cpu(),
"orbit_unitary": U_spin.cpu(), # Store the spin-orbital unitary (block-diagonal)
}

optimized_path = context.folder() / f"optimized_basis_step_{step}.pt"
Expand Down
110 changes: 70 additions & 40 deletions qmp/algorithms/orbit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,60 +48,71 @@ def main(
n_orbit = typing.cast(typing.Any, model).n_qubits // 2
calculator = NaturalOrbitCalculator(n_orbit)

# 2. Calculate 1-RDM and Unitary transformation
rdm = calculator.calculate_rdm(configs, psi)
logging.info("1-RDM calculated. Diagonalizing...")
# 2. Calculate spatial 1-RDM and Unitary transformation
rdm_spatial = calculator.calculate_rdm(configs, psi)
logging.info("Spatial 1-RDM calculated. Diagonalizing...")

# Natural orbitals diagonalize the RDM
eigvals, U = torch.linalg.eigh(rdm)
# U_spatial is now a (n_orbit, n_orbit) matrix for spatial orbitals only
eigvals, U_spatial = torch.linalg.eigh(rdm_spatial)

# 3. Load original integrals
# 3. Load original integrals (these are spatial orbital integrals from FCIDUMP)
(norb, nelec, nspin), e0, h1, h2 = _read_fcidump_tensors(self.src_fcidump)

Copilot AI Mar 27, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This code now assumes a spin-restricted setting (same spatial-orbital rotation applied to both α/β, and RDM built by combining αα/ββ blocks). However, nspin (FCIDUMP MS2) / electron parity isn’t validated before proceeding, so running this on open-shell or otherwise non-restricted inputs will silently apply a potentially incorrect transformation. Consider explicitly checking for the restricted condition (e.g., MS2==0 and even electron count) and either raising a clear error or falling back to a spin-orbital/unrestricted path.

Suggested change
# Validate that the FCIDUMP corresponds to a closed-shell, spin-restricted reference.
# For a restricted calculation we require MS2 == 0 and an even number of electrons.
if nspin != 0 or (nelec % 2) != 0:
raise ValueError(
"Orbit optimization currently assumes a spin-restricted (closed-shell) reference, "
f"but FCIDUMP header reports nelec={nelec}, MS2={nspin}. "
"Support for open-shell / unrestricted inputs is not yet implemented."
)

Copilot uses AI. Check for mistakes.
# Convert to spin-orbital basis (complex128)
device = context.device
U_spatial = U_spatial.to(device=device, dtype=torch.complex128)
h1 = h1.to(device=device, dtype=torch.complex128)
h2 = h2.to(device=device, dtype=torch.complex128)

# 4. Transform SPATIAL orbital integrals first, then expand to spin-orbitals
# This preserves spin symmetry - same transformation for alpha and beta
logging.info("Transforming spatial orbital integrals...")

# Transform h1: h1' = U^† h1 U
h1_opt_spatial = U_spatial.conj().T @ h1 @ U_spatial

# Transform h2 using sequential contractions: h2'_{abcd} = sum U^*_{pa} U^*_{qb} h2_{pqrs} U_{rc} U_{sd}
# This is O(N^5) but correct for 4-index tensor transformation
tmp = torch.einsum("sd,pqrs->pqrd", U_spatial, h2)
tmp = torch.einsum("rc,pqrd->pqcd", U_spatial, tmp)
tmp = torch.einsum("qb,pqcd->pbcd", U_spatial.conj(), tmp)
h2_opt_spatial = torch.einsum("pa,pbcd->abcd", U_spatial.conj(), tmp)

# 5. Expand transformed spatial integrals to spin-orbital basis
# For restricted calculations, alpha and beta use the same spatial orbitals
logging.info("Expanding transformed integrals to spin-orbital basis...")
h1_so = torch.zeros((2 * norb, 2 * norb), dtype=torch.complex128, device=device)
h1_so[0::2, 0::2] = h1.to(device)
h1_so[1::2, 1::2] = h1.to(device)
h1_so[0::2, 0::2] = h1_opt_spatial # alpha-alpha
h1_so[1::2, 1::2] = h1_opt_spatial # beta-beta

h2_so = torch.zeros((2 * norb, 2 * norb, 2 * norb, 2 * norb), dtype=torch.complex128, device=device)
h2_so[0::2, 0::2, 0::2, 0::2] = h2.to(device)
h2_so[0::2, 1::2, 1::2, 0::2] = h2.to(device)
h2_so[1::2, 0::2, 0::2, 1::2] = h2.to(device)
h2_so[1::2, 1::2, 1::2, 1::2] = h2.to(device)

# 4. Transform integrals
logging.info("Transforming integrals to optimized basis...")
h1_opt = U.conj().T @ h1_so @ U

# V'_{abcd} = sum U^*_{pa} U^*_{qb} V_{pqrs} U_{rc} U_{sd}
# Contractions sequential for O(N^5)
# Contract s with U_sd -> d
tmp = torch.einsum("sd,pqrs->pqrd", U, h2_so)
# Contract r with U_rc -> c
tmp = torch.einsum("rc,pqrd->pqcd", U, tmp)
# Contract q with U^*_qb -> b
tmp = torch.einsum("qb,pqcd->pbcd", U.conj(), tmp)
# Contract p with U^*_pa -> a
h2_opt = torch.einsum("pa,pbcd->abcd", U.conj(), tmp)

# 5. Build Hamiltonian dictionary
h2_so[0::2, 0::2, 0::2, 0::2] = h2_opt_spatial # aaaa
h2_so[0::2, 1::2, 1::2, 0::2] = h2_opt_spatial # abba
h2_so[1::2, 0::2, 0::2, 1::2] = h2_opt_spatial # baab
h2_so[1::2, 1::2, 1::2, 1::2] = h2_opt_spatial # bbbb

# Build spin-orbital unitary for storage (block-diagonal: same U for alpha and beta)
U_spin = torch.zeros((2 * norb, 2 * norb), dtype=torch.complex128, device=device)
U_spin[0::2, 0::2] = U_spatial # alpha block
U_spin[1::2, 1::2] = U_spatial # beta block

# 6. Build Hamiltonian dictionary
logging.info("Building Hamiltonian dictionary...")
ham_dict: dict[tuple[tuple[int, int], ...], complex] = {}
ham_dict[()] = complex(e0)

h1_indices = torch.nonzero(torch.abs(h1_opt) > 1e-12)
h1_indices = torch.nonzero(torch.abs(h1_so) > 1e-12)
for p, q in h1_indices:
p, q = p.item(), q.item()
ham_dict[((p, 1), (q, 0))] = h1_opt[p, q].item()
ham_dict[((p, 1), (q, 0))] = h1_so[p, q].item()

h2_indices = torch.nonzero(torch.abs(h2_opt) > 1e-12)
h2_indices = torch.nonzero(torch.abs(h2_so) > 1e-12)
for p, q, r, s in h2_indices:
p, q, r, s = p.item(), q.item(), r.item(), s.item()
key = ((p, 1), (q, 1), (r, 0), (s, 0))
ham_dict[key] = ham_dict.get(key, 0j) + h2_opt[p, q, r, s].item()
ham_dict[key] = ham_dict.get(key, 0j) + h2_so[p, q, r, s].item()

# 6. Save results
# 7. Save results
logging.info("Preparing and saving optimized basis model...")
site, kind, coef = Hamiltonian._prepare(ham_dict)

Expand All @@ -112,7 +123,7 @@ def main(
"n_spins": nspin,
"ref_energy": self.ref_energy if self.ref_energy is not None else model.ref_energy,
"rdm_eigvals": eigvals.cpu(),
"orbit_unitary": U.cpu(), # Store the unitary matrix
"orbit_unitary": U_spin.cpu(), # Store the spin-orbital unitary matrix
}

self.dst_optimized.parent.mkdir(parents=True, exist_ok=True)
Expand Down Expand Up @@ -186,20 +197,39 @@ def __init__(self, n_orbit: int):

def calculate_rdm(self, configs: torch.Tensor, psi: torch.Tensor) -> torch.Tensor:
"""
Calculate the 1-RDM (One-Particle Reduced Density Matrix).
Calculate the spatial 1-RDM (One-Particle Reduced Density Matrix) for restricted calculations.

For restricted (closed-shell) systems, we compute the spatial orbital RDM by summing
contributions from both alpha and beta spin-orbitals, ensuring spin symmetry is preserved.

Comment on lines +200 to +204

Copilot AI Mar 27, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The updated calculate_rdm doc/comments say the spatial RDM is formed by “summing” / D_aa + D_bb, but the implementation actually averages the blocks: (rdm_aa + rdm_bb) / 2. Please make the docstring and inline comments consistent with the chosen convention (sum vs average), and clarify what the resulting eigenvalues represent (per-spin occupations vs spatial occupations).

Copilot uses AI. Check for mistakes.
Returns a RDM of shape (n_orbit, n_orbit) for spatial orbitals only.
"""
n_qubits = self.n_orbit * 2
rdm = torch.zeros((n_qubits, n_qubits), dtype=torch.complex128, device=psi.device)
logging.info("Calculating 1-RDM for %d spin-orbitals", n_qubits)
# Calculate full spin-orbital RDM first
rdm_so = torch.zeros((n_qubits, n_qubits), dtype=torch.complex128, device=psi.device)
logging.info("Calculating spin-orbital 1-RDM for %d spin-orbitals", n_qubits)

for q in range(n_qubits):
if q % 10 == 0:
logging.info("Processing column %d/%d", q, n_qubits)
for p in range(n_qubits):
op = Hamiltonian({((p, 1), (q, 0)): 1.0}, kind="fermi")
val = psi.conj() @ op.apply_within(configs, psi, configs)
rdm[p, q] = val
return rdm
rdm_so[p, q] = val

# Convert to spatial RDM: D_spatial = D_aa + D_bb (for restricted closed-shell)
# Alpha spin-orbitals are at even indices (0, 2, 4, ...), beta at odd indices (1, 3, 5, ...)
# For restricted case, we take the average: D_spatial = (D_aa + D_bb) / 2
# Then the spatial RDM for the transformation should preserve spin symmetry
rdm_aa = rdm_so[0::2, 0::2] # alpha-alpha block
rdm_bb = rdm_so[1::2, 1::2] # beta-beta block

# For spin-restricted case, average the two blocks
# This ensures the resulting natural orbitals are spin-symmetric
rdm_spatial = (rdm_aa + rdm_bb) / 2.0

logging.info("Spatial 1-RDM calculated from spin-orbital RDM (spin-restricted)")
return rdm_spatial


action_dict["orbit"] = OrbitConfig
Loading