demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cec63ff3-4227-410e-8d13-0917d44947e9",
   "metadata": {},
   "source": [
    "# Estimating sandwich bounds on the rate-distortion function of a Gaussian source\n",
    "\n",
    "This notebook demonstrates the method proposed in \n",
    "\n",
    ">  \"Towards Empirical Sandwich Bounds on the Rate-Distortion Function\"\n",
    ">\n",
    ">  Yibo Yang, Stephan Mandt\n",
    ">\n",
    ">  [https://arxiv.org/abs/2111.12166](https://arxiv.org/abs/2111.12166)\n",
    "\n",
    "on a randomly generated 2-dimensional Gaussian distribution.\n",
    "Also see [blog post](https://yiboyang.com/posts/estimating-the-rate-distortion-function-of-real-world-data-part-1/) for an introduction to the topic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8237af16-7e19-4b99-80c1-afa8142cd5bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Basic setup/imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c575d12c-2319-4417-aeba-f1a33d1ce740",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "dtype = np.float32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3524a723-c6e3-4e9f-b1da-870bcd79389c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "be71ee68-a77f-47fa-bce1-fba23e7997cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d1fa3bf2-be81-4c42-afd4-e2f97379031b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_probability as tfp\n",
    "tfd = tfp.distributions\n",
    "tfb = tfp.bijectors\n",
    "\n",
    "seed = 0\n",
    "tf.random.set_seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1364a400-ccbb-47ef-a909-0acd747a95ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "## We keep lambda fixed throughout the notebook; rerunning with different lamb traces out the full UB/LB\n",
    "lamb = 3.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccccdf6f-9724-4c1f-9750-d705a72ea35e",
   "metadata": {},
   "source": [
    "## Create a Gaussian source with a randomly sampled location and scale vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5ff1b580-c531-4e27-90a3-b8255f22b022",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dim = 2\n",
    "np.random.seed(seed)\n",
    "loc = np.random.uniform(-0.5, 0.5, data_dim).astype(dtype)\n",
    "var = np.random.uniform(0, 2, data_dim).astype(dtype)\n",
    "scale = var ** 0.5\n",
    "source = tfd.Normal(loc=loc, scale=scale)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7613fe9-92eb-42ed-b226-bace194ac794",
   "metadata": {},
   "source": [
    "## Estimate a point on the upper bound of the R-D curve, by training a VAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "369f8d7f-4151-4589-a248-18ac76412917",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/yiboyang/extra/envs/tf25/lib/python3.7/site-packages/tensorflow_probability/python/distributions/distribution.py:298: calling MultivariateNormalDiag.__init__ (from tensorflow_probability.python.distributions.mvn_diag) with scale_identity_multiplier is deprecated and will be removed after 2020-01-01.\n",
      "Instructions for updating:\n",
      "`scale_identity_multiplier` is deprecated; please combine it with `scale_diag` directly instead.\n"
     ]
    }
   ],
   "source": [
    "args = {\n",
    "    \"encoder_units\": [],\n",
    "    \"encoder_activation\": None,\n",
    "    \"decoder_units\": [\n",
    "        10\n",
    "    ],\n",
    "    \"decoder_activation\": \"leaky_relu\",\n",
    "    \"posterior_type\": \"gaussian\",\n",
    "    \"prior_type\": \"gmm_1\",  # factorized Gaussian prior with learned mean and scale\n",
    "    \"maf_stacks\": 0,\n",
    "    \"ar_activation\": None,\n",
    "\n",
    "    \"batchsize\": 64,\n",
    "    \"data_dim\": data_dim,\n",
    "    \"latent_dim\": data_dim,\n",
    "\n",
    "    \"lmbda\": lamb,\n",
    "    \"lr\": 5e-4,\n",
    "\n",
    "    \"epochs\": 20,\n",
    "    \"seed\": 0,\n",
    "    \"steps_per_epoch\": 10000,\n",
    "    \"max_validation_steps\": 1000,\n",
    "    \"nats\": True,\n",
    "    \"rpd\": False  # do not scale lambda by the source dimension\n",
    "}\n",
    "args = utils.get_args_as_obj(args)\n",
    "\n",
    "import rdub_mlp\n",
    "ub_model = rdub_mlp.Model.create_model(args)\n",
    "ub_model.compile(\n",
    "    optimizer=tf.keras.optimizers.Adam(learning_rate=args.lr),\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4699f7db-8536-4a99-b920-d287dffe4566",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/20\n",
      "10000/10000 [==============================] - 18s 2ms/step - loss: 2.4748 - rate: 1.3163 - mse: 0.3862\n",
      "Epoch 2/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2366 - rate: 1.2343 - mse: 0.3341\n",
      "Epoch 3/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2357 - rate: 1.2341 - mse: 0.3339\n",
      "Epoch 4/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2375 - rate: 1.2366 - mse: 0.3336\n",
      "Epoch 5/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2371 - rate: 1.2355 - mse: 0.3339\n",
      "Epoch 6/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2367 - rate: 1.2354 - mse: 0.3338\n",
      "Epoch 7/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2380 - rate: 1.2362 - mse: 0.3339\n",
      "Epoch 8/20\n",
      "10000/10000 [==============================] - 18s 2ms/step - loss: 2.2373 - rate: 1.2361 - mse: 0.3337\n",
      "Epoch 9/20\n",
      "10000/10000 [==============================] - 18s 2ms/step - loss: 2.2355 - rate: 1.2342 - mse: 0.3338\n",
      "Epoch 10/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2375 - rate: 1.2361 - mse: 0.3338\n",
      "Epoch 11/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2378 - rate: 1.2367 - mse: 0.3337\n",
      "Epoch 12/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2367 - rate: 1.2373 - mse: 0.3331\n",
      "Epoch 13/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2357 - rate: 1.2348 - mse: 0.3336\n",
      "Epoch 14/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2365 - rate: 1.2353 - mse: 0.3338\n",
      "Epoch 15/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2361 - rate: 1.2371 - mse: 0.3330\n",
      "Epoch 16/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2352 - rate: 1.2337 - mse: 0.3338\n",
      "Epoch 17/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2380 - rate: 1.2353 - mse: 0.3342\n",
      "Epoch 18/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2343 - rate: 1.2368 - mse: 0.3325\n",
      "Epoch 19/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2345 - rate: 1.2359 - mse: 0.3329\n",
      "Epoch 20/20\n",
      "10000/10000 [==============================] - 17s 2ms/step - loss: 2.2357 - rate: 1.2342 - mse: 0.3338\n"
     ]
    }
   ],
   "source": [
    "dataset = tf.data.Dataset.from_tensors([])\n",
    "dataset = dataset.repeat()\n",
    "dataset = dataset.map(lambda _: source.sample(args.batchsize))\n",
    "\n",
    "tmp_save_dir = '/tmp/rdvae'\n",
    "lr_scheduler = rdub_mlp.get_lr_scheduler(args.lr, args.epochs, decay_factor=0.2)\n",
    "hist = ub_model.fit(\n",
    "    dataset.prefetch(tf.data.AUTOTUNE),\n",
    "    epochs=args.epochs,\n",
    "    steps_per_epoch=args.steps_per_epoch,\n",
    "    validation_data=None,\n",
    "    validation_freq=1,\n",
    "    callbacks=[\n",
    "        tf.keras.callbacks.TerminateOnNaN(),\n",
    "        tf.keras.callbacks.experimental.BackupAndRestore(tmp_save_dir),\n",
    "        tf.keras.callbacks.LearningRateScheduler(lr_scheduler),\n",
    "    ],\n",
    "    verbose=True\n",
    ")\n",
    "ub_records = hist.history\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cece7fbf-8bab-4ef4-bf9f-f71e757a681c",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Estimate a point on the R-D curve by running Blahut-Arimoto on the discretized source"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "623abb87-6eb0-44c8-86d9-dd8ad20584d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 100000\n",
    "samples = source.sample(N).numpy()\n",
    "bins = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6889d399-aa11-4000-b908-9bbd38e34148",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f73e02c9bd0>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAT8AAAD4CAYAAACaPOETAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8rg+JYAAAACXBIWXMAAAsTAAALEwEAmpwYAAAjaklEQVR4nO2de4wc13Xmv9PdM90zw5dI6mWRjugVY4CyF/FGkAIoQB7KrqhEMP+RN9TChowIIBaQsDbWC69oA15AgBYWAvgBREZA2EK0jrCUQNsI4VVWkSA7gYPoQTleK3o5Y8kyyUiW+JoZPubR3Wf/6KIy99zD6Ts9NdNdXd8PaKCr+txbt7prztz6zqlzRVVBCCFlo9LvARBCSD+g8yOElBI6P0JIKaHzI4SUEjo/QkgpqfXjoKNS1wYm+nFoQkrNDE6fUNXLl9vu1t+b0JOnWkm2L/507klV3b3swa0xfXF+DUzgJrmlH4cmpNQ8rYfe6qXdiVMtPPfktiTbkat/vrWXY6w1fXF+hJCioWhpu9+DyBU6P0JIVxRAG8P1QASdHyEkiTY48yOElAyFYoG3vYSQsqEAWrztJYSUEWp+hJDSoQBaQ1YBis6PEJLEcCl+dH6EkAQUSs2PEFI+VIGF4fJ9dH6EkBQELUi/B5ErdH6EkK4ogDZnfoSQMsKZHyGkdHSSnOn8CCElQwEs6HDVPqbzI4R0RSFoDVnhdzq/MiPOxTxkD6+T/Ggrb3tdRKQK4AiA46p6e179EkL6DzW/pfkMgFcBbMixT0LIQCBoDZnml8vZiMg2AH8E4Jt59EcIGSw6lZwrSa+ikNfM72sAPg9g/aUMRGQfgH0A0MB4Tocl7+Ppdxar5/Wq79lj5aUTrla/ZMWoCua12u9h5MqK3bSI3A7gXVV9cSk7VT2gqjeo6g0jqK/0sISQNaYNSXoVhTxmfjcD+LiI/CGABoANIvKXqvrJHPomhAwAnYBHcW5pU1jx2ajqflXdpqrXAtgL4Bk6PkKGjU7AI+VVFJjnVwRStDC7L0UD7PVYvbTpRZNkHuLAcDHgMUzk6vxU9YcAfphnn4SQwaDFJGdCSNlQCBZ0uNzFcJ0NIWRVGMaAB50fIaQrCuFtL+kDPQQdpOJdqGGSqrZay+7XNanaflO6jcen7YSZBROh+wYDHoSQ0qGKQqWxpEDnRwjpSifgMVyPt9H5EUKSYMCD5Esvyb9uN6GGpglLbVVGR6N97YXmkv269DC+FFLaSNU5h/n5lM7DbWqHS6IQFjMlhJQTzvwIIaWjs24vnR8hpHQIy9iTVcbRnmweXcoD/1GbS/RtiXRAp43ND5R6WJ9RL8zGHds8RG98Nu/QO7bRMrXdjGy6HTsZ6oLv01m6criivcM1jyWErAqqgrZWkl4piMhuEXldRCZF5D7n87qIPJZ9/pyIXLvos/3Z/tdF5FbTrioi/ygi3+82Bjo/QkgSedXzy1Z6fAjAbQB2AbhTRHYZs7sBnFbV6wB8FcCDWdtd6NQNvR7AbgDfyPq7yMWF1LpC50cI6Uqnnl9uZexvBDCpqm+o6jyAgwD2GJs9AB7J3h8CcIuISLb/oKrOqeqbACaz/pa9kBqdHyEkgWVVct4qIkcWvfaZzq4BcHTR9rFsn2ujqk0AUwC2dGn7NXQWUksSZxnw6JW8KhMbvMTjuNvuCcxu4MQEJqTijEdMsnQzDijYYEW3bddmxLn0ojE7/YyaIgpuQvPyhXm3qEKJAxyWTqpLcrT3hKresIrDiVi8kJqI/G5KGzo/QkhXcn629ziA7Yu2t2X7PJtjIlIDsBHAySXafhzLXEiNt72EkCRyXLT8BQA7RWSHiIyiE8A4bGwOA7gre38HOgujabZ/bxYN3gFgJ4Dne1lIjTM/QkhXOiWt8klyVtWmiNwL4El0NIqHVfVlEbkfwBFVPQzgWwC+LSKTAE6h49CQ2T0O4BUATQD3qKZUkIyh8wPyS2btkmjsFQ+1NrawgNdvpd6I+xkxicZeoVK1CcLdE6ql5lwiVtMzx6pMjMeHttqc02/0fXl6ntE73WRp83tqc6GrjVS8ZO64WZnJs7CBqj4B4Amz70uL3s8C+MQl2j4A4IEl+v4hEhZSo/MjhHSlU9VluFQyOj9CSFc6j7fR+RFCSgdnfoSQkpL49EZhoPMDekpG9uhWfaVSH4kPbSukePr9aMPYOEYmaVhaToKwCTKkBGBQjb8LnQ8DCDI+Zto44zP9yIjzXcyFAQ6blN0xMkGb2bnux4JzLBMESUoc9+h2rQxJonSe0d5Bgc6PEJIEb3sJIaWDa3gQQkqJAmhy5jfAJFQ4TsJrk6IDRomz5j+lsxqZoPsD/1Gyb0LhADiaGppGX9xymWNjkqy9ZOSaOZbVyzwtsWH0u7PnYpsxk7ztJGG7+6KOTHGGhMrSlYbzvduE84RrKdIOU67JvK7bVYa3vYSQ8qG87SWElJCLxUyHCTo/QkgSnPkNMnnpJClFR71inAZbpKDi6Wcj5mF+R9OqrF8f7li/zjmY0dm8HDmrxVntDkBrS9h3ZTYutKDV8I+g3Qj1xcqcUwB13uyrx0VbxWps5y9ENlEbp4gCbJFWjXP4xHzPNncRiDXb9oLTj9VxU/S8XleT6yPLLGZaCIbL+RFCVgWFoOlVuy4wdH6EkCSo+RFCyofytpcQUkKo+TmIyHYA/wvAleh8RwdU9esr7beveCufJVRctiuvVUxib5ToC0QJwm4gxa7o5oxPN0yYwcT6THs07LvdiAMeagT8C1eNRTajUyaQMxcGUmwABAB0wnw3TlCkcsGMue4EiGbOhzu8ZO7pmXDbKc4QBX88m+byA2hSc8ZjcCtLFwA6v5gmgM+p6o9FZD2AF0XkKVV9JYe+CSEDgELQYsAjRFXfBvB29n5GRF5FZxFhOj9ChggGPJZARK4F8DEAzzmf7QOwDwAacHKzCCEDizLgcWlEZB2A7wD4rKpO289V9QCAAwCwQTb3WDlybfCKhdoH1q2e57XzEpajNlbj27A+NjJJuq3Njo25I2muj5OI2yY5ee4yR1MzUlh7NL7gW3WzOpo59uhUXNhgbnN4rOps/N00ToQFHKrnHV3QaJlyciqywaaNoY2jz+r0dHcbU+TBLSIbNTLJ094qenmtFrjGKJ1fjIiMoOP4HlXV7+bRJyFkkGBhgwgREXQWGH5VVb+y8iERQgYRzvxibgbwKQAvichPsn1fyBYlJoQMAapAq03nF6CqPwIKHgYyGoyn09gcPrsYkEdlzOTIjTo5YHafk+fXNgsf6YhX/DLcXFgX99McC3+mhfH4Z7O6oDjqbMXIY62RsM25q+Lx1aeMZroQH7s5YfIQ67HGNmIkPv21KyKb2mmTCzgdF0iIFlDyfk+bK2mLygJRnmZ73lw7vRYqHUBdkNFeQkjpUPC2lxBSShjwIISUFKcsYqGh8yOEJMHb3n7Ro3AcJR57Sae2jfNwum3nXQaRTbR6m3MOJllaa05BAlMUYPbyOMG60gz/LV/YEvdzYatZ1czJ2W2Zrt2Ax9zS/cxvjBvZ6s82SAIALZNQPXIu/n3b1YaxiZOTW+tCGzc12VRu1jNOsrQk/LGba9AGxrzrTVMKJgzYCm+daC+f7SWElBDe9hJCSglvewkhpUMhdH59o0d9wxYk8LDanP8Ae3ftMCpWahNpvYKZhvZ4XJDAJiw3G06CcCPse25jbGO1ueZEZIL5raGGVj3rrPC2NfxOrS5YOxuf5+xWc5yNkQnGfmV0wdG4H62E41l/1Pl9x8Pvq+4lrbfC60k2XxYf672T4Q5PA7TXSoKmnJTAPIBJznne9YrIbgBfR+cP65uq+mXzeR2dIsm/CeAkgD9W1V9kn+0HcDeAFoD/oqpPikgDwN8BqKPj1w6p6v9YagzFcX6EkP6hgOb0eJuIVAE8BODfAzgG4AUROWwKIN8N4LSqXiciewE8COCPRWQXgL0ArgfwAQBPi8ivA5gD8PuqejYrtPIjEflrVX32UuMYrvANIWTVUJWkVwI3AphU1TdUdR7AQQB7jM0eAI9k7w8BuCUrorIHwEFVnVPVNwFMArhRO5zN7Eey15KTVTo/QkgSqmmvBK4BcHTR9rFsn2ujqk0AUwC2LNVWRKpZcZV3ATylqlFR5cWU77bXy5eyJOkrcT9iCxkYzU/H4vy85pawqvX8Jk/zC4/laX7nrzDFBT4Ya0+1mVCfal4W58ihEZ77+FVnI5PLJsLCARfmw/M8eWZd1KZ9xhSGaDp5fua8xt+Obepnwr+ume3x97XhrbAAgS0MAQCyYHIyz87GNnWTs3chLpAQaX5mW+Dk+SXIgoOg8S1mmc/2bhWRI4u2D2TFjFcVVW0B+A0R2QTgeyLyEVX9p0vZl8/5EUKWjwJId34nVPWGJT4/DmD7ou1t2T7P5piI1ABsRCfw0bWtqp4RkR8A2A3gks6Pt72EkCRyvO19AcBOEdkhIqPoBDAOG5vDAO7K3t8B4BlV1Wz/XhGpi8gOADsBPC8il2czPojIGDrBlNeWGgRnfoSQBCS3aK+qNkXkXgBPopPq8rCqviwi9wM4oqqH0akO/20RmQRwCh0HiczucXRWh2wCuEdVWyJyNYBHskhyBcDjqvr9pcZB50cISSPHRL+s0vsTZt+XFr2fBfCJS7R9AMADZt9P0Vk5Mpnhd35JVXONIO4ERaJV1rzKv3a1tpoRw52iBZaR6YVo37krw0BKcywywYUrTeJxy0mEvioMBHz42rcjmw+Mhw/4Ty/EBzs1G+77yNawn3cn4tXljm8Is5rPvBVnOesHw6BD+714idPZTeF5bXrDCdoYWhPxb1V7xwQ47G8FREnNlfXxedkV3mxQxEuyj5LoPS3NVhdvxtfFmqJ8vI0QUlZY2IAQUk448yOElJHBSj1cMUPv/KJipgmFDvyOzH89L6ZvbUwhA69QaasR/gQL62LtyRYkmL3cObRJnK1uixNyt205HWxvnzgT2VTMvc3vbY6zBf75wpXB9mUjYdLzr0/8Kmrzf+Y/Emy3ts1ENtVq+Nd15kNxUvj4L8MvY+pD8SXcOBWew9jJOKu4efWmYLv2XjyeKEn9vLMKnC1earXgeUerMzq0Wzy33xqfZXl5foVg6J0fISQfWMyUEFJO6PwIIaWEt719IqcFXKJFhTy8fu0D6zbvD4gWI2qbhXTUKYbZrofn1a7FNvNxellEc1Ooa21dF+tTv7nlWLD9/IkPRjaf3v4PwfZHG8cim3839otge0slzJn77sxvRG1uu/qVYPvNC1sjm78/uiPY1pH4d1hYbwRQ57KYM08ijL8T9yNW+2061QbMfZ6MOwmW9l6w1f2ajHRnb5GjXrXpVcRbzKrIFMf5EUL6hwqQ0+NtgwKdHyEkDc78CCGlhM6PEFJK6Pz6RE6rt7kBD9t3Na4ObIVtd/W2ubBwQGU6DDrMX70hamMDHCmicssZXnUiTIqdPhuL8z8/GwYZ/uM1L0Y2W2ph5ebz7fhg22thQnDLPPZ0ZW06anNsfnOwbZOpAaBSCfdV6s53rGFCcPV8ZIJ1x8Pfc25zfJlX50Kbyoa4iIJcCH9POEnOtpCBTITfu07F30XUh7fi24BVcmaSMyGktDDaSwgpJ3R+hJAywplfv8gpydnrx2ou4mkwbfswuvPV2SRmoy9W5xwNy2iSCxPx+Nrmuff2WHze1Vq4b8uGc/H4DDeNvxHtm22H53VzIx7PawvhvhHzV9GoxA/lf3LTC8H2/3zn1sjmAxvDQqo/PxdXcBixEpqX5Lwx/N4n3o4LnopJRtaRuKCETJt2Tadwqi1cYK8BrzBuJSEROmWFt7WGmh8hpHQoeNtLCCkpdH6EkDIiA5Z9s1JycX4ishvA19FZhu6bqvrlPPoNSFqIqPsCQW5+ntHmbIHKTkPzb280LkCJMVPIYDT8emUhPnZrNBxzsxHrKprwKy28G+aXycazkc3l9XDf3577cGSzqxGuHf3S/FRkc7K9Ltj+kMnrO9OKc+b+/ORvB9vvza6LbN46EeYCtudjHW72ivB3WPeL+PuyhV21GtssrA9/v9r0XGRjNT5dcAqMWu3XXl/edZtwnUbRhZR+Vjs3cMhmfitetDxbJ/MhALcB2AXgThHZtdJ+CSGDg2j6qyis2PkBuBHApKq+oarzAA4C2JNDv4SQQUIl7VUQ8nB+1wA4umj7WLYvQET2icgRETmyAOcWgxAy2GjiqyCsWcBDVQ8AOAAAG2Rzgb4iQghQrFvaFPJwfscBbF+0vS3bt/Y4gm+0MlaCjc7PxzZ2la65ePYqE6HQbx+Mb49NRG0aJ0OblhNssUERrcQT9rkt4XmdPhsHHd4eCwsrbBqJqwJsMpUCJiT+Lix/eyGswDxns7IBTDfDYND5Zmwz1giPtTAbX54jU+Ft1UJcKwJ1E6MZmYmTk9sj4Xco804Cc0rVb2OjF8KgiF09EIiDbm7V5pTgxVoWP9Dhi/bmcdv7AoCdIrJDREYB7AVwOId+CSGDBG97Q1S1KSL3AngSnVSXh1X15RWPjBAyWBTIsaWQi+anqk8AeCKPvgghgwk1v4IRFS1wdBxtJug0C6EmVNngLKlmih/AaHPiFDZYuMIUv3QSciummae9SDNsd+F0I7I53tgYbJ+ajXXBq7eFgtmXf3lbZPPRjf8SbG+shQU933WWm/v74x8KtrdMxIUXpk6HmujIsXpkYxO+6+9FJqgu2ATh+K+2ds5JWLZU4uugK1Wjz847SfXm+tK2ozeudQJzCRl650cIyQnO/AghpWMIo710foSQNDjz6xM5FTN1c6rsoUadwpbmAXZ1FrOJiqA2jGZVjRcVGj0T5raNTMea39S/CfW76mxsM/ZOuO/cuthm+mzYz0IjPs+/+PlNwXatGn/HL+EDwfZ4rbt+NlEPz/Otd7bERtNh7l/FkcKqCQ8HjZy12mv8XVTN4lJoOdfSuVCXFOcabJ9dumispx+3TR6pqzE3EzTJNUQwfAGPPPL8CCFlIMc8PxHZLSKvi8ikiNznfF4Xkceyz58TkWsXfbY/2/+6iNya7dsuIj8QkVdE5GUR+Uy3MdD5EUK6k2NVl8RKUHcDOK2q1wH4KoAHs7a70HmQ4noAuwF8I+uvCeBzqroLwG8BuKdbdSk6P0JIGu3EV3dSKkHtAfBI9v4QgFtERLL9B1V1TlXfBDAJ4EZVfVtVfwwAqjoD4FU4BVYWQ+dHCEliGTO/rRcrOGWvfaarlEpQ79uoahPAFIAtKW2zW+SPAXhuqfMpTsCj1yRP084TlyOcoEhUAdqrCO1Vdw7axP1WTBJsa8x54P+ksWnEP1vL1ENo/Ets0zwTnvvZsXg8I1eHhQ3Wj89GNq/98qpge92G0GZuIT72nEm6tknZANA4Ef4vrsZ1FzDxTjjmtnMF186F31dtKj4HmTPRFK9Ks12JzUlAt0nNFrWru7lGBckhSQ94nFDVG1ZxJJdERNYB+A6Az6qqXesvgDM/Qkh3UoMdaQ4ypRLU+zYiUgOwEcDJpdqKyAg6ju9RVf1ut0HQ+RFCksixjH1KJajDAO7K3t8B4BlV1Wz/3iwavAPATgDPZ3rgtwC8qqpfSRlEcW57CSH9Jac8v0tVghKR+wEcUdXD6Diyb4vIJIBT6DhIZHaPA3gFnQjvParaEpHfBvApAC+JyE+yQ30hK7riMlzOLyEROiXJ2dXzrHZYjx+6hy1EaoshzMaFQduwYl38k9RPh7qRLcQJAOeuDPc1nAf+5zaHmlXbW6TuZ+GqaicujwskVM+F2uHMjNEpHWmsshDubLwbn4Opj4CKV0fV9L3+aKypVewqeVa7A4Ca0X4dzS/S65pO1rVd4c1eXwl6XtI1OQDk+XibVwlKVb+06P0sgE9cou0DAB4w+34E98q7NMPl/Aghq0PBCpWmQOdHCOmKYJnTqgJA50cISYMzP0JIGRm2wgbD5fw8cTmhIm5K4nOU5Owlr54Ps3Kjfk+cjppULgurK3s/iJogSG02PofL/jkU3s9fGSdLV+fCG5fKXHwjs2CKMI8di0dkAyW1X4Xn2QxjJh0bU/xk3fH4L6lZD8czdio+T3vuXkXm6hkTOak51VjqpoLM9NnIJqoA7QROouvCC7p1o8frtmub1Hap0PkRQkoHi5kSQkoLZ36EkDJCza8MpGgwXrML4QP0sj4Uv7QVJ8nKTKg1eUdpV8JE45EpJ1l6JNTdqrNOcQYjWdWnYg1r/D2bpBuPZ35D2E6NtLn+aNyoZXLC7ViAWONrnIzPszJvktYdPU+sVjcTV93GmXCVOvUSmI3GF+l78BKUY5ue6EWrW+0CCXR+hJAywpkfIaR8KFILlRYGOj9CSFeGcQGj4Xd+OT1YXjG5drrg6HdGyNI5o1l1KXwJwH3AvjJldK7R+GerzIZa08RsPD5bBHV+k9PPQvhdqFPAs25KRKpZHa12Pta9qiY/rzUe51aOng6/r1Y9tolXXXN+u6Y5/rmlV1jrdOysoHbB0QoN4qwMF/ThXVv2mlzt/Ly8oPMjhJSRKJBUcOj8CCHdYVUXQkhZoeZHCCklfLytaKQ8IJ4gLrdn58JuvWIIdl/Cim/R+MxxAEQCftXrpxFmEbtBEbNSnNZisd4mEVcTAiethk2wjsdXmwnPq7LgFF6YCgtDVJ1CAnLGFCCoOZfwjLVxghlzCb+n+W20mbASWy8MYnDDgzM/QkjpSF+cqDDQ+RFC0qDzI4SUDSY5F5FV0lO85Fadt0nNRkfy9CnTRqrxamlRcrSnT50LE3KlGS/NJvOhftc4OxvZtBthu/Z4rM3ZIp/1d8Mk4spM3K/Ww3MfOedom3NGU2s72qbNNTMFZIFYz4NzKJjkY+2xIEHXBPmi6HkJSEFWmUtl+J0fIWTlMM+PEFJWhi3VpYcFB/4VEflTEXlNRH4qIt8TkU05jYsQMmho4qsgrHTm9xSA/araFJEHAewH8N9XPqw+k/CguVfYsttCSJEWBQB1k593Nn4IP+r3Qqyp2cKpcPrB2Fi47RTwrBjdTebrkU3NFmxI0ILkvCkS4D0nWjF5dU5hARkxGmQrno5ERSe84gNWi3Mkv7btp9dFhLrZFEQXHLaAx4pmfqr6N6p68Qp5FsC2lQ+JEDJwKDr/sFJeBSFPze9PADx2qQ9FZB+AfQDQwHiOhyWErAXDpvl1dX4i8jSAq5yPvqiqf5XZfBFAE8Cjl+pHVQ8AOAAAG2Rzcf49EELKmeenqn+w1Oci8mkAtwO4RbVAc15CSDoFu6VNYUW3vSKyG8DnAfyOqsbZpkWlx+rP2g4F8ihQ4YnhNujgJEK37apwo07i8fRMaFOPAxV6+kz3fswFLl41YxOYwLxJTnYKCUR/OE6gwgYmvACSmsIP7XNOkrMJRlVGnO/UBDO8pPXo+M6Sc1HgKyV4UZAAh6V0M78u/BmAOoCnpFOB41lV/c8rHhUhZPCg8/tXVPW6vAZCCBlsOPMjhJQPhb9YVIGh8+sVT7exxS+NHuRqWFYzcrTESI9ykpOTLkujJ1r9DEBcTNUTuW0/NunaW6XOGXO3fl2MVuh9p3afl5DerQ3QY5LzEDNsM78VJTkTQkpEjknOIrJbRF4XkUkRuc/5vC4ij2WfPyci1y76bH+2/3URuXXR/odF5F0R+aeUMdD5EUKSEE17de1HpArgIQC3AdgF4E4R2WXM7gZwOosrfBXAg1nbXQD2ArgewG4A38j6A4C/yPYlQedHCOlOalGDtInfjQAmVfUNVZ0HcBDAHmOzB8Aj2ftDAG6RTkrJHgAHVXVOVd8EMJn1B1X9OwCnUk+Jml+OWG0uzgFzHrDv0ofbT0rhBUfPk8q83RH3Y/Q6t1hDin4XHzzc9vSzhH67Fg/1+nbO037Pkb7n9ePQi75YRASApAc8torIkUXbB7InvC5yDYCji7aPAbjJ9PG+TVY4ZQrAlmz/s6btNakDWwydHyEkCUl/wuOEqt6wmmPJA972EkK6k+9t73EA2xdtb8v2uTYiUgOwEcDJxLZJ0PkRQhJIjPSmzQ5fALBTRHaIyCg6AYzDxuYwgLuy93cAeCarHXAYwN4sGrwDwE4Az/dyRnR+hJAk8or2ZjVA7wXwJIBXATyuqi+LyP0i8vHM7FsAtojIJID/CuC+rO3LAB4H8AqA/wvgHlVtAYCI/G8A/wDgwyJyTETuXmoc1PxypFtgwhXDI5v52CR6wN6pXmwCAeLVFrA2Faef+XCMioXIJmpjk7kTigR4gYvIphkfu6fEce8vMqeE5WENcLjkWNVFVZ8A8ITZ96VF72cBfOISbR8A8ICz/87ljIHOjxDSHV1WtLcQ0PkRQtIYLt9H50cISWMZqS6FgM4vT3JYlStFw+q2Styljh3pUyn9OAnC3TS+NG3TselhdbSUY5W9IEFu0PkRQkqHAhiy/yF0foSQrgiUt72EkJLSHq6pH51fnnTTlrxCAraLhLyxJJ3LO3xCrp09By9fMG5i+knQCX1N0uxIOKdLDKh7P9QBlwdvewkhZYW3vYSQckLnRwgpH1y0nBBSRrh6G1kReYnsKZWcvbhJSvKvTSLOqXJyEintevkOGdzIBWp+hJByQudHCCkdCndN6SJD50cISYABDzII9KphpbTLQx9bzfGxaEH/oPMjhJQOBdAarn80dH6EkAR06GbZdH6EkDR420tID+SVwzdks4/CwGgvIaS0cOZHCCkldH6EkNKhCgzZGsV0foSQNDjzI6QHVitQwSrNawedHyGkfOjQRXt7rD0UIiKfExEVka159EcIGTAUUG0nvYrCimd+IrIdwH8A8MuVD4cQMrAM2eNtecz8vgrg8+ikQRKytmg7fpH8Ue0sXZnyKggrmvmJyB4Ax1X1/4lIN9t9APYBQAPjKzksIaQflC3gISJPA7jK+eiLAL6Azi1vV1T1AIADALBBNg/Xt0hICdACzepS6Or8VPUPvP0i8lEAOwBcnPVtA/BjEblRVd/JdZSEkD7DYqbvo6ovAbji4raI/ALADap6IodxEUIGCRY2IISUEQWgQ/Z4Wy55fgCgqtdy1kfIkKLqR9Z7jLaLyG4ReV1EJkXkPufzuog8ln3+nIhcu+iz/dn+10Xk1tQ+Lbk5P0LIcKNtTXp1Q0SqAB4CcBuAXQDuFJFdxuxuAKdV9Tp00ukezNruArAXwPUAdgP4hohUE/sMoPMjhKSR38zvRgCTqvqGqs4DOAhgj7HZA+CR7P0hALdIJ7K6B8BBVZ1T1TcBTGb9pfQZ0BfNbwanTzyth97qsflWAEW/veY5DAZFP4dexv9rvRxoBqeffFoPpT6+2hCRI4u2D2Spbhe5BsDRRdvHANxk+njfRlWbIjIFYEu2/1nT9prsfbc+A/ri/FT18l7bisgRVb0hz/GsNTyHwaDo57CW41fV3WtxnLWEt72EkLXmOIDti7a3ZftcGxGpAdgI4OQSbVP6DKDzI4SsNS8A2CkiO0RkFJ0AxmFjcxjAXdn7OwA8o6qa7d+bRYN3ANgJ4PnEPgOKmOd3oLvJwMNzGAyKfg6FHH+m4d0L4EkAVQAPq+rLInI/gCOqehjAtwB8W0QmAZxCx5khs3scwCsAmgDuUdUWAHh9LjUO0SF7ZIUQQlLgbS8hpJTQ+RFCSkmhnV+Ry+eLyJ+KyGsi8lMR+Z6IbOr3mFJY7iNEg4aIbBeRH4jIKyLysoh8pt9j6pXsyYZ/FJHv93ssRaSwzm8Iyuc/BeAjqvpvAfwMwP4+j6crvTxCNIA0AXxOVXcB+C0A9xTwHC7yGQCv9nsQRaWwzg8FL5+vqn+jqs1s81l08pIGnWU/QjRoqOrbqvrj7P0MOs7jmqVbDR4isg3AHwH4Zr/HUlQK6fwWl8/v91hy4k8A/HW/B5GA91hS4RzHRbJKIR8D8Fyfh9ILX0Pnn/9wlVdeQwY2zy+v8vn9ZKlzUNW/ymy+iM6t2KNrObayIyLrAHwHwGdVdbrf41kOInI7gHdV9UUR+d0+D6ewDKzzG4by+Zc6h4uIyKcB3A7gFi1GwuWyHyEaRERkBB3H96iqfrff4+mBmwF8XET+EEADwAYR+UtV/WSfx1UoCp/kXNTy+SKyG8BXAPyOqr7X7/GkkD1j+TMAt6Dj9F4A8J+6ZdIPEllZpEcAnFLVz/Z5OCsmm/n9N1W9vc9DKRyF1PyGhD8DsB7AUyLyExH5834PqBtZgObiI0SvAni8SI4v42YAnwLw+9n3/pNsBkVKRuFnfoQQ0guc+RFCSgmdHyGklND5EUJKCZ0fIaSU0PkRQkoJnR8hpJTQ+RFCSsn/BwE1sGm6T749AAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 360x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "## Plot the discretized source\n",
    "fig, ax = plt.subplots(figsize=(5, 4))\n",
    "h = ax.hist2d(samples[:, 0], samples[:, 1], bins=bins, cmap='viridis', weights=np.ones(N) / N) # plot freq: https://stackoverflow.com/a/16399202\n",
    "fig.colorbar(h[3], ax=ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6687d441-7a46-4d2f-be96-a3833dc08cd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Get the PMF of the discretized source, and set up the discretized alphabets\n",
    "from ba import bin_edges_to_grid_pts, vectorized_mse, blahut_arimoto\n",
    "\n",
    "n = samples.shape[1]\n",
    "hist_res = np.histogramdd(samples, bins=bins)  # , range=(xlim, ylim))\n",
    "counts = hist_res[0]  # joint counts\n",
    "bin_edges = hist_res[1]  # length-n list of arrays, each array has length bins+1\n",
    "\n",
    "grid_axes = [bin_edges_to_grid_pts(edges) for edges in bin_edges]\n",
    "\n",
    "# Enumerate grid points corresponding to the histogram (using the center of each bin).\n",
    "meshgrid = np.meshgrid(*grid_axes, indexing='ij')  # length-n list, one 'mesh' for each data dimension\n",
    "grid_pts = np.dstack(meshgrid)  # each grid point (n-tuple) resides in the inner-most dimension\n",
    "grid_pts_flat = np.reshape(grid_pts, [-1, n])  # preserve the inner-most dim while flatten the rest\n",
    "counts_flat = counts.ravel()\n",
    "\n",
    "\n",
    "good_pts_ind = (counts_flat != 0)\n",
    "src_alphabet = grid_pts_flat[good_pts_ind]  # remove bins with zero samples from the source alphabet\n",
    "src_dist = counts_flat[good_pts_ind]\n",
    "src_dist /= src_dist.sum()\n",
    "\n",
    "rep_alphabet = grid_pts_flat  # use all bins from the histogram for the reproduction alphabet\n",
    "Rho = vectorized_mse(src_alphabet, rep_alphabet)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "160105b1-d993-40be-8e70-2b3f2ac1c3a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'step': 0, 'ub': 3.702002082010963, 'D': 0.33313611002238447, 'R': 2.7025937519438097}\n",
      "{'step': 10, 'ub': 2.2365682997117156, 'D': 0.333159628843132, 'R': 1.2370894131823196}\n",
      "{'step': 20, 'ub': 2.236505584319688, 'D': 0.33320517472767885, 'R': 1.2368900601366515}\n",
      "{'step': 30, 'ub': 2.236474444165019, 'D': 0.3331970825557947, 'R': 1.2368831964976348}\n",
      "Tolerance reached, terminating after 33 steps\n",
      "{'step': 33, 'ub': 2.2364675408042825, 'D': 0.3331950760606714, 'R': 1.2368823126222686}\n"
     ]
    }
   ],
   "source": [
    "ba_records, log_Q, log_q_y = blahut_arimoto(Rho=Rho, p_x=src_dist, steps=100, lamb=lamb, verbose=True, tol=1e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c819646d-5704-4057-9f83-cccd2e22d1f0",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Estimate a linear lower bound of the R-D curve, by training a neural network \"-log u\" function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b114f1ef-18a3-480c-9232-f16a0d9ab4ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nn_models import make_mlp\n",
    "lb_model = make_mlp(units=[20, 10, 1], activation='selu', name='mlp', input_shape=[data_dim])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "fbad77cc-42a9-428b-8c21-cf1cee2bc4d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_lb(log_u_model, nn_optimizer, train_dataset_iter, args):\n",
    "    # Simplified version of rdlb.train\n",
    "    records = []\n",
    "    from rdlb import optimize_y, compute_Ckobj\n",
    "    M = args.num_Ck_samples\n",
    "    log_M = np.log(float(M))\n",
    "\n",
    "    for step in range(-1, args.last_step):\n",
    "        # Run M many global optimization runs to draw M samples of C_k, in order to estimate the training objective\n",
    "        opt_ys = []\n",
    "        log_Ck_samples = []\n",
    "        x_M_batch = []  # a batch containing M minibatches of data samples, each minibatch consisting of k samples\n",
    "        for j in range(M):  # embarrassingly parallelizable\n",
    "            x = next(train_dataset_iter)\n",
    "            x_M_batch.append(x)\n",
    "\n",
    "            res = optimize_y(log_u_fun=log_u_model, x=x, lamb=args.lamb, num_steps=args.y_steps,\n",
    "                             lr=args.y_lr, tol=args.y_tol, init=args.y_init,\n",
    "                             quick_topn=args.y_quick_topn, verbose=False, chunksize=args.chunksize)\n",
    "            opt_y = res['opt_y']\n",
    "            opt_ys.append(opt_y)\n",
    "            log_Ck_samples.append(res['opt_log_supobj'])\n",
    "\n",
    "        if step == -1:  # the initial run is just to set the log expansion point alpha\n",
    "            # 'log_avg_Ck' here is just the avg of the M samples of Ck in log domain, for numerical reasons\n",
    "            # (and as a R.V., its expected value underestimates the true log E[C_k], just like in IWAE)\n",
    "            log_avg_Ck = tf.reduce_logsumexp(log_Ck_samples) - log_M\n",
    "            log_alpha = prev_log_avg_Ck = log_avg_Ck  # for next iter\n",
    "            continue\n",
    "\n",
    "        # Update log_alpha for next iter\n",
    "        beta = args.beta  # should be in [0, 1)\n",
    "        if abs(prev_log_avg_Ck - log_alpha) <= 10.0:  # heuristic, skips update if the new value would be too extreme\n",
    "            if beta == 0:\n",
    "                log_alpha = prev_log_avg_Ck\n",
    "            else:  # retain beta fraction of its current value, and update alpha with (1-beta) fraction of prev_log_Ck\n",
    "                # alpha = beta * alpha + (1-beta) * prev E[C_k]\n",
    "                log_alpha = tf.reduce_logsumexp(\n",
    "                    [log_alpha + tf.math.log(beta), prev_log_avg_Ck + tf.math.log(1 - beta)])\n",
    "                \n",
    "        # Estimate the RD LB objective and do gradient update\n",
    "        with tf.GradientTape() as tape:\n",
    "            log_Ck_samples = []\n",
    "            log_us = []\n",
    "            for j in tf.range(M):\n",
    "                x = x_M_batch[j]\n",
    "                opt_y = opt_ys[j]\n",
    "                log_Ck, log_u = compute_Ckobj(log_u_fun=log_u_model, x=x, y=opt_y, lamb=lamb)\n",
    "                log_Ck_samples.append(log_Ck)\n",
    "                log_us.append(log_u)  # each is a length k tensor\n",
    "\n",
    "            log_avg_Ck = tf.reduce_logsumexp(log_Ck_samples) - log_M\n",
    "            log_us = tf.concat(log_us, axis=0)\n",
    "            E_log_u = tf.reduce_mean(log_us)\n",
    "            log_E_Ck_est = tf.math.exp(\n",
    "                log_avg_Ck - log_alpha) + log_alpha - 1  # overestimator of log(E[C_k]) by linearization\n",
    "            loss = E_log_u + log_E_Ck_est\n",
    "\n",
    "        prev_log_avg_Ck = log_avg_Ck  # for next iter\n",
    "        trainable_vars = log_u_model.trainable_variables\n",
    "        grads = tape.gradient(loss, trainable_vars)\n",
    "        nn_optimizer.apply_gradients(zip(grads, trainable_vars))\n",
    "        step_rcd = dict(log_alpha=log_alpha, log_avg_Ck=log_avg_Ck, E_log_u=E_log_u,\n",
    "                        log_E_Ck_est=log_E_Ck_est, loss=loss)\n",
    "\n",
    "        print_to_console = (10 * step) % args.last_step == 0\n",
    "        if print_to_console:\n",
    "            str_to_print = f\"step {step}:\\t\\tloss = {loss:.4g}, log_alpha = {log_alpha:.4g}, log_avg_Ck = {log_avg_Ck:.4g}, log_E_Ck_est = {log_E_Ck_est:.4g}, \"\n",
    "            str_to_print += f\"E_log_u = {E_log_u:.4}\"\n",
    "            print(str_to_print)\n",
    "        step_rcd['step'] = step\n",
    "\n",
    "        records.append(step_rcd)\n",
    "\n",
    "        finished = (step + 1 >= args.last_step)\n",
    "        if finished:\n",
    "            break\n",
    "        \n",
    "    return records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "02040570-05d3-42aa-a701-2bd58dcb4649",
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 1024\n",
    "M = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "33fa8b87-ca23-4c03-8f24-2d0e7a8591a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "args = {\n",
    "    \"data_dim\": data_dim,\n",
    "    \"lamb\": lamb,\n",
    "    \"batchsize\": k,\n",
    "    \"num_Ck_samples\": M,\n",
    "\n",
    "    # for the global optimization inner loop\n",
    "    \"y_init\": \"quick\",\n",
    "    \"y_quick_topn\": 3,\n",
    "    \"y_steps\": 1000,\n",
    "    \"y_tol\": 1e-6,\n",
    "    \"y_lr\": 0.01,\n",
    "    \"chunksize\": None,\n",
    "\n",
    "    \"lr\": 5e-4,\n",
    "    \"last_step\": 400,\n",
    "    \"beta\": 0.2,\n",
    "}\n",
    "\n",
    "args = utils.get_args_as_obj(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2d4e309b-1ea4-4014-aa9a-a4490e12354b",
   "metadata": {},
   "outputs": [],
   "source": [
    "nn_optimizer = tf.keras.optimizers.Adam(learning_rate=args.lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "4b294a7e-d8b8-40d5-9d51-f54fb23161c1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 0:\t\tloss = -1.379, log_alpha = -1.49, log_avg_Ck = -1.474, log_E_Ck_est = -1.474, E_log_u = 0.09497\n",
      "step 40:\t\tloss = -1.538, log_alpha = -1.788, log_avg_Ck = -1.815, log_E_Ck_est = -1.814, E_log_u = 0.2759\n",
      "step 80:\t\tloss = -1.647, log_alpha = -1.999, log_avg_Ck = -2.043, log_E_Ck_est = -2.042, E_log_u = 0.3956\n",
      "step 120:\t\tloss = -1.725, log_alpha = -2.195, log_avg_Ck = -2.209, log_E_Ck_est = -2.209, E_log_u = 0.4838\n",
      "step 160:\t\tloss = -1.737, log_alpha = -2.303, log_avg_Ck = -2.317, log_E_Ck_est = -2.317, E_log_u = 0.5798\n",
      "step 200:\t\tloss = -1.762, log_alpha = -2.436, log_avg_Ck = -2.38, log_E_Ck_est = -2.379, E_log_u = 0.617\n",
      "step 240:\t\tloss = -1.845, log_alpha = -2.438, log_avg_Ck = -2.506, log_E_Ck_est = -2.503, E_log_u = 0.6581\n",
      "step 280:\t\tloss = -1.831, log_alpha = -2.538, log_avg_Ck = -2.551, log_E_Ck_est = -2.551, E_log_u = 0.7204\n",
      "step 320:\t\tloss = -1.936, log_alpha = -2.619, log_avg_Ck = -2.656, log_E_Ck_est = -2.656, E_log_u = 0.7195\n",
      "step 360:\t\tloss = -1.897, log_alpha = -2.666, log_avg_Ck = -2.672, log_E_Ck_est = -2.672, E_log_u = 0.7742\n"
     ]
    }
   ],
   "source": [
    "dataset = tf.data.Dataset.from_tensors([])\n",
    "dataset = dataset.repeat()\n",
    "dataset = dataset.map(lambda _: source.sample(args.batchsize))\n",
    "\n",
    "lb_records = train_lb(lb_model, nn_optimizer, iter(dataset), args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "392321bf-bfbd-41b0-bc4d-0e101bab7682",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Run full global optimization procedure to obtain final estimate of the lower bound intercept\n",
    "# (takes a while; can be parallelized)\n",
    "from rdlb import est_R_\n",
    "from copy import copy\n",
    "test_args = copy(args)\n",
    "test_args.y_init = 'exhaustive'\n",
    "test_args.verbose = False\n",
    "lb_result = est_R_(lb_model, lamb, iter(dataset), test_args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "778ae83d-5943-4b30-bc4f-b868451e36f2",
   "metadata": {},
   "source": [
    "## Compare the results with the ground-truth R-D function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "71672925-d2b1-44dd-9d67-bc24b06ac09f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, 'Rate (nats per sample)')"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAX8AAAEKCAYAAAD6q1UVAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8rg+JYAAAACXBIWXMAAAsTAAALEwEAmpwYAABG60lEQVR4nO3dd3hUZfbA8e+bHlJIIPSWgPRAQkBAaSKKBZBiXUVExbaua9n1J7quIiugK6uuiquu4ooi4rrKWlgsiAqCKCXJ0EJLIDQhgUAaaXN+f9zJkEDKBGYyKefzPPMwc+eWcyfDee+89973GBFBKaVU4+Lj7QCUUkrVPk3+SinVCGnyV0qpRkiTv1JKNUKa/JVSqhHS5K+UUo2QnydXboxJA7KBEqBYRAZ4cntKKaVc49Hk7zBSRDJqYTtKKaVcpN0+SinVCBlP3uFrjEkFjgECvC4ib1Qwz53AnQAhISH9e/Tocdbb23Ywm9AgP9pHBp/1OpRSqj5Zv359hoi0qOlynk7+7URkvzGmJfA1cJ+I/FDZ/AMGDJB169ad9fZG/e07urcO49Wb+p/1OpRSqj4xxqw/m/OpHu32EZH9jn8PA58AAz25vdAgf3IKSjy5CaWUahA8lvyNMSHGmLDS58BoYJOntgcQGuhLbkGxJzehlFINgiev9mkFfGKMKd3O+yKyzIPbIzTQj4zsPE9uQimlGgSPJX8R2Q3EeWr9FQkJ9CNHj/xVPVRUVMS+ffs4efKkt0NRdVRQUBDt27fH39/fLeurjev8a01ooB+5hZr8Vf2zb98+wsLCiI6OxvFrWSknESEzM5N9+/YRExPjlnU2qOv8QwP9yDlZjBaoUfXNyZMnad68uSZ+VSFjDM2bN3frL8MGlfxDAv0otgsFxXZvh6JUjWniV1Vx9/ejQSX/sCCrF0v7/ZVSqmoNKvmHB1knQk7kF3k5EqWUqtsaVPJv2sRK/lma/JVSqkoNKvlHBFvJ/3ieJn+laiorK4tXX33VY+t//fXXad26NXFxcXTp0oUFCxY438vPz2fEiBGUlJQ4523Tpg3x8fHExcVx7bXXkpqaCkBhYSHDhw+nuFi7d89Fw0r+TQIAyMov9HIkStU/lSV/EcFuP/eLKGw2GzNmzCApKYlFixbx0EMPOd+bP38+kyZNwtfX1znvzJkzSUxMJCkpiVGjRjFp0iREhICAAEaNGsXixYvPOabGrEFd51965J+lR/6qHnvqs81sOXDCrevs1TacJ8f1rnKe6dOns2vXLuLj4+natSvJyckMGjSI9evXs3TpUkSEsWPHsmmTNUrL3LlzycnJYcaMGbz33nu89NJLFBYWMmjQIF599VVnIi+VnJzM1VdfDUBMTAwBAQHO9xYuXMj7779fbt6bbrrJ+fruu+/mqaeeIj09nY4dOzJhwgQeffTRcvOommlQR/7hwf4Yo8lfqbPxzDPP0KVLFxITE3nuuefYsWMHv/3tb9m8eTOdOnWqdLmtW7eyePFifvzxRxITE/H19WXhwoVnzGez2ejevTsiwiuvvMKsWbMAqxtn9+7dREdHO+fdtGkTvXuXb6yCg4M5duwYALGxsfzyyy9u2OvGq0Ed+fv6GMKD/MnK024fVX9Vd4ReWzp16sTgwYOrnW/58uWsX7+e888/H7D671u2bFlunvT0dLKzs7nyyivZv38/ffv2ZcaMGQBkZGQQERFRbt7Q0FDCw8Od04qKijh48CCdO3cGwNfXl4CAALKzswkLCzvHPW2cGlTyB4ho4q9X+yjlBiEhIeVe+/n5lev7L73bVES45ZZbmDNnTqXrstlsDB8+nG+//ZZjx44RGxvLmjVruPDCCwkODi5356rNZiM2Nrbc8m+//TYXX3xxuURfUFBAUFDQOe1jY9agun3A6vfXbh+lai4sLIzs7OxK32/VqhWHDx8mMzOTgoICPv/8cwBGjRrFRx99xOHDhwE4evQoe/bsKbdscnIy/fr1AyAyMpIbb7yRL774wvm6pKTE2QAkJyeXS/5fffUVc+bMYe7cuc5pmZmZREVFuW2Qs8aowSX/pk0C9MhfqbPQvHlzhgwZQmxsLA8//PAZ7/v7+/PEE08wcOBALr30UkpLrvbq1Yunn36a0aNH07dvXy699FIOHjxYblmbzeZM/gDjxo1j6dKlztejR49m1apVznkXLlxI//79SUhI4J133mHZsmX07NnTOf+KFSsYM2aMW/e/sfFoGceaOtcyjgC/X7SR5H1ZfPfwSDdFpZTnbd26tVxya2w2bNjACy+8wLvvvuvS/JMmTeKZZ56hW7duHo6sbqnoe1Inyzh6g/b5K1X/JCQkMHLkSOdNXlUpLCxkwoQJjS7xu1vDO+Eb7M/x/CLsdsHHR0dJVKq+uO2221yaLyAggClTpng4moavAR75ByACx/XoXymlKtXgkn9UWCAAGTkFXo5EKaXqrgaX/FuEWsn/SLYmf6WUqkzDS/6OI/8jeuSvlFKVarjJX4/8lVKqUg0u+YcH+RHg56PJXymlqtDgkr8xhhahgdrto1QjZ7PZaN26NTabzduh1EkNLvmDdcWPHvkr1bjNnj2b1atXM3v2bG+HUic1uJu8wLriZ9+xPG+HoZTyokWLFpX7V5XXII/8W4QF6nX+SilVhQab/DNzCykuOfe6o0qpumnGjBnlhnkuVVWh+Iq4Wjy+oRWOb5DJv03TIETgV+33V8qt3FXM3ZOqKhRfEVeLx/v7+zeowvENNvkDHMzK93IkStUfaWlp9OjRg5tuuomePXtyzTXXkJeXR1paGt27d2fKlCnExsaSnp7O888/T2xsLLGxsbz44otVLg9UOD9Abm4uY8aMIS4ujtjYWGdife+99xg4cCDx8fHcddddzqPyWbNm0a1bN4YOHUpKSkqF+5GcnEz37t2BMwvFV2ThwoWMHz++3PJli8ncfffdHDp0iPT0dCZMmFBhfeJ6SUTqzKN///7iDtsPnZBOj3wuSzbuc8v6lPK0LVu2eDsESU1NFUBWrVolIiK33nqrPPfcc5KamirGGFmzZo2IiKxbt05iY2MlJydHsrOzpVevXrJhw4ZKl69sfhGRjz76SKZNm+aMISsrS7Zs2SJjx46VwsJCERG555575J133nGuJzc3V44fPy5dunSR55577oz9iIiIkP3794vdbpcnnnhC5s+fX+k+FxQUSKtWrcpNi4yMlOPHj5ebFhMTI4mJiVJcXCxRUVE1/WjdpqLvCbBOziLfNsirfdpEBANwIOtkNXMqVQc98AAkJrp3nfHxUOaIuzIdOnRgyJAhAEyePJmXXnqJa665plwx91WrVjFx4kRnjd9JkyaxcuVKrrrqqgqX9/f3r3D+fv360adPH/7whz/wyCOPMHbsWIYNG8a7775bYUH4o0ePMnHiRJo0aQLAVVdddUb8VRWKr0hNi8c3pMLxDbLbJzTQj/AgPw4e124fpWrCGFPh69OLudd0+cp069aNDRs20KdPHx5//HFmzpzpLAifmJhIYmIiKSkpVSbwskoLxScmJrJ9+3a2bdvGmjVrnO+PGzeu3PxnUzy+oRSOb5BH/gBtI4I5oH3+qj5y4QjdU/bu3cuaNWu44IILeP/99xk6dOgZ8wwbNoypU6cyffp0RIRPPvnEWX6xouWrmv/AgQM0a9aMyZMnExERwZtvvsns2bMZP348Dz74oPOIPzs7m+HDhzN16lQeffRRiouL+eyzz7jrrrvKxVZZofgLL7yQtLQ0oqOjy81ftnh8UFBQpcXjS+sNN6TC8R4/8jfG+BpjNhpjPvf0tspqFxHMfu32UapGunfvzrx58+jZsyfHjh3jnnvuOWOehIQEpk6dysCBAxk0aBDTpk1zJtyKlq9qfpvN5jyx+9RTT/H4449XWhA+ISGB66+/nri4OK644gpnt1BZVRWK37BhAwkJCWcsU5Pi8Q2qcPzZnCioyQN4CHgf+Ly6ed11wldE5E+fJEvfGV+6bX1KeVJdOeHbu3dvry3vaY899pgkJSWdMX39+vUyefJkl9YxceJESUlJcXdoLqs3J3yNMe2BMcAsRyNQazpENuF4fhHH84toGlz/f6Ippc7Nhg0byMzMxM/Pj2uvvZYRI0YA5YvHl17rX5GGVjje033+LwL/B1R6WtwYcydwJ0DHjh3dtuFOza0TVHsz8+jTvqnb1qtUQxUdHc2mTZu8tryn/e9//6v0PVeKxze0wvEe6/M3xowFDovI+qrmE5E3RGSAiAxo0aKF27YfHWVdDpaameu2dSqlVEPhyRO+Q4CrjDFpwAfAxcaY9zy4vXI6NbOO/PdkaPJXSqnTeSz5i8ijItJeRKKBG4BvRWSyp7Z3uuAAX1qHB5GWqUM7K6XU6RrkTV6loqOakKbdPkopdYZaSf4i8p2IjK2NbZUV3TyEPZr8lVLqDA38yD+EjJxCsk8WeTsUpZSqUxp28m9uXfGTlqH9/kopVVaDTv5dWoQCsOtIjpcjUUqpuqVBJ//oqBD8fQ0pv2Z7OxSllKpTGnTy9/f1oUuLULYf0uSvVGNjs9lo3bo1NpvN26HUSdUmf2NMS2PMRGPMvcaY24wxA40x9abR6NYqTI/8lWqEZs+ezerVq5k9e7a3Q6mTKh3bxxgzEpgONAM2AoeBIGAC0MUY8xHwNxE5UQtxnrXurcP4NOkAOQXFhAY22PIFSqnTLFq0qNy/qryqsuGVwB0isvf0N4wxfsBY4FLgPx6KzS26tbLGlNvxazb9OkZ6ORqllKobKu2+EZGHK0r8jveKRWSJiNTpxA/Q3ZH8U7TfX6kGZcaMGcydO/eM6a+//jqtW7cmLi6OLl26sGDBgmrXlZ+fz4gRIygpKXGuo02bNsTHxxMXF8e1115LamoqYA3tPHz4cIqLi927Q7XMlT7/VsaYt4wxyxyvexljbvd8aO7RPjKYJgG+bD1Yp3unlKoXRAS73e7tMKpks9mYMWMGSUlJLFq0iIceqr6UyPz585k0aZJzPH+bzcbMmTNJTEwkKSmJUaNGMWnSJESEgIAARo0axeLFiz29Kx7lyonbfwFfAm0cr7cDD3goHrfz8THEtm2Kbf9xb4eilFvsz8rnlW93lFbKQ0R45dsd7D/HmtVpaWn06NGDm266iZ49e3LNNdeQl5dHWloa3bt3Z8qUKcTGxpKens7zzz9PbGwssbGxvOioOVzZ8kCF8wPk5uYyZswY4uLiiI2NdSbU9957z1ne8a677nIekc+aNYtu3boxdOhQUlJSKtyP5ORkunfvDkBMTAwBAQHV7vvChQsZP358uXWUreV79913c+jQIdLT0wGYMGECCxcudPGTrZtcSf5RIvIhYAerywco8WhUbta3fVM2HzhBUUndPmJRyhWfbNjH3K+2M/PzLYgIMz/fwtyvtvPJhn3nvO6UlBR++9vfsnXrVsLDw3n11VcB2LFjB7/97W/ZvHkzGRkZvP3226xdu5affvqJf/7zn2zcuLHS5devX1/p/MuWLaNt27YkJSWxadMmLr/8crZu3crixYv58ccfSUxMxNfXl4ULF7J+/Xo++OADEhMTWbp0Kb/88kuF+2Cz2ejevbvVKL7yCrNmzapynwsLC9m9e3e54u6bNm2id+/e5eYLDg7m2LFjAMTGxla6/frCleSfa4xpDgiAMWYwUK8Oo/u0b0pBsZ3tesmnagDuHXketw6J5u0f04h5dClv/5jGrUOiuXfkeee87g4dOjBkyBAAJk+e7Cxs3qlTJwYPHgzAqlWrmDhxIiEhIYSGhjJp0iRWrlxZ6fJVzd+nTx++/vprHnnkEVauXEnTpk1Zvnw569ev5/zzzyc+Pp7ly5eze/duVq5cycSJE2nSpAnh4eFcddVVZ8Sfnp5OdnY2V155JS1btmTVqlVMnTq1yn3OyMggIiKi3DpCQ0MJDw93TisqKuLgwYN07twZAF9fXwICAsjOrr85xZXk/xDwKdblnT8CC4D7PBqVm8W1jwDAtq9etVlKVcgYwxNje5Wb9sTYXhhj3LLuil6HhISc0/KV6datGxs2bKBPnz48/vjjzJw5ExHhlltuITExkcTERFJSUpgxY4ZL27fZbAwfPpzExES2b9/Otm3bWLNmDWCd1L3++uud806YMAGwjuhPnjxZbh1lu3wA3n77bS6++GLCwk5VpC0oKCAoKMiluOqiapO/iGwARgAXAncBvUUk2dOBuVOn5k0ID/IjSZO/agBKu3rKKu0COld79+51Jsv333+foUOHnjHPsGHDWLJkCXl5eeTm5vLJJ58wbNiwSpevav4DBw7QpEkTJk+ezMMPP8yGDRsYNWoUH330EYcPHwbg6NGj7Nmzh+HDh7NkyRLy8/PJzs7ms88+OyO25ORk+vXrB0BkZCQ33ngjX3zxBQBJSUnO9woKCggMDHTOV1JS4mwATu/v/+qrr5gzZ065K4syMzOJiorC39//bD9qr6vqJq9JlbzVzRiDiHzsoZjczhhD3/YRJO/L8nYoSp2zeSt2Ort6nhjbi5mfb+HtH9NoHhLA7y7uek7r7t69O/PmzeO2226jV69e3HPPPc4kXCohIYGpU6cycOBAAKZNm0a/fv2cJ4ZPX75JkyYVzg/WUfbDDz+Mj48P/v7+/OMf/6BXr148/fTTjB49Grvdjr+/P/PmzWPw4MFcf/31xMXF0bJlS84///wz4rfZbFxxxRXO1+PGjeP+++9n1qxZzq6k0vn69OnjnG/06NGsWrWKSy65BJvNxnfffcfy5csREXr27MmyZcucJ5EBVqxYwZgxY87ps/Y6EanwAbxdxWN+Zcudy6N///7iKc/+b6t0efQLyS8s9tg2lDpbW7ZscXnefcfy5OXl28Vut4uIiN1ul5eXb5d9x/LOKYbU1FTp3bu315b3tAceeEB27NghIiJPPvmkrF692vne+vXrZfLkyS6va+LEiZKSkuL2GKtT0fcEWCdnkW8rPfIXkVs93/TUnvgOERTbBdv+45wf3czb4Sh11tpFBJc7wjfGnPMRf2MwYcIEHnzwQaKjo2nWrBkXXHCB872EhARGjhxJSUmJ81r/yhQWFjJhwgS6devm6ZA9ykg1/YSOK32eBIZiXfGzCpgpIpnuDmbAgAGybt06d68WgKO5hST85Wsevqy7W66KUMqdtm7dSs+ePb0dhqrjKvqeGGPWi8iAmq7Llat9PgCOAFcD1zie17tb25qFBNC9VRhrU496OxSllPI6V5J/GxH5i4ikOh5PA608HZgnDIxpxvq0oxTrzV5KqUbOleT/lTHmBmOMj+NxHdZwD/XOwJhm5BaWsEXH+VFKNXKuJP87gPeBQsfjA+AuY0y2MaZeZdFBMdaJ3rW7tetHKdW4uXKTV5iI+IiIn+Ph45gWJiLh1S1fl7QMDyImKoS1qW4/V62UUvWKS6WtjDF9geiy80s9usmrrAu7NGfJxv0UFtsJ8Ks31SiVUsqtXBnPfz4wH+tqn3GOx1gPx+UxI7q1ILewhPV7jnk7FKWU8hpXjvwHi0iv6merHy48Lwo/H8P3249wQZfm3g5HKaW8wpV+jzXGmAaT/EMD/RgQHcn32494OxSlVANjs9lo3bo1NpvN26FUy5XkvwCrAUgxxiQbY2zGmHo1qufphndrwdaDJ/j1xMnqZ1ZKKRfNnj2b1atXM3v2bG+HUi1Xkv9bwM3A5Zzq7x/nyaA8bUS3FgB8l3K4mjmVUsp1ixYtonPnzixatMjboVTLleR/REQ+ddzdu6f04fHIPKhXm3DaRwazbNMhb4eilFJe4Ury32iMed8Y8xtjzKTSh8cj8yBjDJf3bs2POzM5cbLI2+EoVWf4+voSHx9PXFwcCQkJrF69GoDQ0NAql0tLSzuj+lV1srKynDWCK7NkyRKMMWzbtq3K+S688MIabftcYqrI66+/TuvWrYmLi6NLly4sWLDA+V5+fj4jRoxwFqF//fXXadOmjfNzvvbaa0lNTQWsEUOHDx9OcXGxW/anKq4k/2CgABhNA7jUs9Tlsa0pLLGzYpt2/ShVKjg4mMTERJKSkpgzZw6PPvqox7blSqJdtGgRQ4cOrbQbRUSw2+3ORqo2YqqIzWZjxowZJCUlsWjRIh566CHne/Pnz2fSpEnOoaJtNhszZ850fs6jRo1i0qRJiAgBAQGMGjWKxYs9P3amK3f43lrB47bqljPGBBljfjbGJBljNhtjnnJPyO6R0DGSlmGB2vWj6p+sdPjhOSgdjl3Eep2V7tbNnDhxgsjIyDOmT5gwgf79+9O7d2/eeOMN5/SSkhLuuOMOevfuzejRo8nPzwfO/FUwd+5cZsyYwfTp09m1axfx8fE8/PDDZ2wnJyeHVatW8dZbb/HBBx84p5dWDJsyZQqxsbHOguul7/Xo0YOpU6fSrVs3brrpJr755huGDBlC165d+fnnn53ref7554mNjSU2NpYXX3wRoMKYKprvdMnJyc5KXzExMQQEBDjfW7hwIePHjy83b9nP4+677+bQoUOkp6c7P9+FCxdWuB23qq7aCxAE3Au8inWz13xcqOQFGCDU8dwfWIt1z4BXKnlV5PFPbNLj8f9JXoFW91LeVZNKXvL9X0WeDBdZ+oiI3W79+2S4Nf0c+fj4SFxcnHTv3l3Cw8Nl3bp1IiISEhLinCczM1NERPLy8qR3796SkZEhqamp4uvrKxs3bhQRkWuvvVbeffddETmzwtdzzz0nTz75ZLWVv9577z257bbbRETkggsucMaSmpoqxhhZs2aNc97S+ErjSE5OlpKSEklISJBbb71V7Ha7LFmyRMaPHy8iIuvWrZPY2FjJycmR7Oxs6dWrl2zYsOGMmCqb73QRERGyf/9+sdvt8sQTT8j8+fNFRKSgoEBatWpVbt7IyEg5fvx4uWkxMTGSmJgoIiLFxcUSFRVV4WfizkpernT7vAu0Bi4DvgfaA9kuNCoiIjmOl/6Ox7lXmHajK2Jbk19Uwrfa9aPqk2F/hEH3wNp/wFMR1r+D7rGmn6PSbp9t27axbNkypkyZckZh+Jdeeom4uDgGDx5Meno6O3bsAKwj3vj4eAD69+9PWlraOcWyaNEibrjhBgBuuOGGcl0/nTp1YvDgwRUuFxMTQ58+ffDx8aF3796MGjUKYwx9+vRxxrRq1SomTpxISEgIoaGhTJo0iZUrV56xLlfmS09PJzs7myuvvJKWLVuyatUqpk6dCkBGRgYRERHl5g0NDSU8/NSwaEVFRRw8eJDOnTsD1nmXgIAAsrOrTbPnxJXkf56I/BnIFZF3gDHAIFdWbozxNcYkAoeBr0Vk7VlH6gGDOjenVXggn2zc5+1QlHKdMXD5nPLTLp9jTXejCy64gIyMDI4cOXVD5Hfffcc333zDmjVrSEpKol+/fpw8ad0vExgY6JzP19fXedLSz88Pu/1UDY3S+U83b9484uPjiY+P58CBA3z77bdMmzaN6OhonnvuOT788ENnQxQSElJp3GXj8PHxcb728fHxyIlUm83G8OHDSUxMZPv27Wzbto01a9YAVmNadn9tNtsZJ8bffvttLr74YsLCwpzTCgoKCAoKcnusZbmS/Esvh8kyxsQCTYGWrqxcREpEJB7r18JAx/LlGGPuNMasM8asK/slqw2+PoYJ/drxXcoRMnIKanXbSp01EVh22onYZY+eOgfgJtu2baOkpITmzU8Ng3L8+HEiIyNp0qQJ27Zt46effqp2Pa1ateLw4cNkZmZSUFDA559/DkBYWFi5o9t7772XxMREEhMT+fzzz7n55pvZs2cPaWlppKenExMTU+HR+dkYNmwYS5YsIS8vj9zcXD755BOGDRt2RkyVzVdWcnIy/fr1AyAyMpIbb7yRL774wvm6pKTE2QCc3t//1VdfMWfOHObOneuclpmZSVRUFP7+/m7Z18q4kvzfMMZEAo8DnwJbgL/WZCMikgWswLpR7PT33hCRASIyoEWLFjVZrVtcndCeYrvwWdKBWt+2Umdl5dxTXT1PZp3qAlo5t9pFq5Ofn+88+r7++ut55513yhU0v/zyyykuLqZnz55Mnz690q6Xsvz9/XniiScYOHAgl156KT169ACgefPmDBkyhNjY2DNO+C5atIiJEyeWm3b11Ve77eaphIQEpk6dysCBAxk0aBDTpk2jX79+Z8RU2Xxl2Wy2ctPGjRvH0qVLna9Hjx7NqlWrnPMuXLiQ/v37k5CQwDvvvMOyZcvK1eVdsWIFY8aMcct+VulsThS48gBaABGO58HASmBsVcvU9gnfUmNfWiljX1rplW0rJVLDE77H9lond+1267Xdbr0+ttczwalzsn79epk8ebLL80+cOFFSUlIqfK9WT/gaY+43xoQby5vGmA3GmNEutCttgBWOcYB+werz//xsGylPmpTQDtv+42w7VK8Kk6nGKqIDDH/4VB+/MdbriA7ejUtVKCEhgZEjRzpv8qpKYWEhEyZMoFu3bh6Py5Vun9tE5ATWTV7Nscb5eaa6hUQkWUT6iUhfEYkVkZnnGKvHTIhvR6CfD+/9VK9HrVBK1VG33XZbue6zygQEBDBlypRaiMi15F96CcGVwAIR2VxmWoMQGRLAuLi2fLJhP9k63INSqhFwJfmvN8Z8hZX8vzTGhAH2apapd24e3IncwhI+2bjf26EopZTHuZL8bwemA+eLSB4QANzq0ai8IK5DBHHtm/Lumj1n3NSiVG3Q752qiru/H66M7WMXkQ1iXa6JiGSKSL0u5lKZyYM7seNwDmt2Z3o7FNXIBAUFkZmZqQ2AqpCIkJmZ6dYbv1yp4dtojItry5z/bePNlalc2CXK2+GoRqR9+/bs27eP2r7RUdUfQUFBtG/f3m3r0+RfRpC/L7deGM3fvt7O1oMn6NkmvPqFlHIDf39/YmJivB2GakSq7PZxjM1TdRWFBmbKBdGEBPjy+ve7vB2KUkp5TJXJX0RKgBRjTMdaisfrmjbx58ZBHfks+SDpR/O8HY5SSnmEK1f7RAKbjTHLjTGflj48HZg33T60Mz4GXtOjf6VUA+VKn/+fPR5FHdO6aRDXDejAh+vSuXtEFzo0a+LtkJRSyq1cudTzeyAN8Hc8/wXY4OG4vO6+i7viYwx/X77D26EopZTbuTKw2x3AR8DrjkntgCUejKlOaN00iCkXdOLjDfvYedizFXWUUqq2udLnfy8wBDgBICI7cLGYS313z0XnEezvy/Nfb/d2KEop5VauJP8CESksfWGM8aOO1eL1lGYhAUwb1pmltkP8knbU2+EopZTbuJL8vzfGPAYEG2MuBf4NfObZsOqOu0Z0pk3TIGZ+tgW7vVG0eUqpRsCV5D8dOALYgLuApVglHRuFJgF+TL+iB7b9x/logxZ6V0o1DNVe6ikidmPMO8BarO6eFGlko09dFdeWBWv28NdlKVwR25qwIM8WVlZKKU9z5WqfMcAu4CXgFWCnMeYKTwdWlxhjeHJcLzJyCnjxG730UylV/7nS7fM3YKSIXCQiI4CRwAueDavu6ds+gpsGdeTtH1NJSs/ydjhKKXVOXEn+2SKys8zr3UCjvPD9kSt60CIskOkf2ygqaXDFzJRSjYgryX+dMWapMWaqMeYWrCt9fjHGTDLGTPJwfHVKeJA/M8fHsvXgCf65cre3w1FKqbPmSvIPAn4FRgAXYV35EwyMA8Z6LLI66rLerbm8d2te/GYHO35tlD+AlFINgKlLF+4MGDBA1q1b5+0wqnUku4DLXvyB1uFBLLl3CAF+rrShSinlfsaY9SIyoKbLadY6Cy3CAnlmUh+2HDzBC9/o0A9KqfpHk/9ZGt27NTec34HXvt/FWi34rpSqZzT5n4M/j+1Fx2ZNeHBxIsdyC6tfQCml6ghXbvK63xgTbixvGWM2GGNG10ZwdV1IoB8v/6YfGTmF3L84kRId+0cpVU+4cuR/m4icAEZjlXS8GXjGo1HVI33bR/DkVb34YfsRXv5W7/5VStUPriR/4/j3SuBdEdlcZpoCbhzYkUkJ7fj78h18l3LY2+EopVS1XEn+640xX2El/y+NMWGA3t5ahjGGWRP60L1VGA8sTmRPZq63Q1JKqSq5kvxvxxrW+XwRyQMCgFs9GlU9FBzgy2uT+wNw279+4Xh+kZcjUkqpyrmS/L8WkQ0ikgUgIpk0woHdXBEdFcJrk/uz92ge9y7coOP/KKXqrEqTvzEmyBjTDIgyxkQaY5o5HtFYRdxVBQZ3bs7siX1YtTODJz/dTF26g1oppUpVVczlLuABoC2wnlMneU9gjeuvKnHtgA7szsjlH9/tol1EMPeOPM/bISmlVDmVJn8R+Tvwd2PMfSLyck1XbIzpACwAWmFVAHvDsc5G4eHR3TmYlc9zX6YQ0cSfmwZ18nZISinl5EoZx5eNMbFAL6wRPkunL6hm0WLgDyKywXGF0HpjzNcisuWcIq4nfHwMz10bx4mTxTy+ZBNNg/0Z27ett8NSSinAtTt8nwRedjxGAn8FrqpuORE5KCIbHM+zga00snMF/r4+zLsxgQGdInlwcSLfbz/i7ZCUUgpw7Wqfa4BRwCERuRWIA5rWZCOOk8T9sIrAn/7encaYdcaYdUeONLzkGBzgy5u3nE/XlmHcuWAdK3c0vH1UStU/riT/fBGxA8XGmHDgMNDB1Q0YY0KB/wAPOIaJKEdE3hCRASIyoEWLFq6u9vSVnN1ytaRpsD/vTRtETFQIt7+zTn8BKKW8ztUyjhHAP7Gu+tkArHFl5cYYf6zEv1BEPj7bIKskAiNHwp/+BHX4l0OzkAAW3TGY81qEcseCdazQYSCUUl5UbfIXkd+KSJaIvAZcCtzi6P6pkjHGAG8BW0Xk+XMPtRLZ2dCyJcyZA9HR8NBDcOCAxzZ3LiJDAnj/jkF0axXKXQvW882WX70dklKqkXJpPH9jTDtjzIVARyDCGDPchcWGYI0AerExJtHxuPIcYq1YeDh8+CFs3gxXXw0vvQQxMXDPPZCW5vbNnauIJgEsvH0wPduEcdd76/n3unRvh6SUaoSqreFrjHkWuB7YApQ4JouIVHvFT025pYbv7t3w7LPwr39BSQlMngyPPgrdu7slRnfJKSjm7nfXs2pnBo9c3oO7R3TG+rGklFKuO9savq4k/xSgr4gUnG1wrnJrAff9+2HuXHj9dTh5Eq69Fh57DOLi3LN+NygstvOHfyfxWdIBbh8aw5+u7ImPjzYASinXebKA+27Av+YheVm7dvDCC1bXz/Tp8L//QXw8jBsHa8+44tQrAvx8+Pv18Uy9MJq3VqVy36KN5BeWVL+gUkqdI1eSfx6QaIx53RjzUunD04G5TcuWMHs27NkDTz0Fq1fD4MFw6aXw/fdev0zUx8fw5LhePHZlD5ZuOsh1r6/h0PGTXo1JKdXwuZL8PwX+AqzGutSz9FG/REbCE09YjcBzz4HNBhddBMOGWb8KvNgIGGO4c3gX/nnzAHYfyWH8vFXY9h33WjxKqYav2j7/2uTWPv/q5OfD/PnWyeH0dEhIsO4VmDABfFy6CMojth48wbR31pGZW8Bz18QxLk7HA1JKVc7tff7GmM+MMeMcN2qd/l5nY8xMY8xtNd1gnREcDPfeCzt3wltvwYkT1qWiffrAwoVQXOyVsHq2Cee/vxtCbNum3LdoI099tpnCYi0Ko5Ryr6oOce8AhgHbjDG/GGOWGmO+NcakAq8D60Vkfq1E6UkBAXDbbbB1K7z/PhhjXR7aowe8+SYUFtZ6SFGhgbx/x2BuHRLN2z+mccMbazh4PL/W41BKNVwudfs4BmZrA+QD2x21fN2uVrt9KmO3w6efwtNPw/r10KEDPPwwTJtm/VqoZZ8nH+CRj5IJ8vflpd/0Y8h5UbUeg1Kq7vLkpZ6ISJqIrBGRRE8l/jrDx8fq9//lF1i2DDp1gt//3rpr+K9/tYaTqEVj+7blv78bSrOQACa/tZbnvtymtYGVUufMe2c26zpj4LLLYOVK65LQvn3hkUesxuCpp+Do0VoL5byWoSy5dwjX9m/PvBW7uOYfq0nNyK217SulGh5N/q4YPhy++gp+/tl6PmOG1QhMnw6/1s7gbCGBfvz1mjj+cVMCaZl5jHlpJR/+kq4F4pVSZ8XVgd2CjTF1a3Acbzj/fFiyBJKTYexYqxsoOhruvx/27auVEK7o04ZlDwwjvkME//efZO55bwNHsj0+8oZSqoFxpYzjOCARWOZ4HW+M+dTDcdVtffrAokWwbRv85jfw6qvQuTPccQfs2uXxzbdpGsx7tw/isSt78O22w4x+4Xv+m7hffwUopVzmypH/DGAgkAUgIolAjMciqk+6dbNuFNuxw0r8775rTZs8GbZ4tk69j491V/DS+4cSHRXC/R8kcseCdTo0hFLKJa4k/yIROX2sAT3ELCs6GubNg9RUq5jMkiUQGwvXXAMbN3p00+e1DOOjuy/kz2N7sWpnBpe+8D0f/LwXu13/REqpyrmS/DcbY24EfI0xXY0xL2ON86NO16aNNW5QWpo1VMQ331jDRowZYw0o5yG+Pobbh8aw7P7h9GoTzvSPbVz3+hq2HjyjZLJSSgGuJf/7gN5AAfA+cBy435NB1XtRUfCXv1iDyM2aZV0lNGSIVWt4+XKPDSIXHRXCojsG89dr+rLrSA5jX17FXz7fQk6Bd4aqUErVXa4Uc7lWRP5d3TR3qBN3+HpCbi688YZVXObAARg0CB5/3PpF4KHqXcdyC/nrlyl88MteWoYF8viYXozt20arhSnVwHjyDt9HXZymKhMSAg8+aF0J9Npr1r0B48ZBv35W/eES9xdwiQwJYM6kPnx8z4VEhQZy36KN3PDGT2zar0NFK6WqOPI3xlwBXAlcBywu81Y40EtEBro7mAZ75H+6oiJrELk5cyAlxaov/Nhj1mWj/u4vmlZiF97/eS8vfL2dY3mFTOrXnocv607rpkFu35ZSqnZ54sj/ALAOOEn5Ii6fApedTZDKwd8fbrkFNm+GxYshMNB63a2bVXO4wL03bfn6GG4e3InvHr6IO4d35rOkA4yc+x0vfL2dvEI9H6BUY+RKn7+/iBTVRjCN5sj/dCLw+efWyeG1a6FtW2sk0TvusLqM3Cz9aB7PLNvGF8kHaRUeyO9HdeW6AR3w99XRPpSqb872yN+V5N8VmAP0Apz9BCLSuaYbq06jTf6lRKyrgWbNgu++s64aevBBq+hM06Zu39y6tKPMXrqVDXuz6NS8CQ9e0o1xcW3x9dGTwkrVF5484fs28A+gGBgJLADeq+mGlAuMgUsugRUrYNUqayyhP/3JGkTuz3+GjAy3bm5AdDP+c8+FzJ86gCYBfjywOJEr/76SLzcf0qEilGrgXEn+wSKyHOtXwh4RmQGM8WxYiiFDYOlSq6DMqFFWcZnoaPjjH+HgQbdtxhjDxT1a8cV9Q3nlxn4Uldi56931THh1NR+v3Mjqt6cjdqt+gNjtrPnXdA7t3eG27SulvMOV5F9gjPEBdhhjfmeMmQiEejguVSohAf7zH9i0ySoy88ILVmGZe++1biJzEx8fw9i+bfnqweH89eq+ZGQX8NAXB5i1vT0vPz+DkuIS1r52Fxek/YPUb99y23aVUt7hSp//+cBWIAL4C9AUeFZE1ro7mEbf5++KXbvgmWfgnXescwQ33wyPPgpdu7p1M0Uldj7ZsI/n/7uGQ8WhdDX7uNfvv7Ro3Z4L73kN46Mnh5WqCzx2wreCDfkCN4jIwppurDqa/GsgPd0aR+if/7SKzF93nXWvQJ8+bt1McXEJXzw1hleLx5MiHenUrAn3XNSFiQntCPTzdeu2lFI15/YTvsaYcGPMo8aYV4wxo43ld8BOrBu/lDd16AAvvWQNIvfHP1qXivbte6r+sBuI3c66N+5mvO8a/hfwKK/7P49vzn6mf2xjyDMreHn5Do7lFrplW0qp2lXVb/d3ge6ADZgGrACuBSaKyPhaiE25olUrePZZq///ySfhhx9g4EAYPdp6fg5+WvAYgw9/yE8tr8M8eZSmbTqz3NzLYx1s9G4bzt++3s4FzyznT5/Y2H0kx007pJSqDVUN72ATkT6O577AQaCjiHisWoh2+7hBdrZVWez55+HwYRg61BpEbvToGg8id2jvDlK/fYvBU2ZjfHwQu52fFjxGzMW307pjV7b/ms1bK1P5ZON+iux2RvVoxR3DYhgY00wHkFOqlri9z98Ys0FEEip77Qma/N0oLw/eesuqM7xvHwwYYDUC48aBm0/WHsku4N01abz70x6O5RXRo3UYkwd3YmK/doQE+rl1W0qp8jyR/EuA3NKXQDCQ53guIhJ+lrFWSpO/BxQWwoIF1hVCu3ZZFcYee8w6Qezr3hO2+YUl/DdxPwvW7GHLwROEBfpxdf/2TB7ckfNahrl1W0opS61d7eNJmvw9qLjYGkRu9myrvnDXrjB9ulVvOCDArZsSETbsPca7a/aw1HaIwhI7F3ZpzpQLOnFJz1b46RhCSrlNnUv+xpj5wFjgsIjEurKMJv9aYLdbNYafftqqL9yxI/zf/8Ftt0FwsNs3l5FTwOJf0nl/7V72Z+XTMiyQq/u357oBHYiJcv+gdUo1NnUx+Q8HcoAFmvzrIBFYtsxqBFavtq4a+uMf4e67IdT9N3CX2IXlW3/lg1/S+S7lMHaBgTHNuH5AB67s04bgAL1nQKmzUeeSP4AxJhr4XJN/HSYC339vjST6zTfQrBk88ADcdx9ERHhkk4eOn+Q/G/bx4bp09mTmERbox7j4tlw/oAN92zfVK4WUqoF6m/yNMXcCdwJ07Nix/x43jlejamjtWqsR+OwzCAuD3/3OGlK6RQuPbE5EWJt6lA9/SWfppoOcLLLTvVUYkxLacVV8W9o0dX83lFINTb1N/mXpkX8dkZRknRj+978hKAjuvNMqLtOuncc2eeJkEZ8lHeDf6/aRmJ6FMXBB5+ZM6NeOy2NbEx7k/vKWSjUEmvyV+23bZl0i+t571mWht94KjzxijSrqQWkZuSxJ3M+SjftJy8wj0M+HS3q1YmJ8O4Z3a0GAn14tpFQpTf7Kc1JTrZvF5s+HkhK48UZrJNGePT26WREhad9xlmzcz2dJB8jMLSSyiT9X9mnDmD5tGBjTTC8bVY1enUv+xphFwEVAFPAr8KSIVDkQvCb/Om7/fvjb36wi8/n5cPXVVqWx+HiPb7qoxM6qHRl8snE/32z9lbzCEqJCA7isd2ttCFSjVueS/9nQ5F9PHDkCf/87vPwynDgBY8ZYQ0cMHlwrm88vLOG7lMN8YTvIt9sOa0OgGjVN/qr2ZWXBK69Y1cWOHoWLL7YagYsuqvEgcmersoZgdO/WjO7Vigu6NNe6A6pB0+SvvCcnx+oKmjsXDh2CCy6wGoErrqi1RgAqbghCA/0Y0b0Fo3u14qLuLWkarFcNqYZFk7/yvpMnrZPCzz4Le/dCv37WIHKTJrl9JNFqQykqYfWuDL7e8itfbzlMRk4Bfj6GC7o059Jerbi0Vyu9j0A1CJr8Vd1RWAgLF8KcObBjh3VV0GOPwQ03gF/tD/Fstwsb07P4asshvt78K7szrMFq+7RryiU9W3Fxj5b0bhuOj4/eWazqH03+qu4pKbFuFJs9G2w26NzZGkl0yhQIDPRaWDsP5/D1ll/5asshEtOzEIGo0EAu6t6Ckd1bMrRrlHYPqXpDk7+qu+x2a8iIWbOs+sLt2ll3DN9xBzRp4tXQMnIK+GH7EVakHOGH7Uc4nl+Er4+hf8dILuphNQY9WofpeEOqztLkr+o+EWvwuFmzrMHkWrSAhx6C3/4Wwt1eG6jGikvsJO3LYsW2I6xIOczmAycAaB0exEXdWzCiWwsu7BJF0yb6q0DVHZr8Vf2ycqXVCHz5pTV66O9/bz2aN/d2ZE6/njjJ9ylH+G77YVZuzyC7oBgfA33aRzDsvCiGnBdFQqcIvZRUeZUmf1U//fKLdU5gyRIICbF+BTz0ELRu7e3IyikqsZOYnsXKHRn8uDODxPQsSuxCsL8vgzo3Y+h5UQztGkX3VtpFpGqXJn9Vv23aZDUCixdbZSWnTbMqjHXo4O3IKnTiZBE/7crkx50ZrNyZwe4j1hVEUaGBDD2vOUO7tuCCLs1pF6GXkyrP0uSvGoYdO6yRRBcssG4QmzLFukLovPO8HVmV9mfl8+PODFY5fhlk5hYC0KFZMINjmjO4c3MGa2OgPECTv2pY9u61RhJ9800oKrLuEXjsMejd29uRVctuF7YdymZtaiY/7c5kbepRsvKKAKsxGFTaGHRuRvtI717tpOo/Tf6qYTp0yBpJ9B//gNxcmDjRGkm0f39vR+Yyu11I+TWbn3af2Ri0jwx2NATNGRTTjPaRwXrOQNWIJn/VsGVmnhpJNCsLLr/cagSGDvV2ZDVmtwvbD2fz065Mftp9lLWpmRxzNAatwgMZ0KkZ/TtFcn50M3q2CdMRSlWVNPmrxuHECZg3zxpJ9MgRGDHCGkRu1KhaHUTOnUp/GaxLO8ovacdYv+cY+7PyAWgS4Et8hwgGRDdjQKdI+nWMIExLWqoyNPmrxiUvD954A557Dg4cgIEDrUZg7Nh62wiUdSArn3V7jrE+7Sjr9hxj68ET2AV8DPRoHc6A6Ej6d4pkQHQzPYncyGnyV41TQQG88451hVBqKvTta50YvuYaq+5wA5F9sojE9CzWpR1j3Z6jbNybRV5hCWB1FcV3iCC+QyTxHSLo274pIYG1P4Ce8g5N/qpxKy6GRYusewW2bYNu3aw6wzfdBP4Nr5ukuMTOtkNWV1FiehaJ6VmkZeYB1q+Dbq3C6NcxwtkonNcyFF8dtbRB0uSvFFgjiX7yCTz9NCQlQadO8MgjcOutEBTk7eg86mhuIUnpWWx0NAZJ6Vkcz7dOJIcG+tGnXVPiO0bQr0ME8R0jaBnWsD+PxkKTv1JlicDSpVYj8NNP0KYN/PGPcNdd1jASjYCIkJqRy8a9Wc5fB1sPnqDYbv2fb9M0iD7tmlqP9ta/zUO9N9S2Ojua/JWqiAisWGE1AitWQFQUPPAA/O530LSpt6OrdSeLSth84Dgb92aRvO84m/Yfdxa3AWjbNMjZEMS20wahPtDkr1R11qyxRhL94gtrCOn77rMagqgob0fmVSdOFrF5/wk27T+OzfFILdMgtIsIJrZduOMXQgR92jWlWUiAFyNWZWnyV8pVGzdajcDHH0NwMNx9N/zhD9C2rbcjqzNKGwTb/ixs+09g23fqhDJYDULPNuH0ahtOrzbh9G4brncne4kmf6VqautWq87w++9bl4Xedpt1cjg62tuR1UnH84vYfOC44xfCCbYePMHuIzk4TiEQFujnbBB6tgmjV5umdG0VSpB/w7nkti7S5K/U2dq9G559Ft5+2zpHcNNN1mWi3bt7O7I6L7+whJRfs9l68ARbDpxgy0GrUSi9B8HXx3Bei1CrMWgbTq82TenZJkzPI7iRJn+lztW+fTB3rnXn8MmTcN111g1jfft6O7J6xW4X9h7NY8tpDcLB4yed87QKD6RH63C6tw6je6swurcO47yW+ivhbGjyV8pdDh+2xg6aNw+ys2HcOGvoiIEDvR1ZvXY0t9D5C2HrwRNsO5TNziM5FBbbAevmtOjmIXRzNAalj07NmujgdlXQ5K+Uux07Bq+8Ai++CEePwiWXWI3A8OENYvyguqC4xE5aZh7bf81m26Fsth/KJuXXbNIycylNTQF+PnRtGer8hdDN8WuhTdMgPcGMJn+lPCc7G157zeoSOnwYhgyxGoHLLtNGwEPyC0vYeTiHlF+zSTl0gpRfc9h+KJtDJ051HYUF+dG9ldVdVPbRtmkwPo1oKAtN/kp5Wn4+vPWWVWEsPd0qKPOnP8H48eCj3RK1ISuvkO2/5jgahGy2H8ph55EcjjrKZoI1DHaXFuUbhK4tQ+nYQLuPNPkrVVsKC+G996xB5HbtskpLPvaYdYLYT0fT9IbMnAJ2HrYagp2HTz3KnmQO8PUhOqoJXVuG0aVMoxATFVKvTzRr8leqthUXw4cfWo3A5s3QpYtVbH7KFAjQO2DrguyTRew6klumQchm5+Ec9h7Nc96f4GOgQ7MmnNcilM4tQoiJsv7tHBVCi7DAOn9eQZO/Ut5it8N//2vdNbx+PXToAP/3f3D77dYdxKrOOVlUQmpGbrlfCbuO5JCakUuB4+ojsEZDjYkKcTQKIXRuEUrnKOt5XamZUCeTvzHmcuDvgC/wpog8U9X8mvxVvSYCX35pDSL344/QqhU89BDccw+EhXk7OuUCu104cDyf1Ixcdh/JJTUj19ko7M/Kp2y6bBUeSOeoUGIcvxJKfzV0iAyu1XMLdS75G2N8ge3ApcA+4BfgNyKypbJlNPmrBkEEfvjB+iXw9dcQGQn33w+//731XNVLJ4tK2JOZx+4jOex2Ng7W86y8Iud8fj6Gjs2b0DkqhI7NQoiOakLHZk3o1DyE9pHB+Lu5YaiLyf8CYIaIXOZ4/SiAiMypbBlN/qrB+flnqxH49FPr6P/ee+HBB6FlS29HptzoWG4huzNy2H0kl90ZuaQ6fjXsOZrLyaJT3Ui+Poa2EUF0ahZCx+ZNiG7ehI7NQujUvAk9Woed1fmFupj8rwEuF5Fpjtc3A4NE5HeVLaPJXzVYycnWieEPP7Qqit15p1Vcpn17b0emPEhEOJJdwJ6jeaRl5LL3aB5pmXnszcxlz9E85y+GZiEBbPjzpWe1jbNN/l4/Y2GMuRO4E6Bjx45ejkYpD+nbFz74AGbOtIrNz5sHr75qlZd85BHo3NnbESoPMMbQMjyIluFBnB/d7Iz3j+cVsedo+W6j2uLJsxL7gQ5lXrd3TCtHRN4QkQEiMqBFixYeDEepOqBbN5g/H3buhGnT4F//sqbdfDNsqfR0mGqgmjbxp2/7CIZ3q/3c58nk/wvQ1RgTY4wJAG4APvXg9pSqPzp1so78U1Otk8EffwyxsXDNNVaxGaU8zGPJX0SKgd8BXwJbgQ9FZLOntqdUvdS2Lfztb7Bnj3WX8NdfQ0ICjBkDq1d7OzrVgNWpm7yMMUeAPWe5eBSQ4cZw6gLdp/pB96l+aKj7FCIiNe43qlPJ/1wYY9adzRnvukz3qX7QfaofdJ/Ka3hD3CmllKqWJn+llGqEGlLyf8PbAXiA7lP9oPtUP+g+ldFg+vyVUkq5riEd+SullHKRJn+llGqE6l3yN8ZcboxJMcbsNMZMr+D9QGPMYsf7a40x0V4Is0Zc2KeHjDFbjDHJxpjlxphO3oizJqrbpzLzXW2MEWNMnb8Ez5V9MsZc5/hbbTbGvF/bMdaUC9+9jsaYFcaYjY7v35XeiNNVxpj5xpjDxphNlbxvjDEvOfY32RiTUNsx1pQL+3STY19sxpjVxpg4l1YsIvXmgVUUZhfQGQgAkoBep83zW+A1x/MbgMXejtsN+zQSaOJ4fk9D2CfHfGHAD8BPwABvx+2Gv1NXYCMQ6Xjd0ttxu2Gf3gDucTzvBaR5O+5q9mk4kABsquT9K4H/AQYYDKz1dsxu2KcLy3znrnB1n+rbkf9AYKeI7BaRQuADYPxp84wH3nE8/wgYZep2Ec5q90lEVohInuPlT1iD5NVlrvydAP4CPAucrOC9usaVfboDmCcixwBE5HAtx1hTruyTAOGO502BA7UYX42JyA/A0SpmGQ8sEMtPQIQxpk3tRHd2qtsnEVld+p2jBvmhviX/dkB6mdf7HNMqnEes8YWOA81rJbqz48o+lXU71pFLXVbtPjl+bncQkS9qM7Bz4MrfqRvQzRjzozHmJ0cZ07rMlX2aAUw2xuwDlgL31U5oHlPT/2/1jcv5wevj+SvXGWMmAwOAEd6O5VwYY3yA54GpXg7F3fywun4uwjr6+sEY00dEsrwZ1Dn6DfAvEfmbozrfu8aYWBGxV7egql3GmJFYyX+oK/PXtyN/V2oEOOcxxvhh/VTNrJXozo5LdQ+MMZcAfwKuEpGCWortbFW3T2FALPCdMSYNq+/10zp+0teVv9M+4FMRKRKRVKwa1l1rKb6z4co+3Q58CCAia4AgrMHE6iuX/r/VN8aYvsCbwHgRcSnf1bfk70qNgE+BWxzPrwG+FceZkDqq2n0yxvQDXsdK/HW9Hxmq2ScROS4iUSISLSLRWP2UV4lIXa7h6cp3bwnWUT/GmCisbqDdtRhjTbmyT3uBUQDGmJ5Yyf9IrUbpXp8CUxxX/QwGjovIQW8HdS6MMR2Bj4GbRWS7ywt6+0z2WZz5vhLriGoX8CfHtJlYyQOsL+e/gZ3Az0Bnb8fshn36BvgVSHQ8PvV2zOe6T6fN+x11/GofF/9OBqs7awtgA27wdsxu2KdewI9YVwIlAqO9HXM1+7MIOAgUYf0Sux24G7i7zN9onmN/bfXke1fdPr0JHCuTH9a5sl4d3kEppRqh+tbto5RSyg00+SulVCOkyV8ppRohTf5KKdUIafJXSqlGSJN/I2KMKTHGJDpGnEwyxvzBcbctxpgBxpiXqlg22hhz41lss9xy1W2nhus2xphvjTHh1c/dOBljcry03W+MMZHe2LZyjSb/xiVfROJFpDdwKdYIgE8CiMg6Efl9FctGAzVK/o47rMst58J2auJKIElETrhpffWC43Otje34VvW6kmWM44DiXawRdlVd5e0bGPRRew8g57TXnbGGvjBYd6Z+7pg+glM3jGzEGo7hJ6xB8hKBB7Fupnsb60aZjcBIx7JTse6i/Bb4voLlym6nGdZdscmO+fo6ps8A5mPd/LUb+H0l+/M+cJHjeTSwDfgX1k1LC4FLsG5Q2gEMdMwX4lj3z464x5dZfiWwwfG40DH9IkccHznWvxBH+dPTYvk91s1dycAHjmnNga+AzVg34uzBGhohmjLD8wJ/BGY4nt+BdedtEvAfTg3l/S/gNWAt1o1kXYBlwHpH3D0c88UAaxx/l6c57W9eZpuTHZ9BItbd476l3xHgb47tD63g9UPAJsfjgTKfXQqwwLGvnYBIKhmCWB914+H1APRRi3/sChIBkAW0onxS/gwY4ngeijVgmfN9x/Q/APMdz3tgDQMQhJX89wHNHO+dvlzZ7bwMPOl4fjGQ6Hg+A1gNBDqSZSbgX0Hse4Awx/NooBjog/WLdj1WkjdYw/guccw3G5jseB6B1VCEAE2AIMf0rjjuknTEexxrDBgfR2IdWkEsB4DA0vU6/n0JeMLxfAzW8MjVJf/mZaY/DdzneP4v4HNOJenlQFfH80FYw5iAY/gCx/N7K/mb93T8jf0dr18ts4wA15WZ1/ka6I/VqIRgfS82A/0c+2MHBp+2nR1l90cfdeuh3T6qIj8Czxtjfo+VyIormGco8B6AiGzDSsTdHO99LSJVjaledh3vOtbxLdC8TP/9FyJSICIZwGGsBup0zUQku8zrVBGxiTXi5GZguVhZyIaVoABGA9ONMYlYR/RBQEfAH/inMcaGNTxIrzLr/VlE9jnWm1hmXWUlAwsdI6+Wfl7DOfUZfYF1C351Yo0xKx1x3AT0LvPev0WkxBgTilXA49+O/XgdKB2TfgjWcADg+GwrMAorkf/iWH4U1q9AgBKsXxxU8Hoo8ImI5IpIDtZ4MsMc7+0Ra3z8sg4DbavfZeUNOqRzI2aM6Yz1n/sw1tEgACLyjDHmC6w+9R+NMZfVcNW5bgiv7MilJVT8XS02xvjIqeGFyy5jL/PaXmZ5A1wtIillV2SMmYE1flIc1hF+2QIzrsQyBivZjwP+ZIzpU/FuWXFT/nxbUJnn/wImiEiSMWYqjoHiHEo/Vx8gS0TiK1l/dWO2GOAdEXm0gvdOikhJFa8rU9HfPAjId2FZ5QV65N9IGWNaYPUhv+I4Oi77XhfHEfSzWP3PPYBsrL7/UiuxjkwxxnTDOnoul1AdTl+urLLruAjIkJqdvE3h1BGrq74E7iut7uYYMRWsob8POhqSm7FKHLrEcYKzg4isAB5xrCsUq0TljY55rsDqBwerkWlpjGlujAkExpZZXRhw0Bjjj+OzOZ3jM0o1xlzrWLcpU7f1R6zROalseawuo2uMMS0dyzczrtWFXglMMMY0McaEABMd087g+HxbA2kurFd5gSb/xiW49FJPrJFCvwKeqmC+B4wxm4wxyVgjCf4Pq1ujxHGJ6INY/cQ+ju6JxcBUqbjOwOnLlTUD6O/YzjOcGorbVV9Q/sjYFX/B6uJJdnwOf3FMfxW4xRiThNXY1eTXiy/wnuOz2Ai8JFYBl6eA4Y7tTMI6L4KIFGGNnPkz8DXWieRSf8Y6qfvjadNPdxNwuyPezZwqv3g/cK8jlgorVInIFuBx4CvHZ/81p7qNKiUiG7B+mfzsiPFNEdlYyez9gZ8q6TJUdYCO6qnqLWPVXl0gIpd6OxZXOArXDHCcx2jQjDF/xxp6fLm3Y1EV0yN/VW+JVYTjn3qTV520SRN/3aZH/kop1Qjpkb9SSjVCmvyVUqoR0uSvlFKNkCZ/pZRqhDT5K6VUI/T/9LO8CVLpJQwAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "figsize = (6, 4)\n",
    "fig = plt.figure(figsize=figsize)\n",
    "\n",
    "gDs, gRs = utils.diag_gaussian_rdf(var, num_points=1000)\n",
    "plt.plot(gDs, gRs, label='true $R(D)$')\n",
    "\n",
    "results = {}\n",
    "results['proposed $\\\\hat R_U(D)$'] = {'R': ub_records['rate'][-1], 'D': ub_records['mse'][-1]}\n",
    "results['Blahut-Arimoto $\\\\hat R(D)$'] = {'R': ba_records[-1]['R'], 'D': ba_records[-1]['D']}\n",
    "results['proposed $\\\\hat R_L(D)$'] = {'lamb': lamb, 'R_': float(lb_result['R_'])}\n",
    "\n",
    "# plot point on the upper bound\n",
    "for key, res in results.items():\n",
    "    if key == 'proposed $\\\\hat R_L(D)$':\n",
    "        xs = np.linspace(0, 0.8, 100)\n",
    "        y = -res['lamb'] * xs + res['R_']  # nats\n",
    "        plt.plot(xs, y, label=key, color='red')\n",
    "        \n",
    "    else:\n",
    "        plt.scatter(res['D'], res['R'], label=key, s=30, marker='x')\n",
    "\n",
    "# plt.xlim(0, 0.6)\n",
    "plt.ylim(-0.1, 5)\n",
    "\n",
    "plt.legend()\n",
    "plt.xlabel('Distortion (mean squared error)')\n",
    "plt.ylabel('Rate (nats per sample)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "1e448736-0ae6-4856-a604-0946b6d0e6dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'proposed $\\\\hat R_U(D)$': {'R': 1.234236478805542, 'D': 0.3338179588317871}, 'Blahut-Arimoto $\\\\hat R(D)$': {'R': 1.2368823126222686, 'D': 0.3331950760606714}, 'proposed $\\\\hat R_L(D)$': {'lamb': 3.0, 'R_': 1.9364495277404785}}\n"
     ]
    }
   ],
   "source": [
    "print(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76fe2b7b-ffae-4701-9457-a61d926fe945",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The points from the proposed upper bound and BA are visually indistinguishable on the plot."
   ]
  }
 ],
 "metadata": {
  "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}