Add basic orbital optimization for a fixed 1-/2-RDM (restricted case)

src/openfermion/__init__.py

@@ -57,6 +57,7 @@
     generate_hamiltonian,
     bose_hubbard,
     fermi_hubbard,
+    optimize_orbitals,
     dual_basis_kinetic,
     dual_basis_potential,
     dual_basis_jellium_model,

src/openfermion/hamiltonians/__init__.py

@@ -28,6 +28,8 @@
 
 from .hubbard import bose_hubbard, fermi_hubbard
 
+from .orbital_optimization import optimize_orbitals
+
 from .jellium import (
     dual_basis_kinetic,
     dual_basis_potential,

src/openfermion/hamiltonians/orbital_optimization.py

@@ -0,0 +1,201 @@
+"""
+Basic orbital optimization for a fixed 1- and 2-particle reduced density
+matrix (RDM), restricted (closed-shell) case.
+
+Motivation (Issue #711): a common post-processing step for a correlated
+calculation (CASCI/FCI in a truncated active space, or a 1-/2-RDM recovered
+from a quantum device via sample-based diagonalization) is to ask whether a
+*different* choice of one-particle orbitals -- expressed as a rotation of the
+orbitals the RDM was computed in -- would lower the total energy, holding the
+RDM itself fixed. This module reuses the existing restricted-orbital
+generator parametrization from `hartree_fock.py` (`rhf_params_to_matrix`,
+an antihermitian kappa matrix restricted to occupied/virtual blocks,
+exponentiated into a unitary) and asks scipy to minimize the resulting
+energy over the rotation parameters.
+
+This is deliberately "basic" (per the issue's own wording): the gradient
+used is scipy's numerical one, not an analytic one. Each trial rotation is
+scored by rotating the *Hamiltonian* integrals into the trial basis (via
+`general_basis_change`, the same utility `HartreeFockFunctional.__init__`
+already uses to change basis) and evaluating that rotated Hamiltonian
+against the *fixed* given RDM -- i.e. holding the CI wavefunction's
+expansion coefficients fixed while asking what energy those same
+coefficients would give if they described occupations of a different,
+rotated one-particle basis instead. This is a physically real question
+with a nontrivial answer, not a change of labels: reusing a wavefunction's
+coefficients under rotated orbitals is generally a *different* state, and
+its energy is generally different from (and, importantly, never lower
+than -- see below) the state the RDM actually came from.
+
+Correctness/scope note verified in the test suite: for a *full* active
+space (all molecular orbitals included, RDM from an untruncated FCI
+calculation), the identity rotation (kappa=0) is provably the *global
+minimum* of this objective -- any rotation keeps the trial state inside
+the same complete N-electron Fock space that full CI already minimizes
+over exactly, so no rotation can score below the FCI energy, and the
+optimizer started away from kappa=0 must converge back down to it (not
+below). The routine's actual use case is the *active-space-truncated*
+case, where the RDM comes from a CI diagonalization over a strict subset
+of orbitals -- there, orbital rotation between the active and excluded
+space is not a symmetry of the truncated problem, and rotating orbitals
+can genuinely recover some of the energy lost to the truncation (this is
+exactly the orbital-rotation step of CASSCF-style methods).
+"""
+
+from typing import Optional
+
+import numpy as np
+import scipy as sp
+from scipy.optimize import OptimizeResult
+
+from openfermion.hamiltonians.hartree_fock import generate_hamiltonian, rhf_params_to_matrix
+from openfermion.ops.representations import general_basis_change
+
+
+def _energy_from_rdms(
+    hamiltonian_one_body: np.ndarray,
+    hamiltonian_two_body: np.ndarray,
+    constant: float,
+    one_rdm: np.ndarray,
+    two_rdm: np.ndarray,
+) -> float:
+    """<H> for a fixed Hamiltonian and a fixed (possibly rotated) RDM pair.
+
+    Uses the same elementwise-sum-product convention as
+    `InteractionRDM.expectation()` (both tensors are assumed to already be
+    expressed in the same orbital-index basis).
+    """
+    energy = constant
+    energy += np.sum(one_rdm * hamiltonian_one_body).real
+    energy += np.sum(two_rdm * hamiltonian_two_body).real
+    return energy
+
+
+def optimize_orbitals(
+    one_body_integrals: np.ndarray,
+    two_body_integrals: np.ndarray,
+    one_rdm: np.ndarray,
+    two_rdm: np.ndarray,
+    n_electrons: int,
+    *,
+    nuclear_repulsion: float = 0.0,
+    initial_guess: Optional[np.ndarray] = None,
+    method: str = 'BFGS',
+    verbose: bool = True,
+    sp_options: Optional[dict] = None,
+) -> OptimizeResult:
+    """Restricted orbital-rotation optimization for a fixed 1-/2-RDM.
+
+    Finds the antihermitian generator kappa (parametrized exactly as in
+    `hartree_fock.rhf_params_to_matrix` -- a rotation restricted to
+    occupied-virtual blocks, using `n_electrons // 2` occupied spatial
+    orbitals) that minimizes
+
+        E(kappa) = sum_pq h_pq(kappa) D_qp + sum_pqrs V_pqrs(kappa) Gamma_qpsr
+
+    where h(kappa)/V(kappa) are `one_body_integrals`/`two_body_integrals`
+    rotated into the trial orbital basis U(kappa) = expm(kappa), and D/Gamma
+    are the *fixed* given `one_rdm`/`two_rdm` (i.e. the CI wavefunction's
+    expansion coefficients are held fixed while the orbitals they refer to
+    are rotated, not re-solved at every step).
+
+    Args:
+        one_body_integrals: spatial-orbital one-body integrals, shape
+            (n_orbitals, n_orbitals), in the same reference basis the RDMs
+            were computed in.
+        two_body_integrals: spatial-orbital two-body integrals, shape
+            (n_orbitals,) * 4, chemist ordering matching
+            `hartree_fock.generate_hamiltonian`.
+        one_rdm: fixed spin-orbital 1-RDM, <a^dagger_p a_q>, shape
+            (2 * n_orbitals,) * 2, in the same reference basis.
+        two_rdm: fixed spin-orbital 2-RDM, <a^dagger_p a^dagger_q a_r a_s>,
+            shape (2 * n_orbitals,) * 4, in the same reference basis.
+        n_electrons: total electron count (used only to split occupied vs.
+            virtual spatial orbitals for the restricted parametrization;
+            the RDM's actual trace need not equal this exactly, e.g. for an
+            active-space RDM computed with frozen core orbitals excluded
+            from `one_body_integrals`/`two_body_integrals` -- pass the
+            electron count for *this* integral set).
+        nuclear_repulsion: constant energy offset added to every evaluation.
+        initial_guess: starting kappa parameter vector. Defaults to zero
+            (start from the reference orbitals, i.e. no rotation).
+        method: scipy.optimize.minimize method. Gradient-free by default
+            (numerical differentiation) -- see module docstring.
+        verbose: passed through as scipy's 'disp' option.
+        sp_options: extra options merged into the scipy optimizer options.
+
+    Returns:
+        scipy.optimize.OptimizeResult. `result.x` is the optimal kappa
+        parameter vector; `result.fun` is the optimized energy.
+    """
+    if one_body_integrals.ndim != 2 or one_body_integrals.shape[0] != one_body_integrals.shape[1]:
+        raise ValueError(
+            f"one_body_integrals must be a square 2D array, got shape "
+            f"{one_body_integrals.shape}"
+        )
+    n_orbitals = one_body_integrals.shape[0]
+    if two_body_integrals.shape != (n_orbitals,) * 4:
+        raise ValueError(
+            f"two_body_integrals must have shape {(n_orbitals,) * 4} to match "
+            f"one_body_integrals (n_orbitals={n_orbitals}), got "
+            f"{two_body_integrals.shape}"
+        )
+    n_spin_orbitals = 2 * n_orbitals
+    if one_rdm.shape != (n_spin_orbitals,) * 2:
+        raise ValueError(
+            f"one_rdm must have shape {(n_spin_orbitals,) * 2} (spin-orbital "
+            f"basis, 2 * n_orbitals with n_orbitals={n_orbitals}), got "
+            f"{one_rdm.shape}"
+        )
+    if two_rdm.shape != (n_spin_orbitals,) * 4:
+        raise ValueError(
+            f"two_rdm must have shape {(n_spin_orbitals,) * 4} (spin-orbital "
+            f"basis, 2 * n_orbitals with n_orbitals={n_orbitals}), got "
+            f"{two_rdm.shape}"
+        )
+    if n_electrons % 2 != 0:
+        raise ValueError(
+            f"optimize_orbitals is restricted (closed-shell) -- n_electrons "
+            f"must be even, got {n_electrons}"
+        )
+    nocc = n_electrons // 2
+    nvirt = n_orbitals - nocc
+    if nocc <= 0 or nvirt <= 0:
+        raise ValueError(
+            f"optimize_orbitals needs at least one occupied and one virtual "
+            f"spatial orbital (got n_orbitals={n_orbitals}, n_electrons={n_electrons})"
+        )
+    occ = list(range(nocc))
+    virt = list(range(nocc, n_orbitals))
+
+    def energy(params: np.ndarray) -> float:
+        kappa = rhf_params_to_matrix(params, n_orbitals, occ, virt)
+        rotation = sp.linalg.expm(kappa)
+        rotated_obi = general_basis_change(one_body_integrals, rotation, (1, 0), transpose=False)
+        rotated_tbi = general_basis_change(
+            two_body_integrals, rotation, (1, 1, 0, 0), transpose=False
+        )
+        hamiltonian = generate_hamiltonian(rotated_obi, rotated_tbi, nuclear_repulsion)
+        return _energy_from_rdms(
+            hamiltonian.one_body_tensor,
+            hamiltonian.two_body_tensor,
+            hamiltonian.constant,
+            one_rdm,
+            two_rdm,
+        )
+
+    if initial_guess is None:
+        init_params = np.zeros(nocc * nvirt)
+    else:
+        init_params = np.asarray(initial_guess).flatten()
+        if init_params.size != nocc * nvirt:
+            raise ValueError(
+                f"initial_guess has {init_params.size} parameters, expected "
+                f"nocc * nvirt = {nocc} * {nvirt} = {nocc * nvirt}"
+            )
+
+    sp_optimizer_options = {'disp': verbose}
+    if sp_options is not None:
+        sp_optimizer_options.update(sp_options)
+
+    return sp.optimize.minimize(energy, init_params, method=method, options=sp_optimizer_options)