diff --git a/python/ffsim/__init__.py b/python/ffsim/__init__.py index 7401ef8b8..806e7b2a4 100644 --- a/python/ffsim/__init__.py +++ b/python/ffsim/__init__.py @@ -90,6 +90,7 @@ HopGateAnsatzOperator, NumNumAnsatzOpSpinBalanced, UCCSDOpRestrictedReal, + UCJAnglesOpSpinBalanced, UCJOpSpinBalanced, UCJOpSpinless, UCJOpSpinUnbalanced, @@ -119,6 +120,7 @@ "SupportsLinearOperator", "SupportsTrace", "UCCSDOpRestrictedReal", + "UCJAnglesOpSpinBalanced", "UCJOpSpinBalanced", "UCJOpSpinUnbalanced", "UCJOpSpinless", diff --git a/python/ffsim/variational/__init__.py b/python/ffsim/variational/__init__.py index 33ca0324a..983162473 100644 --- a/python/ffsim/variational/__init__.py +++ b/python/ffsim/variational/__init__.py @@ -18,6 +18,7 @@ ) from ffsim.variational.num_num import NumNumAnsatzOpSpinBalanced from ffsim.variational.uccsd import UCCSDOpRestrictedReal +from ffsim.variational.ucj_angles_spin_balanced import UCJAnglesOpSpinBalanced from ffsim.variational.ucj_spin_balanced import UCJOpSpinBalanced from ffsim.variational.ucj_spin_unbalanced import UCJOpSpinUnbalanced from ffsim.variational.ucj_spinless import UCJOpSpinless @@ -27,6 +28,7 @@ "HopGateAnsatzOperator", "NumNumAnsatzOpSpinBalanced", "UCCSDOpRestrictedReal", + "UCJAnglesOpSpinBalanced", "UCJOpSpinBalanced", "UCJOpSpinUnbalanced", "UCJOpSpinless", diff --git a/python/ffsim/variational/ucj_angles_spin_balanced.py b/python/ffsim/variational/ucj_angles_spin_balanced.py new file mode 100644 index 000000000..83eda7675 --- /dev/null +++ b/python/ffsim/variational/ucj_angles_spin_balanced.py @@ -0,0 +1,248 @@ +# (C) Copyright IBM 2024. +# +# This code is licensed under the Apache License, Version 2.0. You may +# obtain a copy of this license in the LICENSE.txt file in the root directory +# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0. +# +# Any modifications or derivative works of this code must retain this +# copyright notice, and modified files need to carry a notice indicating +# that they have been altered from the originals. + +"""Spin-balanced (local) UCJ ansatz parameterized by gate rotation angles.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from ffsim import gates, protocols +from ffsim.variational.givens import GivensAnsatzOp +from ffsim.variational.num_num import NumNumAnsatzOpSpinBalanced +from ffsim.variational.ucj_spin_balanced import UCJOpSpinBalanced + + +@dataclass(frozen=True) +class UCJAnglesOpSpinBalanced: + r"""A spin-balanced UCJ operator parameterized by gate rotation angles.""" + + norb: int + num_num_ansatz_ops: list[NumNumAnsatzOpSpinBalanced] + givens_ansatz_ops: list[GivensAnsatzOp] + final_givens_ansatz_op: GivensAnsatzOp | None = None + + def __post_init__(self): + if len(self.num_num_ansatz_ops) != len(self.givens_ansatz_ops): + raise ValueError( + "The number of number-number ansatz operations must equal the number " + "of Givens ansatz operations. " + f"Got {len(self.num_num_ansatz_ops)} and {len(self.givens_ansatz_ops)}." + ) + + @property + def n_reps(self): + """The number of ansatz repetitions.""" + return len(self.num_num_ansatz_ops) + + @staticmethod + def n_params( + norb: int, + n_reps: int, + num_num_interaction_pairs: tuple[list[tuple[int, int]], list[tuple[int, int]]], + givens_interaction_pairs: list[tuple[int, int]], + with_final_givens_ansatz_op: bool = False, + ) -> int: + """Return the number of parameters of an ansatz with given settings.""" + pairs_aa, pairs_ab = num_num_interaction_pairs + n_params_num_num = len(pairs_aa) + len(pairs_ab) + n_params_givens = 2 * len(givens_interaction_pairs) + norb + return ( + n_reps * (n_params_num_num + n_params_givens) + + with_final_givens_ansatz_op * norb**2 + ) + + @staticmethod + def from_parameters( + params: np.ndarray, + *, + norb: int, + n_reps: int, + num_num_interaction_pairs: tuple[list[tuple[int, int]], list[tuple[int, int]]], + givens_interaction_pairs: list[tuple[int, int]], + with_final_givens_ansatz_op: bool = False, + ) -> UCJAnglesOpSpinBalanced: + r"""Initialize the UCJ operator from a real-valued parameter vector.""" + n_params = UCJAnglesOpSpinBalanced.n_params( + norb, + n_reps, + num_num_interaction_pairs=num_num_interaction_pairs, + givens_interaction_pairs=givens_interaction_pairs, + with_final_givens_ansatz_op=with_final_givens_ansatz_op, + ) + if len(params) != n_params: + raise ValueError( + "The number of parameters passed did not match the number expected " + "based on the function inputs. " + f"Expected {n_params} but got {len(params)}." + ) + n_params_num_num = sum(len(pairs) for pairs in num_num_interaction_pairs) + n_params_givens = 2 * len(givens_interaction_pairs) + norb + num_num_ansatz_ops = [] + givens_ansatz_ops = [] + index = 0 + for _ in range(n_reps): + # Givens + n_params = n_params_givens + givens_ansatz_ops.append( + GivensAnsatzOp.from_parameters( + params[index : index + n_params], + norb=norb, + interaction_pairs=givens_interaction_pairs, + ) + ) + index += n_params + # Number-number + n_params = n_params_num_num + num_num_ansatz_ops.append( + NumNumAnsatzOpSpinBalanced.from_parameters( + params[index : index + n_params], + norb=norb, + interaction_pairs=num_num_interaction_pairs, + ) + ) + index += n_params + final_givens_ansatz_op = None + if with_final_givens_ansatz_op: + final_givens_ansatz_op = GivensAnsatzOp.from_parameters( + params[index:], + norb=norb, + interaction_pairs=list(brickwork(norb, norb)), + ) + return UCJAnglesOpSpinBalanced( + norb, + num_num_ansatz_ops=num_num_ansatz_ops, + givens_ansatz_ops=givens_ansatz_ops, + final_givens_ansatz_op=final_givens_ansatz_op, + ) + + def to_parameters(self) -> np.ndarray: + r"""Convert the UCJ operator to a real-valued parameter vector.""" + param_arrays = [] + for givens_ansatz, num_num_ansatz in zip( + self.givens_ansatz_ops, self.num_num_ansatz_ops + ): + param_arrays.append(givens_ansatz.to_parameters()) + param_arrays.append(num_num_ansatz.to_parameters()) + if self.final_givens_ansatz_op is not None: + param_arrays.append(self.final_givens_ansatz_op.to_parameters()) + return np.concatenate(param_arrays) + + @staticmethod + def from_t_amplitudes( + t2: np.ndarray, + *, + t1: np.ndarray | None = None, + n_reps: int | None = None, + interaction_pairs: tuple[ + list[tuple[int, int]] | None, list[tuple[int, int]] | None + ] + | None = None, + tol: float = 1e-8, + ) -> UCJAnglesOpSpinBalanced: + """Initialize the UCJ operator from t2 (and optionally t1) amplitudes.""" + ucj_op = UCJOpSpinBalanced.from_t_amplitudes( + t2, t1=t1, n_reps=n_reps, interaction_pairs=interaction_pairs, tol=tol + ) + return UCJAnglesOpSpinBalanced.from_ucj_op(ucj_op) + + @staticmethod + def from_ucj_op(ucj_op: UCJOpSpinBalanced) -> UCJAnglesOpSpinBalanced: + """Initialize the angles-based UCJ operator from a matrix-based UCJ operator.""" + num_num_ansatz_ops = [ + NumNumAnsatzOpSpinBalanced.from_diag_coulomb_mats(diag_coulomb_mats) + for diag_coulomb_mats in ucj_op.diag_coulomb_mats + ] + givens_ansatz_ops = [ + GivensAnsatzOp.from_orbital_rotation(orbital_rotation) + for orbital_rotation in ucj_op.orbital_rotations + ] + final_givens_ansatz_op = None + if ucj_op.final_orbital_rotation is not None: + final_givens_ansatz_op = GivensAnsatzOp.from_orbital_rotation( + ucj_op.final_orbital_rotation + ) + return UCJAnglesOpSpinBalanced( + norb=ucj_op.norb, + num_num_ansatz_ops=num_num_ansatz_ops, + givens_ansatz_ops=givens_ansatz_ops, + final_givens_ansatz_op=final_givens_ansatz_op, + ) + + def _apply_unitary_( + self, vec: np.ndarray, norb: int, nelec: int | tuple[int, int], copy: bool + ) -> np.ndarray: + if isinstance(nelec, int): + return NotImplemented + if copy: + vec = vec.copy() + current_basis = np.eye(norb) + for num_num_ansatz_op, givens_ansatz_op in zip( + self.num_num_ansatz_ops, self.givens_ansatz_ops + ): + orbital_rotation = givens_ansatz_op.to_orbital_rotation() + vec = gates.apply_orbital_rotation( + vec, + orbital_rotation.T.conj() @ current_basis, + norb=norb, + nelec=nelec, + copy=False, + ) + vec = protocols.apply_unitary( + vec, num_num_ansatz_op, norb=norb, nelec=nelec, copy=False + ) + current_basis = orbital_rotation + if self.final_givens_ansatz_op is None: + vec = gates.apply_orbital_rotation( + vec, current_basis, norb=norb, nelec=nelec, copy=False + ) + else: + final_orbital_rotation = self.final_givens_ansatz_op.to_orbital_rotation() + vec = gates.apply_orbital_rotation( + vec, + final_orbital_rotation @ current_basis, + norb=norb, + nelec=nelec, + copy=False, + ) + return vec + + def _approx_eq_(self, other, rtol: float, atol: float) -> bool: + if isinstance(other, UCJAnglesOpSpinBalanced): + if self.norb != other.norb: + return False + if len(self.num_num_ansatz_ops) != len(other.num_num_ansatz_ops): + return False + if len(self.givens_ansatz_ops) != len(other.givens_ansatz_ops): + return False + if not protocols.approx_eq( + self.final_givens_ansatz_op, other.final_givens_ansatz_op + ): + return False + if not all( + protocols.approx_eq(op1, op2, rtol=rtol, atol=atol) + for op1, op2 in zip(self.num_num_ansatz_ops, other.num_num_ansatz_ops) + ): + return False + if not all( + protocols.approx_eq(op1, op2, rtol=rtol, atol=atol) + for op1, op2 in zip(self.givens_ansatz_ops, other.givens_ansatz_ops) + ): + return False + return True + return NotImplemented + + +def brickwork(norb: int, n_layers: int): + for i in range(n_layers): + for j in range(i % 2, norb - 1, 2): + yield (j, j + 1) diff --git a/tests/python/variational/ucj_angles_spin_balanced_test.py b/tests/python/variational/ucj_angles_spin_balanced_test.py new file mode 100644 index 000000000..7ed7fc247 --- /dev/null +++ b/tests/python/variational/ucj_angles_spin_balanced_test.py @@ -0,0 +1,238 @@ +# (C) Copyright IBM 2024. +# +# This code is licensed under the Apache License, Version 2.0. You may +# obtain a copy of this license in the LICENSE.txt file in the root directory +# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0. +# +# Any modifications or derivative works of this code must retain this +# copyright notice, and modified files need to carry a notice indicating +# that they have been altered from the originals. + +"""Tests for spin-balanced UCJ ansatz parameterized by gate rotation angles.""" + +import itertools + +import numpy as np +import pyscf +import pyscf.cc + +import ffsim + + +def _brickwork(norb: int, n_layers: int): + for i in range(n_layers): + for j in range(i % 2, norb - 1, 2): + yield (j, j + 1) + + +def test_n_params(): + for norb, n_reps, with_final_orbital_rotation in itertools.product( + [1, 2, 3], [1, 2, 3], [False, True] + ): + triu_indices = list(itertools.combinations_with_replacement(range(norb), 2)) + brickwork_indices = list(_brickwork(norb, norb)) + + ucj_op = ffsim.random.random_ucj_op_spin_balanced( + norb, + n_reps=n_reps, + with_final_orbital_rotation=with_final_orbital_rotation, + interaction_pairs=(triu_indices, triu_indices), + ) + operator = ffsim.UCJAnglesOpSpinBalanced.from_ucj_op(ucj_op) + actual = ffsim.UCJAnglesOpSpinBalanced.n_params( + norb, + n_reps, + num_num_interaction_pairs=(triu_indices, triu_indices), + givens_interaction_pairs=brickwork_indices, + with_final_givens_ansatz_op=with_final_orbital_rotation, + ) + expected = len(operator.to_parameters()) + assert actual == expected + + pairs_aa = triu_indices[:norb] + pairs_ab = triu_indices[:norb] + ucj_op = ffsim.random.random_ucj_op_spin_balanced( + norb, + n_reps=n_reps, + with_final_orbital_rotation=with_final_orbital_rotation, + interaction_pairs=(pairs_aa, pairs_ab), + ) + operator = ffsim.UCJAnglesOpSpinBalanced.from_ucj_op(ucj_op) + actual = ffsim.UCJAnglesOpSpinBalanced.n_params( + norb, + n_reps, + num_num_interaction_pairs=(pairs_aa, pairs_ab), + givens_interaction_pairs=brickwork_indices, + with_final_givens_ansatz_op=with_final_orbital_rotation, + ) + expected = len(operator.to_parameters()) + assert actual == expected + + +def test_parameters_roundtrip(): + rng = np.random.default_rng() + norb = 5 + n_reps = 2 + triu_indices = list(itertools.combinations_with_replacement(range(norb), 2)) + pairs_aa = triu_indices[:norb] + pairs_ab = triu_indices[:norb] + brickwork_indices = list(_brickwork(norb, norb)) + + for with_final_orbital_rotation in [False, True]: + ucj_op = ffsim.random.random_ucj_op_spin_balanced( + norb, + n_reps=n_reps, + with_final_orbital_rotation=with_final_orbital_rotation, + interaction_pairs=(pairs_aa, pairs_ab), + seed=rng, + ) + operator = ffsim.UCJAnglesOpSpinBalanced.from_ucj_op(ucj_op) + roundtripped = ffsim.UCJAnglesOpSpinBalanced.from_parameters( + operator.to_parameters(), + norb=norb, + n_reps=n_reps, + num_num_interaction_pairs=(pairs_aa, pairs_ab), + givens_interaction_pairs=brickwork_indices, + with_final_givens_ansatz_op=with_final_orbital_rotation, + ) + assert ffsim.approx_eq(operator, roundtripped) + + +def test_apply_unitary_consistent_with_ucj_op(): + rng = np.random.default_rng(44299) + + norb = 8 + nelec = (5, 5) + n_reps = 3 + + ucj_op = ffsim.random.random_ucj_op_spin_balanced( + norb, n_reps=n_reps, with_final_orbital_rotation=True + ) + ucj_angles_op = ffsim.UCJAnglesOpSpinBalanced.from_ucj_op(ucj_op) + vec = ffsim.random.random_state_vector(ffsim.dim(norb, nelec), seed=rng) + + actual = ffsim.apply_unitary(vec, ucj_angles_op, norb=norb, nelec=nelec) + expected = ffsim.apply_unitary(vec, ucj_op, norb=norb, nelec=nelec) + + np.testing.assert_allclose(actual, expected) + + +def test_from_t_amplitudes_consistent_with_ucj_op(): + rng = np.random.default_rng(6072) + mol = pyscf.gto.Mole() + mol.build( + atom=[["N", (0, 0, 0)], ["N", (0, 0, 1.0)]], + basis="sto-6g", + symmetry="Dooh", + ) + n_frozen = 2 + active_space = range(n_frozen, mol.nao_nr()) + scf = pyscf.scf.RHF(mol).run() + ccsd = pyscf.cc.CCSD( + scf, frozen=[i for i in range(mol.nao_nr()) if i not in active_space] + ).run() + + # Get molecular data and molecular Hamiltonian + mol_data = ffsim.MolecularData.from_scf(scf, active_space=active_space) + norb = mol_data.norb + nelec = mol_data.nelec + assert norb == 8 + assert nelec == (5, 5) + + # Construct UCJ operator + n_reps = 2 + ucj_op = ffsim.UCJOpSpinBalanced.from_t_amplitudes( + ccsd.t2, + t1=ccsd.t1, + n_reps=n_reps, + ) + ucj_angles_op = ffsim.UCJAnglesOpSpinBalanced.from_t_amplitudes( + ccsd.t2, + t1=ccsd.t1, + n_reps=n_reps, + ) + vec = ffsim.random.random_state_vector(ffsim.dim(norb, nelec), seed=rng) + + actual = ffsim.apply_unitary(vec, ucj_angles_op, norb=norb, nelec=nelec) + expected = ffsim.apply_unitary(vec, ucj_op, norb=norb, nelec=nelec) + + np.testing.assert_allclose(actual, expected) + + +def test_t_amplitudes_energy(): + mol = pyscf.gto.Mole() + mol.build( + atom=[["N", (0, 0, 0)], ["N", (0, 0, 1.0)]], + basis="sto-6g", + symmetry="Dooh", + ) + n_frozen = 2 + active_space = range(n_frozen, mol.nao_nr()) + scf = pyscf.scf.RHF(mol).run() + ccsd = pyscf.cc.CCSD( + scf, frozen=[i for i in range(mol.nao_nr()) if i not in active_space] + ).run() + + # Get molecular data and molecular Hamiltonian + mol_data = ffsim.MolecularData.from_scf(scf, active_space=active_space) + norb = mol_data.norb + nelec = mol_data.nelec + assert norb == 8 + assert nelec == (5, 5) + mol_hamiltonian = mol_data.hamiltonian + + # Construct UCJ operator + n_reps = 2 + operator = ffsim.UCJAnglesOpSpinBalanced.from_t_amplitudes( + ccsd.t2, t1=ccsd.t1, n_reps=n_reps + ) + + # Construct the Hartree-Fock state to use as the reference state + n_alpha, n_beta = nelec + reference_state = ffsim.slater_determinant( + norb=norb, occupied_orbitals=(range(n_alpha), range(n_beta)) + ) + + # Apply the operator to the reference state + ansatz_state = ffsim.apply_unitary( + reference_state, operator, norb=norb, nelec=nelec + ) + + # Compute the energy ⟨ψ|H|ψ⟩ of the ansatz state + hamiltonian = ffsim.linear_operator(mol_hamiltonian, norb=norb, nelec=nelec) + energy = np.real(np.vdot(ansatz_state, hamiltonian @ ansatz_state)) + np.testing.assert_allclose(energy, -108.563917) + + +def test_t_amplitudes_restrict_indices(): + # Build an H2 molecule + mol = pyscf.gto.Mole() + mol.build( + atom=[["H", (0, 0, 0)], ["H", (0, 0, 1.8)]], + basis="sto-6g", + symmetry="Dooh", + ) + scf = pyscf.scf.RHF(mol).run() + ccsd = pyscf.cc.CCSD(scf).run() + + # Get molecular data and molecular Hamiltonian (one- and two-body tensors) + mol_data = ffsim.MolecularData.from_scf(scf) + norb = mol_data.norb + + # Construct UCJ operator + n_reps = 2 + pairs_aa = [(p, p + 1) for p in range(norb - 1)] + pairs_ab = [(p, p) for p in range(norb)] + + operator = ffsim.UCJAnglesOpSpinBalanced.from_t_amplitudes( + ccsd.t2, n_reps=n_reps, interaction_pairs=(pairs_aa, pairs_ab) + ) + other_operator = ffsim.UCJAnglesOpSpinBalanced.from_parameters( + operator.to_parameters(), + norb=norb, + n_reps=n_reps, + num_num_interaction_pairs=(pairs_aa, pairs_ab), + givens_interaction_pairs=list(_brickwork(norb, norb)), + ) + + assert ffsim.approx_eq(operator, other_operator, rtol=1e-12)