diff --git a/qmp/algorithms/haar_with_orbit.py b/qmp/algorithms/haar_with_orbit.py index fb68d767..7206745f 100644 --- a/qmp/algorithms/haar_with_orbit.py +++ b/qmp/algorithms/haar_with_orbit.py @@ -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) + + # 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) @@ -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" diff --git a/qmp/algorithms/orbit.py b/qmp/algorithms/orbit.py index dddeb124..606a9ab2 100644 --- a/qmp/algorithms/orbit.py +++ b/qmp/algorithms/orbit.py @@ -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) - # 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) @@ -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) @@ -186,11 +197,17 @@ 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. + + 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: @@ -198,8 +215,21 @@ def calculate_rdm(self, configs: torch.Tensor, psi: torch.Tensor) -> torch.Tenso 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