From 32c38575f26f787481e4187fdd9eee634036cf38 Mon Sep 17 00:00:00 2001
From: "Kevin J. Sung" <kevjsung@umich.edu>
Date: Mon, 6 Nov 2023 17:36:30 -0500
Subject: [PATCH] update entanglement forging notebook to use larger active
 space and brickwork ansatz (#72)

---
 docs/tutorials/05-entanglement-forging.ipynb | 26 ++++++++++++--------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/docs/tutorials/05-entanglement-forging.ipynb b/docs/tutorials/05-entanglement-forging.ipynb
index 740abb417..c7c482112 100644
--- a/docs/tutorials/05-entanglement-forging.ipynb
+++ b/docs/tutorials/05-entanglement-forging.ipynb
@@ -19,7 +19,6 @@
     "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",
@@ -44,7 +43,7 @@
     "hartree_fock.kernel()\n",
     "\n",
     "# Define active space\n",
-    "active_space = [1, 2, 4, 5, 6]\n",
+    "active_space = range(1, mol.nao_nr())\n",
     "\n",
     "# Get molecular data and molecular Hamiltonian (one- and two-body tensors)\n",
     "mol_data = ffsim.MolecularData.from_hartree_fock(\n",
@@ -68,24 +67,31 @@
    "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",
+    "def brickwork(norb: int, n_layers: int):\n",
+    "    for i in range(n_layers):\n",
+    "        for j in range(i % 2, norb - 1, 2):\n",
+    "            yield (j, j + 1)\n",
+    "\n",
+    "\n",
+    "n_layers = norb\n",
+    "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",
     "\n",
     "# Construct ansatz state\n",
-    "reference_occupations_spatial = [(0, 1, 2), (1, 2, 3), (1, 2, 4)]\n",
+    "reference_occupations_spatial = [(0, 1, 2, 3), (1, 2, 3, 4), (0, 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",
+    "    mol_hamiltonian, operator, reference_occupations, norb=norb, nelec=nelec\n",
     ")\n",
     "\n",
     "# Compute the energy ⟨ψ|H|ψ⟩ of the ansatz state\n",
+    "hamiltonian = ffsim.linear_operator(mol_hamiltonian, norb=norb, nelec=nelec)\n",
     "energy = np.real(np.vdot(ansatz_state, hamiltonian @ ansatz_state))\n",
     "print(f\"Energy at initialialization: {energy}\")"
    ]
@@ -111,7 +117,7 @@
     "\n",
     "\n",
     "result = scipy.optimize.minimize(\n",
-    "    fun, x0=operator.thetas, method=\"COBYLA\", options=dict(maxiter=100)\n",
+    "    fun, x0=operator.thetas, method=\"L-BFGS-B\", options=dict(maxfun=100)\n",
     ")\n",
     "\n",
     "print(f\"Number of parameters: {len(result.x)}\")\n",
@@ -135,7 +141,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.10.12"
   },
   "orig_nbformat": 4
  },