Add spin-balanced UCJ operator and refactor spin-unbalanced UCJ (#208)

* add closed-shell UCJ operator * rename open/closed to spin unbalanced/balanced * update docs * rename files * use single interaction_pairs tuple * rename file * rename file * allow specifying different reps for ab and aabb * delete beh test * fix test * update docs and rename files * add ucj spin balanced qiskit gate * reorder classes in file * handle None * docs * fix sampler * fix deprecations
qiskit-community · May 31, 2024 · 16e9030 · 16e9030
1 parent 13ba88f
commit 16e9030
Show file tree

Hide file tree

Showing 28 changed files with 1,521 additions and 547 deletions.
diff --git a/docs/explanations/index.md b/docs/explanations/index.md
@@ -7,5 +7,6 @@ state-vectors-and-gates
 hamiltonians
 orbital-rotation
 double-factorized
+lucj
 qiskit-gate-decompositions
 ```
diff --git a/docs/explanations/lucj.ipynb b/docs/explanations/lucj.ipynb
diff --git a/docs/explanations/qiskit-gate-decompositions.ipynb b/docs/explanations/qiskit-gate-decompositions.ipynb
diff --git a/docs/how-to-guides/lucj.ipynb b/docs/how-to-guides/lucj.ipynb
diff --git a/docs/how-to-guides/qiskit-circuits.ipynb b/docs/how-to-guides/qiskit-circuits.ipynb
diff --git a/docs/how-to-guides/qiskit-sampler.ipynb b/docs/how-to-guides/qiskit-sampler.ipynb
@@ -134,7 +134,7 @@
     "\n",
     "### Sampling from an LUCJ circuit for a closed-shell molecule\n",
     "\n",
-    "The following code cell demonstrates a possible workflow for sampling from an LUCJ circuit for a nitrogen molecule in the 6-31g basis.\n"
+    "The following code cell demonstrates a possible workflow for sampling from a [spin-balanced LUCJ](../explanations/lucj.ipynb) circuit for a nitrogen molecule in the 6-31g basis.\n"
    ]
   },
   {
@@ -146,25 +146,25 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "converged SCF energy = -108.835236570775\n",
-      "norb = 16\n",
-      "nelec = (5, 5)\n",
-      "E(CCSD) = -109.0398256929734  E_corr = -0.204589122198831\n"
+      "converged SCF energy = -108.835236570774\n",
+      "norb = 14\n",
+      "nelec = (3, 3)\n",
+      "E(CCSD) = -108.9630419334855  E_corr = -0.1278053627110058\n"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "{'00000000000111110000000000011111': 9897,\n",
-       " '00000000001101110000000000011111': 27,\n",
-       " '00000000000111110000000000110111': 21,\n",
-       " '10000000000011110000000000011111': 7,\n",
-       " '00010000000101110000000000011111': 5,\n",
-       " '00000000000111110000000001011011': 5,\n",
-       " '00000000000111110001000000010111': 5,\n",
-       " '00000000010110110000000001011011': 3,\n",
-       " '01000000000011110000000000011111': 3,\n",
-       " '00000000000111111000000000001111': 3}"
+       "{'0000000000011100000000000111': 9924,\n",
+       " '0000000000110100000000001101': 16,\n",
+       " '0000000001011000000000010110': 10,\n",
+       " '0000000001110000000000000111': 10,\n",
+       " '0000000000011100000000011100': 10,\n",
+       " '0001000001010000000000000111': 5,\n",
+       " '0000000001011000100000000110': 4,\n",
+       " '0000000000011100100000001100': 3,\n",
+       " '0010000000011000000000010110': 3,\n",
+       " '0000000000110100010000000101': 3}"
       ]
      },
      "execution_count": 3,
@@ -173,19 +173,20 @@
     }
    ],
    "source": [
+    "import pyscf\n",
+    "import pyscf.cc\n",
     "import pyscf.data.elements\n",
-    "from pyscf import cc, gto\n",
     "\n",
     "# Build N2 molecule\n",
-    "mol = gto.Mole()\n",
+    "mol = pyscf.gto.Mole()\n",
     "mol.build(\n",
     "    atom=[[\"N\", (0, 0, 0)], [\"N\", (1.0, 0, 0)]],\n",
     "    basis=\"6-31g\",\n",
     "    symmetry=\"Dooh\",\n",
     ")\n",
     "\n",
     "# Define active space\n",
-    "n_frozen = pyscf.data.elements.chemcore(mol)\n",
+    "n_frozen = 4\n",
     "active_space = range(n_frozen, mol.nao_nr())\n",
     "\n",
     "# Get molecular data and Hamiltonian\n",
@@ -197,26 +198,24 @@
     "print(f\"nelec = {nelec}\")\n",
     "\n",
     "# Get CCSD t2 amplitudes for initializing the ansatz\n",
-    "ccsd = cc.CCSD(scf, frozen=[i for i in range(mol.nao_nr()) if i not in active_space])\n",
-    "_, _, t2 = ccsd.kernel()\n",
-    "\n",
-    "# Construct LUCJ operator using indices for a square lattice mapping\n",
-    "n_reps = 1\n",
-    "alpha_alpha_indices = [(p, p + 1) for p in range(norb - 1)]\n",
-    "alpha_beta_indices = [(p, p) for p in range(norb)]\n",
-    "ucj_op = ffsim.UCJOperator.from_parameters(\n",
-    "    ffsim.UCJOperator.from_t_amplitudes(t2).to_parameters(),\n",
-    "    norb=norb,\n",
-    "    n_reps=n_reps,\n",
-    "    alpha_alpha_indices=alpha_alpha_indices,\n",
-    "    alpha_beta_indices=alpha_beta_indices,\n",
+    "ccsd = pyscf.cc.CCSD(\n",
+    "    scf, frozen=[i for i in range(mol.nao_nr()) if i not in active_space]\n",
+    ").run()\n",
+    "\n",
+    "# Use 2 ansatz layers\n",
+    "n_reps = 2\n",
+    "# Use interactions implementable on a square lattice\n",
+    "pairs_aa = [(p, p + 1) for p in range(norb - 1)]\n",
+    "pairs_ab = [(p, p) for p in range(norb)]\n",
+    "ucj_op = ffsim.UCJOpSpinBalanced.from_t_amplitudes(\n",
+    "    ccsd.t2, n_reps=n_reps, interaction_pairs=(pairs_aa, pairs_ab)\n",
     ")\n",
     "\n",
     "# Construct circuit\n",
     "qubits = QuantumRegister(2 * norb)\n",
     "circuit = QuantumCircuit(qubits)\n",
     "circuit.append(ffsim.qiskit.PrepareHartreeFockJW(norb, nelec), qubits)\n",
-    "circuit.append(ffsim.qiskit.UCJOperatorJW(ucj_op), qubits)\n",
+    "circuit.append(ffsim.qiskit.UCJOpSpinBalancedJW(ucj_op), qubits)\n",
     "circuit.measure_all()\n",
     "\n",
     "# Sample 10,000 shots from the circuit using FfsimSampler\n",
@@ -237,7 +236,7 @@
    "source": [
     "### Sampling from an LUCJ circuit for an open-shell molecule\n",
     "\n",
-    "The following code cell demonstrates a possible workflow for sampling from an LUCJ circuit for a beryllium monohydride molecule in the 6-31g basis."
+    "The following code cell demonstrates a possible workflow for sampling from a [spin-unbalanced LUCJ](../explanations/lucj.ipynb) circuit for a hydroxyl radical in the 6-31g basis."
    ]
   },
   {
@@ -249,21 +248,28 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "converged SCF energy = -15.1172960668632\n",
+      "SCF not converged.\n",
+      "SCF energy = -75.3484557094412\n",
       "norb = 11\n",
-      "nelec = (3, 2)\n",
+      "nelec = (5, 4)\n",
       "\n",
       "WARN: RCCSD method does not support ROHF method. ROHF object is converted to UHF object and UCCSD method is called.\n",
       "\n",
-      "E(UCCSD) = -15.1402344412343  E_corr = -0.02293837437111472\n"
+      "E(UCCSD) = -75.45619739083253  E_corr = -0.1077416813913165\n"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "{'0000000001100000000111': 9998,\n",
-       " '1000000000100000100101': 1,\n",
-       " '0000000010100000100101': 1}"
+       "{'0000000111100000011111': 9991,\n",
+       " '0000100101100000111011': 2,\n",
+       " '0000100110100000111011': 1,\n",
+       " '0000000111100110000111': 1,\n",
+       " '0000010101100001011011': 1,\n",
+       " '1000000101100000111011': 1,\n",
+       " '0101000001100000011111': 1,\n",
+       " '0000010110100001011011': 1,\n",
+       " '0100000110100100001111': 1}"
       ]
      },
      "execution_count": 4,
@@ -275,12 +281,13 @@
     "import pyscf.data.elements\n",
     "from pyscf import cc, gto\n",
     "\n",
-    "# Build BeH molecule\n",
+    "# Build HO molecule\n",
     "mol = gto.Mole()\n",
     "mol.build(\n",
-    "    atom=[[\"H\", (0, 0, 0)], [\"Be\", (0, 0, 1.1)]],\n",
+    "    atom=[[\"H\", (0, 0, 0)], [\"O\", (0, 0, 1.1)]],\n",
     "    basis=\"6-31g\",\n",
     "    spin=1,\n",
+    "    symmetry=\"Coov\",\n",
     ")\n",
     "\n",
     "# Get molecular data and Hamiltonian\n",
@@ -292,27 +299,23 @@
     "print(f\"nelec = {nelec}\")\n",
     "\n",
     "# Get CCSD t2 amplitudes for initializing the ansatz\n",
-    "ccsd = cc.CCSD(scf)\n",
-    "_, _, t2 = ccsd.kernel()\n",
-    "\n",
-    "# Construct LUCJ operator using indices for a square lattice mapping\n",
-    "n_reps = 4\n",
-    "alpha_alpha_indices = [(p, p + 1) for p in range(norb - 1)]\n",
-    "alpha_beta_indices = [(p, p) for p in range(norb)]\n",
-    "beta_beta_indices = [(p, p + 1) for p in range(norb - 1)]\n",
-    "ucj_op = ffsim.UCJOperatorOpenShell.from_t_amplitudes(\n",
-    "    t2,\n",
-    "    n_reps=n_reps,\n",
-    "    alpha_alpha_indices=alpha_alpha_indices,\n",
-    "    alpha_beta_indices=alpha_beta_indices,\n",
-    "    beta_beta_indices=beta_beta_indices,\n",
+    "ccsd = cc.CCSD(scf).run()\n",
+    "\n",
+    "# Use 4 layers from opposite-spin amplitudes and 2 layers from same-spin amplitudes\n",
+    "n_reps = (4, 2)\n",
+    "# Use interactions implementable on a square lattice\n",
+    "pairs_aa = [(p, p + 1) for p in range(norb - 1)]\n",
+    "pairs_ab = [(p, p) for p in range(norb)]\n",
+    "pairs_bb = [(p, p + 1) for p in range(norb - 1)]\n",
+    "ucj_op = ffsim.UCJOpSpinUnbalanced.from_t_amplitudes(\n",
+    "    ccsd.t2, n_reps=n_reps, interaction_pairs=(pairs_aa, pairs_ab, pairs_bb)\n",
     ")\n",
     "\n",
     "# Construct circuit\n",
     "qubits = QuantumRegister(2 * norb)\n",
     "circuit = QuantumCircuit(qubits)\n",
     "circuit.append(ffsim.qiskit.PrepareHartreeFockJW(norb, nelec), qubits)\n",
-    "circuit.append(ffsim.qiskit.UCJOperatorOpenShellJW(ucj_op), qubits)\n",
+    "circuit.append(ffsim.qiskit.UCJOpSpinUnbalancedJW(ucj_op), qubits)\n",
     "circuit.measure_all()\n",
     "\n",
     "# Sample 10,000 shots from the circuit using FfsimSampler\n",

diff --git a/python/ffsim/__init__.py b/python/ffsim/__init__.py
@@ -83,7 +83,8 @@
     HopGateAnsatzOperator,
     RealUCJOperator,
     UCJOperator,
-    UCJOperatorOpenShell,
+    UCJOpSpinBalanced,
+    UCJOpSpinUnbalanced,
     multireference_state,
     multireference_state_prod,
 )
@@ -107,7 +108,8 @@
     "SupportsLinearOperator",
     "SupportsTrace",
     "UCJOperator",
-    "UCJOperatorOpenShell",
+    "UCJOpSpinBalanced",
+    "UCJOpSpinUnbalanced",
     "apply_diag_coulomb_evolution",
     "apply_fsim_gate",
     "apply_givens_rotation",

diff --git a/python/ffsim/qiskit/__init__.py b/python/ffsim/qiskit/__init__.py
@@ -23,7 +23,8 @@
     PrepareSlaterDeterminantJW,
     PrepareSlaterDeterminantSpinlessJW,
     UCJOperatorJW,
-    UCJOperatorOpenShellJW,
+    UCJOpSpinBalancedJW,
+    UCJOpSpinUnbalancedJW,
 )
 from ffsim.qiskit.sampler import FfsimSampler
 from ffsim.qiskit.transpiler_passes import DropNegligible, MergeOrbitalRotations
@@ -53,7 +54,8 @@
     "PrepareSlaterDeterminantJW",
     "PrepareSlaterDeterminantSpinlessJW",
     "UCJOperatorJW",
-    "UCJOperatorOpenShellJW",
+    "UCJOpSpinBalancedJW",
+    "UCJOpSpinUnbalancedJW",
     "ffsim_vec_to_qiskit_vec",
     "pre_init_passes",
     "qiskit_vec_to_ffsim_vec",

diff --git a/python/ffsim/qiskit/gates/__init__.py b/python/ffsim/qiskit/gates/__init__.py
@@ -25,8 +25,8 @@
     PrepareSlaterDeterminantJW,
     PrepareSlaterDeterminantSpinlessJW,
 )
-from ffsim.qiskit.gates.ucj import UCJOperatorJW
-from ffsim.qiskit.gates.ucj_open_shell import UCJOperatorOpenShellJW
+from ffsim.qiskit.gates.ucj import UCJOpSpinBalancedJW, UCJOpSpinUnbalancedJW
+from ffsim.qiskit.gates.ucj_operator import UCJOperatorJW
 
 __all__ = [
     "DiagCoulombEvolutionJW",
@@ -39,5 +39,6 @@
     "PrepareSlaterDeterminantJW",
     "PrepareSlaterDeterminantSpinlessJW",
     "UCJOperatorJW",
-    "UCJOperatorOpenShellJW",
+    "UCJOpSpinBalancedJW",
+    "UCJOpSpinUnbalancedJW",
 ]