Optimize orbitals¶

In this guide, we show how to use ffsim.optimize_orbitals to lower the energy of a correlated state by applying an orbital rotation. The function takes one- and two-body reduced density matrices (1- and 2-RDMs) and a MolecularHamiltonian, and returns an optimized orbital rotation matrix.

The RDMs can come from any source. Here, we’ll use a nitrogen molecule and obtain them from a configuration interaction singles and doubles (CISD) calculation with PySCF.

Build molecule and active space Hamiltonian¶

We’ll use a nitrogen molecule (\(N_2\)) at a stretched bond length (\(R = 2.4\text{ \AA}\)) in the 6-31G basis set. At stretched bond lengths, static correlation is pronounced, making orbital relaxation effects particularly evident. We freeze the core \(1s\) orbitals and construct an active space on the remaining orbitals.

[1]:

import pyscf

import ffsim

# Build N2 molecule at stretched geometry
mol = pyscf.gto.Mole()
mol.build(
    atom=[["N", (0, 0, -1.2)], ["N", (0, 0, 1.2)]],
    basis="6-31g",
    symmetry="Dooh",
)

# Define active space by freezing 1s core orbitals
n_frozen = 2
active_space = range(n_frozen, mol.nao_nr())

# Run RHF and obtain molecular Hamiltonian
scf = pyscf.scf.RHF(mol).run()
mol_data = ffsim.MolecularData.from_scf(scf, active_space=active_space)
mol_hamiltonian = mol_data.hamiltonian
norb = mol_data.norb

print(f"Number of active orbitals: {norb}")
print(f"Number of active electrons: {mol_data.nelec}")

converged SCF energy = -108.139527186846
Number of active orbitals: 16
Number of active electrons: (5, 5)

Compute correlated state and reduced density matrices¶

Next, we run a CISD calculation within the active space and extract the 1- and 2-RDMs. We wrap them in ReducedDensityMatrix, which provides methods to rotate RDMs and compute energy expectation values with molecular Hamiltonians.

We check that the initial energy evaluated using the canonical Hartree-Fock orbitals matches the CISD total energy.

[2]:

import numpy as np
import pyscf.ci

# Run CISD in the active space
cisd = pyscf.ci.CISD(scf, frozen=n_frozen).run()

# Extract 1- and 2-RDMs in the active space
rdm1 = cisd.make_rdm1()[n_frozen:, n_frozen:]
rdm2 = cisd.make_rdm2()[n_frozen:, n_frozen:, n_frozen:, n_frozen:]
rdm = ffsim.ReducedDensityMatrix(rdm1, rdm2)

# Compute energy using canonical Hartree-Fock orbitals
initial_energy = rdm.expectation(mol_hamiltonian)
np.testing.assert_allclose(initial_energy, cisd.e_tot)
print(f"Initial CISD energy: {initial_energy:.8f} Ha")

E(RCISD) = -108.5853221021409  E_corr = -0.4457949152950733
Initial CISD energy: -108.58532210 Ha

Optimize orbitals¶

We now call ffsim.optimize_orbitals to find the orbital rotation matrix \(U\) that minimizes the energy of the RDMs with respect to mol_hamiltonian. The default settings are sufficient for this example.

The RDMs and Hamiltonian stay fixed during optimization; only the orbital rotation varies.

We can compute the optimized energy by rotating the RDMs by \(U\) or, equivalently, rotating the Hamiltonian by \(U^\dagger\). The following code cell demonstrates both methods.

[3]:

orbital_rotation = ffsim.optimize_orbitals(rdm, mol_hamiltonian)

# Compute the energy in the optimized orbital basis
energy_rotated_rdm = rdm.rotated(orbital_rotation).expectation(mol_hamiltonian)
energy_rotated_ham = rdm.expectation(mol_hamiltonian.rotated(orbital_rotation.T.conj()))
np.testing.assert_allclose(energy_rotated_ham, energy_rotated_rdm)

print(f"Optimized energy: {energy_rotated_rdm:.8f} Ha")
print(f"Energy lowering: {energy_rotated_rdm - initial_energy:.8f} Ha")

Optimized energy: -108.58613394 Ha
Energy lowering: -0.00081183 Ha

Configure SciPy optimization¶

