{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "444cdef5-48f1-49a6-9b96-9b63a8ab2663",
   "metadata": {},
   "source": [
    "# QBM Comparison Simulation vs. Annealer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3e3e13aa-be77-4306-9a2e-8effc428a3fd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time: 1.1 s (started: 2022-03-29 12:06:51 +02:00)\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%load_ext autotime\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy.constants import k as k_B, h as h_P\n",
    "\n",
    "k_B /= h_P * 1e9\n",
    "matplotlib.rcParams.update({\"font.size\": 14})\n",
    "\n",
    "from qbm.models import BQRBM\n",
    "from qbm.utils import (\n",
    "    Discretizer,\n",
    "    kl_divergence,\n",
    "    get_project_dir,\n",
    "    get_rng,\n",
    "    lr_exp_decay,\n",
    ")\n",
    "\n",
    "# configure directories\n",
    "project_dir = get_project_dir()\n",
    "example_dir = project_dir / \"example\"\n",
    "models_dir = example_dir / \"artifacts\"\n",
    "plots_dir = example_dir / \"plots\"\n",
    "for dir_ in [example_dir, models_dir, plots_dir]:\n",
    "    if not dir_.exists():\n",
    "        dir_.mkdir(parents=True)\n",
    "\n",
    "# load anneal schedule data\n",
    "qpu_params = {\"region\": \"na-west-1\", \"solver\": \"Advantage_system4.1\"}\n",
    "# the anneal schedules xlsx files can be found here:\n",
    "# https://docs.dwavesys.com/docs/latest/doc_physical_properties.html#anneal-schedules\n",
    "# from the xlsx file, save the processor-annealing-schedule sheet as a csv\n",
    "csv_name = \"09-1263A-A_Advantage_system4_1_annealing_schedule.csv\"\n",
    "df_anneal = pd.read_csv(\n",
    "    project_dir / f\"data/anneal_schedules/csv/{csv_name}\", index_col=\"s\",\n",
    ")\n",
    "if 0.5 not in df_anneal.index:\n",
    "    df_anneal.loc[0.5] = (df_anneal.loc[0.499] + df_anneal.loc[0.501]) / 2\n",
    "df_anneal.sort_index(inplace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e87fdfc1-7d3c-41f9-b789-d89c818172ec",
   "metadata": {},
   "source": [
    "## Train Data Creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4f998341-610b-4900-ae94-c8cca34ec7ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 1, 1, ..., 0, 0, 1],\n",
       "       [0, 1, 0, ..., 1, 1, 1],\n",
       "       [0, 1, 0, ..., 1, 0, 1],\n",
       "       ...,\n",
       "       [1, 0, 0, ..., 0, 0, 0],\n",
       "       [1, 1, 1, ..., 0, 0, 1],\n",
       "       [1, 1, 0, ..., 1, 1, 1]])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time: 10.3 ms (started: 2022-03-29 12:06:52 +02:00)\n"
     ]
    }
   ],
   "source": [
    "seed = 42\n",
    "n_visible = 8\n",
    "n_hidden = 4\n",
    "n_qubits = n_visible + n_hidden\n",
    "\n",
    "rng = get_rng(seed)\n",
    "n_samples = 1500\n",
    "α = 2 / 3\n",
    "N_1 = rng.normal(-2, 1, int(round(n_samples * α, 0)))\n",
    "N_2 = rng.normal(3, 1, int(round(n_samples * (1 - α), 0)))\n",
    "x = np.concatenate((N_1, N_2))\n",
    "\n",
    "df = pd.DataFrame.from_dict({\"x\": x})\n",
    "discretizer = Discretizer(df, n_bits=n_visible)\n",
    "V_train = discretizer.df_to_bit_array(df)\n",
    "V_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "43251d5d-418f-47b1-ab9d-22466a4bf0f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 23556 characters]\n",
      "text/plain": [
       "<Figure size 1440x864 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time: 231 ms (started: 2022-03-29 12:06:52 +02:00)\n"
     ]
    }
   ],
   "source": [
    "data = discretizer.undiscretize_df(discretizer.discretize_df(df))\n",
    "\n",
    "n_bins = 32\n",
    "fig, ax = plt.subplots(dpi=144, figsize=(10, 6))\n",
    "_, bins, _ = ax.hist(data, bins=n_bins, density=True, label=\"Data\")\n",
    "ax.grid(alpha=0.7)\n",
    "plt.tight_layout()\n",
    "plt.savefig(plots_dir / \"histogram_train.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "da319a69-0e5c-4b89-847f-3716253893b7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time: 581 µs (started: 2022-03-29 12:06:53 +02:00)\n"
     ]
    }
   ],
   "source": [
    "def callback(model, sample_state_vectors):\n",
    "    V_train = model.V_train\n",
    "    n_visible = V_train.shape[1]\n",
    "    train_states = (\n",
    "        model._eigen_to_binary(V_train) * 2.0 ** np.arange(n_visible - 1, -1, -1)\n",
    "    ).sum(axis=1)\n",
    "    sample_states = (\n",
    "        model._eigen_to_binary(sample_state_vectors[:, :n_visible])\n",
    "        * 2.0 ** np.arange(n_visible - 1, -1, -1)\n",
    "    ).sum(axis=1)\n",
    "\n",
    "    dkl = kl_divergence(train_states, sample_states, n_bins=32)\n",
    "\n",
    "    return {\"value\": dkl, \"print\": f\"D_KL = {dkl:.3f}\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70b6216f-7998-42b0-ae52-78323311fc8a",
   "metadata": {},
   "source": [
    "## Simulation-based Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ed6bca6d-4486-4312-b849-1de7cefbd647",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[BQRBM] epoch 1: β = 0.939, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.013722\n",
      "D_KL = 0.269\n",
      "[BQRBM] epoch 2: β = 0.848, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.057796\n",
      "D_KL = 0.176\n",
      "[BQRBM] epoch 3: β = 0.762, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.021872\n",
      "D_KL = 0.103\n",
      "[BQRBM] epoch 4: β = 0.709, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.018841\n",
      "D_KL = 0.052\n",
      "[BQRBM] epoch 5: β = 0.671, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.029039\n",
      "D_KL = 0.056\n",
      "[BQRBM] epoch 6: β = 0.644, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.022161\n",
      "D_KL = 0.053\n",
      "[BQRBM] epoch 7: β = 0.619, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.031665\n",
      "D_KL = 0.046\n",
      "[BQRBM] epoch 8: β = 0.602, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.018389\n",
      "D_KL = 0.042\n",
      "[BQRBM] epoch 9: β = 0.585, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.022448\n",
      "D_KL = 0.041\n",
      "[BQRBM] epoch 10: β = 0.577, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.016117\n",
      "D_KL = 0.051\n",
      "[BQRBM] epoch 11: β = 0.561, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.014682\n",
      "D_KL = 0.032\n",
      "[BQRBM] epoch 12: β = 0.551, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.045321\n",
      "D_KL = 0.039\n",
      "[BQRBM] epoch 13: β = 0.541, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.016691\n",
      "D_KL = 0.036\n",
      "[BQRBM] epoch 14: β = 0.530, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.016191\n",
      "D_KL = 0.030\n",
      "[BQRBM] epoch 15: β = 0.525, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.020830\n",
      "D_KL = 0.024\n",
      "[BQRBM] epoch 16: β = 0.522, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.011880\n",
      "D_KL = 0.025\n",
      "[BQRBM] epoch 17: β = 0.520, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.028253\n",
      "D_KL = 0.026\n",
      "[BQRBM] epoch 18: β = 0.517, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.052182\n",
      "D_KL = 0.024\n",
      "[BQRBM] epoch 19: β = 0.513, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.018593\n",
      "D_KL = 0.024\n",
      "[BQRBM] epoch 20: β = 0.516, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.024267\n",
      "D_KL = 0.029\n",
      "[BQRBM] epoch 21: β = 0.511, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.094493\n",
      "D_KL = 0.021\n",
      "[BQRBM] epoch 22: β = 0.510, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.020511\n",
      "D_KL = 0.018\n",
      "[BQRBM] epoch 23: β = 0.506, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.015206\n",
      "D_KL = 0.024\n",
      "[BQRBM] epoch 24: β = 0.505, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.019299\n",
      "D_KL = 0.019\n",
      "[BQRBM] epoch 25: β = 0.502, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.017332\n",
      "D_KL = 0.016\n",
      "[BQRBM] epoch 26: β = 0.500, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.019422\n",
      "D_KL = 0.021\n",
      "[BQRBM] epoch 27: β = 0.498, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.022277\n",
      "D_KL = 0.021\n",
      "[BQRBM] epoch 28: β = 0.499, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.019550\n",
      "D_KL = 0.017\n",
      "[BQRBM] epoch 29: β = 0.503, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.018721\n",
      "D_KL = 0.018\n",
      "[BQRBM] epoch 30: β = 0.503, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.024697\n",
      "D_KL = 0.016\n",
      "[BQRBM] epoch 31: β = 0.502, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.050626\n",
      "D_KL = 0.017\n",
      "[BQRBM] epoch 32: β = 0.503, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.023000\n",
      "D_KL = 0.018\n",
      "[BQRBM] epoch 33: β = 0.503, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.025858\n",
      "D_KL = 0.020\n",
      "[BQRBM] epoch 34: β = 0.497, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.024858\n",
      "D_KL = 0.016\n",
      "[BQRBM] epoch 35: β = 0.499, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.032439\n",
      "D_KL = 0.018\n",
      "[BQRBM] epoch 36: β = 0.497, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.025122\n",
      "D_KL = 0.020\n",
      "[BQRBM] epoch 37: β = 0.496, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.021540\n",
      "D_KL = 0.016\n",
      "[BQRBM] epoch 38: β = 0.495, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.017282\n",
      "D_KL = 0.013\n",
      "[BQRBM] epoch 39: β = 0.496, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.021170\n",
      "D_KL = 0.017\n",
      "[BQRBM] epoch 40: β = 0.499, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.051067\n",
      "D_KL = 0.015\n",
      "[BQRBM] epoch 41: β = 0.499, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.017868\n",
      "D_KL = 0.013\n",
      "[BQRBM] epoch 42: β = 0.503, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.016865\n",
      "D_KL = 0.019\n",
      "[BQRBM] epoch 43: β = 0.499, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.014662\n",
      "D_KL = 0.011\n",
      "[BQRBM] epoch 44: β = 0.499, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.016259\n",
      "D_KL = 0.016\n",
      "[BQRBM] epoch 45: β = 0.500, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.023120\n",
      "D_KL = 0.012\n",
      "[BQRBM] epoch 46: β = 0.499, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.015960\n",
      "D_KL = 0.016\n",
      "[BQRBM] epoch 47: β = 0.501, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.015739\n",
      "D_KL = 0.012\n",
      "[BQRBM] epoch 48: β = 0.503, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.021606\n",
      "D_KL = 0.015\n",
      "[BQRBM] epoch 49: β = 0.503, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.014887\n",
      "D_KL = 0.014\n",
      "[BQRBM] epoch 50: β = 0.508, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:00:02.047007\n",
      "D_KL = 0.015\n",
      "[BQRBM] epoch 51: β = 0.506, learning rate = 9.33e-02, β learning rate = 9.66e-02, epoch duration = 0:00:02.064187\n",
      "D_KL = 0.009\n",
      "[BQRBM] epoch 52: β = 0.501, learning rate = 8.71e-02, β learning rate = 9.33e-02, epoch duration = 0:00:02.051784\n",
      "D_KL = 0.009\n",
      "[BQRBM] epoch 53: β = 0.498, learning rate = 8.12e-02, β learning rate = 9.01e-02, epoch duration = 0:00:02.022658\n",
      "D_KL = 0.012\n",
      "[BQRBM] epoch 54: β = 0.496, learning rate = 7.58e-02, β learning rate = 8.71e-02, epoch duration = 0:00:02.020979\n",
      "D_KL = 0.012\n",
      "[BQRBM] epoch 55: β = 0.496, learning rate = 7.07e-02, β learning rate = 8.41e-02, epoch duration = 0:00:02.016154\n",
      "D_KL = 0.014\n",
      "[BQRBM] epoch 56: β = 0.497, learning rate = 6.60e-02, β learning rate = 8.12e-02, epoch duration = 0:00:02.018136\n",
      "D_KL = 0.012\n",
      "[BQRBM] epoch 57: β = 0.497, learning rate = 6.16e-02, β learning rate = 7.85e-02, epoch duration = 0:00:02.026563\n",
      "D_KL = 0.009\n",
      "[BQRBM] epoch 58: β = 0.499, learning rate = 5.74e-02, β learning rate = 7.58e-02, epoch duration = 0:00:02.017382\n",
      "D_KL = 0.011\n",
      "[BQRBM] epoch 59: β = 0.497, learning rate = 5.36e-02, β learning rate = 7.32e-02, epoch duration = 0:00:02.053702\n",
      "D_KL = 0.012\n",
      "[BQRBM] epoch 60: β = 0.497, learning rate = 5.00e-02, β learning rate = 7.07e-02, epoch duration = 0:00:02.017849\n",
      "D_KL = 0.012\n",
      "[BQRBM] epoch 61: β = 0.498, learning rate = 4.67e-02, β learning rate = 6.83e-02, epoch duration = 0:00:02.017494\n",
      "D_KL = 0.012\n",
      "[BQRBM] epoch 62: β = 0.500, learning rate = 4.35e-02, β learning rate = 6.60e-02, epoch duration = 0:00:02.019530\n",
      "D_KL = 0.010\n",
      "[BQRBM] epoch 63: β = 0.500, learning rate = 4.06e-02, β learning rate = 6.37e-02, epoch duration = 0:00:02.022314\n",
      "D_KL = 0.011\n",
      "[BQRBM] epoch 64: β = 0.498, learning rate = 3.79e-02, β learning rate = 6.16e-02, epoch duration = 0:00:02.025714\n",
      "D_KL = 0.010\n",
      "[BQRBM] epoch 65: β = 0.497, learning rate = 3.54e-02, β learning rate = 5.95e-02, epoch duration = 0:00:02.023247\n",
      "D_KL = 0.012\n",
      "[BQRBM] epoch 66: β = 0.499, learning rate = 3.30e-02, β learning rate = 5.74e-02, epoch duration = 0:00:02.024861\n",
      "D_KL = 0.008\n",
      "[BQRBM] epoch 67: β = 0.500, learning rate = 3.08e-02, β learning rate = 5.55e-02, epoch duration = 0:00:02.020144\n",
      "D_KL = 0.010\n",
      "[BQRBM] epoch 68: β = 0.502, learning rate = 2.87e-02, β learning rate = 5.36e-02, epoch duration = 0:00:02.023129\n",
      "D_KL = 0.010\n",
      "[BQRBM] epoch 69: β = 0.501, learning rate = 2.68e-02, β learning rate = 5.18e-02, epoch duration = 0:00:02.053550\n",
      "D_KL = 0.013\n",
      "[BQRBM] epoch 70: β = 0.500, learning rate = 2.50e-02, β learning rate = 5.00e-02, epoch duration = 0:00:02.022104\n",
      "D_KL = 0.007\n",
      "[BQRBM] epoch 71: β = 0.500, learning rate = 2.33e-02, β learning rate = 4.83e-02, epoch duration = 0:00:02.021212\n",
      "D_KL = 0.008\n",
      "[BQRBM] epoch 72: β = 0.499, learning rate = 2.18e-02, β learning rate = 4.67e-02, epoch duration = 0:00:02.024338\n",
      "D_KL = 0.010\n",
      "[BQRBM] epoch 73: β = 0.498, learning rate = 2.03e-02, β learning rate = 4.51e-02, epoch duration = 0:00:02.025899\n",
      "D_KL = 0.009\n",
      "[BQRBM] epoch 74: β = 0.498, learning rate = 1.89e-02, β learning rate = 4.35e-02, epoch duration = 0:00:02.023541\n",
      "D_KL = 0.011\n",
      "[BQRBM] epoch 75: β = 0.498, learning rate = 1.77e-02, β learning rate = 4.20e-02, epoch duration = 0:00:02.024596\n",
      "D_KL = 0.010\n",
      "[BQRBM] epoch 76: β = 0.498, learning rate = 1.65e-02, β learning rate = 4.06e-02, epoch duration = 0:00:02.021185\n",
      "D_KL = 0.007\n",
      "[BQRBM] epoch 77: β = 0.498, learning rate = 1.54e-02, β learning rate = 3.92e-02, epoch duration = 0:00:02.023632\n",
      "D_KL = 0.008\n",
      "[BQRBM] epoch 78: β = 0.498, learning rate = 1.44e-02, β learning rate = 3.79e-02, epoch duration = 0:00:02.048295\n",
      "D_KL = 0.008\n",
      "[BQRBM] epoch 79: β = 0.497, learning rate = 1.34e-02, β learning rate = 3.66e-02, epoch duration = 0:00:02.019859\n",
      "D_KL = 0.009\n",
      "[BQRBM] epoch 80: β = 0.497, learning rate = 1.25e-02, β learning rate = 3.54e-02, epoch duration = 0:00:02.021562\n",
      "D_KL = 0.009\n",
      "[BQRBM] epoch 81: β = 0.497, learning rate = 1.17e-02, β learning rate = 3.42e-02, epoch duration = 0:00:02.020804\n",
      "D_KL = 0.008\n",
      "[BQRBM] epoch 82: β = 0.496, learning rate = 1.09e-02, β learning rate = 3.30e-02, epoch duration = 0:00:02.047260\n",
      "D_KL = 0.008\n",
      "[BQRBM] epoch 83: β = 0.497, learning rate = 1.02e-02, β learning rate = 3.19e-02, epoch duration = 0:00:02.039120\n",
      "D_KL = 0.009\n",
      "[BQRBM] epoch 84: β = 0.497, learning rate = 9.47e-03, β learning rate = 3.08e-02, epoch duration = 0:00:02.019192\n",
      "D_KL = 0.007\n",
      "[BQRBM] epoch 85: β = 0.498, learning rate = 8.84e-03, β learning rate = 2.97e-02, epoch duration = 0:00:02.050638\n",
      "D_KL = 0.007\n",
      "[BQRBM] epoch 86: β = 0.497, learning rate = 8.25e-03, β learning rate = 2.87e-02, epoch duration = 0:00:02.049080\n",
      "D_KL = 0.007\n",
      "[BQRBM] epoch 87: β = 0.497, learning rate = 7.69e-03, β learning rate = 2.77e-02, epoch duration = 0:00:02.021485\n",
      "D_KL = 0.010\n",
      "[BQRBM] epoch 88: β = 0.497, learning rate = 7.18e-03, β learning rate = 2.68e-02, epoch duration = 0:00:02.047871\n",
      "D_KL = 0.008\n",
      "[BQRBM] epoch 89: β = 0.497, learning rate = 6.70e-03, β learning rate = 2.59e-02, epoch duration = 0:00:02.019209\n",
      "D_KL = 0.008\n",
      "[BQRBM] epoch 90: β = 0.497, learning rate = 6.25e-03, β learning rate = 2.50e-02, epoch duration = 0:00:02.015671\n",
      "D_KL = 0.007\n",
      "[BQRBM] epoch 91: β = 0.497, learning rate = 5.83e-03, β learning rate = 2.41e-02, epoch duration = 0:00:02.019394\n",
      "D_KL = 0.007\n",
      "[BQRBM] epoch 92: β = 0.497, learning rate = 5.44e-03, β learning rate = 2.33e-02, epoch duration = 0:00:02.026408\n",
      "D_KL = 0.008\n",
      "[BQRBM] epoch 93: β = 0.497, learning rate = 5.08e-03, β learning rate = 2.25e-02, epoch duration = 0:00:02.023763\n",
      "D_KL = 0.009\n",
      "[BQRBM] epoch 94: β = 0.497, learning rate = 4.74e-03, β learning rate = 2.18e-02, epoch duration = 0:00:02.020650\n",
      "D_KL = 0.008\n",
      "[BQRBM] epoch 95: β = 0.497, learning rate = 4.42e-03, β learning rate = 2.10e-02, epoch duration = 0:00:02.043108\n",
      "D_KL = 0.011\n",
      "[BQRBM] epoch 96: β = 0.497, learning rate = 4.12e-03, β learning rate = 2.03e-02, epoch duration = 0:00:02.018738\n",
      "D_KL = 0.007\n",
      "[BQRBM] epoch 97: β = 0.496, learning rate = 3.85e-03, β learning rate = 1.96e-02, epoch duration = 0:00:02.055366\n",
      "D_KL = 0.008\n",
      "[BQRBM] epoch 98: β = 0.497, learning rate = 3.59e-03, β learning rate = 1.89e-02, epoch duration = 0:00:02.017947\n",
      "D_KL = 0.008\n",
      "[BQRBM] epoch 99: β = 0.497, learning rate = 3.35e-03, β learning rate = 1.83e-02, epoch duration = 0:00:02.020823\n",
      "D_KL = 0.007\n",
      "[BQRBM] epoch 100: β = 0.497, learning rate = 3.13e-03, β learning rate = 1.77e-02, epoch duration = 0:00:02.026146\n",
      "D_KL = 0.009\n",
      "time: 3min 22s (started: 2022-03-29 12:06:53 +02:00)\n"
     ]
    }
   ],
   "source": [
    "train_model = True\n",
    "model_path = models_dir / \"model_simulation.pkl\"\n",
    "if train_model:\n",
    "    # model params\n",
    "    s_freeze = 1.0\n",
    "    embedding = None\n",
    "    beta_initial = 1\n",
    "    simulation_params = {\"beta\": 0.5}\n",
    "\n",
    "    # training params\n",
    "    n_epochs = 100\n",
    "    n_samples = 10_000\n",
    "    learning_rate = 0.1\n",
    "    mini_batch_size = 10\n",
    "    decay_epoch = 50\n",
    "    decay_period = 10\n",
    "    epochs = np.arange(1, n_epochs + 1)\n",
    "    learning_rates = learning_rate * lr_exp_decay(epochs, decay_epoch=50, period=10)\n",
    "    learning_rates_beta = learning_rate * lr_exp_decay(\n",
    "        epochs, decay_epoch=50, period=20\n",
    "    )\n",
    "\n",
    "    # model init\n",
    "    model_simulation = BQRBM(\n",
    "        V_train=V_train,\n",
    "        n_hidden=n_hidden,\n",
    "        A_freeze=df_anneal.loc[s_freeze, \"A(s) (GHz)\"],\n",
    "        B_freeze=df_anneal.loc[s_freeze, \"B(s) (GHz)\"],\n",
    "        beta_initial=beta_initial,\n",
    "        simulation_params=simulation_params,\n",
    "    )\n",
    "\n",
    "    # model train and save\n",
    "    model_simulation.train(\n",
    "        n_epochs=n_epochs,\n",
    "        n_samples=n_samples,\n",
    "        learning_rate=learning_rates,\n",
    "        learning_rate_beta=learning_rates_beta,\n",
    "        mini_batch_size=mini_batch_size,\n",
    "        callback=callback,\n",
    "    )\n",
    "    model_simulation.save(model_path)\n",
    "else:\n",
    "    model_simulation = BQRBM.load(model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3dcd5188-6e44-4db1-ac02-a0fb27d1fd67",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 216936 characters]\n",
      "text/plain": [
       "<Figure size 3000x1500 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time: 785 ms (started: 2022-03-29 01:29:11 +02:00)\n"
     ]
    }
   ],
   "source": [
    "dkls = [x[\"value\"] for x in model_simulation.callback_history]\n",
    "epochs = np.arange(1, len(dkls) + 1)\n",
    "fig, ax = plt.subplots(1, 2, figsize=(10, 5), dpi=300)\n",
    "ax[0].plot(epochs, dkls, linewidth=1.8)\n",
    "ax[0].set_xlabel(\"Epoch\")\n",
    "ax[0].set_ylabel(r\"$D_{KL}(p_{data} \\ || \\ p_{model})$\")\n",
    "ax[0].set_yticks(np.arange(0, 22.5, 2.5) / 100)\n",
    "ax[1].set_yticks(np.arange(40, 110, 10))\n",
    "ax[0].set_ylim((0, 0.2))\n",
    "ax[1].set_ylim((40, 100))\n",
    "ax[0].grid(alpha=0.7)\n",
    "\n",
    "ax[1].plot(\n",
    "    range(len(model_simulation.beta_history)),\n",
    "    1 / k_B / np.array(model_simulation.beta_history) * 1000,\n",
    "    linewidth=1.8,\n",
    ")\n",
    "ax[1].set_xlabel(\"Epoch\")\n",
    "ax[1].set_ylabel(r\"$\\hat{T}$ [mK]\")\n",
    "ax[1].axhline(\n",
    "    1 / k_B / 0.5 * 1000,\n",
    "    color=\"k\",\n",
    "    linestyle=\"--\",\n",
    "    label=r\"$T_{effective}$\",\n",
    "    linewidth=1.8,\n",
    ")\n",
    "ax[1].grid(alpha=0.7)\n",
    "ax[1].legend(loc=\"lower right\")\n",
    "plt.tight_layout()\n",
    "plt.savefig(plots_dir / \"train_results_simulation.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c845feb-ab7e-48e2-800f-e04df4430e53",
   "metadata": {},
   "source": [
    "## Annealer-based Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d8d90469-70eb-4a60-88ca-a35f4e49058d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[BQRBM] epoch 1: β = 0.484, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:07:08.907931\n",
      "D_KL = 0.232\n",
      "[BQRBM] epoch 2: β = 0.474, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:06:48.725121\n",
      "D_KL = 0.106\n",
      "[BQRBM] epoch 3: β = 0.486, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:04:01.671331\n",
      "D_KL = 0.081\n",
      "[BQRBM] epoch 4: β = 0.486, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:03:56.923078\n",
      "D_KL = 0.109\n",
      "[BQRBM] epoch 5: β = 0.483, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:03:38.070698\n",
      "D_KL = 0.059\n",
      "[BQRBM] epoch 6: β = 0.494, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:03:44.256469\n",
      "D_KL = 0.062\n",
      "[BQRBM] epoch 7: β = 0.485, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:03:01.542740\n",
      "D_KL = 0.059\n",
      "[BQRBM] epoch 8: β = 0.499, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:03:17.857143\n",
      "D_KL = 0.105\n",
      "[BQRBM] epoch 9: β = 0.505, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:45.187334\n",
      "D_KL = 0.069\n",
      "[BQRBM] epoch 10: β = 0.502, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:44.753510\n",
      "D_KL = 0.046\n",
      "[BQRBM] epoch 11: β = 0.502, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:43.701216\n",
      "D_KL = 0.039\n",
      "[BQRBM] epoch 12: β = 0.497, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:47.841487\n",
      "D_KL = 0.040\n",
      "[BQRBM] epoch 13: β = 0.503, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:42.250770\n",
      "D_KL = 0.047\n",
      "[BQRBM] epoch 14: β = 0.505, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:45.582664\n",
      "D_KL = 0.060\n",
      "[BQRBM] epoch 15: β = 0.501, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:41.280884\n",
      "D_KL = 0.035\n",
      "[BQRBM] epoch 16: β = 0.506, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:43.825837\n",
      "D_KL = 0.050\n",
      "[BQRBM] epoch 17: β = 0.500, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:43.196413\n",
      "D_KL = 0.037\n",
      "[BQRBM] epoch 18: β = 0.497, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:42.245617\n",
      "D_KL = 0.041\n",
      "[BQRBM] epoch 19: β = 0.496, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:43.378315\n",
      "D_KL = 0.054\n",
      "[BQRBM] epoch 20: β = 0.503, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:46.798225\n",
      "D_KL = 0.039\n",
      "[BQRBM] epoch 21: β = 0.508, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:45.673851\n",
      "D_KL = 0.061\n",
      "[BQRBM] epoch 22: β = 0.521, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:50.319919\n",
      "D_KL = 0.037\n",
      "[BQRBM] epoch 23: β = 0.526, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:47.534919\n",
      "D_KL = 0.051\n",
      "[BQRBM] epoch 24: β = 0.533, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:43.582384\n",
      "D_KL = 0.027\n",
      "[BQRBM] epoch 25: β = 0.528, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:50.557590\n",
      "D_KL = 0.047\n",
      "[BQRBM] epoch 26: β = 0.538, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:42.965456\n",
      "D_KL = 0.040\n",
      "[BQRBM] epoch 27: β = 0.534, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:43.454228\n",
      "D_KL = 0.047\n",
      "[BQRBM] epoch 28: β = 0.537, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:42.496459\n",
      "D_KL = 0.061\n",
      "[BQRBM] epoch 29: β = 0.534, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:03:53.357594\n",
      "D_KL = 0.093\n",
      "[BQRBM] epoch 30: β = 0.540, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:03:02.790267\n",
      "D_KL = 0.037\n",
      "[BQRBM] epoch 31: β = 0.542, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:03:02.298378\n",
      "D_KL = 0.040\n",
      "[BQRBM] epoch 32: β = 0.540, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:42.451451\n",
      "D_KL = 0.036\n",
      "[BQRBM] epoch 33: β = 0.536, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:47.173624\n",
      "D_KL = 0.039\n",
      "[BQRBM] epoch 34: β = 0.541, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:40.074136\n",
      "D_KL = 0.027\n",
      "[BQRBM] epoch 35: β = 0.547, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:42.367784\n",
      "D_KL = 0.079\n",
      "[BQRBM] epoch 36: β = 0.520, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:41.632255\n",
      "D_KL = 0.052\n",
      "[BQRBM] epoch 37: β = 0.524, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:43.822197\n",
      "D_KL = 0.020\n",
      "[BQRBM] epoch 38: β = 0.521, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:40.941778\n",
      "D_KL = 0.022\n",
      "[BQRBM] epoch 39: β = 0.485, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:04:17.433543\n",
      "D_KL = 0.024\n",
      "[BQRBM] epoch 40: β = 0.473, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:06:44.943720\n",
      "D_KL = 0.019\n",
      "[BQRBM] epoch 41: β = 0.454, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:07:25.917607\n",
      "D_KL = 0.029\n",
      "[BQRBM] epoch 42: β = 0.449, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:06:53.615115\n",
      "D_KL = 0.037\n",
      "[BQRBM] epoch 43: β = 0.440, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:07:00.560991\n",
      "D_KL = 0.021\n",
      "[BQRBM] epoch 44: β = 0.468, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:07:49.345984\n",
      "D_KL = 0.049\n",
      "[BQRBM] epoch 45: β = 0.482, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:03:02.584091\n",
      "D_KL = 0.029\n",
      "[BQRBM] epoch 46: β = 0.485, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:49.108712\n",
      "D_KL = 0.029\n",
      "[BQRBM] epoch 47: β = 0.490, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:43.011319\n",
      "D_KL = 0.025\n",
      "[BQRBM] epoch 48: β = 0.497, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:40.860075\n",
      "D_KL = 0.029\n",
      "[BQRBM] epoch 49: β = 0.507, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:40.831445\n",
      "D_KL = 0.026\n",
      "[BQRBM] epoch 50: β = 0.502, learning rate = 1.00e-01, β learning rate = 1.00e-01, epoch duration = 0:02:45.643364\n",
      "D_KL = 0.047\n",
      "[BQRBM] epoch 51: β = 0.509, learning rate = 9.33e-02, β learning rate = 9.66e-02, epoch duration = 0:02:57.294182\n",
      "D_KL = 0.026\n",
      "[BQRBM] epoch 52: β = 0.510, learning rate = 8.71e-02, β learning rate = 9.33e-02, epoch duration = 0:02:44.202837\n",
      "D_KL = 0.018\n",
      "[BQRBM] epoch 53: β = 0.510, learning rate = 8.12e-02, β learning rate = 9.01e-02, epoch duration = 0:02:40.833022\n",
      "D_KL = 0.018\n",
      "[BQRBM] epoch 54: β = 0.514, learning rate = 7.58e-02, β learning rate = 8.71e-02, epoch duration = 0:02:43.961709\n",
      "D_KL = 0.039\n",
      "[BQRBM] epoch 55: β = 0.523, learning rate = 7.07e-02, β learning rate = 8.41e-02, epoch duration = 0:02:41.420784\n",
      "D_KL = 0.027\n",
      "[BQRBM] epoch 56: β = 0.523, learning rate = 6.60e-02, β learning rate = 8.12e-02, epoch duration = 0:02:43.910451\n",
      "D_KL = 0.019\n",
      "[BQRBM] epoch 57: β = 0.518, learning rate = 6.16e-02, β learning rate = 7.85e-02, epoch duration = 0:02:46.606390\n",
      "D_KL = 0.024\n",
      "[BQRBM] epoch 58: β = 0.517, learning rate = 5.74e-02, β learning rate = 7.58e-02, epoch duration = 0:02:40.529208\n",
      "D_KL = 0.018\n",
      "[BQRBM] epoch 59: β = 0.512, learning rate = 5.36e-02, β learning rate = 7.32e-02, epoch duration = 0:03:00.093450\n",
      "D_KL = 0.031\n",
      "[BQRBM] epoch 60: β = 0.519, learning rate = 5.00e-02, β learning rate = 7.07e-02, epoch duration = 0:02:40.621479\n",
      "D_KL = 0.029\n",
      "[BQRBM] epoch 61: β = 0.520, learning rate = 4.67e-02, β learning rate = 6.83e-02, epoch duration = 0:02:44.035333\n",
      "D_KL = 0.019\n",
      "[BQRBM] epoch 62: β = 0.520, learning rate = 4.35e-02, β learning rate = 6.60e-02, epoch duration = 0:02:43.550385\n",
      "D_KL = 0.034\n",
      "[BQRBM] epoch 63: β = 0.522, learning rate = 4.06e-02, β learning rate = 6.37e-02, epoch duration = 0:03:02.450734\n",
      "D_KL = 0.039\n",
      "[BQRBM] epoch 64: β = 0.522, learning rate = 3.79e-02, β learning rate = 6.16e-02, epoch duration = 0:03:03.080556\n",
      "D_KL = 0.021\n",
      "[BQRBM] epoch 65: β = 0.521, learning rate = 3.54e-02, β learning rate = 5.95e-02, epoch duration = 0:02:41.218441\n",
      "D_KL = 0.011\n",
      "[BQRBM] epoch 66: β = 0.525, learning rate = 3.30e-02, β learning rate = 5.74e-02, epoch duration = 0:02:40.693951\n",
      "D_KL = 0.033\n",
      "[BQRBM] epoch 67: β = 0.522, learning rate = 3.08e-02, β learning rate = 5.55e-02, epoch duration = 0:02:42.756640\n",
      "D_KL = 0.029\n",
      "[BQRBM] epoch 68: β = 0.522, learning rate = 2.87e-02, β learning rate = 5.36e-02, epoch duration = 0:02:47.121776\n",
      "D_KL = 0.016\n",
      "[BQRBM] epoch 69: β = 0.523, learning rate = 2.68e-02, β learning rate = 5.18e-02, epoch duration = 0:02:43.415603\n",
      "D_KL = 0.011\n",
      "[BQRBM] epoch 70: β = 0.517, learning rate = 2.50e-02, β learning rate = 5.00e-02, epoch duration = 0:02:47.585034\n",
      "D_KL = 0.016\n",
      "[BQRBM] epoch 71: β = 0.520, learning rate = 2.33e-02, β learning rate = 4.83e-02, epoch duration = 0:02:44.162901\n",
      "D_KL = 0.016\n",
      "[BQRBM] epoch 72: β = 0.525, learning rate = 2.18e-02, β learning rate = 4.67e-02, epoch duration = 0:02:43.420097\n",
      "D_KL = 0.026\n",
      "[BQRBM] epoch 73: β = 0.516, learning rate = 2.03e-02, β learning rate = 4.51e-02, epoch duration = 0:02:44.361308\n",
      "D_KL = 0.040\n",
      "[BQRBM] epoch 74: β = 0.510, learning rate = 1.89e-02, β learning rate = 4.35e-02, epoch duration = 0:02:40.677690\n",
      "D_KL = 0.016\n",
      "[BQRBM] epoch 75: β = 0.511, learning rate = 1.77e-02, β learning rate = 4.20e-02, epoch duration = 0:02:40.761636\n",
      "D_KL = 0.027\n",
      "[BQRBM] epoch 76: β = 0.512, learning rate = 1.65e-02, β learning rate = 4.06e-02, epoch duration = 0:02:43.153207\n",
      "D_KL = 0.074\n",
      "[BQRBM] epoch 77: β = 0.510, learning rate = 1.54e-02, β learning rate = 3.92e-02, epoch duration = 0:02:42.981059\n",
      "D_KL = 0.022\n",
      "[BQRBM] epoch 78: β = 0.508, learning rate = 1.44e-02, β learning rate = 3.79e-02, epoch duration = 0:02:41.887736\n",
      "D_KL = 0.010\n",
      "[BQRBM] epoch 79: β = 0.505, learning rate = 1.34e-02, β learning rate = 3.66e-02, epoch duration = 0:02:46.777951\n",
      "D_KL = 0.033\n",
      "[BQRBM] epoch 80: β = 0.509, learning rate = 1.25e-02, β learning rate = 3.54e-02, epoch duration = 0:02:41.565834\n",
      "D_KL = 0.029\n",
      "[BQRBM] epoch 81: β = 0.500, learning rate = 1.17e-02, β learning rate = 3.42e-02, epoch duration = 0:03:03.571748\n",
      "D_KL = 0.018\n",
      "[BQRBM] epoch 82: β = 0.498, learning rate = 1.09e-02, β learning rate = 3.30e-02, epoch duration = 0:02:45.323058\n",
      "D_KL = 0.013\n",
      "[BQRBM] epoch 83: β = 0.503, learning rate = 1.02e-02, β learning rate = 3.19e-02, epoch duration = 0:02:42.121055\n",
      "D_KL = 0.032\n",
      "[BQRBM] epoch 84: β = 0.506, learning rate = 9.47e-03, β learning rate = 3.08e-02, epoch duration = 0:02:45.258842\n",
      "D_KL = 0.011\n",
      "[BQRBM] epoch 85: β = 0.512, learning rate = 8.84e-03, β learning rate = 2.97e-02, epoch duration = 0:03:06.646914\n",
      "D_KL = 0.016\n",
      "[BQRBM] epoch 86: β = 0.516, learning rate = 8.25e-03, β learning rate = 2.87e-02, epoch duration = 0:02:59.898895\n",
      "D_KL = 0.034\n",
      "[BQRBM] epoch 87: β = 0.516, learning rate = 7.69e-03, β learning rate = 2.77e-02, epoch duration = 0:02:39.736599\n",
      "D_KL = 0.014\n",
      "[BQRBM] epoch 88: β = 0.519, learning rate = 7.18e-03, β learning rate = 2.68e-02, epoch duration = 0:02:44.552166\n",
      "D_KL = 0.016\n",
      "[BQRBM] epoch 89: β = 0.521, learning rate = 6.70e-03, β learning rate = 2.59e-02, epoch duration = 0:02:39.321136\n",
      "D_KL = 0.033\n",
      "[BQRBM] epoch 90: β = 0.524, learning rate = 6.25e-03, β learning rate = 2.50e-02, epoch duration = 0:02:43.346359\n",
      "D_KL = 0.033\n",
      "[BQRBM] epoch 91: β = 0.526, learning rate = 5.83e-03, β learning rate = 2.41e-02, epoch duration = 0:02:40.517694\n",
      "D_KL = 0.019\n",
      "[BQRBM] epoch 92: β = 0.529, learning rate = 5.44e-03, β learning rate = 2.33e-02, epoch duration = 0:02:42.414572\n",
      "D_KL = 0.039\n",
      "[BQRBM] epoch 93: β = 0.528, learning rate = 5.08e-03, β learning rate = 2.25e-02, epoch duration = 0:02:40.579656\n",
      "D_KL = 0.023\n",
      "[BQRBM] epoch 94: β = 0.529, learning rate = 4.74e-03, β learning rate = 2.18e-02, epoch duration = 0:02:44.102085\n",
      "D_KL = 0.024\n",
      "[BQRBM] epoch 95: β = 0.528, learning rate = 4.42e-03, β learning rate = 2.10e-02, epoch duration = 0:02:48.271655\n",
      "D_KL = 0.023\n",
      "[BQRBM] epoch 96: β = 0.529, learning rate = 4.12e-03, β learning rate = 2.03e-02, epoch duration = 0:02:45.851408\n",
      "D_KL = 0.034\n",
      "[BQRBM] epoch 97: β = 0.529, learning rate = 3.85e-03, β learning rate = 1.96e-02, epoch duration = 0:02:42.188319\n",
      "D_KL = 0.024\n",
      "[BQRBM] epoch 98: β = 0.531, learning rate = 3.59e-03, β learning rate = 1.89e-02, epoch duration = 0:02:38.901239\n",
      "D_KL = 0.029\n",
      "[BQRBM] epoch 99: β = 0.534, learning rate = 3.35e-03, β learning rate = 1.83e-02, epoch duration = 0:02:41.169669\n",
      "D_KL = 0.026\n",
      "[BQRBM] epoch 100: β = 0.534, learning rate = 3.13e-03, β learning rate = 1.77e-02, epoch duration = 0:02:39.609383\n",
      "D_KL = 0.016\n",
      "time: 5h 14min 27s (started: 2022-03-29 01:29:11 +02:00)\n"
     ]
    }
   ],
   "source": [
    "train_model = True\n",
    "model_path = models_dir / \"model_annealer.pkl\"\n",
    "if train_model:\n",
    "    # model params\n",
    "    s_freeze = 1.0\n",
    "    beta_initial = 0.5\n",
    "\n",
    "    # set embedding\n",
    "    # minorminer.find_embedding can be used to generate an embedding:\n",
    "    # https://docs.ocean.dwavesys.com/projects/system/en/stable/reference/generated/minorminer.find_embedding.html\n",
    "    embedding = {\n",
    "        0: [4909],\n",
    "        1: [4939],\n",
    "        2: [4879],\n",
    "        3: [4774],\n",
    "        4: [4789],\n",
    "        5: [4924],\n",
    "        6: [4864],\n",
    "        7: [4894],\n",
    "        8: [895],\n",
    "        9: [850],\n",
    "        10: [880],\n",
    "        11: [865],\n",
    "    }\n",
    "\n",
    "    # set anneal schedule\n",
    "    s_pause = 0.55\n",
    "    t_r = 20\n",
    "    α_quench = 2\n",
    "    t_pause = round(s_pause * t_r, 3)\n",
    "    Δ_quench = round((1 - s_pause) / α_quench, 3)\n",
    "    if s_pause == 1:\n",
    "        anneal_schedule = [(0, 0), (t_pause, s_pause)]\n",
    "    else:\n",
    "        anneal_schedule = [\n",
    "            (0, 0),\n",
    "            (t_pause, s_pause),\n",
    "            (round(t_pause + Δ_quench, 3), 1),\n",
    "        ]\n",
    "\n",
    "    # training params\n",
    "    n_epochs = 100\n",
    "    n_samples = 10_000\n",
    "    learning_rate = 0.1\n",
    "    mini_batch_size = 10\n",
    "    epochs = np.arange(1, n_epochs + 1)\n",
    "    learning_rates = learning_rate * lr_exp_decay(epochs, decay_epoch=50, period=10)\n",
    "    learning_rates_beta = learning_rate * lr_exp_decay(\n",
    "        epochs, decay_epoch=50, period=20\n",
    "    )\n",
    "\n",
    "    # set the anneal params\n",
    "    annealer_params = {\n",
    "        \"schedule\": anneal_schedule,\n",
    "        \"embedding\": embedding,\n",
    "        \"qpu_params\": qpu_params,\n",
    "    }\n",
    "\n",
    "    # skip if model already exists\n",
    "    if model_path.exists():\n",
    "        raise Exception(\"Model already exists\")\n",
    "\n",
    "    # model init\n",
    "    model_annealer = BQRBM(\n",
    "        V_train=V_train,\n",
    "        n_hidden=n_hidden,\n",
    "        A_freeze=df_anneal.loc[s_freeze, \"A(s) (GHz)\"],\n",
    "        B_freeze=df_anneal.loc[s_freeze, \"B(s) (GHz)\"],\n",
    "        beta_initial=beta_initial,\n",
    "        annealer_params=annealer_params,\n",
    "    )\n",
    "\n",
    "    # model train and save\n",
    "    model_annealer.train(\n",
    "        n_epochs=n_epochs,\n",
    "        n_samples=n_samples,\n",
    "        learning_rate=learning_rates,\n",
    "        learning_rate_beta=learning_rates_beta,\n",
    "        mini_batch_size=mini_batch_size,\n",
    "        callback=callback,\n",
    "    )\n",
    "    model_annealer.save(model_path)\n",
    "else:\n",
    "    model_annealer = BQRBM.load(model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e155249c-3907-4a4e-9ec9-c3b70ffa4418",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 398064 characters]\n",
      "text/plain": [
       "<Figure size 3000x1500 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time: 873 ms (started: 2022-03-29 06:43:39 +02:00)\n"
     ]
    }
   ],
   "source": [
    "fig, ax = plt.subplots(1, 2, figsize=(10, 5), dpi=300)\n",
    "dkls_simulation = [d[\"value\"] for d in model_simulation.callback_history]\n",
    "ax[0].set_xlabel(\"Epoch\")\n",
    "ax[0].set_ylabel(r\"$D_{KL}(p_{data} \\ || \\ p_{model})$\")\n",
    "ax[0].set_yticks(np.arange(0, 22.5, 2.5) / 100)\n",
    "ax[1].set_yticks(np.arange(80, 110, 2))\n",
    "ax[0].set_ylim((0, 0.2))\n",
    "ax[1].set_xlabel(\"Epoch\")\n",
    "ax[1].set_ylabel(r\"$\\hat{T}$ [mK]\")\n",
    "dkls_annealer = [d[\"value\"] for d in model_annealer.callback_history]\n",
    "\n",
    "epochs = np.arange(1, len(dkls_simulation) + 1)\n",
    "ax[0].plot(\n",
    "    epochs,\n",
    "    dkls_annealer,\n",
    "    label=fr\"BQRBM Advantage 4.1\",\n",
    "    color=\"tab:blue\",\n",
    "    linewidth=2,\n",
    ")\n",
    "ax[0].plot(\n",
    "    epochs,\n",
    "    dkls_simulation,\n",
    "    label=r\"BQRBM Simulation\",\n",
    "    color=\"k\",\n",
    "    linestyle=\"--\",\n",
    "    linewidth=2,\n",
    ")\n",
    "\n",
    "ax[1].plot(\n",
    "    range(len(model_annealer.beta_history)),\n",
    "    1 / k_B / np.array(model_annealer.beta_history) * 1000,\n",
    "    color=\"tab:blue\",\n",
    "    linewidth=2,\n",
    ")\n",
    "\n",
    "ax[0].grid(alpha=0.7)\n",
    "ax[1].grid(alpha=0.7)\n",
    "ax[0].legend()\n",
    "plt.tight_layout()\n",
    "plt.savefig(plots_dir / \"train_results_comparison.png\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
