jaxpt_vmap.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "aVYWzHNFo-D9"
   },
   "outputs": [],
   "source": [
    "import jax.numpy as np\n",
    "import jax\n",
    "from jax import jit, grad, vmap, random, lax\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 54
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 449,
     "status": "ok",
     "timestamp": 1574887734924,
     "user": {
      "displayName": "Eric Jang",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAMjPmMQfp2nkKuIOsGK7ZRAYfuieou1bC1tdbMMQ=s64",
      "userId": "12717831053666785416"
     },
     "user_tz": 480
    },
    "id": "FwDE5gmGpJ4Q",
    "outputId": "8ce4a7d7-6743-435b-fefd-323ad11b75df"
   },
   "outputs": [],
   "source": [
    "RNG_KEY = random.PRNGKey(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iSnYhAWcr7Og"
   },
   "outputs": [],
   "source": [
    "# object ids\n",
    "OBJ_NONE=0.0\n",
    "OBJ_FLOOR=0.1\n",
    "OBJ_CEIL=.2\n",
    "OBJ_WALL_RD=.3\n",
    "OBJ_WALL_WH=.4\n",
    "OBJ_WALL_GR=.5\n",
    "OBJ_SHORT_BLOCK=.6\n",
    "OBJ_TALL_BLOCK=.7\n",
    "OBJ_LIGHT=1.0\n",
    "OBJ_SPHERE=0.9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EwNCtn3rpQvV"
   },
   "outputs": [],
   "source": [
    "def length(p):\n",
    "  return np.linalg.norm(p)\n",
    "\n",
    "def normalize(p):\n",
    "  return p/length(p)\n",
    "\n",
    "# def sdSphere(p,radius):\n",
    "#   return length(p) - radius\n",
    "\n",
    "def udBox(p, b):\n",
    "  # b = half-widths\n",
    "  return length(np.maximum(np.abs(p)-b,0.0))\n",
    "\n",
    "def rotateX(p,a):\n",
    "  c = np.cos(a); s = np.sin(a);\n",
    "  px,py,pz=p[0],p[1],p[2]\n",
    "  return np.array([px,c*py-s*pz,s*py+c*pz])\n",
    "\n",
    "def rotateY(p,a):\n",
    "  c = np.cos(a); s = np.sin(a);\n",
    "  px,py,pz=p[0],p[1],p[2]\n",
    "  return np.array([c*px+s*pz,py,-s*px+c*pz])\n",
    "\n",
    "def rotateZ(p,a):\n",
    "  c = np.cos(a); s = np.sin(a);\n",
    "  px,py,pz=p[0],p[1],p[2]\n",
    "  return np.array([c*px-s*py,s*px+c*py,pz])\n",
    "\n",
    "@jit\n",
    "def opU(a,b):\n",
    "  condition = np.tile(a[1,None]<b[1,None], [2])\n",
    "  return np.where(condition, a, b)\n",
    "\n",
    "def clamp01(v):\n",
    "  # maybe we should use sigmoid instead of hard thresholding for nicer gradients / soft shadows?\n",
    "  return np.minimum(np.maximum(v,0.0),1.0)\n",
    "\n",
    "def relu(a):\n",
    "  return np.maximum(a,0.)\n",
    "\n",
    "def dot(a,b):\n",
    "  return np.sum(a*b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "HFWTeTalpTlk"
   },
   "outputs": [],
   "source": [
    "def df(obj_id, dist):\n",
    "  return np.array([obj_id, dist])\n",
    "\n",
    "@jit\n",
    "def sdScene(p):\n",
    "  # p is [3,]\n",
    "  px,py,pz=p[0],p[1],p[2]\n",
    "  # floor\n",
    "  obj_floor = df(OBJ_FLOOR, py) # py = distance from y=0\n",
    "  res = obj_floor  \n",
    "  # ceiling\n",
    "  obj_ceil = df(OBJ_CEIL, 4.-py)\n",
    "  res = opU(res,obj_ceil)\n",
    "  # backwall\n",
    "  obj_bwall = df(OBJ_WALL_WH, 4.-pz)\n",
    "\n",
    "  res = opU(res,obj_bwall)\n",
    "  \n",
    "  # leftwall\n",
    "  obj_lwall = df(OBJ_WALL_RD, px-(-2))\n",
    "  res = opU(res,obj_lwall)\n",
    "  \n",
    "  # rightwall\n",
    "  obj_rwall = df(OBJ_WALL_GR, 2-px)\n",
    "  res = opU(res,obj_rwall)\n",
    "  \n",
    "  # light\n",
    "  obj_light = df(OBJ_LIGHT, udBox(p - np.array([0,3.9,2]), np.array([.5,.01,.5])))\n",
    "  res = opU(res,obj_light)\n",
    "  \n",
    "  # tall block\n",
    "  bh = 1.3\n",
    "  p2 = rotateY(p- np.array([-.64,bh,2.6]),.15*np.pi)\n",
    "  d = udBox(p2, np.array([.6,bh,.6]))\n",
    "  obj_tall_block = df(OBJ_TALL_BLOCK, d)\n",
    "  res = opU(res,obj_tall_block)  \n",
    "  \n",
    "  # short block\n",
    "  bw = .6\n",
    "  p2 = rotateY(p- np.array([.65,bw,1.7]),-.1*np.pi)\n",
    "  d = udBox(p2, np.array([bw,bw,bw]))\n",
    "  obj_short_block = df(OBJ_SHORT_BLOCK, d)\n",
    "  res = opU(res,obj_short_block)\n",
    "  \n",
    "  return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dXCXM-lapWM6"
   },
   "outputs": [],
   "source": [
    "# Discussion of cosine-weighted importance sampling for Lambertian BRDFs,\n",
    "# http://www.rorydriscoll.com/2009/01/07/better-sampling/\n",
    "# and implementation here:\n",
    "# https://www.shadertoy.com/view/4tl3z4\n",
    "@jit\n",
    "def sampleCosineWeightedHemisphere(rng_key, n):\n",
    "  rng_key, subkey = random.split(rng_key)\n",
    "  u = random.uniform(subkey,shape=(2,),minval=0,maxval=1)\n",
    "  u1, u2 = u[0], u[1]\n",
    "  uu = normalize(np.cross(n, np.array([0.,1.,1.])))\n",
    "  vv = np.cross(uu,n)\n",
    "  ra = np.sqrt(u2)\n",
    "  rx = ra*np.cos(2*np.pi*u1)\n",
    "  ry = ra*np.sin(2*np.pi*u1)\n",
    "  rz = np.sqrt(1.-u2)\n",
    "  rr = rx*uu+ry*vv+rz*n\n",
    "  return normalize(rr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 213
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 2330,
     "status": "ok",
     "timestamp": 1574887738659,
     "user": {
      "displayName": "Eric Jang",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAMjPmMQfp2nkKuIOsGK7ZRAYfuieou1bC1tdbMMQ=s64",
      "userId": "12717831053666785416"
     },
     "user_tz": 480
    },
    "id": "NvxBgZsdpYDU",
    "outputId": "c68905a4-c512-4e8e-a206-5ada1f25884d"
   },
   "outputs": [],
   "source": [
    "# testing cosine-weighted sphere projection\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "\n",
    "nor = normalize(np.array([[1.,1.,0.]]))\n",
    "nor = np.tile(nor,(1000,1))\n",
    "rng_key = random.split(RNG_KEY, 1000)\n",
    "rd = vmap(sampleCosineWeightedHemisphere)(rng_key, nor)\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(121, projection='3d')\n",
    "ax.scatter(rd[:,0],rd[:,2],rd[:,1])\n",
    "ax.set_xlabel('X')\n",
    "ax.set_ylabel('Z')\n",
    "ax.set_zlabel('Y')\n",
    "ax.set_aspect('equal')\n",
    "\n",
    "ax = fig.add_subplot(122)\n",
    "ax.scatter(rd[:,0],rd[:,1])\n",
    "ax.set_aspect('equal')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 34
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 488,
     "status": "ok",
     "timestamp": 1574887739184,
     "user": {
      "displayName": "Eric Jang",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAMjPmMQfp2nkKuIOsGK7ZRAYfuieou1bC1tdbMMQ=s64",
      "userId": "12717831053666785416"
     },
     "user_tz": 480
    },
    "id": "qKw6HW6o_4WL",
    "outputId": "59f1f188-eda1-41e2-9dd1-0070ecf1d32a"
   },
   "outputs": [],
   "source": [
    "fig.savefig('cos_sampling.png', pad_inches=0, bbox_inches='tight')\n",
    "# %download_file cos_sampling.png"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "k2oyeRL3pZK_"
   },
   "outputs": [],
   "source": [
    "MAX_ITERS=50\n",
    "HORIZON=20.0\n",
    "\n",
    "@jit\n",
    "def raymarch(ro,rd):\n",
    "  t = 0.0\n",
    "  for i in range(MAX_ITERS):\n",
    "    res = sdScene(ro + t*rd)\n",
    "    t += res[1] # t is scalar.\n",
    "  obj_id = res[0]\n",
    "  obj_id = np.where(t>HORIZON, OBJ_NONE, obj_id) \n",
    "  return df(obj_id, t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YiX-aPmIpbqY"
   },
   "outputs": [],
   "source": [
    "nor_light = np.array([0.,-1.,0.])\n",
    "LIGHT_AREA=1.*1.\n",
    "emissive_const = np.array([25, 25, 25]) # Watts\n",
    "emittedRadiance = emissive_const / (np.pi * LIGHT_AREA)\n",
    "lightDiffuseColor = np.array([0.2,0.2,0.2])\n",
    "leftWallColor = np.array([.611, .0555, .062]) * 1.5\n",
    "rightWallColor = np.array([.117, .4125, .115]) * 1.5\n",
    "whiteWallColor = np.array([255, 239, 196]) / 255\n",
    "\n",
    "\n",
    "BRDF_MAP = [\n",
    "  (OBJ_NONE, 0.0),\n",
    "  (OBJ_CEIL, whiteWallColor),\n",
    "  (OBJ_FLOOR, whiteWallColor),\n",
    "  (OBJ_LIGHT, lightDiffuseColor),\n",
    "  (OBJ_SHORT_BLOCK, whiteWallColor),\n",
    "  (OBJ_TALL_BLOCK, whiteWallColor),\n",
    "  (OBJ_WALL_GR, rightWallColor),\n",
    "  (OBJ_WALL_RD, leftWallColor),\n",
    "  (OBJ_WALL_WH, whiteWallColor),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ALjpExVBphlK"
   },
   "outputs": [],
   "source": [
    "# Calculating Normals with Finite-Differencing\n",
    "@jit\n",
    "def calcNormalFiniteDifference(p):\n",
    "  # derivative approximation via midpoint rule\n",
    "  eps = 0.001\n",
    "  dx=np.array([eps,0,0])\n",
    "  dy=np.array([0,eps,0])\n",
    "  dz=np.array([0,0,eps])\n",
    "  # extract just the distance component\n",
    "  nor = np.array([\n",
    "    sdScene(p+dx)[1] - sdScene(p-dx)[1],\n",
    "    sdScene(p+dy)[1] - sdScene(p-dy)[1],\n",
    "    sdScene(p+dz)[1] - sdScene(p-dz)[1],\n",
    "  ])\n",
    "  return normalize(nor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1mIRVmlrpU8-"
   },
   "outputs": [],
   "source": [
    "# Calculating Normals using AutoDiff\n",
    "@jit\n",
    "def dist(p):\n",
    "  # return the distance-component only\n",
    "  return sdScene(p)[1]\n",
    "\n",
    "@jit\n",
    "def calcNormal(p):\n",
    "  return normalize(grad(dist)(p))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "GborOW-vpjMp"
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def trace_direct(rng_key, ro, rd):\n",
    "  # emission and direct lighting contributions to radiance.\n",
    "  res = raymarch(ro,rd)\n",
    "  \n",
    "  # assign lambertian brdfs to intersections\n",
    "  condlist, choicelist = [], []\n",
    "  for obj_idx, obj_color in BRDF_MAP:\n",
    "    condlist.append(res[0] == obj_idx)\n",
    "    choicelist.append(obj_color)\n",
    "  brdf = np.select(condlist, choicelist)\n",
    "\n",
    "  radiance = np.zeros(3)\n",
    "\n",
    "  t = res[1]\n",
    "  eps = 0.0001\n",
    "  p = ro + (t-eps)*rd\n",
    "  nor = calcNormal(p)\n",
    "  did_intersect = res[0] != OBJ_NONE\n",
    "  rng_key, subkey = random.split(rng_key)\n",
    "\n",
    "  # Estimate direct area contribution\n",
    "  # each of the intersect points has some amount of light arriving\n",
    "  # sample point on area light TODO (ejang) - fix hardcoded light position\n",
    "  rng_key, subkey = random.split(rng_key)\n",
    "  p_light_x = random.uniform(subkey, minval=-.5, maxval=.5)\n",
    "  rng_key, subkey = random.split(rng_key)\n",
    "  p_light_z = random.uniform(subkey, minval=2. + -.5, maxval=2. + .5)\n",
    "  p_light = np.array([p_light_x, 3.9, p_light_z])\n",
    "  \n",
    "  wi_light = normalize(p_light - p)\n",
    "  res2 = raymarch(p + 0.001 * nor, wi_light)\n",
    "  # occlusion factor\n",
    "  vis = res2[0] == OBJ_LIGHT\n",
    "  pdf_A = 1./LIGHT_AREA\n",
    "  square_distance = np.sum(np.square(p_light-p))\n",
    "  Li_direct = (brdf * \n",
    "               relu(dot(nor, wi_light)) * \n",
    "               relu(dot(nor_light,-wi_light)) * \n",
    "               emittedRadiance) / (square_distance * pdf_A)\n",
    "  zeros = np.zeros_like(radiance)\n",
    "  isect_and_vis = np.logical_and(did_intersect, vis)\n",
    "  radiance += np.where(isect_and_vis, Li_direct, zeros)\n",
    "  return res, did_intersect, p, nor, brdf, radiance\n",
    "\n",
    "@jit\n",
    "def sample_recursive_ray(rng_key, p, nor):\n",
    "  # To allow JAX to jit this, we don't prune p to get did_intersect items.\n",
    "  ro2 = p + 0.001 * nor # bump along normal.\n",
    "  rd2 = sampleCosineWeightedHemisphere(rng_key, nor)\n",
    "  return ro2, rd2\n",
    "\n",
    "def trace(rng_key,ro,rd,depth):\n",
    "  res, did_intersect, p, nor, brdf, radiance = trace_direct(rng_key, ro, rd)\n",
    "  zeros = np.zeros_like(radiance)\n",
    "\n",
    "  # emitted radiance. This is only counted if it's an eye ray, since\n",
    "  # this contribution is also added in at every bounce.\n",
    "  if depth==0: # is eye ray\n",
    "    is_light = res[0] == OBJ_LIGHT\n",
    "    Li_e = emittedRadiance # Lambertian emitter\n",
    "    radiance += np.where(is_light, Li_e, zeros)\n",
    "\n",
    "  # indirect incoming contribution for intersected points\n",
    "  # instead of RR sampling, we just trace for a fixed number of steps\n",
    "  # to yield a good approximation and allow python to statically unroll.\n",
    "  if depth < 3: # if depth >= 4, hangs on compile?\n",
    "    rng_key, subkey = random.split(rng_key)\n",
    "    ro2, rd2 = sample_recursive_ray(rng_key, p, nor)\n",
    "    Li_indirect = trace(subkey,ro2,rd2,depth+1)\n",
    "    # doing cosweighted sampling cancels out the geom term\n",
    "    radiance += np.where(did_intersect, brdf * Li_indirect, zeros)\n",
    "  return radiance # actually, radiance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Sbz8l78xplyj"
   },
   "outputs": [],
   "source": [
    "# render the image\n",
    "# perspective camera with image plane centered at 0,0,0\n",
    "\n",
    "N=64 # width of image plane\n",
    "xs=np.linspace(0,1,N) # 10 pixels\n",
    "us,vs = np.meshgrid(xs,xs) \n",
    "uv = np.vstack([us.flatten(),vs.flatten()]).T # 10x10 image grid\n",
    "# normalize pixel locations to -1,1\n",
    "p = np.concatenate([-1+2*uv, np.zeros((N*N,1))], axis=1)\n",
    "# Render a pinhole camera.\n",
    "eye = np.tile(np.array([0,2.,-3.5]),(p.shape[0],1))\n",
    "look = np.array([[0,2.0,0]]) # look straight ahead\n",
    "w = vmap(normalize)(look - eye)\n",
    "up = np.array([[0,1,0]]) # up axis of world\n",
    "u = vmap(normalize)(np.cross(w,up))\n",
    "v = vmap(normalize)(np.cross(u,w))\n",
    "d=2.2 # focal distance\n",
    "rd = vmap(normalize)(p[:,0,None]*u + p[:,1,None]*v + d*w)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4RGI1knZDeZE"
   },
   "outputs": [],
   "source": [
    "# # UNCOMMENT TO USE Orthographic projection camera\n",
    "\n",
    "# N=150 # width of image plane\n",
    "# xs=np.linspace(0,1,N) # 10 pixels\n",
    "# us,vs = np.meshgrid(xs,xs) \n",
    "# us = (2*us-1)\n",
    "# vs *= 2\n",
    "# uv = np.vstack([us.flatten(),vs.flatten()]).T # 10x10 image grid\n",
    "# eye = np.concatenate([uv, np.zeros((N*N,1))], axis=1)*2\n",
    "# rd = np.zeros_like(eye) + np.array([[0, 0, 1]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "k8-BR3Mzr_Gi"
   },
   "source": [
    "# Render\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jn3MJME_pohd"
   },
   "outputs": [],
   "source": [
    "FULL_JIT = False\n",
    "DEVICE_PARALLEL = True\n",
    "\n",
    "render_fn = lambda rng_key, ro, rd : trace(rng_key, ro, rd, 0)\n",
    "\n",
    "if FULL_JIT:\n",
    "  trace = jit(trace, static_argnums=(3,))\n",
    "\n",
    "if DEVICE_PARALLEL:\n",
    "  vec_render_fn = jax.soft_pmap(render_fn)\n",
    "else:\n",
    "  vec_render_fn = vmap(render_fn) # will be very slow if we attempt to jit this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4MS39pwf5lI0"
   },
   "outputs": [],
   "source": [
    "rng_key = random.split(RNG_KEY, eye.shape[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 34
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 298,
     "status": "ok",
     "timestamp": 1574887743020,
     "user": {
      "displayName": "Eric Jang",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAMjPmMQfp2nkKuIOsGK7ZRAYfuieou1bC1tdbMMQ=s64",
      "userId": "12717831053666785416"
     },
     "user_tz": 480
    },
    "id": "NANdp3F25wpI",
    "outputId": "428192b2-4ed5-4662-9ced-762199bd35c9"
   },
   "outputs": [],
   "source": [
    "rng_key.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 68
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 93199,
     "status": "ok",
     "timestamp": 1574887836233,
     "user": {
      "displayName": "Eric Jang",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAMjPmMQfp2nkKuIOsGK7ZRAYfuieou1bC1tdbMMQ=s64",
      "userId": "12717831053666785416"
     },
     "user_tz": 480
    },
    "id": "-tl6e3HL1TSt",
    "outputId": "30c62026-47d0-47e4-8cdb-e7accafe37f9"
   },
   "outputs": [],
   "source": [
    "# %%time\n",
    "img = vec_render_fn(rng_key, eye, rd)\n",
    "print('done compiling')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 322
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 425,
     "status": "ok",
     "timestamp": 1574887836675,
     "user": {
      "displayName": "Eric Jang",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAMjPmMQfp2nkKuIOsGK7ZRAYfuieou1bC1tdbMMQ=s64",
      "userId": "12717831053666785416"
     },
     "user_tz": 480
    },
    "id": "seS8xihy_BUQ",
    "outputId": "a9bb130e-a5d1-46c8-d3d2-5b358198aa13"
   },
   "outputs": [],
   "source": [
    "img = np.fliplr(np.flipud(img.reshape((N,N,3))))\n",
    "plot = plt.imshow(img,interpolation='none',vmin = 0, vmax = 1)\n",
    "plt.grid('off')\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LUT-BpM23YkM"
   },
   "outputs": [],
   "source": [
    "# Render the same pixel multiple times.\n",
    "\n",
    "def vec_render_multisample(rng_key, eye, rd, num_samples):\n",
    "  img = vec_render_fn(rng_key, eye, rd)\n",
    "  for i in range(2, num_samples+1):\n",
    "    rng_key = random.split(rng_key[0], eye.shape[0])\n",
    "    sample = vec_render_fn(rng_key, eye, rd)\n",
    "    if i % 10 == 0:\n",
    "      print('Sample %d' % i)\n",
    "    img = (img + sample)\n",
    "  return np.fliplr(np.flipud(img.reshape((N,N,3))))/num_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 472
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 13873,
     "status": "ok",
     "timestamp": 1574887850885,
     "user": {
      "displayName": "Eric Jang",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAMjPmMQfp2nkKuIOsGK7ZRAYfuieou1bC1tdbMMQ=s64",
      "userId": "12717831053666785416"
     },
     "user_tz": 480
    },
    "id": "Rmq_hPwm4aM_",
    "outputId": "0f724ae7-6714-4cae-9354-f85d502c040a"
   },
   "outputs": [],
   "source": [
    "# %%time\n",
    "NUM_SAMPLES=100\n",
    "rng_key = random.split(RNG_KEY, eye.shape[0])\n",
    "img = vec_render_multisample(rng_key, eye, rd, NUM_SAMPLES)\n",
    "plot = plt.imshow(img,interpolation='none',vmin = 0, vmax = 1)\n",
    "plt.grid('off')\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vPjNQFw7z9JP"
   },
   "outputs": [],
   "source": [
    "plot.get_figure().savefig('cornell_box_ortho.png', bbox_inches='tight', pad_inches=0.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 33
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 340,
     "status": "ok",
     "timestamp": 1574745678981,
     "user": {
      "displayName": "Eric Jang",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAMjPmMQfp2nkKuIOsGK7ZRAYfuieou1bC1tdbMMQ=s64",
      "userId": "12717831053666785416"
     },
     "user_tz": 480
    },
    "id": "rwc66JKry8GM",
    "outputId": "63178485-1339-4f7d-9402-c7c8a20236b3"
   },
   "outputs": [],
   "source": [
    "%download_file cornell_box_ortho.png"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6UDUcIGEzLtg"
   },
   "outputs": [],
   "source": [
    "# plot = plt.imshow((img > 1).astype(np.float32),interpolation='none',vmin = 0, vmax = 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FDv-9Wsz-Y70"
   },
   "source": [
    "# Visualizing Normals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 285
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 11353,
     "status": "ok",
     "timestamp": 1574578299539,
     "user": {
      "displayName": "Eric Jang",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAMjPmMQfp2nkKuIOsGK7ZRAYfuieou1bC1tdbMMQ=s64",
      "userId": "12717831053666785416"
     },
     "user_tz": 480
    },
    "id": "OfL6Opbw-YbN",
    "outputId": "f7dc0815-4e2e-4688-a52d-92e8806f24fc"
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def calcNormal(p):\n",
    "  return normalize(grad(dist)(p))\n",
    "\n",
    "# @jit\n",
    "# def calcNormal(p):\n",
    "#   # derivative approximation via midpoint rule\n",
    "#   eps = 0.001\n",
    "#   dx=np.array([eps,0,0])\n",
    "#   dy=np.array([0,eps,0])\n",
    "#   dz=np.array([0,0,eps])\n",
    "#   # exactract just the distance component\n",
    "#   nor = np.array([\n",
    "#     sdScene(p+dx)[1] - sdScene(p-dx)[1],\n",
    "#     sdScene(p+dy)[1] - sdScene(p-dy)[1],\n",
    "#     sdScene(p+dz)[1] - sdScene(p-dz)[1],\n",
    "#   ])\n",
    "#   return normalize(nor)\n",
    "\n",
    "def trace(rng_key,ro,rd,depth):\n",
    "  res = raymarch(ro,rd)\n",
    "  t = res[1]\n",
    "  p = ro + (t-0.0001)*rd # the 0.001 is needed for stable normals.\n",
    "  nor = calcNormal(p)\n",
    "  return nor\n",
    "\n",
    "trace = jit(trace, static_argnums=(3,))\n",
    "render_fn = lambda rng_key, ro, rd : trace(rng_key, ro, rd, 0)\n",
    "vec_render_fn = vmap(render_fn) # will be very slow if we attempt to jit this.\n",
    "\n",
    "img = vec_render_fn(rng_key, eye, rd)\n",
    "\n",
    "img = np.fliplr(np.flipud(img.reshape((N,N,3))))\n",
    "plot = plt.imshow(img,interpolation='none')#,vmin = 0, vmax = 1)\n",
    "plt.grid('off')\n",
    "plt.axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 285
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 389,
     "status": "ok",
     "timestamp": 1574578300286,
     "user": {
      "displayName": "Eric Jang",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAMjPmMQfp2nkKuIOsGK7ZRAYfuieou1bC1tdbMMQ=s64",
      "userId": "12717831053666785416"
     },
     "user_tz": 480
    },
    "id": "tgMmjnLU_dgP",
    "outputId": "678f59af-ea98-4741-fa3c-9accd2a4f56f"
   },
   "outputs": [],
   "source": [
    "\n",
    "img = vec_render_fn(rng_key, eye, rd)\n",
    "\n",
    "img = np.fliplr(np.flipud(img.reshape((N,N,3))))\n",
    "plot = plt.imshow(img,interpolation='none',vmin = 0, vmax = 1)\n",
    "plt.grid('off')\n",
    "plt.axis('off')\n",
    "\n",
    "plt.imshow((img < 0).astype(np.float32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "height": 34
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 331,
     "status": "ok",
     "timestamp": 1574575467265,
     "user": {
      "displayName": "Eric Jang",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mAMjPmMQfp2nkKuIOsGK7ZRAYfuieou1bC1tdbMMQ=s64",
      "userId": "12717831053666785416"
     },
     "user_tz": 480
    },
    "id": "BCHvzS2xAtew",
    "outputId": "a7cd900d-75e3-43da-e838-9e2f92186be9"
   },
   "outputs": [],
   "source": [
    "img.min()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "1W9QMZvNgA4Z"
   },
   "source": [
    "# Bonus: lax.scan stack implementation\n",
    "\n",
    "Implementing path tracing as a for loop is a bit tricker than the typical C++ path tracer implementation, which is to use recursion. This implementation divides 'trace' into two `lax.scan` calls. \n",
    "\n",
    "The first scan call computes the new `ro, rd` for the next loop (e.g. the intersection and reflected eye ray), and the outputs are the emission and direct light contributions to radiance arriving at the intersected point `ro`. We also return BRDF and whether the `rd` rays intersected anything.\n",
    "\n",
    "After we finish the final intersection (`depth == 2`), we proceed backwards from the tail case and accumulate the indirect light contributions to the eye point. Here, we use the BRDF and intersection terms to attenuate the incoming light and compute outgoing light.\n",
    "\n",
    "If you think about this a little bit, this is essentially what a recursive implementation does, except instead of letting the interpreter push recursive `trace()` calls onto a stack (storing each bounce's emission and direct light contributions in the stack frame memory for use when we unwind out of recursion), we simply do this ourselves manually.\n",
    "\n",
    "This ended up crashing my computer because it consumed all available RAM. But I'm leaving this code snippet here because I found it a worthwhile exercise to think about."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QZ5V2dDjHc2Y"
   },
   "outputs": [],
   "source": [
    "# faster-to-jit version implemented with lax.scan\n",
    "# depth is static.\n",
    "# from functools import partial\n",
    "\n",
    "def scatter_eye_rays(carry, depth):\n",
    "  # Scatters eye-to-scene rays elsewhere into the scene, and carries these \n",
    "  # to next iteration. Output contains:\n",
    "  #   - emission + direct light contributions to the intersected point\n",
    "  #   - BRDF of intersected point\n",
    "  #   - did_intersect\n",
    "  rng_key, ro, rd = carry\n",
    "\n",
    "  res = raymarch(ro,rd)\n",
    "  \n",
    "  # assign lambertian brdfs to intersections\n",
    "  condlist, choicelist = [], []\n",
    "  for obj_idx, obj_color in BRDF_MAP:\n",
    "    condlist.append(res[0] == obj_idx)\n",
    "    choicelist.append(obj_color)\n",
    "  brdf = np.select(condlist, choicelist)\n",
    "  radiance = np.zeros(3)\n",
    "  # emitted radiance. This is only counted if it's an eye ray, since\n",
    "  # this contribution is also added in at every bounce.\n",
    "  # if depth==0: # is eye ray\n",
    "  is_light = res[0] == OBJ_LIGHT\n",
    "  is_eye_ray = np.maximum(0, 1.-depth) # hack to avoid 'Abstract value passed to `bool` error'\n",
    "  Li_e = relu(dot(-rd, nor_light)) * emittedRadiance\n",
    "  radiance = np.where(is_light, Li_e, radiance) * is_eye_ray\n",
    "  # end emitted radiance\n",
    "\n",
    "  t = res[1]\n",
    "  p = ro + t*rd\n",
    "  nor = calcNormal(p)\n",
    "  did_intersect = res[0] != OBJ_NONE\n",
    "  rng_key, subkey = random.split(rng_key)\n",
    "\n",
    "  # Estimate direct area contribution\n",
    "  # each of the intersect points has some amount of light arriving\n",
    "  # sample point on area light\n",
    "  rng_key, subkey = random.split(rng_key)\n",
    "  p_light_x = random.uniform(subkey, minval=-.25, maxval=.25)\n",
    "  rng_key, subkey = random.split(rng_key)\n",
    "  p_light_z = random.uniform(subkey, minval=2. + -.25, maxval=2. + .25)\n",
    "  p_light = np.array([p_light_x, 3.9, p_light_z])\n",
    "  \n",
    "  wi_light = normalize(p_light - p)\n",
    "  res2 = raymarch(p + 0.001 * nor, wi_light)\n",
    "  # occlusion factor\n",
    "  vis = res2[0] == OBJ_LIGHT\n",
    "  Li_direct = (brdf * \n",
    "               relu(dot(nor, wi_light)) * \n",
    "               relu(dot(nor_light,-wi_light)) * \n",
    "               emittedRadiance)\n",
    "  \n",
    "  isect_and_vis = np.logical_and(did_intersect, vis)\n",
    "  zeros = np.zeros_like(radiance)\n",
    "  radiance += np.where(isect_and_vis, Li_direct, zeros)\n",
    "  \n",
    "  ro2 = p + 0.001 * nor # bump along normal.\n",
    "  rng_key, subkey = random.split(rng_key)\n",
    "  rd2 = sampleCosineWeightedHemisphere(subkey, nor)\n",
    "\n",
    "  carry = (rng_key, ro, rd)\n",
    "  outputs = (did_intersect, radiance, brdf)\n",
    "  return carry, outputs\n",
    "  # rng_key, subkey = random.split(rng_key)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "weHOAUFXqzx3"
   },
   "source": [
    "In the accumulation step, we accumulate the outgoing radiance for each point by adding the indirect light contribution, starting at $p_3$ (defined to be zero), $p_2$, .. $p_0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "STz-SGcprKid"
   },
   "outputs": [],
   "source": [
    "def accumulate_outgoing(Li_indirect, x):\n",
    "  (did_intersect, radiance, brdf) = x\n",
    "  zeros = np.zeros(3)\n",
    "  carry = radiance + np.where(did_intersect, brdf * Li_indirect, zeros)\n",
    "  return carry, zeros"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "u2bnT54nqAd-"
   },
   "outputs": [],
   "source": [
    "def trace(rng_key,ro,rd,depth):\n",
    "  depths = np.array([0, 1, 2])\n",
    "  initial_value =  (rng_key, ro, rd)\n",
    "  carry, outputs = lax.scan(scatter_eye_rays, initial_value, depths)\n",
    "  # accumulate\n",
    "  initial_value = np.zeros(3) # incoming indirect radiance to the final intersection point is 0. (depth termination)\n",
    "  xs = tuple(np.flip(o,axis=0) for o in outputs)\n",
    "  radiance, _ = lax.scan(accumulate_outgoing, initial_value, xs)\n",
    "  return radiance"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "last_runtime": {
    "build_target": "//learning/deepmind/dm_python:dm_notebook3_tpu",
    "kind": "private"
   },
   "name": "jaxpt-vmap.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}