Improve argument unzipping

Author: eigenvivek

diffdrr/renderers.py

@@ -43,14 +43,12 @@ def forward(self, volume, spacing, source, target):
 # %% ../notebooks/api/01_renderers.ipynb 8
 def _get_alphas(source, target, spacing, dims, eps):
     # Get the CT sizing and spacing parameters
-    dx, dy, dz = spacing
-    nx, ny, nz = dims
+    alphax = torch.arange(dims[0]).to(source) * spacing[0]
+    alphay = torch.arange(dims[1]).to(source) * spacing[1]
+    alphaz = torch.arange(dims[2]).to(source) * spacing[2]
 
     # Get the alpha at each plane intersection
     sx, sy, sz = source[..., 0], source[..., 1], source[..., 2]
-    alphax = torch.arange(nx).to(source) * dx
-    alphay = torch.arange(ny).to(source) * dy
-    alphaz = torch.arange(nz).to(source) * dz
     alphax = alphax.expand(len(source), 1, -1) - sx.unsqueeze(-1)
     alphay = alphay.expand(len(source), 1, -1) - sy.unsqueeze(-1)
     alphaz = alphaz.expand(len(source), 1, -1) - sz.unsqueeze(-1)
@@ -67,9 +65,11 @@ def _get_alphas(source, target, spacing, dims, eps):
     alphas[~good_idxs] = torch.nan
 
     # Sort the alphas by ray, putting nans at the end of the list
-    # Drop indices where alphas for all rays are nan
     alphas = torch.sort(alphas, dim=-1).values
+
+    # Drop indices where alphas for all rays are nan
     alphas = alphas[..., ~alphas.isnan().all(dim=0).all(dim=0)]
+
     return alphas
 
 
@@ -110,7 +110,7 @@ def _get_index(alpha, source, target, spacing, dims, maxidx, eps):
     idxs[idxs > maxidx] = maxidx
     return idxs
 
-# %% ../notebooks/api/01_renderers.ipynb 10
+# %% ../notebooks/api/01_renderers.ipynb 12
 from torch.nn.functional import grid_sample
 
 

notebooks/api/01_renderers.ipynb

@@ -151,14 +151,12 @@
     "#| export\n",
     "def _get_alphas(source, target, spacing, dims, eps):\n",
     "    # Get the CT sizing and spacing parameters\n",
-    "    dx, dy, dz = spacing\n",
-    "    nx, ny, nz = dims\n",
+    "    alphax = torch.arange(dims[0]).to(source) * spacing[0]\n",
+    "    alphay = torch.arange(dims[1]).to(source) * spacing[1]\n",
+    "    alphaz = torch.arange(dims[2]).to(source) * spacing[2]\n",
     "\n",
     "    # Get the alpha at each plane intersection\n",
     "    sx, sy, sz = source[..., 0], source[..., 1], source[..., 2]\n",
-    "    alphax = torch.arange(nx).to(source) * dx\n",
-    "    alphay = torch.arange(ny).to(source) * dy\n",
-    "    alphaz = torch.arange(nz).to(source) * dz\n",
     "    alphax = alphax.expand(len(source), 1, -1) - sx.unsqueeze(-1)\n",
     "    alphay = alphay.expand(len(source), 1, -1) - sy.unsqueeze(-1)\n",
     "    alphaz = alphaz.expand(len(source), 1, -1) - sz.unsqueeze(-1)\n",
@@ -175,9 +173,11 @@
     "    alphas[~good_idxs] = torch.nan\n",
     "\n",
     "    # Sort the alphas by ray, putting nans at the end of the list\n",
-    "    # Drop indices where alphas for all rays are nan\n",
     "    alphas = torch.sort(alphas, dim=-1).values\n",
+    "    \n",
+    "    # Drop indices where alphas for all rays are nan\n",
     "    alphas = alphas[..., ~alphas.isnan().all(dim=0).all(dim=0)]\n",
+    "\n",
     "    return alphas\n",
     "\n",
     "\n",
@@ -219,6 +219,20 @@
     "    return idxs"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Compilier-friendly Siddon"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -238,19 +252,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "ename": "NameError",
-     "evalue": "name 'torch' is not defined",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[1], line 5\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m#| export\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnn\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mfunctional\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m grid_sample\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mTrilinear\u001b[39;00m(\u001b[43mtorch\u001b[49m\u001b[38;5;241m.\u001b[39mnn\u001b[38;5;241m.\u001b[39mModule):\n\u001b[1;32m      6\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m      7\u001b[0m \u001b[38;5;124;03m    Instead of computing the exact line integral over the voxel grid (i.e., Siddon's method), we can sample colors at points along the each ray using trilinear interpolation.\u001b[39;00m\n\u001b[1;32m      8\u001b[0m \u001b[38;5;124;03m    \u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[38;5;124;03m    where $\\mathbf V[\\cdot]$ is the trilinear interpolation function and $M$ is the number of points sampled per ray.\u001b[39;00m\n\u001b[1;32m     14\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[1;32m     15\u001b[0m     \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\n\u001b[1;32m     16\u001b[0m         \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m     17\u001b[0m         near\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0.0\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     20\u001b[0m         eps\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1e-8\u001b[39m,\n\u001b[1;32m     21\u001b[0m     ):\n",
-      "\u001b[0;31mNameError\u001b[0m: name 'torch' is not defined"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "#| export\n",
     "from torch.nn.functional import grid_sample\n",