scripts/_mkl/notebooks/00b - Utils.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# My Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|default_exp utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "import matplotlib.pyplot as plt\n",
    "from   matplotlib.collections import LineCollection\n",
    "import numpy as np\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import genjax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## JAX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "key       = jax.random.PRNGKey(0)\n",
    "logsumexp = jax.scipy.special.logsumexp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def keysplit(key, *ns):\n",
    "    if len(ns) == 0:  \n",
    "        return jax.random.split(key, 1)[0]\n",
    "    elif len(ns) == 1:\n",
    "        n, = ns\n",
    "        if n == 1: return keysplit(key)\n",
    "        else:      return jax.random.split(key, ns[0])\n",
    "    else:\n",
    "        keys = []\n",
    "        for n in ns: keys.append(keysplit(key, n))\n",
    "        return keys\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([ 928981903, 3453687069], dtype=uint32),\n",
       " Array([ 928981903, 3453687069], dtype=uint32),\n",
       " Array([[3668660785,  713825972],\n",
       "        [1185646547, 2092858387],\n",
       "        [4260797006,  129535844],\n",
       "        [ 928977296, 1618649917],\n",
       "        [2708837749, 4129373854],\n",
       "        [ 652965180, 3955248629],\n",
       "        [1312337421, 1285539814],\n",
       "        [2974568872, 3669116123],\n",
       "        [1997906629, 3379841639],\n",
       "        [4278014892, 1203387755]], dtype=uint32),\n",
       " [Array([ 928981903, 3453687069], dtype=uint32),\n",
       "  Array([[3668660785,  713825972],\n",
       "         [1185646547, 2092858387],\n",
       "         [4260797006,  129535844],\n",
       "         [ 928977296, 1618649917],\n",
       "         [2708837749, 4129373854],\n",
       "         [ 652965180, 3955248629],\n",
       "         [1312337421, 1285539814],\n",
       "         [2974568872, 3669116123],\n",
       "         [1997906629, 3379841639],\n",
       "         [4278014892, 1203387755]], dtype=uint32)])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(\n",
    "    keysplit(key),\n",
    "    keysplit(key, 1),\n",
    "    keysplit(key, 10),\n",
    "    keysplit(key, 1, 10),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def bounding_box(arr, pad=0):\n",
    "    \"\"\"Takes a euclidean-like arr (`arr.shape[-1] == 2`) and returns its bounding box.\"\"\"\n",
    "    return jnp.array([\n",
    "        [jnp.min(arr[...,0])-pad, jnp.min(arr[...,1])-pad],\n",
    "        [jnp.max(arr[...,0])+pad, jnp.max(arr[...,1])+pad]\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def argmax_axes(a, axes=None):\n",
    "    \"\"\"Argmax along specified axes\"\"\"\n",
    "    if axes is None: return jnp.argmax(a)\n",
    "    \n",
    "    n = len(axes)        \n",
    "    axes_  = set(range(a.ndim))\n",
    "    axes_0 = axes\n",
    "    axes_1 = sorted(axes_ - set(axes_0))    \n",
    "    axes_  = axes_0 + axes_1\n",
    "\n",
    "    b = jnp.transpose(a, axes=axes_)\n",
    "    c = b.reshape(np.prod(b.shape[:n]), -1)\n",
    "\n",
    "    I = jnp.argmax(c, axis=0)\n",
    "    I = jnp.array([jnp.unravel_index(i, b.shape[:n]) for i in I]).reshape(b.shape[n:] + (n,))\n",
    "\n",
    "    return  I"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5, 9, 2)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_shape = (3, 99, 5, 9)\n",
    "a = jnp.arange(np.prod(test_shape)).reshape(test_shape)\n",
    "\n",
    "I = argmax_axes(a, axes=[0,1])\n",
    "I.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Poses and Geometry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def cam_to_screen(x): return jnp.array([x[0]/x[2], x[1]/x[2], jnp.linalg.norm(x)])\n",
    "def screen_to_cam(y): return y[2]*jnp.array([y[0], y[1], 1.0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def rot2d(hd): return jnp.array([\n",
    "    [jnp.cos(hd), -jnp.sin(hd)], \n",
    "    [jnp.sin(hd),  jnp.cos(hd)]\n",
    "    ]);\n",
    "\n",
    "def pack_2dpose(x,hd): \n",
    "    return jnp.concatenate([x,jnp.array([hd])])\n",
    "\n",
    "def apply_2dpose(p, ys): \n",
    "    return ys@rot2d(p[2] - jnp.pi/2).T + p[:2]\n",
    "\n",
    "def unit_vec(hd): \n",
    "    return jnp.array([jnp.cos(hd), jnp.sin(hd)])\n",
    "\n",
    "def adjust_angle(hd):\n",
    "    \"\"\"Adjusts angle to lie in the interval [-pi,pi).\"\"\"\n",
    "    return (hd + jnp.pi)%(2*jnp.pi) - jnp.pi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GenJAX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "from genjax.incremental import UnknownChange, NoChange, Diff\n",
    "\n",
    "\n",
    "def argdiffs(args, other=None):\n",
    "    return tuple(map(lambda v: Diff(v, UnknownChange), args))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "from builtins import property as _property, tuple as _tuple\n",
    "from typing import Any\n",
    "\n",
    "\n",
    "class Args(tuple):\n",
    "    def __new__(cls, *args, **kwargs):\n",
    "        return _tuple.__new__(cls, list(args) + list(kwargs.values()))\n",
    "    \n",
    "    def __init__(self, *args, **kwargs):\n",
    "        self._d = dict()\n",
    "        for k,v in kwargs.items():\n",
    "            self._d[k] = v\n",
    "            setattr(self, k, v)\n",
    "\n",
    "    def __getitem__(self, k: str) -> Any:\n",
    "        return self._d[k]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "__main__.Args"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(Args(x=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "# \n",
    "# Monkey patching `sample` for `BuiltinGenerativeFunction`\n",
    "# \n",
    "cls = genjax._src.generative_functions.static.static_gen_fn.StaticGenerativeFunction\n",
    "\n",
    "def genjax_sample(self, key, *args, **kwargs):\n",
    "    tr = self.simulate(key, args)\n",
    "    return tr.get_retval()\n",
    "\n",
    "setattr(cls, \"sample\", genjax_sample)\n",
    "\n",
    "\n",
    "# \n",
    "# Monkey patching `sample` for `DeferredGenerativeFunctionCall`\n",
    "# \n",
    "cls = genjax._src.generative_functions.supports_callees.SugaredGenerativeFunctionCall\n",
    "\n",
    "def deff_gen_func_call(self, key, **kwargs):\n",
    "    return self.gen_fn.sample(key, *self.args, **kwargs)\n",
    "\n",
    "def deff_gen_func_logpdf(self, x, **kwargs):\n",
    "    return self.gen_fn.logpdf(x, *self.args, **kwargs)\n",
    "\n",
    "setattr(cls, \"__call__\", deff_gen_func_call)\n",
    "setattr(cls, \"sample\", deff_gen_func_call)\n",
    "setattr(cls, \"logpdf\", deff_gen_func_logpdf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bayes3d",
   "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.11.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}