Variational ansatzes

This page explains how ffsim represents variational ansatzes.

Design

In ffsim, a variational ansatz is a parameterized unitary operator. ffsim includes several classes for representing variational ansatzes. Each ansatz class stores the parameters of the ansatz in NumPy arrays, and implements the SupportsApplyUnitary protocol so that it can be passed to ffsim.apply_unitary to apply the ansatz operator to a state vector.

To facilitate variational optimization, each ansatz class implements:

• n_params: a static method that returns the number of real-valued parameters for a given set of ansatz settings.

• from_parameters: a static method that constructs the ansatz from a flat NumPy array of real-valued parameters.

• to_parameters: an instance method that converts the ansatz back to a flat NumPy array of real-valued parameters.

These methods make it straightforward to use the ansatzes with standard numerical optimizers. For a guide on variational optimization, see Simulate the variational quantum eigensolver (VQE).

Ansatz classes

The rest of this page gives an overview of the variational ansatz classes in ffsim.

import numpy as np

import ffsim

rng = np.random.default_rng(12345)

# Set example values for closed- and open-shell systems
norb = 4
nelec_closed = (2, 2)
nelec_open = (3, 1)

nocc, _ = nelec_closed
nvrt = norb - nocc

nocc_a, nocc_b = nelec_open
nvrt_a = norb - nocc_a
nvrt_b = norb - nocc_b

# Initialize the reference state as the Hartree-Fock state
reference_state_closed = ffsim.hartree_fock_state(norb, nelec_closed)
reference_state_open = ffsim.hartree_fock_state(norb, nelec_open)
reference_state_spinless = ffsim.hartree_fock_state(norb, nocc)

UCJ ansatzes

The unitary cluster Jastrow (UCJ) ansatz has the form

\[\lvert \Psi \rangle = \prod_{k = 1}^{L} \mathcal{U}_k e^{i \mathcal{J}_k} \mathcal{U}_k^\dagger \lvert \Phi_0 \rangle\]

where \(\lvert \Phi_0 \rangle\) is a reference state, each \(\mathcal{U}_k\) is an orbital rotation, and each \(\mathcal{J}_k\) is a diagonal Coulomb operator of the form

\[\begin{split}\mathcal{J} = \frac{1}{2}\sum_{\substack{ij \\ \sigma \tau}} \mathbf{J}^{\sigma \tau}_{ij} n_{i\sigma} n_{j\tau}.\end{split}\]

The number of terms \(L\) is referred to as the number of ansatz repetitions. An optional final orbital rotation can be appended at the end to support variational optimization of the orbital basis.

The local UCJ (LUCJ) ansatz is implemented using the same classes as UCJ, by specifying additional arguments. For more detail on the UCJ ansatz, including its connection to CCSD and LUCJ, see The local unitary cluster Jastrow (LUCJ) ansatz.

ffsim provides three UCJ classes that differ in their spin symmetry assumptions. The diagonal Coulomb matrices and orbital rotations are stored as NumPy arrays. The parameter vector encodes the independent entries of the diagonal Coulomb matrices (upper-triangular entries, since the matrices are symmetric) and the entries of the logarithms of the orbital rotations.

UCJOpSpinBalanced

The UCJOpSpinBalanced class represents the spin-balanced UCJ operator, appropriate for closed-shell systems. It imposes the constraints \(\mathbf{J}^{\alpha\alpha} = \mathbf{J}^{\beta\beta}\) and \(\mathbf{J}^{\alpha\beta} = \mathbf{J}^{\beta\alpha}\), so each diagonal Coulomb operator is described by two symmetric matrices: \(\mathbf{J}^{\alpha\alpha}\) and \(\mathbf{J}^{\alpha\beta}\). The same orbital rotation is applied to both spin sectors.

The data stored by this class consists of:

• diag_coulomb_mats: An array of shape (n_reps, 2, norb, norb) containing the \(\mathbf{J}^{\alpha\alpha}\) and \(\mathbf{J}^{\alpha\beta}\) matrices for each repetition.

• orbital_rotations: An array of shape (n_reps, norb, norb) containing the orbital rotation for each repetition.

• final_orbital_rotation: An optional array of shape (norb, norb) containing a final orbital rotation.

n_reps = 2

