pangolin-jags/Example-vmap.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d3819450",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pangolin import Given, d, t, I, IID, vmap, sample, E, P, var, std, cov, corr, jags_code\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "np.set_printoptions(formatter={'float': '{:6.2f}'.format}) # print nicely"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a8c47888",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "warning: no model code - (OK in principle...)\n",
      "E(y) [  2.00   2.00   2.00]\n"
     ]
    }
   ],
   "source": [
    "# you can \"vectorize\" constants\n",
    "\n",
    "x = 2\n",
    "y = IID(x,3)\n",
    "print('E(y)',E(y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9fa71b4f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "E(x)     [  3.08   8.04]\n",
      "expected [  3.00   8.00]\n"
     ]
    }
   ],
   "source": [
    "# very simple vectorization\n",
    "\n",
    "def f(a,b):\n",
    "    return d.norm(a*b,1)\n",
    "\n",
    "A = np.array([1.,2.])\n",
    "B = np.array([3.,4.])\n",
    "\n",
    "x= vmap(f)(A,B)\n",
    "print('E(x)    ',E(x))\n",
    "expected = A*B\n",
    "print('expected',expected)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "424cf7b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "E(x)     [  1.05   2.02]\n",
      "expected [  1.00   2.00]\n",
      "E(y)     [  2.97   8.04]\n",
      "expected [  3.00   8.00]\n"
     ]
    }
   ],
   "source": [
    "def f(a,b):\n",
    "    u = a*b\n",
    "    #s = a+b\n",
    "    return [d.norm(a,b),d.norm(u,1)]\n",
    "\n",
    "A = np.array([1.,2.])\n",
    "B = np.array([3.,4.])\n",
    "\n",
    "x,y = vmap(f)(A,B)\n",
    "print('E(x)    ',E(x))\n",
    "expected = A\n",
    "print('expected',expected)\n",
    "print('E(y)    ',E(y))\n",
    "expected = np.array(A)*np.array(B)\n",
    "print('expected',expected)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "30a81703",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "traced version\n",
      "[  0.54   1.42   2.55]\n",
      "untraced version\n",
      "[0.5466128421399999, 1.47939783778]\n"
     ]
    }
   ],
   "source": [
    "# Tracing that uses IID\n",
    "def f(a,b):\n",
    "    x = a*b\n",
    "    y = a+b\n",
    "    return [d.norm(x,y)]\n",
    "\n",
    "A = IID(d.exp(1),3)\n",
    "B = [.5,1.5,2.5]\n",
    "#x,=RV.trace(f,A,B)\n",
    "x, = vmap(f)(A,B)\n",
    "print('traced version')\n",
    "print(E(x))\n",
    "\n",
    "print('untraced version')\n",
    "x0 = A[0]*B[0]\n",
    "y0 = A[0]+B[0]\n",
    "x1 = A[1]*B[1]\n",
    "y1 = A[1]+B[1]\n",
    "x0 = d.norm(x0,y0)\n",
    "x1 = d.norm(x1,y1)\n",
    "print(E([x0,x1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a6dfdba0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "E(x)   [  0.98   2.27   3.01]\n",
      "std(x) [  3.98   4.78   6.15]\n"
     ]
    }
   ],
   "source": [
    "# here's a realistic-ish example: you might want to create a normal\n",
    "# parameterized by location and scale\n",
    "\n",
    "def norm_loc_scale(loc,scale):\n",
    "    return d.norm(loc,1/scale**2)\n",
    "\n",
    "vec_norm_loc_scale = vmap(norm_loc_scale)\n",
    "\n",
    "x = vec_norm_loc_scale([1,2,3],[4,5,6])\n",
    "print('E(x)  ',E(x))\n",
    "print('std(x)',std(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "926cdeac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([  3.00,   6.00])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# vmap and indexing\n",
    "x_train = [1,2]\n",
    "\n",
    "def get_yi(xi,w):\n",
    "    return xi*w[0]\n",
    "\n",
    "# vectorized_version\n",
    "get_y = vmap(get_yi,[True,False])\n",
    "\n",
    "#w = RV.Constant([3,3,3])\n",
    "w = [3,4,5]\n",
    "y = get_y(x_train,w)\n",
    "E(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "31861f55",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([  0.98,   2.04]), array([  3.01,   3.03])]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# a vmap still works even if one variable flows straight to output\n",
    "def f(a,b):\n",
    "    return [d.norm(a,1),d.norm(b,1)]\n",
    "\n",
    "E(vmap(f,[True,False])([1,2],3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "699d67c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "E(x)   [array([  1.97,   3.04]), array([  0.00,   0.00]), array([  1.00,   1.00])]\n",
      "var(x) [array([  1.00,   0.96]), array([  0.00,   0.00]), array([  0.00,   0.00])]\n"
     ]
    }
   ],
   "source": [
    "# vectorization works even if some things are constants\n",
    "# or are even constants\n",
    "def f(a):\n",
    "    return [d.norm(a,1),a,1] #,a\n",
    "\n",
    "x = vmap(f)([2,3])\n",
    "print('E(x)  ',E(x))\n",
    "print('var(x)',var(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9fc9bbf9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "E(x)\n",
      "[[  1.09  -3.06  -1.41]\n",
      " [ -1.92   0.24  -0.23]\n",
      " [  0.17  -1.99  -0.45]\n",
      " [  0.01   0.47  -1.48]\n",
      " [ -2.94   1.40  -0.04]]\n",
      "expected\n",
      "[[  1.09  -3.06  -1.41]\n",
      " [ -1.92   0.24  -0.23]\n",
      " [  0.17  -1.99  -0.45]\n",
      " [  0.01   0.47  -1.48]\n",
      " [ -2.94   1.40  -0.04]]\n"
     ]
    }
   ],
   "source": [
    "# vmap with indexing\n",
    "\n",
    "def f(a,b):\n",
    "    return a[0]*b\n",
    "\n",
    "A = np.random.randn(5,2)\n",
    "B = np.random.randn(5,3)\n",
    "x = vmap(f)(A,B)\n",
    "print('E(x)')\n",
    "print(E(x))\n",
    "print('expected')\n",
    "print(A[:,0,None]*B)"
   ]
  }
 ],
 "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}