From df54023889a03c674368324c94b464cecd584f41 Mon Sep 17 00:00:00 2001 From: "Kevin J. Sung" Date: Sun, 13 Oct 2024 09:58:16 -0400 Subject: [PATCH] Add real-valued UCCSD ansatz (#327) * add uccsd * remove TODO and fix doc * fix tests by using complex dtype for orbital rotations * fix random uccsd complex t1 generation --- python/ffsim/__init__.py | 2 + python/ffsim/random/__init__.py | 2 + python/ffsim/random/random.py | 38 ++++ python/ffsim/variational/__init__.py | 2 + python/ffsim/variational/uccsd.py | 252 +++++++++++++++++++++++++ python/ffsim/variational/util.py | 29 +-- tests/python/variational/uccsd_test.py | 119 ++++++++++++ tests/python/variational/ucj_test.py | 12 +- 8 files changed, 441 insertions(+), 15 deletions(-) create mode 100644 python/ffsim/variational/uccsd.py create mode 100644 tests/python/variational/uccsd_test.py diff --git a/python/ffsim/__init__.py b/python/ffsim/__init__.py index fd9690c53..33ba14a31 100644 --- a/python/ffsim/__init__.py +++ b/python/ffsim/__init__.py @@ -95,6 +95,7 @@ HopGateAnsatzOperator, NumNumAnsatzOpSpinBalanced, RealUCJOperator, + UCCSDOpRestrictedReal, UCJOperator, UCJOpSpinBalanced, UCJOpSpinless, @@ -126,6 +127,7 @@ "SupportsFermionOperator", "SupportsLinearOperator", "SupportsTrace", + "UCCSDOpRestrictedReal", "UCJOperator", "UCJOpSpinBalanced", "UCJOpSpinUnbalanced", diff --git a/python/ffsim/random/__init__.py b/python/ffsim/random/__init__.py index 85685de45..9ea9ed26c 100644 --- a/python/ffsim/random/__init__.py +++ b/python/ffsim/random/__init__.py @@ -25,6 +25,7 @@ random_statevector, random_t2_amplitudes, random_two_body_tensor, + random_uccsd_restricted, random_ucj_op_spin_balanced, random_ucj_op_spin_unbalanced, random_ucj_op_spinless, @@ -47,6 +48,7 @@ "random_state_vector", "random_t2_amplitudes", "random_two_body_tensor", + "random_uccsd_restricted", "random_ucj_operator", "random_ucj_op_spin_balanced", "random_ucj_op_spin_unbalanced", diff --git a/python/ffsim/random/random.py b/python/ffsim/random/random.py index a6215fe03..a5d4c05bb 100644 --- a/python/ffsim/random/random.py +++ b/python/ffsim/random/random.py @@ -344,6 +344,44 @@ def random_ucj_operator( ) +def random_uccsd_restricted( + norb: int, + nocc: int, + *, + with_final_orbital_rotation: bool = False, + real: bool = False, + seed=None, +) -> variational.UCCSDOpRestrictedReal: + """Sample a random UCCSD operator. + + Args: + norb: The number of spatial orbitals. + nocc: The number of spatial orbitals that are occupied by electrons. + with_final_orbital_rotation: Whether to include a final orbital rotation + in the operator. + real: Whether to sample a real-valued object rather than a complex-valued one. + seed: A seed to initialize the pseudorandom number generator. + Should be a valid input to ``np.random.default_rng``. + + Returns: + The sampled UCCSD operator. + """ + rng = np.random.default_rng(seed) + dtype = float if real else complex + nvrt = norb - nocc + t1: np.ndarray = rng.standard_normal((nocc, nvrt)).astype(dtype, copy=False) + if not real: + t1 += 1j * rng.standard_normal((nocc, nvrt)) + t2 = random_t2_amplitudes(norb, nocc, seed=rng, dtype=dtype) + final_orbital_rotation = None + if with_final_orbital_rotation: + unitary_func = random_orthogonal if real else random_unitary + final_orbital_rotation = unitary_func(norb, seed=rng) + return variational.UCCSDOpRestrictedReal( + t1=t1, t2=t2, final_orbital_rotation=final_orbital_rotation + ) + + def random_ucj_op_spin_balanced( norb: int, *, diff --git a/python/ffsim/variational/__init__.py b/python/ffsim/variational/__init__.py index da611d525..b330adcb9 100644 --- a/python/ffsim/variational/__init__.py +++ b/python/ffsim/variational/__init__.py @@ -17,6 +17,7 @@ multireference_state_prod, ) from ffsim.variational.num_num import NumNumAnsatzOpSpinBalanced +from ffsim.variational.uccsd import UCCSDOpRestrictedReal from ffsim.variational.ucj import RealUCJOperator, UCJOperator from ffsim.variational.ucj_spin_balanced import UCJOpSpinBalanced from ffsim.variational.ucj_spin_unbalanced import UCJOpSpinUnbalanced @@ -28,6 +29,7 @@ "HopGateAnsatzOperator", "NumNumAnsatzOpSpinBalanced", "RealUCJOperator", + "UCCSDOpRestrictedReal", "UCJOperator", "UCJOpSpinBalanced", "UCJOpSpinUnbalanced", diff --git a/python/ffsim/variational/uccsd.py b/python/ffsim/variational/uccsd.py new file mode 100644 index 000000000..40824a37c --- /dev/null +++ b/python/ffsim/variational/uccsd.py @@ -0,0 +1,252 @@ +# (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. + +"""Unitary coupled cluster, singles and doubles ansatz.""" + +from __future__ import annotations + +import itertools +from dataclasses import InitVar, dataclass +from typing import cast + +import numpy as np +import scipy.linalg +import scipy.sparse.linalg + +from ffsim import gates, hamiltonians, linalg, protocols +from ffsim.variational.util import ( + orbital_rotation_from_parameters, + orbital_rotation_to_parameters, +) + + +@dataclass(frozen=True) +class UCCSDOpRestrictedReal: + """Real-valued restricted unitary coupled cluster, singles and doubles operator.""" + + t1: np.ndarray # shape: (nocc, nvrt) + t2: np.ndarray # shape: (nocc, nocc, nvrt, nvrt) + final_orbital_rotation: np.ndarray | None = None # shape: (norb, norb) + validate: InitVar[bool] = True + rtol: InitVar[float] = 1e-5 + atol: InitVar[float] = 1e-8 + + def __post_init__(self, validate: bool, rtol: float, atol: float): + if np.iscomplexobj(self.t1): + raise TypeError( + "UCCSDOpRestricted only accepts real-valued t1 amplitudes. " + "Please pass a t1 amplitudes tensor with a real-valued data type." + ) + if np.iscomplexobj(self.t2): + raise TypeError( + "UCCSDOpRestricted only accepts real-valued t2 amplitudes. " + "Please pass a t2 amplitudes tensor with a real-valued data type." + ) + if self.final_orbital_rotation is not None and np.iscomplexobj( + self.final_orbital_rotation + ): + raise TypeError( + "UCCSDOpRestricted only accepts a real-valued final " + "orbital rotation. Please pass a final orbital rotation with a " + "real-valued data type." + ) + if validate: + # Validate shapes + nocc, nvrt = self.t1.shape + norb = nocc + nvrt + if self.t2.shape != (nocc, nocc, nvrt, nvrt): + raise ValueError( + "t2 shape not consistent with t1 shape. " + f"Expected {(nocc, nocc, nvrt, nvrt)} but got {self.t2.shape}" + ) + if self.final_orbital_rotation is not None: + if self.final_orbital_rotation.shape != (norb, norb): + raise ValueError( + "Final orbital rotation shape not consistent with t1 shape. " + f"Expected {(norb, norb)} but got " + f"{self.final_orbital_rotation.shape}" + ) + if not linalg.is_unitary( + self.final_orbital_rotation, rtol=rtol, atol=atol + ): + raise ValueError("Final orbital rotation was not unitary.") + + @property + def norb(self): + """The number of spatial orbitals.""" + nocc, nvrt = self.t1.shape + return nocc + nvrt + + @staticmethod + def n_params( + norb: int, nocc: int, *, with_final_orbital_rotation: bool = False + ) -> int: + """Return the number of parameters of an ansatz with given settings. + + Args: + norb: The number of spatial orbitals. + nocc: The number of spatial orbitals that are occupied by electrons. + with_final_orbital_rotation: Whether to include a final orbital rotation + in the operator. + + Returns: + The number of parameters of the ansatz. + """ + nvrt = norb - nocc + # Number of occupied-virtual pairs + n_pairs = nocc * nvrt + # t1 has n_pairs parameters + # t2 has n_pairs * (n_pairs + 1) // 2 parameters + # Final orbital rotation has norb * (norb - 1) // 2 parameters + return ( + n_pairs + + n_pairs * (n_pairs + 1) // 2 + + with_final_orbital_rotation * norb * (norb - 1) // 2 + ) + + @staticmethod + def from_parameters( + params: np.ndarray, + *, + norb: int, + nocc: int, + with_final_orbital_rotation: bool = False, + ) -> UCCSDOpRestrictedReal: + r"""Initialize the UCCSD operator from a real-valued parameter vector. + + Args: + params: The real-valued parameter vector. + norb: The number of spatial orbitals. + nocc: The number of spatial orbitals that are occupied by electrons. + with_final_orbital_rotation: Whether to include a final orbital rotation + in the operator. + + Returns: + The UCCSD operator constructed from the given parameters. + + Raises: + ValueError: The number of parameters passed did not match the number + expected based on the function inputs. + """ + n_params = UCCSDOpRestrictedReal.n_params( + norb, nocc, with_final_orbital_rotation=with_final_orbital_rotation + ) + 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)}." + ) + nvrt = norb - nocc + t1 = np.zeros((nocc, nvrt)) + t2 = np.zeros((nocc, nocc, nvrt, nvrt)) + occ_vrt_pairs = list(itertools.product(range(nocc), range(nocc, norb))) + index = 0 + # t1 + for i, a in occ_vrt_pairs: + t1[i, a - nocc] = params[index] + index += 1 + # t2 + for (i, a), (j, b) in itertools.combinations_with_replacement(occ_vrt_pairs, 2): + t2[i, j, a - nocc, b - nocc] = params[index] + t2[j, i, b - nocc, a - nocc] = params[index] + index += 1 + # Final orbital rotation + final_orbital_rotation = None + if with_final_orbital_rotation: + final_orbital_rotation = orbital_rotation_from_parameters( + params[index:], norb, real=True + ) + return UCCSDOpRestrictedReal( + t1=t1, t2=t2, final_orbital_rotation=final_orbital_rotation + ) + + def to_parameters(self) -> np.ndarray: + r"""Convert the UCCSD operator to a real-valued parameter vector. + + Returns: + The real-valued parameter vector. + """ + nocc, nvrt = self.t1.shape + norb = nocc + nvrt + n_params = UCCSDOpRestrictedReal.n_params( + norb, + nocc, + with_final_orbital_rotation=self.final_orbital_rotation is not None, + ) + params = np.zeros(n_params) + occ_vrt_pairs = list(itertools.product(range(nocc), range(nocc, norb))) + index = 0 + # t1 + for i, a in occ_vrt_pairs: + params[index] = self.t1[i, a - nocc] + index += 1 + # t2 + for (i, a), (j, b) in itertools.combinations_with_replacement(occ_vrt_pairs, 2): + params[index] = self.t2[i, j, a - nocc, b - nocc] + index += 1 + # Final orbital rotation + if self.final_orbital_rotation is not None: + params[index:] = orbital_rotation_to_parameters(self.final_orbital_rotation) + return params + + 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() + + nocc, _ = self.t1.shape + assert nelec == (nocc, nocc) + + one_body_tensor = np.zeros((norb, norb)) + two_body_tensor = np.zeros((norb, norb, norb, norb)) + one_body_tensor[:nocc, nocc:] = self.t1 + one_body_tensor[nocc:, :nocc] = -self.t1.T + two_body_tensor[nocc:, :nocc, nocc:, :nocc] = self.t2.transpose(2, 0, 3, 1) + two_body_tensor[:nocc, nocc:, :nocc, nocc:] = -self.t2.transpose(0, 2, 1, 3) + + linop = protocols.linear_operator( + hamiltonians.MolecularHamiltonian( + one_body_tensor=one_body_tensor, two_body_tensor=two_body_tensor + ), + norb=norb, + nelec=nelec, + ) + vec = scipy.sparse.linalg.expm_multiply(linop, vec, traceA=0.0) + + if self.final_orbital_rotation is not None: + vec = gates.apply_orbital_rotation( + vec, self.final_orbital_rotation, norb=norb, nelec=nelec, copy=False + ) + + return vec + + def _approx_eq_(self, other, rtol: float, atol: float) -> bool: + if isinstance(other, UCCSDOpRestrictedReal): + if not np.allclose(self.t1, other.t1, rtol=rtol, atol=atol): + return False + if not np.allclose(self.t2, other.t2, rtol=rtol, atol=atol): + return False + if (self.final_orbital_rotation is None) != ( + other.final_orbital_rotation is None + ): + return False + if self.final_orbital_rotation is not None: + return np.allclose( + cast(np.ndarray, self.final_orbital_rotation), + cast(np.ndarray, other.final_orbital_rotation), + rtol=rtol, + atol=atol, + ) + return True + return NotImplemented diff --git a/python/ffsim/variational/util.py b/python/ffsim/variational/util.py index 3c017a995..5ccf04ed8 100644 --- a/python/ffsim/variational/util.py +++ b/python/ffsim/variational/util.py @@ -32,18 +32,23 @@ def orbital_rotation_to_parameters(orbital_rotation: np.ndarray) -> np.ndarray: The list of real numbers parameterizing the orbital rotation. """ norb, _ = orbital_rotation.shape - triu_indices = list(itertools.combinations_with_replacement(range(norb), 2)) triu_indices_no_diag = list(itertools.combinations(range(norb), 2)) mat = scipy.linalg.logm(orbital_rotation) - params = np.zeros(norb**2) + params = np.zeros( + norb**2 if np.iscomplexobj(orbital_rotation) else norb * (norb - 1) // 2 + ) # real part params[: len(triu_indices_no_diag)] = mat[tuple(zip(*triu_indices_no_diag))].real # imaginary part - params[len(triu_indices_no_diag) :] = mat[tuple(zip(*triu_indices))].imag + if np.iscomplexobj(orbital_rotation): + triu_indices = list(itertools.combinations_with_replacement(range(norb), 2)) + params[len(triu_indices_no_diag) :] = mat[tuple(zip(*triu_indices))].imag return params -def orbital_rotation_from_parameters(params: np.ndarray, norb: int) -> np.ndarray: +def orbital_rotation_from_parameters( + params: np.ndarray, norb: int, real: bool = False +) -> np.ndarray: """Construct an orbital rotation from parameters. Converts a real-valued parameter vector to an orbital rotation. The parameter vector @@ -54,18 +59,20 @@ def orbital_rotation_from_parameters(params: np.ndarray, norb: int) -> np.ndarra params: The real-valued parameters. norb: The number of spatial orbitals, which gives the width and height of the orbital rotation matrix. + real: Whether to construct a real-valued orbital rotation Returns: The orbital rotation. """ - triu_indices = list(itertools.combinations_with_replacement(range(norb), 2)) triu_indices_no_diag = list(itertools.combinations(range(norb), 2)) - generator = np.zeros((norb, norb), dtype=complex) - # imaginary part - vals = 1j * params[len(triu_indices_no_diag) :] - rows, cols = zip(*triu_indices) - generator[rows, cols] = vals - generator[cols, rows] = vals + generator = np.zeros((norb, norb), dtype=float if real else complex) + if not real: + # imaginary part + triu_indices = list(itertools.combinations_with_replacement(range(norb), 2)) + vals = 1j * params[len(triu_indices_no_diag) :] + rows, cols = zip(*triu_indices) + generator[rows, cols] = vals + generator[cols, rows] = vals # real part vals = params[: len(triu_indices_no_diag)] rows, cols = zip(*triu_indices_no_diag) diff --git a/tests/python/variational/uccsd_test.py b/tests/python/variational/uccsd_test.py new file mode 100644 index 000000000..1c6a58ba1 --- /dev/null +++ b/tests/python/variational/uccsd_test.py @@ -0,0 +1,119 @@ +# (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 UCCSD ansatz.""" + +import dataclasses + +import numpy as np + +import ffsim + + +def test_norb(): + """Test norb property.""" + rng = np.random.default_rng(4878) + norb = 5 + nocc = 3 + operator = ffsim.random.random_uccsd_restricted(norb, nocc, real=True, seed=rng) + assert operator.norb == norb + + +def test_n_params(): + """Test computing number of parameters.""" + rng = np.random.default_rng(4878) + norb = 5 + nocc = 3 + for with_final_orbital_rotation in [False, True]: + operator = ffsim.random.random_uccsd_restricted( + norb, + nocc, + with_final_orbital_rotation=with_final_orbital_rotation, + real=True, + seed=rng, + ) + actual = ffsim.UCCSDOpRestrictedReal.n_params( + norb, nocc, with_final_orbital_rotation=with_final_orbital_rotation + ) + expected = len(operator.to_parameters()) + assert actual == expected + + +def test_parameters_roundtrip(): + """Test parameters roundtrip.""" + rng = np.random.default_rng(4623) + norb = 5 + nocc = 3 + for with_final_orbital_rotation in [False, True]: + operator = ffsim.random.random_uccsd_restricted( + norb, + nocc, + with_final_orbital_rotation=with_final_orbital_rotation, + real=True, + seed=rng, + ) + roundtripped = ffsim.UCCSDOpRestrictedReal.from_parameters( + operator.to_parameters(), + norb=norb, + nocc=nocc, + with_final_orbital_rotation=with_final_orbital_rotation, + ) + np.testing.assert_allclose(roundtripped.t1, operator.t1) + np.testing.assert_allclose(roundtripped.t2, operator.t2) + if with_final_orbital_rotation: + np.testing.assert_allclose( + roundtripped.final_orbital_rotation, operator.final_orbital_rotation + ) + + +def test_approx_eq(): + """Test approximate equality.""" + rng = np.random.default_rng(4623) + norb = 5 + nocc = 3 + for with_final_orbital_rotation in [False, True]: + operator = ffsim.random.random_uccsd_restricted( + norb, + nocc, + with_final_orbital_rotation=with_final_orbital_rotation, + real=True, + seed=rng, + ) + roundtripped = ffsim.UCCSDOpRestrictedReal.from_parameters( + operator.to_parameters(), + norb=norb, + nocc=nocc, + with_final_orbital_rotation=with_final_orbital_rotation, + ) + assert ffsim.approx_eq(operator, roundtripped) + assert not ffsim.approx_eq( + operator, dataclasses.replace(operator, t1=2 * operator.t1) + ) + assert not ffsim.approx_eq( + operator, dataclasses.replace(operator, t2=2 * operator.t2) + ) + + +def test_apply_unitary(): + """Test unitary.""" + rng = np.random.default_rng(4623) + norb = 5 + nocc = 3 + vec = ffsim.random.random_state_vector(ffsim.dim(norb, (nocc, nocc)), seed=rng) + for with_final_orbital_rotation in [False, True]: + operator = ffsim.random.random_uccsd_restricted( + norb, + nocc, + with_final_orbital_rotation=with_final_orbital_rotation, + real=True, + seed=rng, + ) + result = ffsim.apply_unitary(vec, operator, norb=norb, nelec=(nocc, nocc)) + np.testing.assert_allclose(np.linalg.norm(result), 1.0) diff --git a/tests/python/variational/ucj_test.py b/tests/python/variational/ucj_test.py index 2c3de272b..e6ca26bbe 100644 --- a/tests/python/variational/ucj_test.py +++ b/tests/python/variational/ucj_test.py @@ -35,9 +35,11 @@ def test_n_params(): ): diag_coulomb_mats_alpha_alpha = np.zeros((n_reps, norb, norb)) diag_coulomb_mats_alpha_beta = np.zeros((n_reps, norb, norb)) - orbital_rotations = np.stack([np.eye(norb) for _ in range(n_reps)]) + orbital_rotations = np.stack( + [np.eye(norb, dtype=complex) for _ in range(n_reps)] + ) - final_orbital_rotation = np.eye(norb) + final_orbital_rotation = np.eye(norb, dtype=complex) operator = ffsim.UCJOperator( diag_coulomb_mats_alpha_alpha=diag_coulomb_mats_alpha_alpha, diag_coulomb_mats_alpha_beta=diag_coulomb_mats_alpha_beta, @@ -254,9 +256,11 @@ def test_real_ucj_n_params(): ): diag_coulomb_mats_alpha_alpha = np.zeros((n_reps, norb, norb)) diag_coulomb_mats_alpha_beta = np.zeros((n_reps, norb, norb)) - orbital_rotations = np.stack([np.eye(norb) for _ in range(n_reps)]) + orbital_rotations = np.stack( + [np.eye(norb, dtype=complex) for _ in range(n_reps)] + ) - final_orbital_rotation = np.eye(norb) + final_orbital_rotation = np.eye(norb, dtype=complex) operator = ffsim.RealUCJOperator( diag_coulomb_mats_alpha_alpha=diag_coulomb_mats_alpha_alpha, diag_coulomb_mats_alpha_beta=diag_coulomb_mats_alpha_beta,