CustomizingVHA.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Customizing VQE\n",
    "\n",
    "I end up needing to rewrite most of the `SwapNetworkTrotterAnsatz` from OpenFermion to get more control over our initialization strategy. [1811.04476](https://arxiv.org/abs/1811.04476) says initial parameters turned out to be very important for reaching the ground state, and they were able to get 100% overlap on the 2x2 Hubbard model, so I'll try their strategies. \n",
    "\n",
    "Refer to the docstrings for more info on implementation, but broadly, the initialization strategies are: \n",
    "- Adiabatically turn on the interacting term \n",
    "- Adiabatically turn on the tunneling and interaction terms \n",
    "- Trotter expansion where all parameters are the same \n",
    "- Adiabatically turn on interaction for shorter time. \n",
    "\n",
    "We also try to implement the \"full optimization\" method outlined in [1506.05135](https://arxiv.org/abs/1506.05135). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A simple callback function we can pass into the optimizer to tell us what iteration it's on\n",
    "# Nevermind: this doesn't work with study.save()\n",
    "def gen_callback():\n",
    "    i = 0 \n",
    "    def callback(xk):\n",
    "        nonlocal i\n",
    "        print(\"Iteration: \", i)\n",
    "        i += 1\n",
    "    return callback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomHubbard(SwapNetworkTrotterHubbardAnsatz):\n",
    "    # Redefining my own Hubbard ansatz to override default_initial_params(self) \n",
    "    def __init__(self,\n",
    "                 x_dim: float,\n",
    "                 y_dim: float,\n",
    "                 tunneling: float,\n",
    "                 coulomb: float,\n",
    "                 initial_param_method: str, # This is the only change! \n",
    "                 periodic: bool=True,\n",
    "                 iterations: int=1,from \n",
    "                 adiabatic_evolution_time=None,\n",
    "                 qubits=None\n",
    "                 ) -> None:\n",
    "        \n",
    "        possible_initial_param_methods = [\n",
    "            'adiabatic', 'adiabatic_tunneling', 'trotter_constant', 'adiabatic_short'\n",
    "        ]\n",
    "        if initial_param_method not in possible_initial_param_methods:\n",
    "            raise ValueError('initial_param_method must be one of {}'.format(possible_initial_param_methods))\n",
    "        self.initial_param_method = initial_param_method\n",
    "        \n",
    "        super().__init__(x_dim, y_dim, tunneling, coulomb, periodic=periodic, iterations=iterations, \n",
    "                         adiabatic_evolution_time=adiabatic_evolution_time, qubits=qubits)\n",
    "        \n",
    "    def default_initial_params(self):\n",
    "        total_time = self.adiabatic_evolution_time\n",
    "        step_time = total_time / self.iterations \n",
    "        \n",
    "        params = [] \n",
    "        # I'm going to keep the format of the original default_initial_params (the -self.tunneling and \n",
    "        # division by pi) even though I don't understand why they do this yet. \n",
    "        if self.initial_param_method == 'adiabatic': \n",
    "            \"\"\"\n",
    "            Approximate evolution by H(t) = T + (t/A)V.\n",
    "            Sets the parameters so that the ansatz circuit consists of a sequence\n",
    "            of second-order Trotter steps approximating the dynamics of the\n",
    "            time-dependent Hamiltonian H(t) = T + (t/A)V, where T is the one-body\n",
    "            term and V is the two-body term of the Hamiltonian used to generate the\n",
    "            ansatz circuit, and t ranges from 0 to A, where A is equal to\n",
    "            `self.adibatic_evolution_time`. The number of Trotter steps\n",
    "            is equal to the number of iterations in the ansatz. This choice is\n",
    "            motivated by the idea of state preparation via adiabatic evolution.\n",
    "            The dynamics of H(t) are approximated as follows. First, the total\n",
    "            evolution time of A is split into segments of length A / r, where r\n",
    "            is the number of Trotter steps. Then, each Trotter step simulates H(t)\n",
    "            for a time length of A / r, where t is the midpoint of the\n",
    "            corresponding time segment. As an example, suppose A is 100 and the\n",
    "            ansatz has two iterations. Then the approximation is achieved with two\n",
    "            Trotter steps. The first Trotter step simulates H(25) for a time length\n",
    "            of 50, and the second Trotter step simulates H(75) for a time length\n",
    "            of 50.\n",
    "\n",
    "            The above docstring was copied from OpenFermion. \n",
    "            \"\"\"\n",
    "            for param, scale_factor in zip(self.params(),\n",
    "                                           self.param_scale_factors()):\n",
    "                if param.letter == 'Th' or param.letter == 'Tv':\n",
    "                    params.append(_canonicalize_exponent(\n",
    "                        -self.tunneling * step_time / np.pi, 4) / scale_factor)\n",
    "                elif param.letter == 'V':\n",
    "                    i, = param.subscripts\n",
    "                    # Use the midpoint of the time segment\n",
    "                    interpolation_progress = 0.5 * (2 * i + 1) / self.iterations\n",
    "                    params.append(_canonicalize_exponent(\n",
    "                        -0.5 * self.coulomb * interpolation_progress *\n",
    "                        step_time / np.pi, 2) / scale_factor)\n",
    "\n",
    "        elif self.initial_param_method == 'adiabatic_tunneling': \n",
    "            \"\"\"\n",
    "            Instead of only adiabatically turning on the interaction term, we also \n",
    "            adiabatically turn on the tunneling term. \n",
    "            \n",
    "            Inspired by 1811.04476. \n",
    "            \"\"\"\n",
    "            for param, scale_factor in zip(self.params(),\n",
    "                                           self.param_scale_factors()):\n",
    "                i, = param.subscripts\n",
    "                # Use the midpoint of the time segment\n",
    "                interpolation_progress = 0.5 * (2 * i + 1) / self.iterations\n",
    "\n",
    "                if param.letter == 'Th' or param.letter == 'Tv':\n",
    "                    params.append(_canonicalize_exponent(\n",
    "                        -self.tunneling * interpolation_progress * \n",
    "                        step_time / np.pi, 4) / scale_factor)\n",
    "                elif param.letter == 'V':\n",
    "                    params.append(_canonicalize_exponent(\n",
    "                        -0.5 * self.coulomb * interpolation_progress *\n",
    "                        step_time / np.pi, 2) / scale_factor)\n",
    "        \n",
    "        elif self.initial_param_method == 'trotter_constant': \n",
    "            \"\"\"\n",
    "            This is just Trotter decomposition for a simulation of length total_time. \n",
    "            I don't use _canonicalize_exponent(), so check what changes and if anything breaks. \n",
    "            \n",
    "            Inspired by 1811.04476.\n",
    "            \"\"\"\n",
    "            for param, scale_factor in zip(self.params(),\n",
    "                                           self.param_scale_factors()):\n",
    "                params.append(step_time)\n",
    "\n",
    "        elif self.initial_param_method == 'adiabatic_short':\n",
    "            \"\"\"\n",
    "            OpenFermion sets adiabatic_evolution_time = 0.1*abs(coulomb)*iterations, if we don't \n",
    "            provide one ourselves. This initialization method will divide that by iterations. \n",
    "            \n",
    "            Inspired by 1811.04476. \n",
    "            \"\"\"\n",
    "            total_time /= 5\n",
    "            step_time /= 5\n",
    "            for param, scale_factor in zip(self.params(),  self.param_scale_factors()):\n",
    "                # Copied 'adiabatic'\n",
    "                if param.letter == 'Th' or param.letter == 'Tv':\n",
    "                    params.append(_canonicalize_exponent(\n",
    "                        -self.tunneling * step_time / np.pi, 4) / scale_factor)\n",
    "                elif param.letter == 'V':\n",
    "                    i, = param.subscripts\n",
    "                    # Use the midpoint of the time segment\n",
    "                    interpolation_progress = 0.5 * (2 * i + 1) / self.iterations\n",
    "                    params.append(_canonicalize_exponent(\n",
    "                        -0.5 * self.coulomb * interpolation_progress *\n",
    "                        step_time / np.pi, 2) / scale_factor)\n",
    "                    \n",
    "        else: \n",
    "            raise ValueError(\"Don't know how to interpret initial parameter method {}\".format(self.initial_param_method))\n",
    "\n",
    "        return np.array(params)\n",
    "    \n",
    "def _canonicalize_exponent(exponent: float, period: int) -> float:\n",
    "    # Shift exponent into [-p/2, +p/2).\n",
    "    # They chose period = bounds (4 for T, 2 for V)\n",
    "    exponent += period / 2\n",
    "    exponent %= period\n",
    "    exponent -= period / 2\n",
    "    # Prefer (-p/2, +p/2] over [-p/2, +p/2).\n",
    "    if exponent <= -period / 2:\n",
    "        exponent += period  # coverage: ignore\n",
    "    return exponent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openfermioncirq import HamiltonianObjective, VariationalStudy\n",
    "from openfermioncirq.optimization import ScipyOptimizationAlgorithm, OptimizationParams\n",
    "from datetime import datetime\n",
    "\n",
    "# Optimizes and saves our VariationalStudy\n",
    "def run_ansatz(index_orbital_energies_combs, initialization_strat):\n",
    "    steps = 10\n",
    "    \n",
    "    obj = HamiltonianObjective(dc_hub) # Define objective function as Hamiltonian averaging\n",
    "    \n",
    "#     ansatz = SwapNetworkTrotterHubbardAnsatz(x_n, y_n, 1., 2., periodic=False, iterations=steps)\n",
    "    ansatz = CustomHubbard(x_n, y_n, 1., 2., initialization_strat, periodic=False, iterations=steps, \n",
    "                           adiabatic_evolution_time=50.)\n",
    "    \n",
    "    prep_circ = Circuit(\n",
    "        prepare_gaussian_state(\n",
    "            ansatz.qubits, \n",
    "            QuadraticHamiltonian(dc_hub.one_body), \n",
    "            occupied_orbitals=orbital_energies_combs[index_orbital_energies_combs]\n",
    "        ))\n",
    "    time = datetime.now().strftime(\"%m.%d.%y-%H:%M:%S\")\n",
    "    study = VariationalStudy(\n",
    "        'Hubbard-VHA-{}-{}-{}'.format(index_orbital_energies_combs, initialization_strat, time), \n",
    "        ansatz, \n",
    "        obj, \n",
    "        preparation_circuit=prep_circ, \n",
    "        target=genergy, \n",
    "        datadir='SavedVariationalStudy'\n",
    "    )\n",
    "    \n",
    "    algorithm = ScipyOptimizationAlgorithm(\n",
    "        kwargs={'method': 'L-BFGS-B', \n",
    "#                 'tol': 1e-16\n",
    "               }, \n",
    "#         options={'maxiter': 100}, \n",
    "        uses_bounds=True)\n",
    "    \n",
    "    optimizaton_params = OptimizationParams(algorithm=algorithm)\n",
    "    \n",
    "    # Optimize\n",
    "    result = study.optimize(optimizaton_params)\n",
    "    \n",
    "    study.save()\n",
    "    return result.optimal_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in top_three:\n",
    "    for init_strat in ['adiabatic', 'adiabatic_tunneling', 'trotter_constant', 'adiabatic_short']:\n",
    "        print(\"For tunneling ground state with index {} and initialization strategy {},the Hubbard ground state energy is {}\".format(\n",
    "            i, init_strat, run_ansatz(i, init_strat)))\n",
    "\n",
    "# We don't see that much improvement..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "path = 'SavedVariationalStudy/' + os.listdir('SavedVariationalStudy')[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# So once we do that optimization, we have decent overlap and we have our parameters. \n",
    "\"\"\"\n",
    "It seems they do an iterated VQE: do VQE for H(1/s) using ground state of H(0), then do VQE for H(2/s) using \n",
    "ground state of H(1/s), ..., until they solve VQE for H(1). They call this the result of \"sequential optimization\"\n",
    "and it isn't *that* close to the ground state (~70% overlap for 3 iterations with 2x2). \n",
    "Then, they used these resulting parameters as the initial parameters for an optimization method they call the \n",
    "\"global vairational\" method, calling this final result \"full optimziation\". \n",
    "Here's what the global variational method does: \n",
    "1. Choose 6 random points near the initial parameters\n",
    "2. For each point, do greedy noisy search for 150 iterations: \n",
    "    a. slightly perturb point\n",
    "        i. If number of acceptances in last 30 trials was large, increase step size. Else, decrease it. \n",
    "    b. keep the new value if it reduces energy\n",
    "3. Do Powell on each point until convergence. \n",
    "4. Keep the point with lowest energy. \n",
    "5. Alternate greedy noisy search and Powell until neither finds improvement. \n",
    "\"\"\"\n",
    "\n",
    "# study = VariationalStudy.load('SavedVariationalStudy/Hubbard-VHA-0-adiabatic-5.13.20-09:39:03')\n",
    "study = VariationalStudy.load(path)\n",
    "params = study.trial_results[0].optimal_parameters \n",
    "\n",
    "def init_6(params):\n",
    "    \"\"\"Initialize 6 'points' that we optimize. Step 1 above. \"\"\"\n",
    "    new_points = [] \n",
    "    for i in range(6):\n",
    "        gaussian_noise = np.random.normal(size=params.size) \n",
    "        new_points.append(params + gaussian_noise)\n",
    "    return new_points\n",
    "\n",
    "def greedy_noisy_search(point, iterations=150):\n",
    "    \"\"\"Do greedy noisy search as described above. \"\"\"\n",
    "    min_point = point\n",
    "    values = []\n",
    "    # Standard deviation of noise - this might not work \n",
    "    step_size = 0.5\n",
    "    for i in range(iterations):\n",
    "        gaussian_noise = np.random.normal(scale=step_size, size=point.size)\n",
    "        # Get value of new parameters\n",
    "        new_val = study.value_of(point + gaussian_noise)\n",
    "        values.append(new_val) \n",
    "        if new_val == min(values): \n",
    "            step_size /= 0.8 \n",
    "            min_point = new_val\n",
    "            print('Minimum is now {}'.format(min_point))\n",
    "        elif np.random.randint(20) == 0: # Do this 1/20 times\n",
    "            step_size *= 0.8\n",
    "        print(i, step_size)\n",
    "    return min_point\n",
    "\n",
    "def powell(point):\n",
    "    res = scipy.optimize.minimize(study.value_of, point, method='Powell', tol=1e-10)\n",
    "    return res.x.real\n",
    "\n",
    "def full_optimization(params):\n",
    "    six_starter_points = init_6(params)\n",
    "    for i in range(len(six_starter_points)):\n",
    "        six_starter_points[i] = greedy_noisy_search(six_starter_points[i])\n",
    "#         six_starter_points[i] = powell(six_starter_points[i])\n",
    "        print(six_starter_points[i])\n",
    "    vals = [study.value_of(point) for point in six_starter_points]\n",
    "    \n",
    "    try:\n",
    "        best = six_starter_points[np.argmin(vals)]\n",
    "        best_val = study.value_of(best)\n",
    "        curr = 'greedy'\n",
    "        while True: \n",
    "            old_best = best\n",
    "            if curr == 'greedy':\n",
    "                best = greedy_noisy_search(best)\n",
    "                curr = 'powell'\n",
    "            else: \n",
    "                best = powell(best)\n",
    "                curr = 'greedy'\n",
    "            if np.abs(study.value_of(best) - study.value_of(old_best)) < 1e-10:\n",
    "                break \n",
    "    except:\n",
    "        return best\n",
    "    return best\n",
    "            \n",
    "            \n",
    "full_optimization(params)\n",
    "# Errors when I run Powell... but either way not seeing much improvement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# So now we have an error of about 0.12. I suspect this might be the Trotter error, which means Jan-Michael's \n",
    "# decomposition will fix it. If it does, then I am justified in working on measurement precision. \n",
    "# For now, let's check the overlap with our true ground state\n",
    "\n",
    "from cirq import resolve_parameters\n",
    "from openfermioncirq.variational import variational_black_box\n",
    "\n",
    "def get_optimal_ground_state(study): \n",
    "    # Modified variational_black_box.UNITARY_SIMULATE.evaluate_noiseless \n",
    "    black_box = study._black_box_type(\n",
    "        study.ansatz, \n",
    "        study.objective, \n",
    "        study._preparation_circuit, \n",
    "        study.initial_state)\n",
    "    \n",
    "    circuit = resolve_parameters(\n",
    "        black_box.preparation_circuit + black_box.ansatz.circuit, \n",
    "        black_box.ansatz.param_resolver(black_box.ansatz.default_initial_params()))\n",
    "    final_state = circuit.final_wavefunction(\n",
    "        black_box.initial_state, \n",
    "        qubit_order=black_box.ansatz.qubit_permutation(black_box.ansatz.qubits))\n",
    "    return final_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LOAD STUDY WITH BEST VALUE AND CHECK FIDELITY\n",
    "study = VariationalStudy.load(path)\n",
    "\n",
    "opt_ground = get_optimal_ground_state(study)\n",
    "fidelity(opt_ground, v_hub[:,0]).real\n",
    "# 97.7% fidelity is great!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How can we implement the above exponentials in a quantum circuit? Hm... \n",
    "\n",
    "OF has a `SwapNetworkTrotterHubbard` ansatz. How does it work? \n",
    "\n",
    "Well, the SwapNetwork uses only `ISWAP`, `PhasedISWAP`, `CZ` and `Z` gates. What does it do? \n",
    "\n",
    "It was proposed in arxiv: 1711.04789. It allows us to simulate a Trotter step in $N$ depth and $N^2 / 2$ two-qubit entangling gates, and lets us prepare arbitrary Slater determinants in $N/2$ depth, all assuming only a linearly connected architecture. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Btw, this our ansatz\n",
    "print('Created a variational study with {} qubits and {} parameters'.format(\n",
    "    len(study.ansatz.qubits), study.num_params))\n",
    "\n",
    "print(\"The value of the objective with default initial parameters is {}\".format(\n",
    "    study.value_of(ansatz.default_initial_params())))\n",
    "\n",
    "print(\"The circuit of the study is:\")\n",
    "print(study.circuit.to_text_diagram(transpose=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that the circuit is made up of only 2 main gates: `ISWAPPowGate` and `CZPowGate`. The Cirq documentation tells us that the matrix decomposition for these gates is: \n",
    "$$\\texttt{ISWAPPowGate}(t) = \\begin{bmatrix} \n",
    "1 & 0 & 0 & 0 \\\\ \n",
    "0 & \\cos(\\frac{\\pi t}{2}) & i \\sin \\frac{\\pi t}{2} & 0 \\\\\n",
    "0 & i \\sin(\\frac{\\pi t}{2}) & \\cos(\\frac{\\pi t}{2}) & 0 \\\\\n",
    "0 & 0 & 0 & 1\n",
    "\\end{bmatrix} \\qquad \\texttt{CZPowGate}(t) = \\begin{bmatrix}\n",
    "1 & 0 & 0 & 0 \\\\\n",
    "0 & 1 & 0 & 0 \\\\\n",
    "0 & 0 & 1 & 0 \\\\\n",
    "0 & 0 & 0 & \\exp(i \\pi t) \n",
    "\\end{bmatrix} $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We have 60 parameters. What are they? \n",
    "# They're in format (U/T/W/V, p, q, i) where p,q are qubits and i is iteration\n",
    "# Seems 12 parameters per iteration. Per iteration, 4 interaction terms like we'd expect, \n",
    "# and then 8 tunneling terms like we'd expect. \n",
    "\n",
    "# How can I make these fewer? YO, SwapNetworkTrotterHubbard has only 3 parameters per iteration!\n",
    "list(ansatz.params())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# What does this mean?\n",
    "obj.variance_bound"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checking commutator relations of JW for Trotterization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openfermion.utils import commutator \n",
    "from openfermion.ops import FermionOperator \n",
    "from openfermion.transforms import jordan_wigner \n",
    "\n",
    "tunneling_01 = FermionOperator('0^ 1') + FermionOperator('1^ 0')\n",
    "tunneling_12 = FermionOperator('1^ 2') + FermionOperator('2^ 1')\n",
    "tunneling_23 = FermionOperator('2^ 3') + FermionOperator('3^ 2')\n",
    "\n",
    "print('Commutator of 01 and 12 is: ', commutator(tunneling_01, tunneling_12))\n",
    "print('\\nJW matrix of above commutator is: ', commutator(jordan_wigner(tunneling_01), jordan_wigner(tunneling_12)))\n",
    "print('\\nCommutator of 01 and 23 is: ', commutator(tunneling_01, tunneling_23))\n",
    "print('\\nJW matrix of above commutator is: ', commutator(jordan_wigner(tunneling_01), jordan_wigner(tunneling_23)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jordan_wigner(FermionOperator('1^')) * jordan_wigner(FermionOperator('0^'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "-0.25 * NKron(X, X) + 0.25 * NKron(X, Y) + 0.25 * NKron(Y, X) + 0.25 * NKron(Y, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hor_tun = FermionOperator('0^ 1') + FermionOperator('1^ 0') + FermionOperator('2^ 3') + FermionOperator('3^ 2')\n",
    "ver_tun = FermionOperator('0^ 2') + FermionOperator('2^ 0') + FermionOperator('1^ 3') + FermionOperator('3^ 1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "commutator(jordan_wigner(hor_tun), jordan_wigner(ver_tun))\n",
    "# Wait so horizontal and vertical terms commute... and the individual horizontal and vertical terms commute for \n",
    "# 2x2 case. Then I should have no error at all! It's just that I need more steps... But other papers didn't need \n",
    "# more steps. I legit just have to optimize longer then?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trying to reduce Trotter error\n",
    "\n",
    "The `HubbardSquareLattice` class has a useful method: `site_pair_iter(edge_type)`. We'll use the 'horizontal_neighbor' and 'vertical_neighbor' edge types. But we also need to differentiate between \"even\" and \"odd\" horizontal and vertical neighbors: even horizontal neighbors will be horizontal neighbors whose leftmost site has even index, and likewise for horizontal odd, vertical even, and vertical odd neighbors. To get this added specificity, we'll need to examine each item in the iterable generated by `site_pair_iter()`. We need the even and odd terms because then when we Trotterize the Hubbard Hamiltonian, we don't introduce any Trotter error, because the our four sets of tunneling terms (even horizontal, even vertical, odd horizontal, odd vertical) commute with each other. \n",
    "\n",
    "Actually, it will be much easier to just subclass `HubbardSquareLattice` and define our own `site_pair_iter()` that allows us to specify 'even' or 'odd'. I'll call this class `DecomposedHubbardSquareLattice`. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "class DecomposedHubbardSquareLattice(HubbardSquareLattice):\n",
    "    @property \n",
    "    def edge_types(self):\n",
    "        # Overriding edge_types property so we can define additional ones \n",
    "        return ('onsite', 'neighbor', 'diagonal_neighbor', 'horizontal_neighbor', 'vertical_neighbor', \n",
    "                'hor_even_neighbor', 'hor_odd_neighbor', 'ver_even_neighbor', 'ver_odd_neighbor')\n",
    "        \n",
    "    def site_pairs_iter(self, edge_type, ordered=True):\n",
    "        # `ordered` parameter just flips the order: if True, for each a,b -> (a,b), (b,a); if False, for each \n",
    "        # a,b -> (a,b), so we only get it once and it doesn't flip order. \n",
    "        # We WANT ordered=False, because there's a helper function tunneling_operator(i, j, coeff) that takes in \n",
    "        # two sites i,j and does coeff*(i^j j^i). \n",
    "        \n",
    "        # Overriding site_pairs_iter() to add functionality for additional edge_types \n",
    "        if edge_type == 'onsite':\n",
    "            return ((i, i) for i in self.site_indices)\n",
    "        elif edge_type == 'neighbor':\n",
    "            return self.neighbors_iter(ordered)\n",
    "        elif edge_type == 'horizontal_neighbor':\n",
    "            return self.horizontal_neighbors_iter(ordered)\n",
    "        elif edge_type == 'vertical_neighbor':\n",
    "            return self.vertical_neighbors_iter(ordered)\n",
    "        elif edge_type == 'diagonal_neighbor':\n",
    "            return self.diagonal_neighbors_iter(ordered)\n",
    "        # Above was copied from utils._lattice.py so we handle old edge_types correctly. \n",
    "        # Below is added functionality for new edge_types. \n",
    "        elif edge_type == 'hor_even_neighbor':\n",
    "            return self.hv_eo_neighbors(lambda x,y: 1-x%2, lambda x,y: (x+1, y), ordered) \n",
    "        elif edge_type == 'hor_odd_neighbor':\n",
    "            return self.hv_eo_neighbors(lambda x,y: x%2, lambda x,y: (x+1, y), ordered)\n",
    "        elif edge_type == 'ver_even_neighbor':\n",
    "            return self.hv_eo_neighbors(lambda x,y: 1-y%2, lambda x,y: (x, y+1), ordered)\n",
    "        elif edge_type == 'ver_odd_neighbor':\n",
    "            return self.hv_eo_neighbors(lambda x,y: y%2, lambda x,y: (x, y+1), ordered)\n",
    "        raise ValueError('Edge type {} is not valid.'.format(edge_type))\n",
    "        \n",
    "    def hv_eo_neighbors(self, filter_xy, map_next_xy, ordered=True):\n",
    "        for i in range(self.x_dimension):\n",
    "            for j in range(self.y_dimension):\n",
    "                if filter_xy(i, j):\n",
    "                    # Get indices for next site  \n",
    "                    k, l = map_next_xy(i, j)\n",
    "                    # Make sure next site isn't out of bounds \n",
    "                    if k >= self.x_dimension or l >= self.y_dimension: continue \n",
    "                    \n",
    "                    site_a = self.to_site_index((i, j))\n",
    "                    site_b = self.to_site_index((k, l))\n",
    "                    \n",
    "                    yield (site_a, site_b)\n",
    "                    if ordered: yield (site_b, site_a)\n",
    "        \n",
    "    # Function OpenFermion uses for 'horizontal_neightbor' edge_type.  \n",
    "    # PROBLEM: This loops around! Notice the % self.x_dimension. \n",
    "    def horizontal_neighbors_iter(self, ordered=True):\n",
    "        n_horizontal_edges_per_y = (\n",
    "            self.x_dimension - (self.x_dimension <= 2 or not self.periodic))\n",
    "        for x in range(n_horizontal_edges_per_y):\n",
    "            for y in range(self.y_dimension):\n",
    "                i = self.to_site_index((x, y))\n",
    "                j = self.to_site_index(((x + 1) % self.x_dimension, y))\n",
    "                yield (i, j)\n",
    "                if ordered:\n",
    "                    yield (j, i)\n",
    "\n",
    "    # Function OpenFermion uses for 'vertical_neighbor' edge_type\n",
    "    def vertical_neighbors_iter(self, ordered=True):\n",
    "        n_vertical_edges_per_x = (self.y_dimension -\n",
    "                                  (self.y_dimension <= 2 or not self.periodic))\n",
    "        for y in range(n_vertical_edges_per_x):\n",
    "            for x in range(self.x_dimension):\n",
    "                i = self.to_site_index((x, y))\n",
    "                j = self.to_site_index((x, (y + 1) % self.y_dimension))\n",
    "                yield (i, j)\n",
    "                if ordered:\n",
    "                    yield (j, i)\n",
    "    \n",
    "    # I'm changing this function so that it uses my functions, so I can compare the spectrums, since \n",
    "    # FermiHubbardModel.hamiltonian() will call this to generate the iterable. \n",
    "    def neighbors_iter(self, ordered=True):\n",
    "        return itertools.chain(\n",
    "            #self.horizontal_neighbors_iter(ordered),\n",
    "            #self.vertical_neighbors_iter(ordered)\n",
    "            \n",
    "            # itertools.chain('ABC', 'DEF') = Iterable('ABCDEF') \n",
    "            self.site_pairs_iter('hor_even_neighbor', ordered), \n",
    "            self.site_pairs_iter('hor_odd_neighbor', ordered), \n",
    "            self.site_pairs_iter('ver_even_neighbor', ordered), \n",
    "            self.site_pairs_iter('ver_odd_neighbor', ordered), \n",
    "        )\n",
    "\n",
    "\n",
    "lattice = DecomposedHubbardSquareLattice(x_n, y_n, n_dofs=n_dofs, periodic=periodic, spinless=spinless)\n",
    "hubbard = FermiHubbardModel(lattice , tunneling_parameters=tunneling, interaction_parameters=interaction, \n",
    "                            potential_parameters=potential, magnetic_field=mag_field, \n",
    "                            particle_hole_symmetry=particle_hole_sym)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openfermion.ops import FermionOperator\n",
    "\n",
    "def tunneling_operator(i, j, coefficient=1.):\n",
    "    # Copied from hamiltonians/_lattice.py in OpenFermion\n",
    "    return (FermionOperator(((i, 1), (j, 0)), coefficient) + FermionOperator(\n",
    "        ((j, 1), (i, 0)), coefficient.conjugate()))\n",
    "def tunneling_terms_hor_even(hor, even, model):\n",
    "    # Mostly copied from FermiHubbardMode.tunneling_terms() \n",
    "    terms = FermionOperator()\n",
    "    for param in model.tunneling_parameters:\n",
    "        a, aa = param.dofs \n",
    "        # We don't use param.edge_type because it's 'neighbor' and we need to be more specific\n",
    "        if hor and even:\n",
    "            site_pairs = model.lattice.site_pairs_iter('hor_even_neighbor', a != aa)\n",
    "        elif hor and not even: \n",
    "            site_pairs = model.lattice.site_pairs_iter('hor_odd_neighbor', a != aa)\n",
    "        elif not hor and even: \n",
    "            site_pairs = model.lattice.site_pairs_iter('ver_even_neighbor', a != aa)\n",
    "        elif not hor and not even:\n",
    "            site_pairs = model.lattice.site_pairs_iter('ver_odd_neighbor', a != aa)\n",
    "\n",
    "        for r, rr in site_pairs: \n",
    "            for spin_index in model.lattice.spin_indices:\n",
    "                i = model.lattice.to_spin_orbital_index(r, a, spin_index)\n",
    "                j = model.lattice.to_spin_orbital_index(rr, aa, spin_index)\n",
    "                terms += tunneling_operator(i, j, -param.coefficient)\n",
    "    return terms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hamiltonians = {\n",
    "    'hub': hubbard.hamiltonian(), \n",
    "    'non_interacting': hubbard.tunneling_terms(), \n",
    "    'hor_even': tunneling_terms_hor_even(True, True, hubbard), \n",
    "    'hor_odd': tunneling_terms_hor_even(True, False, hubbard), \n",
    "    'ver_even': tunneling_terms_hor_even(False, True, hubbard), \n",
    "    'ver_odd': tunneling_terms_hor_even(False, False, hubbard)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check to make sure summing up the parts gives us the whole of non_interacting term \n",
    "print('Sum of horizontal/vertical even/odd terms gives non-interacting term: ', \n",
    "      (hamiltonians['hor_even'] + hamiltonians['hor_odd'] + hamiltonians['ver_even'] + \n",
    "       hamiltonians['ver_odd'] == hamiltonians['non_interacting']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyzing ground states of tunneling term"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Recall we calculated w_tun and v_tun as the 16 degenerate lowest eigenvalues and eigenvectors of the \n",
    "# tunneling term\n",
    "v_tun.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tol = 1e-5\n",
    "def decompose_HF(state):\n",
    "    \"\"\"Decomposes an eigenstate (I can generalize to any state later) as a tensor product of |0> and |1> \"\"\"\n",
    "    tot = 0\n",
    "    for i in range(len(state)):\n",
    "        if np.abs(state[i]) > tol:\n",
    "            tot += 1\n",
    "    return tot\n",
    "        \n",
    "\n",
    "\n",
    "# Why do they all have 196 nonzero coefficients?\n",
    "for j in range(16):\n",
    "    print(decompose_HF(v_tun[:,j]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Oh, right, I guess the ground state *doesn't have to be* a pure state in the computational basis. It has to first and foremost be an eigenvector. For the Hadamard matrix, the eigenvectors are $\\ket{+}$ and $\\ket{-}$, where $\\ket{-}$ is the ground state, for example. \n",
    "Okay, so I need to move this to a basis where the Hamiltonian is diagonal. That's just the Bogoliubov transform. Then each eigenvector should be a separate basis state. Before I do that, let's just verify that the 16 degenerate eigenvectors are orthogonal. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def check_orthogonal(m):\n",
    "    \"\"\"Checks that each vector m[:,i] is orthogonal to the other vectors in m\"\"\"\n",
    "    for i in range(m.shape[1]):\n",
    "        vec1 = m[:,i]\n",
    "        for j in range(i+1, m.shape[1]):\n",
    "            vec2 = m[:,j]\n",
    "            overlap = fidelity(vec1, vec2)\n",
    "            assert np.abs(overlap) < tol, \"Inner product was {} for vectors {} and {}.\".format(overlap, i, j)\n",
    "            \n",
    "check_orthogonal(v_tun)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Right, the eigenvectors of a Hermitian matrix are orthogonal if the eigenvalues are distinct. All these vectors have the same eigenvalues, so the eigenvectors don't have to be orthogonal. \n",
    "\n",
    "Well, we could make them orthogonal with Gram-Schmidt, but wouldn't that change the eigenvalues? No, consider\n",
    "$$ H ( \\ket{\\psi_0} + \\ket{\\psi_1} ) = E_1 \\ket{\\psi_0} + ? \\ket{\\psi_1} $$\n",
    "where $\\ket{\\psi_0}$ is the first eigenvector, and we decomposed the second eigenvector as $\\ket{\\psi_0} + \\ket{\\psi_1}$. Then, by definition of eigenvector, the eigenvalue of $\\ket{\\psi_1}$ must be $E_1$. So, Gram-Schmidt preserves eigenvalues. \n",
    "\n",
    "Simple enough, I'll do GS on the degenerate ground states. Fine, but then what? After the Bogoliubov transform, the eigenvectors no longer represent distinct pure states in HF space. So how can I analyze this?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hubbard ground state is a superposition of 16 HF states\n",
    "decompose_HF(v_hub)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For 2 x 1 lattice, `v_hub` is a superposition of 4 states, and for a 2 x 2 lattice, `v_hub` is a superposition for 16 states. I'm guessing it's a function $2^{\\mathrm{\\# sites}}$, but why? \n",
    "\n",
    "Hypothesis 1: up and down spins are symmetric here since we don't really have different behavior for them. So, suppose the ground state is defined for $x$ fermions where $x = N / 2$, ie half-filling. The $2^x$ occurs because we can choose to put each fermion in either UP spin or DOWN spin, without changing the energy of the system. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}