JSAIT'24_source_code_Variational_Entropy_Constrained_Vector_Quantizers.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ezgimez/neural-distributed-compressor-JSAIT2024/blob/main/JSAIT'24_source_code_Variational_Entropy_Constrained_Vector_Quantizers.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Code definitions"
      ],
      "metadata": {
        "id": "W9ct6--c_JXP"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IqR2PQG4ZaZ0"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "from flax import linen as nn\n",
        "from flax.training import train_state\n",
        "import jax.numpy as jnp\n",
        "import jax\n",
        "from jax import random\n",
        "import numpy as np\n",
        "import optax\n",
        "from google import colab\n",
        "import os\n",
        "from argparse import Namespace\n",
        "from collections import defaultdict\n",
        "from typing import Callable\n",
        "import numpy as np\n",
        "\n",
        "from tensorflow_probability.substrates import jax as tfp\n",
        "tfpd = tfp.distributions\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VboeEquG3dyF"
      },
      "outputs": [],
      "source": [
        "class MLP(nn.Module):\n",
        "  \"\"\"Used for neural decoder  (see Eqs. (11) and (12) in the paper) function. \"\"\"\n",
        "  output_dims: int\n",
        "  activation_fn: Callable\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, *x):\n",
        "    x = jnp.concatenate(x, -1)\n",
        "    x = nn.Dense(100, name=\"fc1\")(x)\n",
        "    x = self.activation_fn(x)\n",
        "    x = nn.Dense(100, name=\"fc2\")(x)\n",
        "    x = self.activation_fn(x)\n",
        "    x = nn.Dense(self.output_dims, name=\"fc3\")(x)\n",
        "    return x\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def sample_xyn_conditional(num_samples, x, var_n, n_rng):\n",
        "  dist_n = tfp.distributions.Normal(loc=0., scale=var_n ** .5)\n",
        "  n = dist_n.sample(seed=n_rng, sample_shape=(num_samples, 1, 1))\n",
        "  y = x + n # corr_pattern: x = y + n, quadratic-Gaussian setup.\n",
        "  y  = y / (1+2*var_n)\n",
        "  # conditional sampling at the encoder, see Eqs. (14) and (15) in the paper.\n",
        "  return y\n",
        "\n",
        "def sample_yxn_conditional(num_samples, x, var_n, n_rng):\n",
        "  dist_n = tfp.distributions.Normal(loc=0., scale=var_n ** .5)\n",
        "  n = dist_n.sample(seed=n_rng, sample_shape=(num_samples, 1, 1))\n",
        "  y = x + n # corr_pattern: y = x + n, quadratic-Gaussian setup.\n",
        "  # conditional sampling at the encoder, see Eqs. (14) and (15) in the paper.\n",
        "  return y\n",
        "\n",
        "def sample_laplacian_conditional(x):\n",
        "  signed_y = jnp.sign(x) # x : Laplacian(0, 1) and y : sign(x)\n",
        "  # conditional sampling at the encoder, see Eqs. (14) and (15) in the paper.\n",
        "  return signed_y[None, :, :]\n",
        "\n",
        "def sample_xyn(num_samples, var_n, xy_rng):\n",
        "  dist = tfp.distributions.Normal(loc=0., scale=[1., var_n ** .5])\n",
        "  yn = dist.sample(seed=xy_rng, sample_shape=(num_samples, 1))\n",
        "  y = yn[..., 0]\n",
        "  x = y + yn[..., 1] # x = y + n, quadratic-Gaussian setup\n",
        "  # (see Section IV.A in the paper).\n",
        "  return x, y\n",
        "\n",
        "def sample_yxn(num_samples, var_n, xy_rng):\n",
        "  dist = tfp.distributions.Normal(loc=0., scale=[1., var_n ** .5])\n",
        "  xn = dist.sample(seed=xy_rng, sample_shape=(num_samples, 1))\n",
        "  x = xn[..., 0]\n",
        "  y = x + xn[..., 1] # y = x + n, quadratic-Gaussian setup.\n",
        "  # (see Section IV.A in the paper)\n",
        "  return x, y\n",
        "\n",
        "def sample_signed_laplacian(num_samples, var_n, xy_rng):\n",
        "  dist = tfp.distributions.Laplace(loc=0., scale=var_n)\n",
        "  x = dist.sample(seed=xy_rng, sample_shape=(num_samples, 1))\n",
        "  y = jnp.sign(x) # 'signed' Laplacian experimental setup.\n",
        "  # (see Section IV.B in the paper)\n",
        "  return x, y\n",
        "\n",
        "def sample_xy(num_samples, xy_rng):\n",
        "  if corr_pattern == \"y=x+n\":\n",
        "    return sample_yxn(num_samples, var_n, xy_rng)\n",
        "  elif corr_pattern == \"x=y+n\":\n",
        "    return sample_xyn(num_samples, var_n, xy_rng)\n",
        "  elif corr_pattern == \"signed_laplacian\":\n",
        "    return sample_signed_laplacian(num_samples, var_n, xy_rng)\n",
        "  assert False\n",
        "\n",
        "def sample_conditional(repeat_num, init_x, sampling_rng):\n",
        "  if corr_pattern == \"y=x+n\":\n",
        "    # shape: (N, B, 1)\n",
        "    return sample_yxn_conditional(repeat_num, init_x, var_n, sampling_rng)\n",
        "  elif corr_pattern == \"x=y+n\":\n",
        "    # shape: (N, B, 1)\n",
        "    return sample_xyn_conditional(repeat_num, init_x, var_n, sampling_rng) #\n",
        "  elif corr_pattern == \"signed_laplacian\":\n",
        "    # shape: (1, B, 1)\n",
        "    # basically, N=1 because there doesn't need to be sampling\n",
        "    # since it is deterministic correlation pattern in this case.\n",
        "    return sample_laplacian_conditional(init_x)\n",
        "  assert False\n"
      ],
      "metadata": {
        "id": "6NCYZ5dnhema"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ps3hxStOB9aX"
      },
      "outputs": [],
      "source": [
        "class WynerZivModel(nn.Module):\n",
        "  source_dims: int\n",
        "  latent_dims: int\n",
        "  var_n: float\n",
        "  num_y_samples: int\n",
        "  lmbda: float\n",
        "  corr_pattern: str\n",
        "\n",
        "  def setup(self):\n",
        "      # this is 'marginal' entropy bottleneck, just learnable parameters.\n",
        "    self.logits = self.param(\n",
        "          \"logits\", jax.nn.initializers.uniform(1e-1), (self.latent_dims,))\n",
        "\n",
        "    self.g = MLP(self.source_dims, nn.leaky_relu) # this is the decoder.\n",
        "\n",
        "\n",
        "  def encode(self, x, encoder_rng):\n",
        "    \"\"\"Implements the marginal encoder function (see Eq. (18) in the paper.\n",
        "    See Section IV.A in the paper for discussion on this compressor.\"\"\"\n",
        "    # shape: (B, 1)\n",
        "    B = x.shape[0]\n",
        "    K = self.latent_dims\n",
        "    N = self.num_y_samples\n",
        "\n",
        "    if self.corr_pattern == \"y=x+n\":\n",
        "      # shape: (N, B, 1)\n",
        "      y_given_x = sample_yxn_conditional(N, x, self.var_n, encoder_rng)\n",
        "\n",
        "    elif self.corr_pattern == \"x=y+n\":\n",
        "      # shape: (N, B, 1)\n",
        "      y_given_x = sample_xyn_conditional(N, x, self.var_n, encoder_rng)\n",
        "\n",
        "    elif self.corr_pattern == \"signed_laplacian\":\n",
        "      # shape: (1, B, 1)\n",
        "      y_given_x = sample_laplacian_conditional(x)\n",
        "\n",
        "    else:\n",
        "      raise ValueError(\"Invalid correlation type. Only `y=x+n`, `x=y+n` and \"\n",
        "                       \"`signed_laplacian` correlation patterns are supported.\")\n",
        "\n",
        "    # shape: (K,)\n",
        "    rates = jax.nn.log_softmax(self.logits) / -jnp.log(2.)\n",
        "\n",
        "    # shape: (K, K)\n",
        "    all_possible_indices = jnp.eye(K)\n",
        "    # shape: (B*N, K, K)\n",
        "    all_possible_indices_broadcasted = jnp.broadcast_to(\n",
        "        all_possible_indices, (B*N, K, K))\n",
        "    # shape: (B*N*K, K)\n",
        "    all_possible_indices_reshaped = all_possible_indices_broadcasted.reshape(\n",
        "        (B*N*K, K))\n",
        "\n",
        "    # shape: (N, B, K, 1)\n",
        "    y_given_x_broadcasted = jnp.broadcast_to(\n",
        "        y_given_x[:, :, None, :], (N, B, K, 1))\n",
        "    # shape: (B*N*K, 1)\n",
        "    y_given_x_reshaped = y_given_x_broadcasted.reshape((B*N*K, 1))\n",
        "\n",
        "    # shape: (B*N*K, 1)\n",
        "    reconstructed_values = self.g(\n",
        "        all_possible_indices_reshaped, y_given_x_reshaped)\n",
        "    # shape: (N, B, K, 1)\n",
        "    reconstructed_values_reshaped = reconstructed_values.reshape(\n",
        "        (N, B, K, 1))\n",
        "    # shape: (N, B, K, 1)\n",
        "    squared_dist = jnp.square(x[:, None, :] - reconstructed_values_reshaped)\n",
        "    # shape: (N, B, K)\n",
        "    summed_dist = jnp.sum(squared_dist, axis=-1)\n",
        "    # shape: (B, K)\n",
        "    distortion_estimated = jnp.mean(summed_dist, axis=0)\n",
        "\n",
        "    # shape: (B, K)\n",
        "    all_rd = rates + self.lmbda * distortion_estimated\n",
        "    # shape: (B,)\n",
        "    z = jnp.argmin(all_rd, axis=-1)\n",
        "\n",
        "    distortion_estimated_passed = jnp.choose(\n",
        "        z, distortion_estimated.T, mode='wrap')\n",
        "\n",
        "    return z, rates[z], distortion_estimated_passed\n",
        "\n",
        "  def decode(self, z, y):\n",
        "    \"\"\"Implements the neural decoder function in the paper.\"\"\"\n",
        "    z = nn.one_hot(z, self.latent_dims, axis=-1)\n",
        "    return self.g(z, y)\n",
        "\n",
        "  def decode_visualization(self, z, y):\n",
        "    return self.g(z, y)\n",
        "\n",
        "  def __call__(self, x, y, encoder_rng, training):\n",
        "    z, rates, distortion = self.encode(x, encoder_rng)\n",
        "\n",
        "    rate = jnp.mean(rates)\n",
        "\n",
        "    if not training:\n",
        "      x_hat = self.decode(z, y)\n",
        "      distortion = jnp.sum(jnp.square(x - x_hat), axis=-1)\n",
        "\n",
        "    distortion = jnp.mean(distortion)\n",
        "\n",
        "    return dict(distortion=distortion, rate=rate)\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "@jax.jit\n",
        "def pretrain_step(state, init_x, var_n, sampling_rng, init_z):\n",
        "  # init_y_noisy shape: (N, B, 1)\n",
        "  init_y_noisy = sample_conditional(repeat_num, init_x, sampling_rng)\n",
        "  # init_y_noisy shape: (B*N, 1)\n",
        "  init_y_noisy = init_y_noisy.reshape((init_x.shape[0]*repeat_num, 1))\n",
        "\n",
        "  def loss_fn(params):\n",
        "    m = model()\n",
        "    x_hat = m.apply({'params': params}, method=m.decode,\n",
        "                    z=init_z, y=init_y_noisy)\n",
        "    init_x_repeated = jnp.tile(init_x, (repeat_num, 1))\n",
        "    loss = jnp.square(init_x_repeated - x_hat).sum(axis=-1).mean()\n",
        "    return loss, dict(distortion=loss)\n",
        "\n",
        "  grads, aux = jax.grad(loss_fn, has_aux=True)(state.params)\n",
        "  return state.apply_gradients(grads=grads), aux"
      ],
      "metadata": {
        "id": "d4WtOAV2_ukr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "@jax.jit\n",
        "def train_step(state, data_rng, encoder_rng):\n",
        "  x, y = sample_xy(batch_size, data_rng)\n",
        "  def loss_fn(params):\n",
        "    m = model()\n",
        "    result = m.apply(\n",
        "        {'params': params}, x=x, y=y, encoder_rng=encoder_rng, training=True)\n",
        "    r = Namespace(**result)\n",
        "    loss = r.rate + m.lmbda * r.distortion\n",
        "    result.update(loss=loss)\n",
        "    return loss, result\n",
        "  grads, aux = jax.grad(loss_fn, has_aux=True)(state.params)\n",
        "  return state.apply_gradients(grads=grads), aux\n"
      ],
      "metadata": {
        "id": "YF14g7zVmyk4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "@jax.jit\n",
        "def eval_step(state, x, y):\n",
        "  m = model()\n",
        "  result = m.apply(\n",
        "      {'params': state.params}, x=x, y=y, training=False,\n",
        "      encoder_rng=random.PRNGKey(0))\n",
        "  r = Namespace(**result)\n",
        "  loss = r.rate + m.lmbda * r.distortion\n",
        "  return dict(val_rate=r.rate, val_distortion=r.distortion, val_loss=loss)\n"
      ],
      "metadata": {
        "id": "vE7Jhzjlm0P6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "@jax.jit\n",
        "def encoder_behaviour(state, x, encoder_rng):\n",
        "  return model().apply({'params': state.params}, method=\"encode\", x=x, encoder_rng=encoder_rng)\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def decoder_behaviour(state, z, y):\n",
        "  return model().apply({'params': state.params}, method=\"decode_visualization\", z=z, y=y)\n"
      ],
      "metadata": {
        "id": "0hiyiAn46fZV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def plot_history(history):\n",
        "  _, axs = plt.subplots(1, 4, figsize=(25, 5))\n",
        "\n",
        "  for k in [\"rate\", \"val_rate\"]:\n",
        "    axs[0].plot(history[k], label=k)\n",
        "  axs[0].legend(loc=\"best\")\n",
        "  axs[0].grid()\n",
        "  axs[0].set_title(\"entropy\")\n",
        "  axs[0].set_yscale(\"log\")\n",
        "\n",
        "  for k in [\"distortion\", \"val_distortion\"]:\n",
        "    axs[1].plot(history[k], label=k)\n",
        "  axs[1].legend(loc=\"best\")\n",
        "  axs[1].grid()\n",
        "  axs[1].set_title(\"distortion\")\n",
        "  axs[1].set_yscale(\"log\")\n",
        "\n",
        "  for k in [\"loss\", \"val_loss\"]:\n",
        "    axs[2].plot(history[k], label=k)\n",
        "  axs[2].legend(loc=\"best\")\n",
        "  axs[2].grid()\n",
        "  axs[2].set_title(\"loss\")\n",
        "  axs[2].set_yscale(\"log\")\n",
        "\n",
        "  for k in [\"lr\"]:\n",
        "    axs[3].plot(history[k], label=k)\n",
        "  axs[3].legend(loc=\"best\")\n",
        "  axs[3].grid()\n",
        "  axs[3].set_yscale(\"log\")\n"
      ],
      "metadata": {
        "id": "Epc7lJ2eY0HN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Experiment"
      ],
      "metadata": {
        "id": "XI1O6JGN_QUg"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KM4yiWiL28de"
      },
      "outputs": [],
      "source": [
        "corr_pattern = \"signed_laplacian\"\n",
        "# options for `corr_pattern` are : {y=x+n, x=y+n, signed_laplacian}\n",
        "# (see Section V in the paper for discussion on different correlation patterns)\n",
        "\n",
        "var_n = 1.0\n",
        "num_latents = 32\n",
        "lmbda = 42.\n",
        "num_y_samples = 16\n",
        "\n",
        "batch_size = 512\n",
        "num_epochs =  100 * 2\n",
        "steps_per_epoch = 1000\n",
        "validation_size = 10 * 1024\n",
        "test_steps = 1024\n",
        "\n",
        "# repeat_num = 2 # recommended for quadratic-Gaussian case.\n",
        "repeat_num = 1 # recommended for signed Laplacian case.\n",
        "\n",
        "init_num_bins = repeat_num * num_latents\n",
        "\n",
        "total_steps = num_epochs * steps_per_epoch\n",
        "start_training_step = steps_per_epoch * 30\n",
        "lr_schedule = optax.piecewise_constant_schedule(\n",
        "    init_value=1e-3,\n",
        "    boundaries_and_scales={\n",
        "        int(7/10 * total_steps): 1e-1,\n",
        "    },\n",
        ")\n",
        "\n",
        "state_dict = dict()\n",
        "mod_epoch = 10\n",
        "\n",
        "def model():\n",
        "  return WynerZivModel(\n",
        "      source_dims=1, latent_dims=num_latents, var_n=var_n,\n",
        "      num_y_samples=num_y_samples, lmbda=lmbda, corr_pattern=corr_pattern)\n",
        "\n",
        "seed, = np.frombuffer(os.getrandom(4), dtype=np.int32)\n",
        "rng = random.PRNGKey(seed)\n",
        "\n",
        "rng, init_rng, init_xy_rng, init_z_rng = random.split(rng, 4)\n",
        "init_x, init_y = sample_xy(num_latents, init_xy_rng)\n",
        "init_z = jax.random.randint(init_z_rng, (init_num_bins,), 0, num_latents)\n",
        "state = train_state.TrainState.create(\n",
        "    apply_fn=model().apply,\n",
        "    params=model().init(\n",
        "        init_rng, x=init_x, y=init_y, training=False,\n",
        "        encoder_rng=random.PRNGKey(0))['params'],\n",
        "    tx=optax.adam(lr_schedule),\n",
        ")\n",
        "\n",
        "rng, data_rng = random.split(rng)\n",
        "validation_set = sample_xy(validation_size, data_rng)\n",
        "\n",
        "history = defaultdict(lambda: np.full((num_epochs,), float(\"nan\")))\n",
        "\n",
        "for epoch in range(num_epochs):\n",
        "\n",
        "  if epoch%mod_epoch==0:\n",
        "    state_dict[epoch] = state\n",
        "\n",
        "  results = defaultdict(lambda: 0, lr=lr_schedule(state.step))\n",
        "\n",
        "  for _ in range(steps_per_epoch):\n",
        "    rng, sampling_rng = random.split(rng, 2)\n",
        "    if state.step < start_training_step:\n",
        "      state, result = pretrain_step(state, init_x, var_n=var_n, sampling_rng=sampling_rng, init_z=init_z)\n",
        "    else:\n",
        "      rng, data_rng, encoder_rng = random.split(rng, 3)\n",
        "      state, result = train_step(state, data_rng=data_rng, encoder_rng=encoder_rng)\n",
        "    for k in result:\n",
        "      results[k] += result[k]\n",
        "  for k in results:\n",
        "    results[k] /= steps_per_epoch\n",
        "\n",
        "  results.update(eval_step(state, *validation_set))\n",
        "\n",
        "  for k in results:\n",
        "    history[k][epoch] = results[k]\n",
        "\n",
        "  colab.output.clear(wait=True)\n",
        "  plot_history(history)\n",
        "  plt.show()\n",
        "  print(f\"epoch {epoch:4}\")\n",
        "  print(f\"train      entropy {history['rate'][epoch]:6.4f}, distortion {history['distortion'][epoch]:7.4f} ({10*np.log10(history['distortion'][epoch]):6.3f} dB), loss {history['loss'][epoch]:7.4f}\")\n",
        "  print(f\"validation entropy {history['val_rate'][epoch]:6.4f}, distortion {history['val_distortion'][epoch]:7.4f} ({10*np.log10(history['val_distortion'][epoch]):6.3f} dB), loss {history['val_loss'][epoch]:7.4f}\", flush=True)\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "metrics = defaultdict(lambda: 0)\n",
        "for _ in range(test_steps):\n",
        "  rng, data_rng = random.split(rng)\n",
        "  result = eval_step(state, *sample_xy(batch_size, data_rng))\n",
        "  for k in result:\n",
        "    metrics[k] += result[k]\n",
        "for k in metrics:\n",
        "  metrics[k] /= test_steps\n",
        "print(f\"test rate {metrics['val_rate']:6.4f}, distortion {metrics['val_distortion']:7.4f}\"\n",
        "      \"({10*np.log10(metrics['val_distortion']):6.3f} dB), loss {metrics['val_loss']:6.3f}\")\n",
        "\n",
        "print(\"pgfplots:\", metrics['val_rate'], 10*np.log10(metrics['val_distortion']))\n"
      ],
      "metadata": {
        "id": "OB-9QEXKB1GQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Visualizing the Neural Compressor"
      ],
      "metadata": {
        "id": "tvvdrBK2r2Ll"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def plot_binning(state):\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  x = np.linspace(-10, 10, 10000)[:, None]\n",
        "\n",
        "  y = jnp.repeat(0, len(x)).reshape((len(x), 1))\n",
        "  # above, repeating 0. to be the same dimension as `x`\n",
        "  # in order to get obtain quantization boundaries\n",
        "  z, _, _ = encoder_behaviour(state, x, encoder_rng=(0,0))\n",
        "\n",
        "  boundaries = jnp.nonzero(z[1:] != z[:-1])[0]\n",
        "  print(boundaries)\n",
        "\n",
        "  categories_used = []\n",
        "  for i in range(len(boundaries)):\n",
        "    categories_used.append(z[boundaries[i]])\n",
        "  categories_used.append(z[boundaries[-1] +1])\n",
        "  categories_used = np.asarray(categories_used)\n",
        "\n",
        "  unique_categories = jnp.unique(categories_used)\n",
        "\n",
        "  if corr_pattern == \"y=x+n\" or corr_pattern == \"x=y+n\":\n",
        "    y_at_decoder_varying = np.linspace(-10, 10, 10000)[:, None]\n",
        "\n",
        "  elif corr_pattern == \"signed_laplacian\":\n",
        "    y_at_decoder_varying = np.sign(x)\n",
        "\n",
        "  else:\n",
        "    raise ValueError(\"Invalid correlation type. Only `y=x+n`, `x=y+n` and \"\n",
        "                     \"`signed_laplacian` correlation patterns are supported.\")\n",
        "\n",
        "  reconstructed = []\n",
        "  for i in range(len(unique_categories)):\n",
        "    ex_z = nn.one_hot(\n",
        "        unique_categories[i], num_latents, axis=-1, dtype='f')[None, :]\n",
        "    ex_z = np.tile(ex_z, (10000,1))\n",
        "    intermed_reconstructed = decoder_behaviour(\n",
        "        state, ex_z, y_at_decoder_varying)\n",
        "    intermed_reconstructed = jnp.squeeze(np.array(intermed_reconstructed), -1)\n",
        "    reconstructed.append((intermed_reconstructed, unique_categories[i]))\n",
        "\n",
        "  return reconstructed, boundaries, categories_used, unique_categories\n"
      ],
      "metadata": {
        "id": "tNV4EXqZ6VgC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "reconstructed, boundaries, categories_used, _ = plot_binning(state)\n",
        "print(categories_used)\n",
        "print(len(categories_used))\n",
        "print(np.unique(categories_used))\n",
        "print(len(np.unique(categories_used)))"
      ],
      "metadata": {
        "id": "MHvjqLsUsT1X"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "plt.figure(figsize=(16, 14))\n",
        "\n",
        "ys = np.linspace(-10, 10, 10000)\n",
        "\n",
        "xs = np.linspace(-10, 10, 10000)\n",
        "\n",
        "for ind_boundary in boundaries:\n",
        "  # you can reach the latest boundary\n",
        "  # since it doesn't change anyway when you vary 'y')\n",
        "  plt.axhline(xs[ind_boundary], -10, 10, linestyle='--', linewidth=2.5)\n",
        "\n",
        "plt.ylim(-7, 7)\n",
        "plt.xlim(-10, 10)\n",
        "\n",
        "\n",
        "# getting random colours\n",
        "vals = np.linspace(0,1,num_latents)\n",
        "np.random.shuffle(vals)\n",
        "#cmap = plt.cm.colors.ListedColormap(plt.cm.gist_rainbow(vals))\n",
        "cmap = plt.cm.colors.ListedColormap(plt.cm.nipy_spectral(vals))\n",
        "\n",
        "points = []\n",
        "for ind_reconstructed, ind_category in reconstructed:\n",
        "  plt.plot(ys, ind_reconstructed, color=cmap(ind_category.item()), linewidth=3)\n",
        "\n",
        "plt.ylabel(\"learned quantization boundaries and $\\hat{x}$\", fontsize=30)\n",
        "plt.grid()\n",
        "\n",
        "for i in range(len(categories_used)):\n",
        "  if i==0:\n",
        "    plt.fill_between(xs, np.repeat(xs[boundaries[i]], len(xs)), np.repeat(-10, len(xs)),alpha=0.15, color=cmap(categories_used[i]))\n",
        "  elif i==len(categories_used)-1:\n",
        "    plt.fill_between(xs, np.repeat(xs[boundaries[i]], len(xs)), np.repeat(10, len(xs)),alpha=0.15, color=cmap(categories_used[i]))\n",
        "  else:\n",
        "    plt.fill_between(xs, np.repeat(xs[boundaries[i-1]], len(xs)), np.repeat(xs[boundaries[i]], len(xs)),alpha=0.15, color=cmap(categories_used[i]))\n",
        "\n",
        "\n",
        "plt.yticks(fontsize=20)\n",
        "plt.xticks([-3, 3], [\"$y=-1$\", \"$y=+1$\"], fontsize=30) # valid for signed Laplacian setup.\n",
        "\n",
        "\n",
        "#plt.savefig('binning_plot.pdf', bbox_inches='tight')"
      ],
      "metadata": {
        "id": "zQ5P6UYS6uXs"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}