demo/demo-eig.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c8a2f179-3175-442f-96f3-667948ad4113",
   "metadata": {},
   "source": [
    "# Question-asking based on Expected Information Gain (EIG)\n",
    "\n",
    "**Inspired by:** Rothe, A., Lake, B. M., & Gureckis, T. M. (2018). _Do people ask good questions?._ Computational Brain & Behavior, 1, 69-89.\n",
    "\n",
    "Bob rolls a red die and a blue die. Alice gets to ask one yes-no question about the sum. What is the most informative question she could ask, in order to learn the most about the two die rolls? For example, is it better to ask if the sum is a perfect square, or if the sum is prime?\n",
    "\n",
    "We'll compute the EIG of various questions..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4ee47cfd-0371-489f-b1c2-a277007eea27",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "48.1 μs ± 892 ns per loop (mean ± std. dev. of 10 runs, 100 loops each)\n",
      "EIG     Question\n",
      "---     ---\n",
      "0.69315 n % 2 == 1,  # odd?\n",
      "0.69315 n % 2 == 0,  # even?\n",
      "0.67919 n > 6\n",
      "0.67919 is_prime[n]\n",
      "0.63651 n % 3 == 0\n",
      "0.59084 n > 5\n",
      "0.59084 n > 8\n",
      "0.56233 n % 4 == 0\n",
      "0.56233 is_pow_2[n]\n",
      "0.49260 n % 5 == 0\n",
      "0.49260 is_square[n]\n",
      "0.45056 n == 7\n",
      "0.28684 n > 10\n",
      "0.12693 n == 12\n"
     ]
    }
   ],
   "source": [
    "from memo import memo\n",
    "import jax\n",
    "import jax.numpy as np\n",
    "\n",
    "is_prime  = np.array([0, 0, 1, 1, 0, 1, 0, 1, 0, 0, 0, 1, 0])\n",
    "is_square = np.array([1, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0])\n",
    "is_pow_2  = np.array([0, 1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0])\n",
    "Qs = [\n",
    "    lambda n: n == 7,\n",
    "    lambda n: n == 12,\n",
    "    lambda n: n > 10,\n",
    "    lambda n: n > 8,\n",
    "    lambda n: n > 6,\n",
    "    lambda n: n > 5,\n",
    "    lambda n: n % 2 == 0,  # even??\n",
    "    lambda n: n % 2 == 1,  # odd??\n",
    "    lambda n: n % 3 == 0,\n",
    "    lambda n: n % 4 == 0,\n",
    "    lambda n: n % 5 == 0,\n",
    "    lambda n: is_prime[n],\n",
    "    lambda n: is_square[n],\n",
    "    lambda n: is_pow_2[n],\n",
    "]\n",
    "\n",
    "N = np.arange(1, 6 + 1)  # single die's outcomes\n",
    "Q = np.arange(len(Qs))   # questions\n",
    "A = np.array([0, 1])     # answers (yes/no)\n",
    "\n",
    "@jax.jit\n",
    "def respond(q, a, n):\n",
    "    return np.array([q_(n) for q_ in Qs])[q] == a\n",
    "\n",
    "@memo\n",
    "def eig[q: Q]():\n",
    "    alice: knows(q)\n",
    "    alice: thinks[\n",
    "        # bob rolls dice...\n",
    "        bob: chooses(n_red in N, wpp=1),\n",
    "        bob: chooses(n_blu in N, wpp=1),\n",
    "\n",
    "        # bob answers question...\n",
    "        bob: knows(q),\n",
    "        bob: chooses(a in A, wpp=respond(q, a, n_red + n_blu))\n",
    "    ]\n",
    "    alice: snapshots_self_as(future_self)\n",
    "\n",
    "    return alice[ imagine[\n",
    "        # if I were to get the answer...\n",
    "        future_self: observes [bob.a] is bob.a,\n",
    "        # EIG = entropy minus conditional entropy\n",
    "        H[bob.n_red, bob.n_blu] - E[future_self[ H[bob.n_red, bob.n_blu] ]]\n",
    "    ] ]\n",
    "\n",
    "z = eig()\n",
    "%timeit -r 10 -n 100 eig().block_until_ready()\n",
    "\n",
    "## print questions and EIGs in sorted order\n",
    "print('EIG     Question')\n",
    "print('---     ---')\n",
    "import inspect\n",
    "q_names = [inspect.getsource(q_).strip()[10:-1] for q_ in Qs]\n",
    "for eig_, q_ in reversed(sorted(list(zip(z, q_names)))):\n",
    "    print(f'{eig_:0.5f}', q_)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}