# Construct from random t2- and t1-amplitudes. The t2-amplitudes are double factorized
# to obtain the diagonal Coulomb matries and orbital rotations. The t1-amplitudes are
# used to set a final orbital rotation.
t2 = ffsim.random.random_t2_amplitudes(norb, nocc, seed=rng, dtype=float)
t1 = rng.standard_normal((nocc, nvrt))

ucj_op = ffsim.UCJOpSpinBalanced.from_t_amplitudes(t2, t1=t1, n_reps=n_reps)

print("diag_coulomb_mats shape:", ucj_op.diag_coulomb_mats.shape)
print("orbital_rotations shape:", ucj_op.orbital_rotations.shape)

# Apply the ansatz to the reference state
final_state = ffsim.apply_unitary(
    reference_state_closed, ucj_op, norb=norb, nelec=nelec_closed
)

diag_coulomb_mats shape: (2, 2, 4, 4)
orbital_rotations shape: (2, 4, 4)

You can also initialize the ansatz from a real-valued parameter vector:

# Construct from a random parameter vector
n_params = ffsim.UCJOpSpinBalanced.n_params(norb, n_reps)
params = rng.standard_normal(n_params)
ucj_op = ffsim.UCJOpSpinBalanced.from_parameters(params, norb=norb, n_reps=n_reps)

The interaction_pairs argument can be used to restrict the orbital interactions in the diagonal Coulomb operators (setting the remaining entries to zero), which yields the local UCJ (LUCJ) variant:

# Restrict to interactions implementable on a square lattice
pairs_aa = [(p, p + 1) for p in range(norb - 1)]
pairs_ab = [(p, p) for p in range(norb)]
interaction_pairs = (pairs_aa, pairs_ab)

lucj_op = ffsim.UCJOpSpinBalanced.from_t_amplitudes(
    t2, t1=t1, n_reps=n_reps, interaction_pairs=interaction_pairs
)

final_state = ffsim.apply_unitary(
    reference_state_closed, lucj_op, norb=norb, nelec=nelec_closed
)

UCJOpSpinUnbalanced

The UCJOpSpinUnbalanced class represents the spin-unbalanced UCJ operator, appropriate for open-shell systems. It allows \(\mathbf{J}^{\alpha\alpha}\) and \(\mathbf{J}^{\beta\beta}\) to differ and does not require \(\mathbf{J}^{\alpha\beta}\) to be symmetric (since \(\mathbf{J}^{\alpha\beta}_{ij} = \mathbf{J}^{\beta\alpha}_{ji}\), the \(\mathbf{J}^{\beta\alpha}\) matrix need not be stored separately). Each diagonal Coulomb operator is therefore described by three matrices: \(\mathbf{J}^{\alpha\alpha}\), \(\mathbf{J}^{\alpha\beta}\), and \(\mathbf{J}^{\beta\beta}\). Separate orbital rotations are applied to the alpha and beta spin sectors.

The data stored by this class consists of:

• diag_coulomb_mats: An array of shape (n_reps, 3, norb, norb) containing the \(\mathbf{J}^{\alpha\alpha}\), \(\mathbf{J}^{\alpha\beta}\), and \(\mathbf{J}^{\beta\beta}\) matrices for each repetition.

• orbital_rotations: An array of shape (n_reps, 2, norb, norb) containing the alpha and beta orbital rotations for each repetition.

• final_orbital_rotation: An optional array of shape (2, norb, norb) containing a final orbital rotation for each spin sector.

t2aa = ffsim.random.random_t2_amplitudes(norb, nocc_a, seed=rng, dtype=float)
t2ab = rng.standard_normal((nocc_a, nocc_b, nvrt_a, nvrt_b))
t2bb = ffsim.random.random_t2_amplitudes(norb, nocc_b, seed=rng, dtype=float)
t1a = rng.standard_normal((nocc_a, nvrt_a))
t1b = rng.standard_normal((nocc_b, nvrt_b))

ucj_op_unbalanced = ffsim.UCJOpSpinUnbalanced.from_t_amplitudes(
    (t2aa, t2ab, t2bb), t1=(t1a, t1b), n_reps=(2, 2)
)

print("diag_coulomb_mats shape:", ucj_op_unbalanced.diag_coulomb_mats.shape)
print("orbital_rotations shape:", ucj_op_unbalanced.orbital_rotations.shape)