Additional keyword arguments to ffsim.optimize_orbitals are passed to scipy.optimize.minimize. Use method to choose the optimizer, options to configure it, and callback to monitor progress. The default method is L-BFGS-B.

Setting return_optimize_result=True returns both the orbital rotation and a scipy.optimize.OptimizeResult. Its fields include the final energy (fun), termination status (success and message), iteration count (nit), and function evaluation count (nfev).

Here we repeat the optimization with an explicit method, a limit of 10 iterations, and a callback that prints the energy after each iteration. The callback argument is named intermediate_result so SciPy passes an OptimizeResult containing the current parameters and energy.

[4]:

def callback(intermediate_result):
    print(f"Energy = {intermediate_result.fun:.8f} Ha")


orbital_rotation_configured, result = ffsim.optimize_orbitals(
    rdm,
    mol_hamiltonian,
    return_optimize_result=True,
    method="L-BFGS-B",
    options=dict(maxiter=10),
    callback=callback,
)

energy_configured = rdm.rotated(orbital_rotation_configured).expectation(
    mol_hamiltonian
)
np.testing.assert_allclose(energy_configured, result.fun)

print(f"Optimized energy:     {result.fun:.8f} Ha")
print(f"Termination:          {result.message}")
print(f"Optimizer success:    {result.success}")
print(f"Iterations:           {result.nit}")
print(f"Function evaluations: {result.nfev}")

Energy = -108.58599926 Ha
Energy = -108.58610334 Ha
Energy = -108.58612508 Ha
Energy = -108.58612954 Ha
Energy = -108.58613168 Ha
Energy = -108.58613372 Ha
Energy = -108.58613394 Ha
Optimized energy:     -108.58613394 Ha
Termination:          CONVERGENCE: RELATIVE REDUCTION OF F <= FACTR*EPSMCH
Optimizer success:    True
Iterations:           7
Function evaluations: 9

For both real and complex optimization, the returned rotation is \(U = \exp(K)U_0\), where \(U_0\) is the initial rotation and result.x parameterizes the additional rotation \(\exp(K)\). The initial rotation stays fixed throughout optimization, and the generator parameters start at zero. This keeps all rotation directions available at initialization, including when \(U_0\) has several eigenvalues equal to \(-1\). Bounds, constraints, and callbacks use these incremental generator parameters. For example, in two real orbitals, an initial rotation of \(0.1\) radians and an additional rotation of \(0.2\) radians give a total rotation of \(0.3\) radians.

Real vs. complex orbital rotations¶

For \(N\) spatial orbitals, the real keyword argument controls whether the optimization uses real orthogonal rotations or general complex unitary rotations:

When real=True, the rotation is parameterized by \(N(N - 1) / 2\) skew-symmetric generator parameters and stays in the determinant component of the initial rotation. The initial rotation defaults to identity (determinant \(+1\)); an initial rotation with determinant \(-1\) is also honored.
When real=False, the rotation is parameterized by \(N^2\) anti-Hermitian generator parameters.
When real=None (the default), ffsim.optimize_orbitals automatically detects whether the RDMs, Hamiltonian, and initial orbital rotation are all real-valued, and defaults to real optimization if so.

With one orbital, real optimization raises a ValueError because there are no continuous parameters. Complex optimization (real=False) is supported; its single phase parameter does not change the energy, but SciPy can still enforce constraints on it.

If we initialize with a complex orbital rotation (or pass real=False), the optimizer optimizes over the full unitary group:

[5]:

import scipy.linalg

# Generate a small complex perturbation to the initial orbital rotation
rng = np.random.default_rng(12345)
perturbation = 1e-4 * ffsim.random.random_antihermitian(norb, seed=rng)
initial_orbital_rotation = scipy.linalg.expm(perturbation)

# Optimize orbitals allowing complex rotations
orbital_rotation_complex, result_complex = ffsim.optimize_orbitals(
    rdm,
    mol_hamiltonian,
    initial_orbital_rotation=initial_orbital_rotation,
    return_optimize_result=True,
)

energy_complex = rdm.rotated(orbital_rotation_complex).expectation(mol_hamiltonian)

print(f"Optimized energy (complex): {energy_complex:.8f} Ha")
print(
    f"Number of parameters:       {len(result_complex.x)} "
    f"(vs {norb * (norb - 1) // 2} real)"
)
print(f"Is complex rotation matrix: {np.iscomplexobj(orbital_rotation_complex)}")

