demo/demo-schelling.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "489aee96-05b8-426e-bafe-5750cd9a73e6",
   "metadata": {},
   "source": [
    "# Schelling game\n",
    "\n",
    "**Inspired by:** Schelling, T. C. (1980). _The Strategy of Conflict._ Harvard university press.\n",
    "\n",
    "Alice and Bob agree to meet at a bar on Sunday, but they forget to decide on a bar to meet at. One bar is slightly more popular than the other. Where do they go? We can model this with recursive reasoning — Alice thinks about where Bob thinks Alice might go."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "50196ddd-e76d-4193-ba2e-6cdfff225e5b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "alice(1) = [0.59900993 0.40099007]; bob(1) = [0.64611655 0.35388345]\n",
      "alice(2) = [0.6905481  0.30945188]; bob(2) = [0.73171747 0.26828253]\n",
      "alice(3) = [0.76923996 0.23076005]; bob(3) = [0.8029279 0.1970721]\n",
      "alice(4) = [0.832767   0.16723298]; bob(4) = [0.8588822 0.1411178]\n",
      "alice(5) = [0.88149947 0.11850046]; bob(5) = [0.90091014 0.09908987]\n",
      "alice(6) = [0.91743904 0.08256097]; bob(6) = [0.9314207  0.06857934]\n",
      "alice(7) = [0.9431811  0.05681884]; bob(7) = [0.9530266  0.04697341]\n",
      "alice(8) = [0.9612362  0.03876386]; bob(8) = [0.96805906 0.031941  ]\n",
      "alice(9) = [0.9737139 0.0262862]; bob(9) = [0.97838986 0.02161017]\n"
     ]
    }
   ],
   "source": [
    "from memo import memo\n",
    "import jax\n",
    "import jax.numpy as np\n",
    "\n",
    "Bar = np.arange(2)\n",
    "@jax.jit\n",
    "def prior(b): return np.array([0.55, 0.45])[b]\n",
    "\n",
    "@memo\n",
    "def alice[b: Bar](depth):\n",
    "    alice: thinks[ bob: chooses(b in Bar, wpp=bob[b](depth - 1)) ]\n",
    "    alice: chooses(b in Bar, wpp=prior(b) * Pr[b == bob.b])\n",
    "    return Pr[alice.b == b]\n",
    "\n",
    "@memo\n",
    "def bob[b: Bar](depth):\n",
    "    bob: thinks[ alice: chooses(b in Bar, wpp=alice[b](depth) if depth > 0 else 1) ]\n",
    "    bob: chooses(b in Bar, wpp=prior(b) * Pr[b == alice.b])\n",
    "    return Pr[bob.b == b]\n",
    "\n",
    "for i in range(1, 10):\n",
    "    print(f'alice({i}) = {alice(i)}; bob({i}) = {bob(i)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2533c2d1-9e3b-414f-b7f8-e751f03843e9",
   "metadata": {},
   "source": [
    "Notice the rapid convergence to the more popular bar.\n",
    "\n",
    "How confident is Alice in meeting Bob? How confident is an observer that Alice will meet Bob? Which do you expect to be higher?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2b23102d-3a45-4742-8295-ee085913cf87",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Alice is 0.510 confident, observer is 0.529 confident.\n",
      "Alice is 0.556 confident, observer is 0.588 confident.\n",
      "Alice is 0.625 confident, observer is 0.663 confident.\n",
      "Alice is 0.702 confident, observer is 0.739 confident.\n",
      "Alice is 0.774 confident, observer is 0.806 confident.\n",
      "Alice is 0.835 confident, observer is 0.860 confident.\n",
      "Alice is 0.882 confident, observer is 0.902 confident.\n",
      "Alice is 0.918 confident, observer is 0.932 confident.\n",
      "Alice is 0.943 confident, observer is 0.953 confident.\n"
     ]
    }
   ],
   "source": [
    "@memo\n",
    "def alice_confidence(depth):\n",
    "    alice: thinks[ bob: chooses(b in Bar, wpp=bob[b](depth - 1)) ]\n",
    "    alice: chooses(b in Bar, wpp=prior(b) * Pr[b == bob.b])\n",
    "    return E[alice[Pr[b == bob.b]]]\n",
    "\n",
    "@memo\n",
    "def obs_confidence(depth):\n",
    "    alice: chooses(b in Bar, wpp=alice[b](depth))\n",
    "    bob: chooses(b in Bar, wpp=bob[b](depth))\n",
    "    return Pr[alice.b == bob.b]\n",
    "\n",
    "for i in range(1, 10):\n",
    "    print(f'Alice is {alice_confidence(i):.3f} confident, observer is {obs_confidence(i):.3f} confident.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7b09ce2-6df4-410d-9c9a-4ba05fcbd745",
   "metadata": {},
   "source": [
    "The observer is always slightly more confident, because they think both Alice and Bob are thinking at level $i$. Alice on the other hand thinks Bob is thinking at level $i-1$.\n",
    "\n",
    "The code below reproduces the scaled-up experiment shown in the paper (100 bars, 100 levels of recursion)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "95af5623-f5c5-4a5d-b225-138f08c424ef",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.06 ms ± 14.4 μs per loop (mean ± std. dev. of 10 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "Bar = np.arange(100)\n",
    "@jax.jit\n",
    "def prior(b): return 1\n",
    "\n",
    "@memo\n",
    "def alice[b: Bar](depth):\n",
    "    alice: thinks[ bob: chooses(b in Bar, wpp=bob[b](depth - 1)) ]\n",
    "    alice: chooses(b in Bar, wpp=prior(b) * Pr[b == bob.b])\n",
    "    return Pr[alice.b == b]\n",
    "\n",
    "@memo\n",
    "def bob[b: Bar](depth):\n",
    "    bob: thinks[ alice: chooses(b in Bar, wpp=alice[b](depth) if depth > 0 else 1) ]\n",
    "    bob: chooses(b in Bar, wpp=prior(b) * Pr[b == alice.b])\n",
    "    return Pr[bob.b == b]\n",
    "\n",
    "alice(100)\n",
    "%timeit -r 10 -n 100 alice(100).block_until_ready()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}