{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Custom Entropy Models With `constriction`\n",
    "\n",
    "- **Author:** Robert Bamler, University of Tuebingen\n",
    "- **Initial Publication Date:** Jan 4, 2022\n",
    "\n",
    "This is an interactive jupyter notebook.\n",
    "You can read this notebook [online](https://github.com/bamler-lab/constriction/blob/main/examples/python/02-custom-entropy-models.ipynb) but if you want to execute any code, we recommend to [download](https://raw.githubusercontent.com/bamler-lab/constriction/main/examples/python/02-custom-entropy-models.ipynb) it.\n",
    "\n",
    "More examples, tutorials, and reference materials are available at <https://bamler-lab.github.io/constriction/>."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Constriction\n",
    "\n",
    "Before you start, install `constriction` and `scipy` by executing the following cell, then restart your jupyter kernel:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install constriction~=0.3.5 scipy  # (this will automatically also install numpy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Don't forget to restart your jupyter kernel.**\n",
    "Then test if you can import `constriction`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import constriction # This cell should produce no output (in particular, no error messages)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We further recommend that you read through the simple [\"Hello, World\"](https://github.com/bamler-lab/constriction/blob/main/examples/python/01-hello-world.ipynb) example before continuing with this example."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 1: Fixed Custom Entropy model\n",
    "\n",
    "Constriction provides a selection of [predefined entropy models and model families](https://bamler-lab.github.io/constriction/apidoc/python/stream/model.html).\n",
    "However, sometimes these do not cover everyone's need.\n",
    "Therefore, `constriction` also provides two escape hatches:\n",
    "\n",
    "- You can use any univariate distribution from the popular `scipy` package and wrap it in a [`ScipyModel`](https://bamler-lab.github.io/constriction/apidoc/python/stream/model.html#constriction.stream.model.ScipyModel).\n",
    "  This turns the model into exactly invertible fixed point arithmetic so that it can be used for entropy coding with `constriction`'s stream codes.\n",
    "- If your model is not available in `scipy`, then you can use a [`CustomModel`](https://bamler-lab.github.io/constriction/apidoc/python/stream/model.html#constriction.stream.model.CustomModel) and explicitly provide your own cumulative distribution function (and its approximate inverse) as an arbitrary python function (that can, optionally, take additional model parameters as in Example 2 below).\n",
    "\n",
    "In this first example, we'll use `scipy`'s Cauchy distribution to define a *fixed* (i.e., fully parametrized) entropy model.\n",
    "In Example 2 below, we'll show how we can construct custom models that accept additional model parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original message: [  3   2   6 -51 -19   5  87]\n",
      "compressed representation: [1831109982  151859470]\n",
      "Reconstructed message: [  3   2   6 -51 -19   5  87]\n"
     ]
    }
   ],
   "source": [
    "import constriction\n",
    "import numpy as np\n",
    "import scipy.stats\n",
    "\n",
    "# Get a scipy model with fixed parameters (`loc` and `scale`). See below\n",
    "# for an example of a model *family* with variable model parameters.\n",
    "scipy_model = scipy.stats.cauchy(loc=10.2, scale=30.9)\n",
    "\n",
    "# Quantize the model to bins of size 1 centered at integers from -100 to 100:\n",
    "model = constriction.stream.model.ScipyModel(scipy_model, -100, 100)\n",
    "\n",
    "# Encode and decode some message using the above model for all model parameters:\n",
    "message = np.array([3, 2, 6, -51, -19, 5, 87], dtype=np.int32)\n",
    "print(f\"Original message: {message}\")\n",
    "encoder = constriction.stream.queue.RangeEncoder()\n",
    "encoder.encode(message, model)\n",
    "compressed = encoder.get_compressed()\n",
    "print(f\"compressed representation: {compressed}\")\n",
    "\n",
    "decoder = constriction.stream.queue.RangeDecoder(compressed)\n",
    "reconstructed = decoder.decode(model, 7) # (decodes 7 symbols)\n",
    "print(f\"Reconstructed message: {reconstructed}\")\n",
    "assert np.all(reconstructed == message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using `CustomModel` Instead of `ScipyModel`\n",
    "\n",
    "In the above example, we used the `ScipyModel` wrapper.\n",
    "We could instead also have defined the model explicitly through a cumulative distribution function (CDF) and its (approximate) inverse (percent point function or PPF) by replacing:\n",
    "\n",
    "```python\n",
    "model = constriction.stream.model.ScipyModel(scipy_model, -100, 100)\n",
    "```\n",
    "\n",
    "with:\n",
    "\n",
    "```python\n",
    "model = constriction.stream.model.CustomModel(scipy_model.cdf, scipy_model.ppf, -100, 100)\n",
    "```\n",
    "\n",
    "The `CustomModel` wrapper allows you to easily define arbitrary custom entropy models:\n",
    "\n",
    "```python\n",
    "model = constriction.stream.model.CustomModel(\n",
    "    lambda x: ... TODO ...,  # define your CDF here\n",
    "    lambda xi: ... TODO ..., # define an (approximate) inverse of the CDF here\n",
    "    -100, 100) # (or on whichever range you want to define your model)\n",
    "```\n",
    "\n",
    "Only the first parameter (the CDF) must be carefully designed.\n",
    "The second parameter (inverse CDF) does not need to be very accurate since it is only used to speed up an internal *exact* inversion of a fixed-point approximation of the CDF.\n",
    "For details, please refer to the [API documentation](https://bamler-lab.github.io/constriction/apidoc/python/stream/model.html#constriction.stream.model.CustomModel)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 2: *Parameterized* Custom Entropy Model\n",
    "\n",
    "Example 1 above encoded and decoded a message of i.i.d. symbols, i.e., we used the same custom defined model for each symbol.\n",
    "We can also use `ScipyModel` or `CustomModel` to define model *families* which still have some free parameters that we can then define for each symbol individually upon encoding or decoding:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original message: [  3   2   6 -51 -19   5  87]\n",
      "compressed representation: [1123914182  166883892]\n",
      "Reconstructed message: [  3   2   6 -51 -19   5  87]\n"
     ]
    }
   ],
   "source": [
    "import constriction\n",
    "import numpy as np\n",
    "import scipy.stats\n",
    "\n",
    "# Get a scipy model *family*, i.e., a model with free parameters \n",
    "scipy_model_family = scipy.stats.cauchy\n",
    "\n",
    "# Turn it into a a family of (eventually) quantized entropy models:\n",
    "model_family = constriction.stream.model.ScipyModel(scipy_model_family, -100, 100)\n",
    "\n",
    "# Provide model parameters for each symbol:\n",
    "message = np.array([3,    2,   6,   -51,   -19,   5,   87  ], dtype=np.int32) # (same message as in Example 1)\n",
    "locs    = np.array([7.2, -1.4, 9.1, -60.1,   3.9, 8.1, 63.2], dtype=np.float64)\n",
    "scales  = np.array([4.3,  5.1, 6,    14.2,  31.9, 7.2, 10.7], dtype=np.float64)\n",
    "\n",
    "# Encode and decode the message using individual model parameters for each symbol:\n",
    "print(f\"Original message: {message}\")\n",
    "encoder = constriction.stream.queue.RangeEncoder()\n",
    "encoder.encode(message, model_family, locs, scales)\n",
    "compressed = encoder.get_compressed()\n",
    "print(f\"compressed representation: {compressed}\")\n",
    "\n",
    "decoder = constriction.stream.queue.RangeDecoder(compressed)\n",
    "reconstructed = decoder.decode(model_family, locs, scales)\n",
    "print(f\"Reconstructed message: {reconstructed}\")\n",
    "assert np.all(reconstructed == message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using `CustomModel` Instead of `ScipyModel`\n",
    "\n",
    "Similar to the discussion of fully parameterized model under Example 1 above, we could again also define our model *family* wiwth a `CustomModel` instead of a `ScipyModel`.\n",
    "Our provided CDF and PPF functions just have to accept additional arguments for the model parameters:\n",
    "\n",
    "```python\n",
    "model_family = constriction.stream.model.CustomModel(\n",
    "    scipy_model_family.cdf, scipy_model_family.ppf, -100, 100)\n",
    "```\n",
    "\n",
    "or, more generally:\n",
    "\n",
    "```python\n",
    "model_family = constriction.stream.model.CustomModel(\n",
    "    lambda x, model_param1, model_param2: ... TODO ...,  # CDF\n",
    "    lambda xi, model_param1, model_param2: ... TODO ..., # approximate inverse CDF\n",
    "    -100, 100) # (or on whichever range you want to define your model)\n",
    "```\n",
    "\n",
    "The number of model parameters is arbitrary but must match the number of parameter-arrays provided when encoding or decoding."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Further Reading\n",
    "\n",
    "See Python API documentation for [`ScipyModel`](https://bamler-lab.github.io/constriction/apidoc/python/stream/model.html#constriction.stream.model.ScipyModel) and [`CustomModel`](https://bamler-lab.github.io/constriction/apidoc/python/stream/model.html#constriction.stream.model.CustomModel)."
   ]
  }
 ],
 "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.12.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "ead1b95f633dc9c51826328e1846203f51a198c6fb5f2884a80417ba131d4e82"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}