add hop gate ansatz and apply_unitary protocol (#49)

* add hop gate ansatz and apply_unitary protocol * remove type annotation syntax that only works on python 3.10+ * Revert "remove type annotation syntax that only works on python 3.10+" This reverts commit a1cd941. * from __future__ import annotations * format * lint * fix docs
qiskit-community · Oct 20, 2023 · af3eb3a · af3eb3a
1 parent eba3a32
commit af3eb3a
Show file tree

Hide file tree

Showing 16 changed files with 402 additions and 72 deletions.
diff --git a/docs/tutorials/04-lucj.ipynb b/docs/tutorials/04-lucj.ipynb
@@ -80,7 +80,7 @@
     "\n",
     "In ffsim, the UCJ ansatz operator $\\prod_{k = 1}^L \\mathcal{W_k} e^{i \\mathcal{J}_k} \\mathcal{W_k^\\dagger}$ is represented by the `UCJOperator` class, which is just a dataclass that stores the diagonal Coulomb matrices and orbital rotations. A constructor method is provided to initialize the operator from a truncated double factorization of t2 amplitudes (e.g. from CCSD or MP2).\n",
     "\n",
-    "In the code cell below, we run CCSD to get the t2 amplitudes for initializing the ansatz. We'll create an ansatz operator with 2 repetitions ($L = 2$). For our reference state, we'll use the Hartree-Fock state. We use the function `apply_ucj_operator` to apply the ansatz operator to the reference state to obtain the ansatz state. Finally, we compute the energy of the ansatz state."
+    "In the code cell below, we run CCSD to get the t2 amplitudes for initializing the ansatz. We'll create an ansatz operator with 2 repetitions ($L = 2$). For our reference state, we'll use the Hartree-Fock state. Since `UCJOperator` defines a unitary effect, we can use the function `apply_unitary` to apply the ansatz operator to the reference state to obtain the ansatz state. Finally, we compute the energy of the ansatz state."
    ]
   },
   {
@@ -110,9 +110,7 @@
     ")\n",
     "\n",
     "# Apply the operator to the reference state\n",
-    "ansatz_state = ffsim.apply_ucj_operator(\n",
-    "    reference_state, operator, norb=norb, nelec=nelec\n",
-    ")\n",
+    "ansatz_state = ffsim.apply_unitary(reference_state, operator, norb=norb, nelec=nelec)\n",
     "\n",
     "# Compute the energy ⟨ψ|H|ψ⟩ of the ansatz state\n",
     "hamiltonian = ffsim.linear_operator(mol_hamiltonian, norb=norb, nelec=nelec)\n",
@@ -142,7 +140,7 @@
     "    # Initialize the ansatz operator from the parameter vector\n",
     "    operator = ffsim.UCJOperator.from_parameters(x, norb=norb, n_reps=n_reps)\n",
     "    # Apply the ansatz operator to the reference state\n",
-    "    final_state = ffsim.apply_ucj_operator(\n",
+    "    final_state = ffsim.apply_unitary(\n",
     "        reference_state, operator, norb=norb, nelec=(n_alpha, n_beta)\n",
     "    )\n",
     "    # Return the energy ⟨ψ|H|ψ⟩ of the ansatz state\n",
@@ -197,7 +195,7 @@
     "        alpha_alpha_indices=alpha_alpha_indices,\n",
     "        alpha_beta_indices=alpha_beta_indices,\n",
     "    )\n",
-    "    final_state = ffsim.apply_ucj_operator(\n",
+    "    final_state = ffsim.apply_unitary(\n",
     "        reference_state, operator, norb=norb, nelec=(n_alpha, n_beta)\n",
     "    )\n",
     "    return np.real(np.vdot(final_state, hamiltonian @ final_state))\n",

diff --git a/docs/tutorials/05-entanglement-forging.ipynb b/docs/tutorials/05-entanglement-forging.ipynb
@@ -0,0 +1,147 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Entanglement forging\n",
+    "\n",
+    "In this tutorial, we show how to simulate entanglement forging for a water molecule at equilibrium."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pyscf\n",
+    "import pyscf.mcscf\n",
+    "import ffsim\n",
+    "import math\n",
+    "import numpy as np\n",
+    "\n",
+    "# Build a water molecule\n",
+    "radius_1 = 0.958  # position for the first H atom\n",
+    "radius_2 = 0.958  # position for the second H atom\n",
+    "bond_angle_deg = 104.478  # bond angles.\n",
+    "\n",
+    "h1_x = radius_1\n",
+    "h2_x = radius_2 * math.cos(math.pi / 180 * bond_angle_deg)\n",
+    "h2_y = radius_2 * math.sin(math.pi / 180 * bond_angle_deg)\n",
+    "\n",
+    "mol = pyscf.gto.Mole()\n",
+    "mol.build(\n",
+    "    atom=[\n",
+    "        [\"O\", (0, 0, 0)],\n",
+    "        [\"H\", (h1_x, 0, 0)],\n",
+    "        [\"H\", (h2_x, h2_y, 0)],\n",
+    "    ],\n",
+    "    basis=\"sto-6g\",\n",
+    "    symmetry=\"c2v\",\n",
+    ")\n",
+    "hartree_fock = pyscf.scf.RHF(mol)\n",
+    "hartree_fock.kernel()\n",
+    "\n",
+    "# Define active space\n",
+    "active_space = [1, 2, 4, 5, 6]\n",
+    "norb = len(active_space)\n",
+    "n_electrons = int(sum(hartree_fock.mo_occ[active_space]))\n",
+    "n_alpha = (n_electrons + hartree_fock.mol.spin) // 2\n",
+    "n_beta = (n_electrons - hartree_fock.mol.spin) // 2\n",
+    "nelec = (n_alpha, n_beta)\n",
+    "\n",
+    "# Compute FCI energy\n",
+    "cas = pyscf.mcscf.CASCI(hartree_fock, ncas=len(active_space), nelecas=nelec)\n",
+    "mo = cas.sort_mo(active_space, base=0)\n",
+    "energy_fci = cas.kernel(mo)[0]\n",
+    "\n",
+    "# Get molecular data and molecular Hamiltonian (one- and two-body tensors)\n",
+    "mol_data = ffsim.MolecularData.from_hartree_fock(\n",
+    "    hartree_fock, active_space=active_space\n",
+    ")\n",
+    "mol_hamiltonian = mol_data.hamiltonian"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "n_reps = 1\n",
+    "\n",
+    "# Construct ansatz operator\n",
+    "interaction_pairs = [(0, 1), (3, 4), (1, 4), (0, 2), (3, 4)]\n",
+    "thetas = np.zeros(n_reps * len(interaction_pairs))\n",
+    "operator = ffsim.HopGateAnsatzOperator(interaction_pairs, thetas=thetas)\n",
+    "\n",
+    "# Construct ansatz state\n",
+    "reference_occupations_spatial = [(0, 1, 2), (1, 2, 3), (1, 2, 4)]\n",
+    "reference_occupations = list(\n",
+    "    zip(reference_occupations_spatial, reference_occupations_spatial)\n",
+    ")\n",
+    "hamiltonian = ffsim.linear_operator(mol_hamiltonian, norb=norb, nelec=nelec)\n",
+    "ansatz_state = ffsim.multireference_state(\n",
+    "    hamiltonian, operator, reference_occupations, norb=norb, nelec=nelec\n",
+    ")\n",
+    "\n",
+    "# Compute the energy ⟨ψ|H|ψ⟩ of the ansatz state\n",
+    "energy = np.real(np.vdot(ansatz_state, hamiltonian @ ansatz_state))\n",
+    "print(f\"Energy at initialialization: {energy}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import scipy.optimize\n",
+    "\n",
+    "\n",
+    "def fun(x):\n",
+    "    # Initialize the ansatz operator from the parameter vector\n",
+    "    operator = ffsim.HopGateAnsatzOperator(interaction_pairs, x)\n",
+    "    # Compute ansatz state\n",
+    "    ansatz_state = ffsim.multireference_state(\n",
+    "        hamiltonian, operator, reference_occupations, norb=norb, nelec=nelec\n",
+    "    )\n",
+    "    # Compute the energy ⟨ψ|H|ψ⟩ of the ansatz state\n",
+    "    return np.real(np.vdot(ansatz_state, hamiltonian @ ansatz_state))\n",
+    "\n",
+    "\n",
+    "result = scipy.optimize.minimize(\n",
+    "    fun, x0=operator.thetas, method=\"COBYLA\", options=dict(maxiter=100)\n",
+    ")\n",
+    "\n",
+    "print(f\"Number of parameters: {len(result.x)}\")\n",
+    "print(result)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "ffsim-a58AE6yt",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.13"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/docs/tutorials/05-fermion-operator.ipynb → docs/tutorials/06-fermion-operator.ipynb b/docs/tutorials/05-fermion-operator.ipynb → docs/tutorials/06-fermion-operator.ipynb
diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst
@@ -11,4 +11,5 @@ ffsim tutorials.
    02-orbital-rotation
    03-double-factorized
    04-lucj
-   05-fermion-operator
+   05-entanglement-forging
+   06-fermion-operator
diff --git a/python/ffsim/__init__.py b/python/ffsim/__init__.py
@@ -39,9 +39,11 @@
 from ffsim.hamiltonians import DoubleFactorizedHamiltonian, MolecularHamiltonian
 from ffsim.molecular_data import MolecularData
 from ffsim.protocols import (
+    SupportsApplyUnitary,
     SupportsApproximateEquality,
     SupportsLinearOperator,
     SupportsTrace,
+    apply_unitary,
     approx_eq,
     linear_operator,
     trace,
@@ -57,4 +59,4 @@
     simulate_qdrift_double_factorized,
     simulate_trotter_double_factorized,
 )
-from ffsim.variational import UCJOperator, apply_ucj_operator
+from ffsim.variational import HopGateAnsatzOperator, UCJOperator, multireference_state
diff --git a/python/ffsim/linalg/__init__.py b/python/ffsim/linalg/__init__.py
@@ -16,10 +16,7 @@
     modified_cholesky,
 )
 from ffsim.linalg.givens import apply_matrix_to_slices, givens_decomposition