final_state = ffsim.apply_unitary(
    reference_state_open, ucj_op_unbalanced, norb=norb, nelec=nelec_open
)

diag_coulomb_mats shape: (4, 3, 4, 4)
orbital_rotations shape: (4, 2, 4, 4)

UCJOpSpinless

The UCJOpSpinless class represents the spinless UCJ operator, for systems without a spin degree of freedom. The diagonal Coulomb operator takes the form

\[\mathcal{J} = \frac{1}{2}\sum_{ij} \mathbf{J}_{ij} n_i n_j\]

where \(\mathbf{J}\) is a single real symmetric matrix.

The data stored by this class consists of:

• diag_coulomb_mats: An array of shape (n_reps, norb, norb) containing the diagonal Coulomb matrix for each repetition.

• orbital_rotations: An array of shape (n_reps, norb, norb) containing the orbital rotation for each repetition.

• final_orbital_rotation: An optional array of shape (norb, norb) containing a final orbital rotation.

ucj_op_spinless = ffsim.UCJOpSpinless.from_t_amplitudes(t2, t1=t1, n_reps=n_reps)

print("diag_coulomb_mats shape:", ucj_op_spinless.diag_coulomb_mats.shape)
print("orbital_rotations shape:", ucj_op_spinless.orbital_rotations.shape)

final_state_spinless = ffsim.apply_unitary(
    reference_state_spinless, ucj_op_spinless, norb=norb, nelec=nocc
)

diag_coulomb_mats shape: (2, 4, 4)
orbital_rotations shape: (2, 4, 4)

UCCSD ansatzes

The unitary coupled cluster singles and doubles (UCCSD) ansatz has the form

\[\lvert \Psi \rangle = e^{T - T^\dagger} \lvert \Phi_0 \rangle\]

where

\[T = T_1 + T_2, \quad T_1 = \sum_{ia} t_{ia} a^\dagger_a a_i, \quad T_2 = \sum_{ijab} t_{ijab} a^\dagger_a a^\dagger_b a_j a_i.\]

Here \(i, j\) index occupied orbitals, \(a, b\) index virtual orbitals, and the \(t\)-amplitudes \(t_{ia}\) (singles) and \(t_{ijab}\) (doubles) are the variational parameters.

UCCSDOpRestrictedReal

The UCCSDOpRestrictedReal class represents the restricted UCCSD operator with real-valued \(t\)-amplitudes. “Restricted” means that the same spatial orbitals are used for both spin sectors, and the same \(t\)-amplitudes apply to both. This is the most common case when starting from a restricted Hartree-Fock (RHF) reference.

The data stored by this class consists of:

• t1: The real-valued \(t_1\)-amplitudes, an array of shape (nocc, nvrt).

• t2: The real-valued \(t_2\)-amplitudes, an array of shape (nocc, nocc, nvrt, nvrt).

• final_orbital_rotation: An optional array of shape (norb, norb) (may be complex-valued).

t1 = rng.standard_normal((nocc, nvrt))
t2 = ffsim.random.random_t2_amplitudes(norb, nocc, seed=rng, dtype=float)

uccsd_op = ffsim.UCCSDOpRestrictedReal(t1, t2)

print("t1 shape:", uccsd_op.t1.shape)
print("t2 shape:", uccsd_op.t2.shape)

final_state = ffsim.apply_unitary(
    reference_state_closed, uccsd_op, norb=norb, nelec=nelec_closed
)

t1 shape: (2, 2)
t2 shape: (2, 2, 2, 2)

UCCSDOpRestricted

The UCCSDOpRestricted class is just like UCCSDOpRestrictedReal except its \(t\)-amplitudes are allowed to take complex values, which doubles the number of real parameters.

t1 = rng.standard_normal((nocc, nvrt)) + 1j * rng.standard_normal((nocc, nvrt))
t2 = ffsim.random.random_t2_amplitudes(norb, nocc, seed=rng, dtype=complex)

uccsd_op = ffsim.UCCSDOpRestricted(t1, t2)

print("t1 shape:", uccsd_op.t1.shape)
print("t2 shape:", uccsd_op.t2.shape)

