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,167 @@
+"""
+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.
+    """
+    n_orbitals = one_body_integrals.shape[0]
+    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()
+
+    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)

src/openfermion/hamiltonians/orbital_optimization_test.py

@@ -0,0 +1,189 @@
+import itertools
+
+import numpy as np
+import pytest
+import scipy as sp
+
+from openfermion.chem import MolecularData
+from openfermion.config import DATA_DIRECTORY
+from openfermion.hamiltonians.hartree_fock import generate_hamiltonian, rhf_params_to_matrix
+from openfermion.hamiltonians.orbital_optimization import optimize_orbitals
+from openfermion.linalg import expectation, get_ground_state, get_sparse_operator
+from openfermion.ops.operators import FermionOperator
+from openfermion.ops.representations import general_basis_change
+
+
+def _fci_ground_state_rdms(hamiltonian, n_qubits):
+    """1-/2-RDM of the exact ground state, computed directly by expectation
+    value against the ground-state vector (not via a measured qubit
+    operator) -- the same construction `measurements.get_interaction_rdm`
+    uses, but starting from `linalg.get_ground_state` instead of a real
+    measurement, appropriate for a known-answer test."""
+    sparse_h = get_sparse_operator(hamiltonian, n_qubits=n_qubits)
+    energy, state = get_ground_state(sparse_h)
+
+    one_rdm = np.zeros((n_qubits, n_qubits))
+    for p, q in itertools.product(range(n_qubits), repeat=2):
+        op = get_sparse_operator(FermionOperator(((p, 1), (q, 0))), n_qubits=n_qubits)
+        one_rdm[p, q] = expectation(op, state).real
+
+    two_rdm = np.zeros((n_qubits,) * 4)
+    for p, q, r, s in itertools.product(range(n_qubits), repeat=4):
+        op = get_sparse_operator(
+            FermionOperator(((p, 1), (q, 1), (r, 0), (s, 0))), n_qubits=n_qubits
+        )
+        two_rdm[p, q, r, s] = expectation(op, state).real
+
+    return energy, one_rdm, two_rdm
+
+
+def _load_h2(bond_length='0.7414', basis_suffix='sto-3g'):
+    m = MolecularData(filename=f"{DATA_DIRECTORY}/H2_{basis_suffix}_singlet_{bond_length}.hdf5")
+    m.load()
+    return m
+
+
+def test_optimize_orbitals_at_identity_reproduces_full_ci_energy():
+    """A necessary correctness check: with the RDM taken from a full
+    (untruncated) FCI calculation in the reference orbitals, evaluating the
+    objective at kappa=0 (the identity rotation) must reproduce the FCI
+    energy exactly -- this is just re-checking energy() is wired up to the
+    same accounting `InteractionRDM.expectation` uses, nothing about
+    optimization yet."""
+    m = _load_h2()
+    hamiltonian = m.get_molecular_hamiltonian()
+    fci_energy, one_rdm, two_rdm = _fci_ground_state_rdms(hamiltonian, m.n_qubits)
+
+    result = optimize_orbitals(
+        m.one_body_integrals,
+        m.two_body_integrals,
+        one_rdm,
+        two_rdm,
+        m.n_electrons,
+        nuclear_repulsion=m.nuclear_repulsion,
+        initial_guess=np.zeros((m.n_electrons // 2) * (m.n_orbitals - m.n_electrons // 2)),
+        method='Nelder-Mead',
+        verbose=False,
+        sp_options={'maxiter': 1},  # don't actually move -- just evaluate near kappa=0
+    )
+    assert np.isclose(result.fun, fci_energy, atol=1e-6)
+
+
+def test_full_space_fci_rdm_is_never_beaten_by_any_rotation():
+    """Physical correctness check, not a code-behavior tautology: when the
+    RDM comes from a full-space FCI calculation, no orbital rotation can
+    produce a state with LOWER energy than the FCI value, because a
+    rotation among all M orbitals stays inside the same complete
+    N-electron Fock space that full CI already minimizes over exactly.
+    kappa=0 must therefore be a global minimum of energy(kappa) -- the
+    optimizer, started away from kappa=0, must converge back down to (not
+    below) the FCI energy, and any explicit nonzero kappa must score
+    >= the FCI energy."""
+    m = _load_h2()
+    hamiltonian = m.get_molecular_hamiltonian()
+    fci_energy, one_rdm, two_rdm = _fci_ground_state_rdms(hamiltonian, m.n_qubits)
+
+    # explicit nonzero rotations must not beat the FCI floor
+    n_orbitals = m.n_orbitals
+    nocc = m.n_electrons // 2
+    occ = list(range(nocc))
+    virt = list(range(nocc, n_orbitals))
+    for scale in (0.3, -0.7, 1.2):
+        params = np.full(nocc * (n_orbitals - nocc), scale)
+        kappa = rhf_params_to_matrix(params, n_orbitals, occ, virt)
+        rotation = sp.linalg.expm(kappa)
+        rotated_obi = general_basis_change(m.one_body_integrals, rotation, (1, 0), transpose=False)
+        rotated_tbi = general_basis_change(
+            m.two_body_integrals, rotation, (1, 1, 0, 0), transpose=False
+        )
+        rotated_hamiltonian = generate_hamiltonian(rotated_obi, rotated_tbi, m.nuclear_repulsion)
+        energy = rotated_hamiltonian.constant
+        energy += np.sum(one_rdm * rotated_hamiltonian.one_body_tensor).real
+        energy += np.sum(two_rdm * rotated_hamiltonian.two_body_tensor).real
+        assert energy >= fci_energy - 1e-8, (
+            f"rotation with params={scale} scored below the FCI floor -- "
+            f"got {energy}, floor is {fci_energy}"
+        )
+
+    # the optimizer, started away from kappa=0, must converge back to the floor
+    rng = np.random.default_rng(1234)
+    init = rng.normal(scale=0.4, size=nocc * (n_orbitals - nocc))
+    result = optimize_orbitals(
+        m.one_body_integrals,
+        m.two_body_integrals,
+        one_rdm,
+        two_rdm,
+        m.n_electrons,
+        nuclear_repulsion=m.nuclear_repulsion,
+        initial_guess=init,
+        verbose=False,
+    )
+    assert np.isclose(result.fun, fci_energy, atol=1e-5)
+    assert result.fun >= fci_energy - 1e-6
+
+
+def test_optimize_orbitals_improves_a_truncated_active_space():
+    """The realistic use case: a CASCI-style active-space-truncated RDM
+    (computed via canonical/reference orbitals, which are not generally
+    CASSCF-optimal) should either be improved by orbital rotation or, at
+    worst, left unchanged -- never made worse than the untruncated
+    (kappa=0) starting point."""
+    m = _load_h2(bond_length='0.75', basis_suffix='6-31g')
+    # 4 spatial orbitals total; restrict the active CI space to the lowest 2
+    # (drop the top 2 virtuals from the CI problem, but keep them in the
+    # one-/two-body integral tensors that optimize_orbitals rotates over --
+    # this is exactly the "orbital rotation between active and excluded
+    # space is not a symmetry" scenario orbital optimization targets).
+    active_indices = [0, 1]
+    active_hamiltonian = m.get_molecular_hamiltonian(active_indices=active_indices)
+    n_active_qubits = 2 * len(active_indices)
+    active_energy, active_one_rdm, active_two_rdm = _fci_ground_state_rdms(
+        active_hamiltonian, n_active_qubits
+    )
+
+    # Pad the active-space RDM back out to the full 4-orbital (8 spin-orbital)
+    # tensor shape optimize_orbitals expects, with the excluded orbitals'
+    # entries left at zero (unoccupied in this trial density).
+    n_orbitals = m.n_orbitals
+    n_spin_orbitals = 2 * n_orbitals
+    n_active_spin = n_active_qubits
+    one_rdm = np.zeros((n_spin_orbitals, n_spin_orbitals))
+    one_rdm[:n_active_spin, :n_active_spin] = active_one_rdm
+    two_rdm = np.zeros((n_spin_orbitals,) * 4)
+    two_rdm[:n_active_spin, :n_active_spin, :n_active_spin, :n_active_spin] = active_two_rdm
+
+    baseline = optimize_orbitals(
+        m.one_body_integrals,
+        m.two_body_integrals,
+        one_rdm,
+        two_rdm,
+        n_electrons=2 * len(active_indices),
+        nuclear_repulsion=m.nuclear_repulsion,
+        initial_guess=np.zeros((len(active_indices)) * (n_orbitals - len(active_indices))),
+        method='Nelder-Mead',
+        verbose=False,
+        sp_options={'maxiter': 1},
+    )
+    optimized = optimize_orbitals(
+        m.one_body_integrals,
+        m.two_body_integrals,
+        one_rdm,
+        two_rdm,
+        n_electrons=2 * len(active_indices),
+        nuclear_repulsion=m.nuclear_repulsion,
+        verbose=False,
+    )
+    assert np.isclose(baseline.fun, active_energy, atol=1e-6)
+    assert optimized.fun <= baseline.fun + 1e-8
+
+
+def test_optimize_orbitals_rejects_degenerate_orbital_split():
+    """No occupied or no virtual spatial orbitals -- nothing to rotate."""
+    obi = np.zeros((2, 2))
+    tbi = np.zeros((2, 2, 2, 2))
+    one_rdm = np.zeros((4, 4))
+    two_rdm = np.zeros((4, 4, 4, 4))
+    with pytest.raises(ValueError):
+        optimize_orbitals(obi, tbi, one_rdm, two_rdm, n_electrons=4)  # all occupied, no virtuals
+    with pytest.raises(ValueError):
+        optimize_orbitals(obi, tbi, one_rdm, two_rdm, n_electrons=0)  # all virtual, no occupied