Optimized energy (complex): -108.58613389 Ha
Number of parameters:       256 (vs 120 real)
Is complex rotation matrix: True

Choose the determinant component¶

With real=True, the returned rotation has the same determinant sign as initial_orbital_rotation. The default initial rotation is identity, with determinant \(+1\).

For a correlated state, flipping one orbital’s sign can change relative configuration signs and thus the energy. This differs from a basis change in which configuration coefficients are transformed consistently to preserve the physical state.

To optimize in the determinant \(-1\) component, supply an initial reflection:

[6]:

initial_reflection = np.eye(norb)
initial_reflection[0, 0] = -1

orbital_rotation_reflected, result_reflected = ffsim.optimize_orbitals(
    rdm,
    mol_hamiltonian,
    initial_orbital_rotation=initial_reflection,
    real=True,
    return_optimize_result=True,
    options=dict(maxiter=10),
)

energy_reflected = rdm.rotated(orbital_rotation_reflected).expectation(mol_hamiltonian)
np.testing.assert_allclose(energy_reflected, result_reflected.fun, atol=1e-8, rtol=0)
np.testing.assert_allclose(np.linalg.det(orbital_rotation_reflected), -1)

print(f"Optimized energy (negative determinant): {energy_reflected:.8f} Ha")
print(f"Determinant: {np.linalg.det(orbital_rotation_reflected):+.0f}")

Optimized energy (negative determinant): -108.53256623 Ha
Determinant: -1

Update molecular orbital coefficients¶

To express the optimized orbitals in the atomic orbital basis, apply the orbital rotation matrix \(U\) to the active-space molecular orbital coefficient matrix \(C\):

\[C_{\text{opt}} = C_{\text{active}} U\]

where \(C_{\text{active}}\) contains the active-space columns of the original Hartree-Fock MO coefficients.

[7]:

# Transform molecular orbital coefficients for downstream calculations
mo_coeff_active = scf.mo_coeff[:, active_space]
mo_coeff_optimized = mo_coeff_active @ orbital_rotation

print(f"Original active MO coefficients shape:  {mo_coeff_active.shape}")
print(f"Optimized active MO coefficients shape: {mo_coeff_optimized.shape}")

Original active MO coefficients shape:  (18, 16)
Optimized active MO coefficients shape: (18, 16)

Updating the coefficients with \(C_{\text{active}}U\) transforms both the one-body and two-body integrals, giving mol_hamiltonian.rotated(U.conj().T).

For the real rotation above, we rebuild the integrals with PySCF and check that they match the rotated Hamiltonian. Using the original RDMs with the rebuilt Hamiltonian should reproduce the optimized energy.

[8]:

import pyscf.ao2mo
import pyscf.mcscf

# Rebuild the integrals from the updated MO coefficients
mo_coeff_full = np.array(scf.mo_coeff, copy=True)
mo_coeff_full[:, active_space] = mo_coeff_optimized
cas = pyscf.mcscf.CASCI(scf, norb, mol_data.nelec)
one_body, core_energy = cas.get_h1cas(mo_coeff_full)
two_body = pyscf.ao2mo.restore(1, cas.get_h2cas(mo_coeff_full), norb)
rebuilt_hamiltonian = ffsim.MolecularHamiltonian(one_body, two_body, core_energy)

# Compare with the Hamiltonian rotated by the adjoint of U
expected_hamiltonian = mol_hamiltonian.rotated(orbital_rotation.conj().T)
np.testing.assert_allclose(one_body, expected_hamiltonian.one_body_tensor, atol=1e-8)
np.testing.assert_allclose(two_body, expected_hamiltonian.two_body_tensor, atol=1e-8)
np.testing.assert_allclose(core_energy, expected_hamiltonian.constant, atol=1e-8)

# Keep the original RDMs fixed when evaluating the rebuilt Hamiltonian
energy_from_mo_coeff = rdm.expectation(rebuilt_hamiltonian)
np.testing.assert_allclose(energy_from_mo_coeff, energy_rotated_rdm)

print(f"Energy from updated MO coefficients: {energy_from_mo_coeff:.8f} Ha")
print(f"Optimized energy:                   {energy_rotated_rdm:.8f} Ha")

Energy from updated MO coefficients: -108.58613394 Ha
Optimized energy:                   -108.58613394 Ha