final_state = ffsim.apply_unitary(
    reference_state_closed, uccsd_op, norb=norb, nelec=nelec_closed
)

t1 shape: (2, 2)
t2 shape: (2, 2, 2, 2)

UCCSDOpUnrestrictedReal

The UCCSDOpUnrestrictedReal class represents the unrestricted UCCSD operator with real-valued \(t\)-amplitudes. “Unrestricted” means that separate \(t\)-amplitudes are used for different spin combinations, suitable for open-shell systems described by an unrestricted Hartree-Fock (UHF) reference.

The data stored by this class consists of:

• t1: A tuple (t1_a, t1_b) of real-valued arrays with shapes (nocc_a, nvrt_a) and (nocc_b, nvrt_b).

• t2: A tuple (t2_aa, t2_ab, t2_bb) of real-valued arrays with shapes (nocc_a, nocc_a, nvrt_a, nvrt_a), (nocc_a, nocc_b, nvrt_a, nvrt_b), and (nocc_b, nocc_b, nvrt_b, nvrt_b).

• final_orbital_rotation: An optional array of shape (2, norb, norb) containing a final orbital rotation for each spin sector.

t1a = rng.standard_normal((nocc_a, nvrt_a))
t1b = rng.standard_normal((nocc_b, nvrt_b))
t2aa = ffsim.random.random_t2_amplitudes(norb, nocc_a, seed=rng, dtype=float)
t2ab = rng.standard_normal((nocc_a, nocc_b, nvrt_a, nvrt_b))
t2bb = ffsim.random.random_t2_amplitudes(norb, nocc_b, seed=rng, dtype=float)

uccsd_op = ffsim.UCCSDOpUnrestrictedReal((t1a, t1b), (t2aa, t2ab, t2bb))

print("t1_a shape:", uccsd_op.t1[0].shape)
print("t1_b shape:", uccsd_op.t1[1].shape)
print("t2_aa shape:", uccsd_op.t2[0].shape)
print("t2_ab shape:", uccsd_op.t2[1].shape)
print("t2_bb shape:", uccsd_op.t2[2].shape)

reference_state_open = ffsim.hartree_fock_state(norb, nelec_open)
final_state = ffsim.apply_unitary(
    reference_state_open, uccsd_op, norb=norb, nelec=nelec_open
)

t1_a shape: (3, 1)
t1_b shape: (1, 3)
t2_aa shape: (3, 3, 1, 1)
t2_ab shape: (3, 1, 1, 3)
t2_bb shape: (1, 1, 3, 3)

UCCSDOpUnrestricted

The UCCSDOpUnrestricted class is just like UCCSDOpUnrestrictedReal except its \(t\)-amplitudes are allowed to take complex values, which doubles the number of real parameters.

t1a = rng.standard_normal((nocc_a, nvrt_a)) + 1j * rng.standard_normal((nocc_a, nvrt_a))
t1b = rng.standard_normal((nocc_b, nvrt_b)) + 1j * rng.standard_normal((nocc_b, nvrt_b))
t2aa = ffsim.random.random_t2_amplitudes(norb, nocc_a, seed=rng, dtype=complex)
t2ab = rng.standard_normal((nocc_a, nocc_b, nvrt_a, nvrt_b)) + 1j * rng.standard_normal(
    (nocc_a, nocc_b, nvrt_a, nvrt_b)
)
t2bb = ffsim.random.random_t2_amplitudes(norb, nocc_b, seed=rng, dtype=complex)

uccsd_op = ffsim.UCCSDOpUnrestricted((t1a, t1b), (t2aa, t2ab, t2bb))

print("t1_a shape:", uccsd_op.t1[0].shape)
print("t1_b shape:", uccsd_op.t1[1].shape)
print("t2_aa shape:", uccsd_op.t2[0].shape)
print("t2_ab shape:", uccsd_op.t2[1].shape)
print("t2_bb shape:", uccsd_op.t2[2].shape)

final_state = ffsim.apply_unitary(
    reference_state_open, uccsd_op, norb=norb, nelec=nelec_open
)

t1_a shape: (3, 1)
t1_b shape: (1, 3)
t2_aa shape: (3, 3, 1, 1)
t2_ab shape: (3, 1, 1, 3)
t2_bb shape: (1, 1, 3, 3)