code-examples.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Code Examples From the Paper \"Understanding Entropy Coding With Asymmetric Numeral Systems (ANS): a Statistician's Perspective\"\n",
    "\n",
    "**Author:** [Robert Bamler](https://robamler.github.io/)<br>\n",
    "**Date:** 22 December 2021\n",
    "\n",
    "This notebook contains all code examples from the paper \"Understanding Entropy Coding With Asymmetric Numeral Systems (ANS): a Statistician's Perspective\" by Robert Bamler.\n",
    "\n",
    "The final streaming ANS coder is implemented in Listing 7 below.\n",
    "While this coder implementation works, it is really only intended as an educational demo.\n",
    "For readers who're looking for an entropy coder, we recommend installing the [constriction package](https://bamler-lab.github.io/constriction/) with `pip install constriction`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Listing 1: `UniformCoder`\n",
    "\n",
    "Use Listings 2 and 3 below to test this implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "class UniformCoder: # optimal stream code for uniformly distributed symbols\n",
    "    def __init__(self, number=0):   # constructor with an optional argument\n",
    "        self.number = number\n",
    "\n",
    "    def push(self, symbol, base):   # Encodes a symbol ∈{0, . . . , base −1}.\n",
    "        self.number = self.number * base + symbol\n",
    "\n",
    "    def pop(self, base):            # Decodes a symbol ∈{0, . . . , base −1}.\n",
    "        symbol = self.number % base # “%”  denotes modulo.\n",
    "        self.number //= base        # “//” denotes integer division.\n",
    "        return symbol"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Listing 2: Simple Usage Example of `UniformCoder`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Encoded number: 365\n",
      "In binary: 101101101\n",
      "Decoded '5'.\n",
      "Decoded '6'.\n",
      "Decoded '3'.\n"
     ]
    }
   ],
   "source": [
    "coder = UniformCoder() # Defined in Listing 1.\n",
    "\n",
    "# Encode the message x = (3, 6, 5):\n",
    "coder.push(3, base=10) # base=10 means that the alphabet is {0, ..., 9}.\n",
    "coder.push(6, base=10)\n",
    "coder.push(5, base=10)\n",
    "print(f\"Encoded number: {coder.number}\")  # Prints: “Encoded number: 365”\n",
    "print(f\"In binary: {coder.number:b}\")     # Prints: “In binary: 101101101”\n",
    "\n",
    "# Decode the encoded symbols (in reverse order):\n",
    "print(f\"Decoded '{coder.pop(base=10)}'.\") # Prints: “Decoded '5'.”\n",
    "print(f\"Decoded '{coder.pop(base=10)}'.\") # Prints: “Decoded '6'.”\n",
    "print(f\"Decoded '{coder.pop(base=10)}'.\") # Prints: “Decoded '3'.”"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Listing 3: Using `UniformCoder` With Inhomogeneous Alphabets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Binary: 10000001011100\n",
      "Decoded '4'.\n",
      "Decoded '12'.\n",
      "Decoded '6'.\n",
      "Decoded '3'.\n"
     ]
    }
   ],
   "source": [
    "coder = UniformCoder() # Defined in Listing 1.\n",
    "\n",
    "# Encode a sequence of symbols from alphabets of varying sizes:\n",
    "coder.push( 3, base=10) # base=10 means that the alphabet is {0, ..., 9}.\n",
    "coder.push( 6, base=10)\n",
    "coder.push(12, base=15) # Setting base=15 switches to alphabet {0, ..., 14}.\n",
    "coder.push( 4, base=15)\n",
    "print(f\"Binary: {coder.number:b}\") # Prints: “Binary: 10000001011100”\n",
    "\n",
    "# For decoding, use the same sequence of bases but in *reverse* order:\n",
    "print(f\"Decoded '{coder.pop(base=15)}'.\") # Prints: “Decoded '4'.”\n",
    "print(f\"Decoded '{coder.pop(base=15)}'.\") # Prints: “Decoded '12'.”\n",
    "print(f\"Decoded '{coder.pop(base=10)}'.\") # Prints: “Decoded '6'.”\n",
    "print(f\"Decoded '{coder.pop(base=10)}'.\") # Prints: “Decoded '3'.”"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Listing 4: `SlowAnsCoder`\n",
    "\n",
    "Use Listing 5 below to test this implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SlowAnsCoder: # Has near-optimal bitrates but high runtime cost.\n",
    "    def __init__(self, precision, compressed=0):\n",
    "        self.n = 2**precision # See Eq. 5 (“**” denotes exponentiation).\n",
    "        self.stack = UniformCoder(compressed) # Defined in Listing 1.\n",
    "\n",
    "    def push(self, symbol, m): # Encodes one symbol.\n",
    "        z = self.stack.pop(base=m[symbol]) + sum(m[0:symbol])\n",
    "        self.stack.push(z, base=self.n)\n",
    "\n",
    "    def pop(self, m):          # Decodes one symbol.\n",
    "        z = self.stack.pop(base=self.n)\n",
    "        # Find the unique symbol that satisfies z ∈ Z_i(symbol) (real\n",
    "        # deployments should use a more efficient method than linear search):\n",
    "        for symbol, m_symbol in enumerate(m):\n",
    "            if z >= m_symbol:\n",
    "                z -= m_symbol\n",
    "            else:\n",
    "                break\n",
    "        self.stack.push(z, base=m_symbol)\n",
    "        return symbol\n",
    "\n",
    "    def get_compressed(self):\n",
    "        return self.stack.number"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Listing 5: Usage Example for `SlowAnsCoder` (and for `AnsCoder`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify an approximated entropy model via precision and m_i(x_i) from Eq. 5:\n",
    "precision = 4 # For demonstration; deployments should use higher precision.\n",
    "m = [7, 3, 6] # Sets Q_i(X_i=0) = 7/16, Q_i(X_i=1) = 3/16, and Q_i(X_i=2) = 6/16.\n",
    "\n",
    "# Encode a message in reversed order so we can decode it in forward order:\n",
    "example_message = [2, 0, 2, 1, 0]\n",
    "encoder = SlowAnsCoder(precision) # Also works with AnsCoder (Listing 7).\n",
    "for symbol in reversed(example_message):\n",
    "    encoder.push(symbol, m) # We could use a different m for each symbol.\n",
    "compressed = encoder.get_compressed()\n",
    "\n",
    "# We could actually reuse the encoder for decoding, but let's pretend that\n",
    "# decoding occurs on a different machine that receives only “compressed”.\n",
    "decoder = SlowAnsCoder(precision, compressed)\n",
    "reconstructed = [decoder.pop(m) for _ in range(5)]\n",
    "assert reconstructed == example_message # Verify correctness."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Listing 6: Self-Contained Reimplementation of `SlowAnsCoder`\n",
    "\n",
    "Use Listing 5 above to test this implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SlowAnsCoder: # Equivalent to Listing 4, just more self-contained.\n",
    "    def __init__(self, precision, compressed=0):\n",
    "        self.n = 2**precision # See Eq. 5 (“**” denotes exponentiation).\n",
    "        self.compressed = compressed # (== stack.number in SlowAnsCoder)\n",
    "\n",
    "    def push(self, symbol, m):        # Encodes one symbol.\n",
    "        z = self.compressed % m[symbol] + sum(m[0:symbol])\n",
    "        self.compressed //= m[symbol] # “//” denotes integer division.\n",
    "        self.compressed = self.compressed * self.n + z\n",
    "\n",
    "    def pop(self, m): # Decodes one symbol.\n",
    "        z = self.compressed % self.n\n",
    "        self.compressed //= self.n # “//” denotes integer division.\n",
    "        for symbol, m_symbol in enumerate(m):\n",
    "            if z >= m_symbol:\n",
    "                z -= m_symbol\n",
    "            else:\n",
    "                break # We found the symbol that satisfies z ∈ Z_i(symbol).\n",
    "        self.compressed = self.compressed * m_symbol + z\n",
    "        return symbol\n",
    "\n",
    "    def get_compressed(self):\n",
    "        return self.compressed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Listing 7: Streaming ANS Coder\n",
    "\n",
    "The following code illustrates a full implementation of a streaming ANS coder.\n",
    "While this coder works, it is really only intended as an educational demo.\n",
    "Readers who want to use an entropy coder for a real project are advised to install [`constriction`](https://bamler-lab.github.io/constriction/) with `pip install constriction`.\n",
    "\n",
    "To test this implementation, replace `SlowAnsCoder` with `AnsCoder` in Listing 5 above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AnsCoder:\n",
    "    def __init__(self, precision, compressed=[]):\n",
    "        self.precision = precision\n",
    "        self.mask = (1 << precision) - 1 # (a string of `precision`` one-bits)\n",
    "        self.bulk = compressed.copy() # (We will mutate `bulk`` below.)\n",
    "        self.head = 0\n",
    "        # Establish invariant (ii):\n",
    "        while len(self.bulk) != 0 and (self.head >> precision) == 0:\n",
    "            self.head = (self.head << precision) | self.bulk.pop()\n",
    "\n",
    "    def push(self, symbol, m): # Encodes one symbol.\n",
    "        # Check if encoding directly onto head would violate invariant (i):\n",
    "        if (self.head >> self.precision) >= m[symbol]:\n",
    "            # Transfer one word of compressed data from head to bulk:\n",
    "            self.bulk.append(self.head & self.mask) # (“&” is bitwise `and`)\n",
    "            self.head >>= self.precision\n",
    "            # At this point, invariant (ii) is definitely violated,\n",
    "            # but the operations below will restore it.\n",
    "\n",
    "        z = self.head % m[symbol] + sum(m[0:symbol])\n",
    "        self.head //= m[symbol]\n",
    "        self.head = (self.head << self.precision) | z # (This is\n",
    "            # equivalent to “ self.head * n + z”, just slightly faster.)\n",
    "\n",
    "    def pop(self, m): # Decodes one symbol.\n",
    "        z = self.head & self.mask # (same as “ self.head % n” but faster)\n",
    "        self.head >>= self.precision # (same as “ //= n” but faster)\n",
    "        for symbol, m_symbol in enumerate(m):\n",
    "            if z >= m_symbol:\n",
    "                z -= m_symbol\n",
    "            else:\n",
    "                break # We found the symbol that satisfies z ∈ Z_i(symbol).\n",
    "        self.head = self.head * m_symbol + z\n",
    "\n",
    "        # Restore invariant (ii) if it is violated (which happens exactly\n",
    "        # if the encoder transferred data from head to bulk at this point):\n",
    "        if (self.head >> self.precision) == 0 and len(self.bulk) != 0:\n",
    "            # Transfer data back from bulk to head (“ |” is bitwise or):\n",
    "            self.head = (self.head << self.precision) | self.bulk.pop()\n",
    "\n",
    "        return symbol\n",
    "\n",
    "    def get_compressed(self):\n",
    "        compressed = self.bulk.copy() # (We will mutate `compressed`` below.)\n",
    "        head = self.head\n",
    "        # Chop `head`` into `precision`-sized words and append to `compressed``:\n",
    "        while head != 0:\n",
    "            compressed.append(head & self.mask)\n",
    "            head >>= self.precision\n",
    "        return compressed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Listing 8: Random-Access Decoding with ANS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SeekableAnsCoder(AnsCoder): # Adds random-access decoding to Listing 7.\n",
    "    def __init__(self, precision, compressed=[]):\n",
    "        return super(SeekableAnsCoder, self).__init__(precision, compressed)\n",
    "\n",
    "    # Inherits push, pop, and get_compressed from Listing 7.\n",
    "\n",
    "    def checkpoint(self): # Records a point to which we can seek later.\n",
    "        return (len(self.bulk), self.head)\n",
    "\n",
    "    def seek(self, checkpoint): # Jumps to a previously taken checkpoint.\n",
    "        position, head = checkpoint\n",
    "        if position > len(self.bulk): # “raise” throws an exception.\n",
    "            raise \"This simple demo can only seek forward.\"\n",
    "        self.bulk = self.bulk[0:position] # Truncates bulk.\n",
    "        self.head = head\n",
    "\n",
    "# Usage example:\n",
    "precision = 4 # For demonstration; deployments should use higher precision.\n",
    "m = [7, 3, 6] # Same demo model as in Listing 5.\n",
    "coder = SeekableAnsCoder(precision)\n",
    "message = [2, 0, 2, 1, 0, 1, 2, 2, 2, 1, 0, 2, 1, 2, 0, 0, 1, 1, 1, 2]\n",
    "\n",
    "for symbol in reversed(message[10:20]): # Encode second half of message.\n",
    "    coder.push(symbol, m)\n",
    "checkpoint = coder.checkpoint()         # Record a checkpoint.\n",
    "for symbol in reversed(message[0:10]):  # Encode first half of message.\n",
    "    coder.push(symbol, m)\n",
    "\n",
    "assert coder.pop(m) == message[0]       # Decode first symbol.\n",
    "assert coder.pop(m) == message[1]       # Decode second symbol.\n",
    "coder.seek(checkpoint)                  # Jump to 11th symbol.\n",
    "assert [coder.pop(m) for _ in range(10)] == message[10:20] # Decode rest."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Listing 9: Demonstration of Non-Local Effect of Entropy Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "case1 = [0, 1, 0, 2]\n",
      "case2 = [1, 1, 2, 0]\n"
     ]
    }
   ],
   "source": [
    "precision = 4 # For demonstration; deployments should use higher precision.\n",
    "m_orig = [7, 3, 6] # Same demo entropy model as in Listing 5.\n",
    "m_mod = [6, 4, 6]  # (Slightly) modified entropy model compared to m_orig.\n",
    "compressed = [0b1001, 0b1110, 0b0110, 0b1110] # Some example bit string.\n",
    "\n",
    "# Case 1: decode 4 symbols using entropy model m_orig for all symbols:\n",
    "decoder = AnsCoder(precision, compressed) # AnsCoder defined in Listing 7.\n",
    "case1 = [decoder.pop(m_orig) for _ in range(4)]\n",
    "\n",
    "# Case 2: change the entropy model, but *only* for the first symbol:\n",
    "decoder = AnsCoder(precision, compressed) # “compressed” hasn't changed.\n",
    "case2 = [decoder.pop(m_mod)] + [decoder.pop(m_orig) for _ in range(3)]\n",
    "\n",
    "print(f\"case1 = {case1}\") # Prints: “case1 = [0, 1, 0, 2]”\n",
    "print(f\"case2 = {case2}\") # Prints: “case2 = [1, 1, 2, 0]”\n",
    "                          #                   ↑     ↑  ↑\n",
    "                          #   We changed only the | But that affected both\n",
    "                          # model for this symbol.| of these symbols too."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Listing 10: Experimental New Coder That Overcomes the Issue Demonstrated in Listing 9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ChainCoder: # Prevents the non-local effect shown in Listing 9.\n",
    "    def __init__(self, precision, compressed, remainders=[]):\n",
    "        \"\"\"Initializes a ChainCoder for decoding from `compressed`.\"\"\"\n",
    "        self.precision = precision\n",
    "        self.mask = (1 << precision) - 1\n",
    "        self.compressed = compressed.copy() # pop decodes from here.\n",
    "        self.remainders = remainders.copy() # pop encodes onto here.\n",
    "        self.remainders_head = 0\n",
    "        # Establish invariant (ii):\n",
    "        while len(self.remainders) != 0 and \\\n",
    "                (self.remainders_head >> precision) == 0:\n",
    "            self.remainders_head <<= self.precision\n",
    "            self.remainders_head |= self.remainders.pop()\n",
    "\n",
    "    def pop(self, m): # Decodes one symbol.\n",
    "        z = self.compressed.pop() # Always read a full word from compressed.\n",
    "        for symbol, m_symbol in enumerate(m):\n",
    "            if z >= m_symbol:\n",
    "                z -= m_symbol\n",
    "            else:\n",
    "                break # We found the symbol that satisfies z ∈Zi(symbol).\n",
    "\n",
    "        self.remainders_head = self.remainders_head * m_symbol + z\n",
    "        if (self.remainders_head >> (2 * self.precision)) != 0:\n",
    "            # Invariant (i) is violated. Flush one word to remainders.\n",
    "            self.remainders.append(self.remainders_head & self.mask)\n",
    "            self.remainders_head >>= self.precision\n",
    "            # It can easily be shown that invariant (i) is restored here.\n",
    "\n",
    "        return symbol\n",
    "\n",
    "    def push(self, symbol, m): # Encodes one symbol.\n",
    "        if len(self.remainders) != 0 and \\\n",
    "                self.remainders_head < (m[symbol] << self.precision):\n",
    "            self.remainders_head <<= self.precision\n",
    "            self.remainders_head |= self.remainders.pop()\n",
    "            # Invariant (i) is now violated but will be restored below.\n",
    "\n",
    "        z = self.remainders_head % m[symbol] + sum(m[0:symbol])\n",
    "        self.remainders_head //= m[symbol]\n",
    "        self.compressed.append(z)\n",
    "\n",
    "    def get_compressed(self):\n",
    "        return self.compressed.copy()\n",
    "\n",
    "    def get_remainders(self):\n",
    "        remainders = self.bulk.copy() # (We will mutate `remainders`` below.)\n",
    "        remainders_head = self.remainders_head\n",
    "        # Chop `remainders_head`` into `precision`-sized words and append to `remainders``:\n",
    "        while remainders_head != 0:\n",
    "            remainders.append(remainders_head & self.mask)\n",
    "            remainders_head >>= self.precision\n",
    "        return compressed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Usage Example\n",
    "\n",
    "Let's replace `AnsCoder` by `ChainCoder` in the example from Listing 9:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "case1 = [2, 0, 2, 1]\n",
      "case2 = [2, 0, 2, 1]\n"
     ]
    }
   ],
   "source": [
    "precision = 4 # For demonstration; deployments should use higher precision.\n",
    "m_orig = [7, 3, 6] # Same demo entropy model as in Listing 5.\n",
    "m_mod = [6, 4, 6]  # (Slightly) modified entropy model compared to m_orig.\n",
    "compressed = [0b1001, 0b1110, 0b0110, 0b1110] # Some example bit string.\n",
    "\n",
    "# Case 1: decode 4 symbols using entropy model m_orig for all symbols:\n",
    "decoder = ChainCoder(precision, compressed) # ChainCoder defined in Listing 10.\n",
    "case1 = [decoder.pop(m_orig) for _ in range(4)]\n",
    "\n",
    "# Case 2: change the entropy model, but *only* for the first symbol:\n",
    "decoder = ChainCoder(precision, compressed) # “compressed” hasn't changed.\n",
    "case2 = [decoder.pop(m_mod)] + [decoder.pop(m_orig) for _ in range(3)]\n",
    "\n",
    "print(f\"case1 = {case1}\") # Prints: “case1 = [2, 0, 2, 1]”\n",
    "print(f\"case2 = {case2}\") # Prints: “case2 = [2, 0, 2, 1]”"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's no surprise that we obtain a different sequence of decoded symbols than with an `AnsCoder` because we used a different entropy coding algorithm.\n",
    "The important point is that `case1` and `case` agree on the second, third, and fourth symbol.\n",
    "(In the above example, they *happen* to also agree on the first symbol, but that's a coincidence.)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "97cc609b13305c559618ec78a438abc56230b9381f827f22d070313b9a1f3777"
  },
  "kernelspec": {
   "display_name": "Python 3.10.0 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}