-from ffsim.linalg.linalg import (
-    expm_multiply_taylor,
-    lup,
-)
+from ffsim.linalg.linalg import expm_multiply_taylor, lup, reduced_matrix
 from ffsim.linalg.predicates import (
     is_antihermitian,
     is_hermitian,

diff --git a/python/ffsim/linalg/linalg.py b/python/ffsim/linalg/linalg.py
@@ -12,6 +12,8 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
+
 import numpy as np
 import scipy.sparse.linalg
 
@@ -47,3 +49,18 @@ def lup(mat: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     ell *= d
     u /= d[:, None]
     return u.T, ell.T, p.T
+
+
+def reduced_matrix(mat: scipy.sparse.linalg.LinearOperator, vecs: Sequence[np.ndarray]):
+    r"""Compute reduced matrix within a subspace spanned by some vectors.
+
+    Given a linear operator :math:`A` and a list of vectors :math:`\{v_i}\}`,
+    return the matrix M where :math:`M_{ij} = v_i^\dagger A v_j`.
+    """
+    dim = len(vecs)
+    result = np.zeros((dim, dim), dtype=complex)
+    for j, state_j in enumerate(vecs):
+        mat_state_j = mat @ state_j
+        for i, state_i in enumerate(vecs):
+            result[i, j] = np.vdot(state_i, mat_state_j)
+    return result
diff --git a/python/ffsim/molecular_data.py b/python/ffsim/molecular_data.py
@@ -103,7 +103,7 @@ def from_hartree_fock(
 
         active_space = list(active_space)
         norb = len(active_space)
-        n_electrons = int(np.sum(hartree_fock.mo_occ[active_space]))
+        n_electrons = int(sum(hartree_fock.mo_occ[active_space]))
         n_alpha = (n_electrons + hartree_fock.mol.spin) // 2
         n_beta = (n_electrons - hartree_fock.mol.spin) // 2
 

diff --git a/python/ffsim/protocols/__init__.py b/python/ffsim/protocols/__init__.py
@@ -10,6 +10,7 @@
 
 """Protocols."""
 
+from ffsim.protocols.apply_unitary import SupportsApplyUnitary, apply_unitary
 from ffsim.protocols.approximate_equality import SupportsApproximateEquality, approx_eq
 from ffsim.protocols.linear_operator import SupportsLinearOperator, linear_operator
 from ffsim.protocols.trace import SupportsTrace, trace
diff --git a/python/ffsim/protocols/apply_unitary.py b/python/ffsim/protocols/apply_unitary.py
@@ -0,0 +1,70 @@
+# (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.
+
+"""Apply unitary protocol."""
+
+from __future__ import annotations
+
+from typing import Any, Protocol
+
+import numpy as np
+
+
+class SupportsApplyUnitary(Protocol):
+    """An object that can apply a unitary transformation to a vector."""
+
+    def _apply_unitary_(
+        self, vec: np.ndarray, norb: int, nelec: tuple[int, int], copy: bool
+    ) -> np.ndarray:
+        """Apply a unitary transformation to a vector.
+
+        Args:
+            vec: The vector to apply the unitary transformation to.
+            norb: The number of spatial orbitals.
+            nelec: The number of alpha and beta electrons.
+            copy: Whether to copy the vector before operating on it.
+                - If ``copy=True`` then this function always returns a newly allocated
+                vector and the original vector is left untouched.
+                - If ``copy=False`` then this function may still return a newly
+                allocated vector, but the original vector may have its data overwritten.
+                It is also possible that the original vector is returned,
+                modified in-place.
+
+        Returns:
+            The transformed vector.
+        """
+
+
+def apply_unitary(
+    vec: np.ndarray, obj: Any, norb: int, nelec: tuple[int, int], copy: bool = True
+) -> np.ndarray:
+    """Apply a unitary transformation to a vector.
+
+    Args:
+        vec: The vector to apply the unitary transformation to.
+        obj: The object with a unitary effect.
+        norb: The number of spatial orbitals.
+        nelec: The number of alpha and beta electrons.
+        copy: Whether to copy the vector before operating on it.
+            - If ``copy=True`` then this function always returns a newly allocated
+            vector and the original vector is left untouched.
+            - If ``copy=False`` then this function may still return a newly
+            allocated vector, but the original vector may have its data overwritten.
+            It is also possible that the original vector is returned,
+            modified in-place.
+
+    Returns:
+        The transformed vector.
+    """
+    method = getattr(obj, "_apply_unitary_", None)
+    if method is not None:
+        return method(vec, norb=norb, nelec=nelec, copy=copy)
+
+    raise TypeError(f"Object of type {type(obj)} has no _apply_unitary_ method.")
diff --git a/python/ffsim/variational/__init__.py b/python/ffsim/variational/__init__.py
@@ -10,7 +10,6 @@
 
 """Variational ansatzes."""
 
-from ffsim.variational.lucj import (
-    UCJOperator,
-    apply_ucj_operator,
-)
+from ffsim.variational.hopgate import HopGateAnsatzOperator
+from ffsim.variational.lucj import UCJOperator
+from ffsim.variational.multireference import multireference_state
diff --git a/python/ffsim/variational/hopgate.py b/python/ffsim/variational/hopgate.py
@@ -0,0 +1,42 @@
+# (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.
+
+"""Hop gate ansatz."""
+
+from __future__ import annotations
+
+import itertools
+from dataclasses import dataclass
+
+import numpy as np
+
+from ffsim.gates import apply_hop_gate
+
+
+@dataclass
+class HopGateAnsatzOperator:
+    r"""A hop gate ansatz operator."""
+
+    interaction_pairs: list[tuple[int, int]]
+    thetas: np.ndarray
+
+    def _apply_unitary_(
+        self, vec: np.ndarray, norb: int, nelec: tuple[int, int], copy: bool
+    ) -> np.ndarray:
+        """Apply the operator to a vector."""
+        if copy:
+            vec = vec.copy()
+        for target_orbs, theta in zip(
+            itertools.cycle(self.interaction_pairs), self.thetas
+        ):
+            vec = apply_hop_gate(
+                vec, theta, target_orbs=target_orbs, norb=norb, nelec=nelec, copy=False
+            )
+        return vec