demo/demo-cheryl.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8d8e4df7-475d-4a2f-81a1-becfe3078512",
   "metadata": {},
   "source": [
    "# Cheryl's Birthday\n",
    "\n",
    "> Everyone knows that Cheryl was born in February, March, or April.\n",
    ">\n",
    "> Cheryl separately tells Alice the month and Bob the day. Then they have this dialogue:\n",
    ">\n",
    "> 1. Alice: \"I don't know when Cheryl's birthday is...\"\n",
    "> 2. Alice: \"but I know that Bob doesn't know either.\"\n",
    "> 3. Bob: \"At first I didn’t know when Cheryl's birthday is...\n",
    "> 4. Bob: \"but now I know.\"\n",
    "> 5. Alice: \"Now I know when Cheryl's birthday is.\"\n",
    ">\n",
    "> When is Cheryl's birthday?\n",
    "\n",
    "_(This is actually our in-house variant of the [original puzzle](https://en.wikipedia.org/wiki/Cheryl%27s_Birthday). It is logically the same, but we find this variant more fun and easier to explain to people, because it doesn't rely on positing an arbitrary subset of dates.)_\n",
    "\n",
    "We will progressively build up a model of this scenario this in memo by writing a model for each of the 5 utterances. To model whether someone knows for certain when Cheryl's birthday is, we will check whether according to them the variance over the date is zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a1c25755-6319-4bf6-9f39-fe290c8895d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from memo import memo, domain\n",
    "import jax\n",
    "import jax.numpy as np\n",
    "from enum import IntEnum\n",
    "\n",
    "class Month(IntEnum):\n",
    "    February = 0\n",
    "    March = 1\n",
    "    April = 2\n",
    "\n",
    "Day = np.arange(1, 31 + 1)\n",
    "\n",
    "class U(IntEnum):\n",
    "    DUNNO = 0\n",
    "    KNOWN = 1\n",
    "\n",
    "@jax.jit\n",
    "def possible(m, d):  # 31 days hath...\n",
    "    return d <= np.array([29, 31, 30])[m]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c544a37-d090-4839-bc39-db2ef0b7fed4",
   "metadata": {},
   "source": [
    "Now, let's build up this model step by step.\n",
    "\n",
    "Alice: (after observing month $m$) \"I don't know when Cheryl's birthday is.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fd945542-3e43-411c-b4fb-bb7c7acd4d02",
   "metadata": {},
   "outputs": [],
   "source": [
    "@memo\n",
    "def a_u1[m: Month, u: U]():\n",
    "    a: thinks[\n",
    "        c: chooses(m in Month, wpp=1),\n",
    "        c: chooses(d in Day, wpp=possible(m, d))\n",
    "    ]\n",
    "    a: observes [c.m] is m\n",
    "    return u == a[Var[c.d] == 0]  # note: Alice's variance over Cheryl's d\n",
    "# print(a_u1())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62ab6e22-7ed5-4427-b8aa-b0945d9ab683",
   "metadata": {},
   "source": [
    "Alice: \"...but I know that Bob doesn't know either.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "611ecaf5-d6a0-40b4-84c7-2101cf8ecbd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "@memo\n",
    "def a_u2[m: Month, u: U]():\n",
    "    a: thinks[\n",
    "        c: chooses(m in Month, wpp=1),\n",
    "        c: chooses(d in Day, wpp=possible(m, d)),\n",
    "        b: thinks[\n",
    "            c: chooses(m in Month, wpp=1),\n",
    "            c: chooses(d in Day, wpp=possible(m, d))\n",
    "        ],\n",
    "        b: observes [c.d] is c.d\n",
    "    ]\n",
    "    a: observes [c.m] is m\n",
    "    return u == a[Pr[b[Var[c.m] == 0]] > 0]\n",
    "# print(a_u2())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "127b9390-d44f-4e9b-a231-d2483d9e3fec",
   "metadata": {},
   "source": [
    "Bob: \"At first I didn't know...\" (similar to `a_u1`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0fe09dd8-b817-426b-82ae-773aa00dfb10",
   "metadata": {},
   "outputs": [],
   "source": [
    "@memo\n",
    "def b_u3[d: Day, u: U]():\n",
    "    b: thinks[\n",
    "        c: chooses(m in Month, wpp=1),\n",
    "        c: chooses(d in Day, wpp=possible(m, d))\n",
    "    ]\n",
    "    b: observes [c.d] is d\n",
    "    return u == b[Var[c.m] == 0]\n",
    "# print(b_u3())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52ed3b7a-ff5e-493c-ab69-36d8f6edb886",
   "metadata": {},
   "source": [
    "Bob: \"But now...\" (conditions on result of `a_u1` and `a_u2`!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d20e2fc1-d38f-4f0c-aa17-66f5247396a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "@memo\n",
    "def b_u4[d: Day, u1: U, u2: U, u: U]():\n",
    "    b: thinks[\n",
    "        c: chooses(m in Month, wpp=1),\n",
    "        c: chooses(d in Day, wpp=possible(m, d)),\n",
    "        a: thinks[\n",
    "            c: chooses(m in Month, wpp=1),\n",
    "            c: chooses(d in Day, wpp=possible(m, d)),\n",
    "        ],\n",
    "        a: observes [c.m] is c.m,\n",
    "        a: chooses(u1 in U, wpp=a_u1[c.m, u1]()),\n",
    "        a: chooses(u2 in U, wpp=a_u2[c.m, u2]())\n",
    "    ]\n",
    "    b: observes [c.d] is d\n",
    "    b: observes [a.u1] is u1\n",
    "    b: observes [a.u2] is u2\n",
    "    return u == b[Var[c.m] == 0]\n",
    "# print(b_u4()[:, U.DUNNO, U.DUNNO])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a5e7d5b-4cce-451d-b0cd-39bc51f5622a",
   "metadata": {},
   "source": [
    "Alice: \"Now I know.\" (conditions on `b_u3` and `b_u4`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3a34bf80-3219-46b5-8d36-2c4f253999ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "@memo\n",
    "def a_u5[m: Month, u1: U, u2: U, u3: U, u4: U, u: U]():\n",
    "    a: knows(u1, u2)\n",
    "    a: thinks[\n",
    "        c: chooses(m in Month, wpp=1),\n",
    "        c: chooses(d in Day, wpp=possible(m, d)),\n",
    "        b: thinks[\n",
    "            c: chooses(m in Month, wpp=1),\n",
    "            c: chooses(d in Day, wpp=possible(m, d))\n",
    "        ],\n",
    "        b: knows(u1, u2),\n",
    "        b: observes [c.d] is c.d,\n",
    "        b: chooses(u3 in U, wpp=b_u3[c.d, u3]()),\n",
    "        b: chooses(u4 in U, wpp=b_u4[c.d, u1, u2, u4]()),\n",
    "    ]\n",
    "    a: observes [c.m] is m\n",
    "    a: observes [b.u3] is u3\n",
    "    a: observes [b.u4] is u4\n",
    "    return u == a[Var[c.d] == 0]\n",
    "# a_u5()[:, U.DUNNO, U.DUNNO, U.DUNNO, U.KNOWN]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31123856-0109-4f3d-93ea-099043fcb15b",
   "metadata": {},
   "source": [
    "Putting everything together, we condition on all 5 utterances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ccc80742-2333-49fa-bcfd-8a99574adfa2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "403 μs ± 5.1 μs per loop (mean ± std. dev. of 10 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "@memo\n",
    "def puzzle[m: Month, d: Day, u1: U, u2: U, u3: U, u4: U, u5: U]():\n",
    "    z: thinks[\n",
    "        c: chooses(m in Month, wpp=1),\n",
    "        c: chooses(d in Day, wpp=possible(m, d)),\n",
    "        c: chooses(u1 in U, wpp=a_u1[m, u1]()),\n",
    "        c: chooses(u2 in U, wpp=a_u2[m, u2]()),\n",
    "        c: chooses(u3 in U, wpp=b_u3[d, u3]()),\n",
    "        c: chooses(u4 in U, wpp=b_u4[d, u1, u2, u4]()),\n",
    "        c: chooses(u5 in U, wpp=a_u5[m, u1, u2, u3, u4, u5]()),\n",
    "    ]\n",
    "    z: observes [c.u1] is u1\n",
    "    z: observes [c.u2] is u2\n",
    "    z: observes [c.u3] is u3\n",
    "    z: observes [c.u4] is u4\n",
    "    z: observes [c.u5] is u5\n",
    "    z: knows(m, d)\n",
    "    return z[E[c.m == m and c.d == d]]\n",
    "\n",
    "answer = puzzle()[:, :, U.DUNNO, U.DUNNO, U.DUNNO, U.KNOWN, U.KNOWN]\n",
    "\n",
    "%timeit -r 10 -n 100 out = puzzle().block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d79c2cc8-e317-4806-a4fb-e3d670d4322f",
   "metadata": {},
   "source": [
    "Finally, we extract the answer by finding the nonzero entry in the inferred $(m, d)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "757e3c66-2453-4e58-9d9f-043a4fa6df64",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "April 30\n"
     ]
    }
   ],
   "source": [
    "for m in Month:\n",
    "    for di, d in enumerate(Day):\n",
    "        if answer[m, di]:\n",
    "            print(m.name, d)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}