{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Main component problem"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example we demonstrate solving the Main Component Problem (MCP). This is a subroutine of our decoding algorithm but is quite interesting on its own. The problem reads as follows:\n",
    "$$\n",
    "\\underset{\\phi_i \\in \\mathbb{C}^2}{\\arg \\max }\\left[\\otimes_i\\left\\langle\\phi_i\\right|\\right] \\rho\\left[\\otimes_i\\left|\\phi_i\\right\\rangle\\right],\n",
    "$$\n",
    "which essentially means finding a basis state which contributes most to a given pure state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from tqdm import tqdm\n",
    "from mdopt.optimiser.dephasing_dmrg import DephasingDMRG as deph_dmrg\n",
    "from mdopt.optimiser.dmrg import DMRG as dmrg\n",
    "from mdopt.optimiser.dephasing_dmrg import EffectiveDensityOperator\n",
    "from mdopt.mps.utils import (\n",
    "    create_state_vector,\n",
    "    create_simple_product_state,\n",
    "    create_custom_product_state,\n",
    "    mps_from_dense,\n",
    "    inner_product,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example also serves as a sanity check for the dephasing DMRG optimiser.\n",
    "Here, we solve the problem using exact diagonalisation, DMRG and dephasing DMRG.\n",
    "Next, we compare the solutions which should be exactly the same."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first create random pure state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_sites = 10\n",
    "psi = create_state_vector(num_sites)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us now bump up the main component (which we choose randomly) amplitude and renormalise the state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "index_to_bump = np.random.randint(0, 2**num_sites)\n",
    "psi[index_to_bump] = 10\n",
    "psi /= np.linalg.norm(psi)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we create an exact Matrix Product State (MPS) version of the state and its Matrix Product Density Operator (MPDO)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "mps = mps_from_dense(psi)\n",
    "mpdo = mps.density_mpo()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we find the main component (a computational basis state having the largest overlap) of the density matrix in the dense form."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1024/1024 [00:00<00:00, 5814.55it/s]\n"
     ]
    }
   ],
   "source": [
    "overlaps_exact = []\n",
    "for i in tqdm(range(2**num_sites)):\n",
    "    state_string = np.binary_repr(i, width=num_sites)\n",
    "    overlaps_exact.append(\n",
    "        np.absolute(create_custom_product_state(state_string).dense() @ psi) ** 2\n",
    "    )\n",
    "\n",
    "main_component_exact = np.argmax(overlaps_exact)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's find the main component of the MPDO using our vanilla 2-site DMRG optimiser."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:03<00:00,  3.63s/it]\n"
     ]
    }
   ],
   "source": [
    "mps_start = create_simple_product_state(num_sites, which=\"+\")\n",
    "engine = dmrg(mps_start, mpdo, mode=\"LA\")\n",
    "engine.run()\n",
    "max_excited_mps_from_dmrg = engine.mps\n",
    "\n",
    "overlaps_dmrg = []\n",
    "for i in range(2**num_sites):\n",
    "    state_string = np.binary_repr(i, width=num_sites)\n",
    "    overlaps_dmrg.append(\n",
    "        np.absolute(\n",
    "            inner_product(\n",
    "                max_excited_mps_from_dmrg,\n",
    "                create_custom_product_state(state_string),\n",
    "            )\n",
    "        )\n",
    "        ** 2\n",
    "    )\n",
    "main_component_dmrg = np.argmax(overlaps_dmrg)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we do the same with our built-in Dephasing DMRG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00, 19.11it/s]\n"
     ]
    }
   ],
   "source": [
    "mps_start = create_simple_product_state(num_sites, which=\"+\")\n",
    "dephasing_engine = deph_dmrg(\n",
    "    mps_start,\n",
    "    mps,\n",
    "    mode=\"LA\",\n",
    ")\n",
    "dephasing_engine.run()\n",
    "main_component_mps = dephasing_engine.mps\n",
    "\n",
    "overlaps_dephased = []\n",
    "for i in range(2**num_sites):\n",
    "    state_string = np.binary_repr(i, width=num_sites)\n",
    "    overlaps_dephased.append(\n",
    "        np.absolute(\n",
    "            inner_product(\n",
    "                main_component_mps,\n",
    "                create_custom_product_state(state_string),\n",
    "            )\n",
    "        )\n",
    "        ** 2\n",
    "    )\n",
    "main_component_dephased = np.argmax(overlaps_dephased)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some sanity checks: check that the answer from the Dephasing DMRG is a product state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "mps_product_answer = dephasing_engine.mps\n",
    "assert mps_product_answer.bond_dimensions == [\n",
    "    1 for _ in range(mps_product_answer.num_bonds)\n",
    "]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some sanity checks: check that all the three answers are the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert np.logical_and(\n",
    "    main_component_exact == main_component_dmrg,\n",
    "    main_component_exact == main_component_dephased,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mdopt-ZdbamFdU-py3.10",
   "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.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}