Add final orbital rotation to hop gate ansatz (#85)

* add orbital rotation to hop gate ansatz

* use in final orbital rotation for ucj

* use for rest of ucj

* add test

* doc

* mypy

* update notebook

docs/tutorials/05-entanglement-forging.ipynb

@@ -87,7 +87,7 @@
     "interaction_pairs = list(brickwork(norb, n_layers))\n",
     "rng = np.random.default_rng(1234)\n",
     "thetas = rng.uniform(-np.pi, np.pi, size=len(interaction_pairs))\n",
-    "operator = ffsim.HopGateAnsatzOperator(interaction_pairs, thetas=thetas)\n",
+    "operator = ffsim.HopGateAnsatzOperator(norb, interaction_pairs, thetas)\n",
     "\n",
     "# Compute energy of ansatz state\n",
     "reference_occupations_spatial = [(0, 1, 2, 3), (1, 2, 3, 4), (0, 1, 2, 4)]\n",
@@ -114,10 +114,10 @@
       "  message: STOP: TOTAL NO. of f AND g EVALUATIONS EXCEEDS LIMIT\n",
       "  success: False\n",
       "   status: 1\n",
-      "      fun: -75.68085278176007\n",
+      "      fun: -75.68085225170408\n",
       "        x: [ 2.996e+00 -7.549e-01 ...  2.650e+00  8.012e-01]\n",
       "      nit: 6\n",
-      "      jac: [ 1.759e-03  9.119e-03 ... -1.192e-02  9.436e-04]\n",
+      "      jac: [ 1.756e-03  9.118e-03 ... -1.192e-02  9.521e-04]\n",
       "     nfev: 112\n",
       "     njev: 7\n",
       " hess_inv: <15x15 LbfgsInvHessProduct with dtype=float64>\n"
@@ -130,7 +130,7 @@
     "\n",
     "def fun(x):\n",
     "    # Initialize the ansatz operator from the parameter vector\n",
-    "    operator = ffsim.HopGateAnsatzOperator(interaction_pairs, x)\n",
+    "    operator = ffsim.HopGateAnsatzOperator(norb, interaction_pairs, x)\n",
     "    # Compute energy\n",
     "    energy, _ = ffsim.multireference_state(\n",
     "        mol_hamiltonian, operator, reference_occupations, norb=norb, nelec=nelec\n",

python/ffsim/variational/hopgate.py

@@ -17,15 +17,21 @@
 
 import numpy as np
 
-from ffsim.gates import apply_hop_gate
+from ffsim.gates import apply_hop_gate, apply_orbital_rotation
+from ffsim.variational.util import (
+    orbital_rotation_from_parameters,
+    orbital_rotation_to_parameters,
+)
 
 
 @dataclass(frozen=True)
 class HopGateAnsatzOperator:
     r"""A hop gate ansatz operator."""
 
+    norb: int
     interaction_pairs: list[tuple[int, int]]
     thetas: np.ndarray
+    final_orbital_rotation: np.ndarray | None = None
 
     def _apply_unitary_(
         self, vec: np.ndarray, norb: int, nelec: tuple[int, int], copy: bool
@@ -39,4 +45,45 @@ def _apply_unitary_(
             vec = apply_hop_gate(
                 vec, theta, target_orbs=target_orbs, norb=norb, nelec=nelec, copy=False
             )
+        if self.final_orbital_rotation is not None:
+            vec = apply_orbital_rotation(
+                vec,
+                mat=self.final_orbital_rotation,
+                norb=norb,
+                nelec=nelec,
+                copy=False,
+            )
         return vec
+
+    def to_parameters(self) -> np.ndarray:
+        """Convert the hop gate ansatz operator to a real-valued parameter vector."""
+        num_params = len(self.thetas)
+        if self.final_orbital_rotation is not None:
+            num_params += self.norb**2
+        params = np.zeros(num_params)
+        params[: len(self.thetas)] = self.thetas
+        if self.final_orbital_rotation is not None:
+            params[len(self.thetas) :] = orbital_rotation_to_parameters(
+                self.final_orbital_rotation
+            )
+        return params
+
+    @staticmethod
+    def from_parameters(
+        params: np.ndarray,
+        norb: int,
+        interaction_pairs: list[tuple[int, int]],
+        with_final_orbital_rotation: bool = False,
+    ) -> HopGateAnsatzOperator:
+        final_orbital_rotation = None
+        if with_final_orbital_rotation:
+            final_orbital_rotation = orbital_rotation_from_parameters(
+                params[-(norb**2) :], norb
+            )
+            params = params[: -(norb**2)]
+        return HopGateAnsatzOperator(
+            norb=norb,
+            interaction_pairs=interaction_pairs,
+            thetas=params,
+            final_orbital_rotation=final_orbital_rotation,
+        )

python/ffsim/variational/ucj.py

@@ -22,6 +22,10 @@
 
 from ffsim.gates import apply_diag_coulomb_evolution, apply_orbital_rotation
 from ffsim.linalg import double_factorized_t2
+from ffsim.variational.util import (
+    orbital_rotation_from_parameters,
+    orbital_rotation_to_parameters,
+)
 
 
 @dataclass(frozen=True)
@@ -84,7 +88,7 @@ def from_parameters(
         alpha_alpha_indices: list[tuple[int, int]] | None = None,
         alpha_beta_indices: list[tuple[int, int]] | None = None,
         with_final_orbital_rotation: bool = False,
-    ) -> "UCJOperator":
+    ) -> UCJOperator:
         """Initialize the UCJ operator from a real-valued parameter vector."""
         return UCJOperator(
             *_ucj_from_parameters(
@@ -120,7 +124,7 @@ def from_t_amplitudes(
         t1_amplitudes: np.ndarray | None = None,
         n_reps: int | None = None,
         tol: float = 1e-8,
-    ) -> "UCJOperator":
+    ) -> UCJOperator:
         """Initialize the UCJ operator from t2 (and optionally t1) amplitudes."""
         # TODO maybe allow specifying alpha-alpha and alpha-beta indices
         nocc, _, nvrt, _ = t2_amplitudes.shape
@@ -417,14 +421,13 @@ def _ucj_from_parameters(
         List[Tuple[int, int]],
         list(itertools.combinations_with_replacement(range(norb), 2)),
     )
-    triu_indices_no_diag = list(itertools.combinations(range(norb), 2))
     if alpha_alpha_indices is None:
         alpha_alpha_indices = triu_indices
     if alpha_beta_indices is None:
         alpha_beta_indices = triu_indices
     diag_coulomb_mats_alpha_alpha = np.zeros((n_reps, norb, norb))
     diag_coulomb_mats_alpha_beta = np.zeros((n_reps, norb, norb))
-    orbital_rotation_generators = np.zeros((n_reps, norb, norb), dtype=complex)
+    orbital_rotations = np.zeros((n_reps, norb, norb), dtype=complex)
     index = 0
     # diag coulomb matrices, alpha-alpha
     indices = alpha_alpha_indices
@@ -446,49 +449,16 @@ def _ucj_from_parameters(
             mat[rows, cols] = vals
             mat[cols, rows] = vals
             index += n_params
-    # orbital rotations, imaginary part
-    indices = triu_indices
-    n_params = len(indices)
-    rows, cols = zip(*indices)
-    for mat in orbital_rotation_generators:
-        vals = 1j * params[index : index + n_params]
-        mat[rows, cols] = vals
-        mat[cols, rows] = vals
-        index += n_params
-    # orbital rotations, real part
-    indices = triu_indices_no_diag
-    n_params = len(indices)
-    rows, cols = zip(*indices)
-    for mat in orbital_rotation_generators:
-        vals = params[index : index + n_params]
-        mat[rows, cols] += vals
-        mat[cols, rows] -= vals
-        index += n_params
-    # exponentiate orbital rotation generators
-    orbital_rotations = np.stack(
-        [scipy.linalg.expm(mat) for mat in orbital_rotation_generators]
-    )
+    # orbital rotations
+    for mat in orbital_rotations:
+        mat[:] = orbital_rotation_from_parameters(
+            params[index : index + norb**2], norb
+        )
+        index += norb**2
     # final orbital rotation
     final_orbital_rotation = None
     if with_final_orbital_rotation:
-        final_orbital_rotation_generator = np.zeros((norb, norb), dtype=complex)
-        # final orbital rotation, imaginary part
-        indices = triu_indices
-        n_params = len(indices)
-        rows, cols = zip(*indices)
-        vals = 1j * params[index : index + n_params]
-        final_orbital_rotation_generator[rows, cols] = vals
-        final_orbital_rotation_generator[cols, rows] = vals
-        index += n_params
-        # final orbital rotation, real part
-        indices = triu_indices_no_diag
-        n_params = len(indices)
-        rows, cols = zip(*indices)
-        vals = params[index : index + n_params]
-        final_orbital_rotation_generator[rows, cols] += vals
-        final_orbital_rotation_generator[cols, rows] -= vals
-        # exponentiate final orbital rotation generator
-        final_orbital_rotation = scipy.linalg.expm(final_orbital_rotation_generator)
+        final_orbital_rotation = orbital_rotation_from_parameters(params[index:], norb)
     return (
         diag_coulomb_mats_alpha_alpha,
         diag_coulomb_mats_alpha_beta,
@@ -511,15 +481,13 @@ def _ucj_to_parameters(
         List[Tuple[int, int]],
         list(itertools.combinations_with_replacement(range(norb), 2)),
     )
-    triu_indices_no_diag = list(itertools.combinations(range(norb), 2))
     if alpha_alpha_indices is None:
         alpha_alpha_indices = triu_indices
     if alpha_beta_indices is None:
         alpha_beta_indices = triu_indices
     ntheta = n_reps * (len(alpha_alpha_indices) + len(alpha_beta_indices) + norb**2)
     if final_orbital_rotation is not None:
         ntheta += norb**2
-    orbital_rotation_generators = [scipy.linalg.logm(mat) for mat in orbital_rotations]
     theta = np.zeros(ntheta)
     index = 0
     # diag coulomb matrices, alpha-alpha
@@ -536,28 +504,15 @@ def _ucj_to_parameters(
         for mat in diag_coulomb_mats_alpha_beta:
             theta[index : index + n_params] = mat[tuple(zip(*indices))]
             index += n_params
-    # orbital rotations, imaginary part
-    indices = triu_indices
-    n_params = len(indices)
-    for mat in orbital_rotation_generators:
-        theta[index : index + n_params] = mat[tuple(zip(*indices))].imag
-        index += n_params
-    # orbital rotations, real part
-    indices = triu_indices_no_diag
-    n_params = len(indices)
-    for mat in orbital_rotation_generators:
-        theta[index : index + n_params] = mat[tuple(zip(*indices))].real
-        index += n_params
+    # orbital rotations
+    for orbital_rotation in orbital_rotations:
+        theta[index : index + norb**2] = orbital_rotation_to_parameters(
+            orbital_rotation
+        )
+        index += norb**2
     # final orbital rotation
     if final_orbital_rotation is not None:
-        mat = scipy.linalg.logm(final_orbital_rotation)
-        # final orbital rotation, imaginary part
-        indices = triu_indices
-        n_params = len(indices)
-        theta[index : index + n_params] = mat[tuple(zip(*indices))].imag
-        index += n_params
-        # final orbital rotation, real part
-        indices = triu_indices_no_diag
-        n_params = len(indices)
-        theta[index : index + n_params] = mat[tuple(zip(*indices))].real
+        theta[index : index + norb**2] = orbital_rotation_to_parameters(
+            final_orbital_rotation
+        )
     return theta

python/ffsim/variational/util.py

@@ -0,0 +1,74 @@
+# (C) Copyright IBM 2023.
+#
+# 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.
+
+"""Variational ansatz utilities."""
+
+from __future__ import annotations
+
+import itertools
+
+import numpy as np
+import scipy.linalg
+
+
+def orbital_rotation_to_parameters(orbital_rotation: np.ndarray) -> np.ndarray:
+    """Convert an orbital rotation to parameters.
+
+    Converts an orbital rotation to a real-valued parameter vector. The parameter vector
+    contains non-redundant real and imaginary parts of the elements of the matrix
+    logarithm of the orbital rotation matrix.
+
+    Args:
+        orbital_rotation: The orbital rotation.
+
+    Returns:
+        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)
+    # 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
+    return params
+
+
+def orbital_rotation_from_parameters(params: np.ndarray, norb: int) -> np.ndarray:
+    """Construct an orbital rotation from parameters.
+
+    Converts a real-valued parameter vector to an orbital rotation. The parameter vector
+    contains non-redundant real and imaginary parts of the elements of the matrix
+    logarithm of the orbital rotation matrix.
+
+    Args:
+        params: The real-valued parameters.
+        norb: The number of spatial orbitals, which gives the width and height of the
+            orbital rotation matrix.
+
+    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
+    # real part
+    vals = params[: len(triu_indices_no_diag)]
+    rows, cols = zip(*triu_indices_no_diag)
+    generator[rows, cols] += vals
+    generator[cols, rows] -= vals
+    return scipy.linalg.expm(generator)

tests/variational/hopgate_test.py

@@ -0,0 +1,56 @@
+# (C) Copyright IBM 2023.
+#
+# 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 hop gate ansatz."""
+
+import itertools
+
+import numpy as np
+from scipy.special import comb
+
+import ffsim
+
+
+def test_parameters_roundtrip():
+    norb = 5
+    n_reps = 2
+    rng = np.random.default_rng()
+
+    def ncycles(iterable, n):
+        "Returns the sequence elements n times"
+        return itertools.chain.from_iterable(itertools.repeat(tuple(iterable), n))
+
+    pairs = itertools.combinations(range(norb), 2)
+    interaction_pairs = list(ncycles(pairs, n_reps))
+    thetas = rng.uniform(-np.pi, np.pi, size=len(interaction_pairs))
+    final_orbital_rotation = ffsim.random.random_unitary(norb)
+
+    operator = ffsim.HopGateAnsatzOperator(
+        norb=norb,
+        interaction_pairs=interaction_pairs,
+        thetas=thetas,
+        final_orbital_rotation=final_orbital_rotation,
+    )
+    assert len(operator.to_parameters()) == n_reps * comb(norb, 2) + norb**2
+    roundtripped = ffsim.HopGateAnsatzOperator.from_parameters(
+        operator.to_parameters(),
+        norb=norb,
+        interaction_pairs=interaction_pairs,
+        with_final_orbital_rotation=True,
+    )
+
+    np.testing.assert_allclose(
+        roundtripped.thetas,
+        operator.thetas,
+    )
+    np.testing.assert_allclose(
+        roundtripped.final_orbital_rotation,
+        operator.final_orbital_rotation,
+    )

tests/variational/ucj_test.py

@@ -8,6 +8,8 @@
 # copyright notice, and modified files need to carry a notice indicating
 # that they have been altered from the originals.
 
+"""Tests for unitary cluster Jastrow ansatz."""
+
 import numpy as np
 import pyscf
 import scipy.linalg