Simplify patch rendering

Author: eigenvivek

diffdrr/drr.py

@@ -194,16 +194,12 @@ def render(
             **kwargs,
         )
     else:
-        n_points = target.shape[1] // self.n_patches
         partials = []
-        for idx in range(self.n_patches):
-            partial = self.renderer(
-                density,
-                source,
-                target[:, idx * n_points : (idx + 1) * n_points],
-                img[..., idx * n_points : (idx + 1) * n_points],
-                **kwargs,
-            )
+        for t, i in zip(
+            target.chunk(self.n_patches, dim=1),
+            img.chunk(self.n_patches, dim=-1),
+        ):
+            partial = self.renderer(density, source, t, i, **kwargs)
             partials.append(partial)
         img = torch.cat(partials, dim=-1)
 

notebooks/api/00_drr.ipynb

@@ -313,16 +313,12 @@
     "            **kwargs,\n",
     "        )\n",
     "    else:\n",
-    "        n_points = target.shape[1] // self.n_patches\n",
     "        partials = []\n",
-    "        for idx in range(self.n_patches):\n",
-    "            partial = self.renderer(\n",
-    "                density,\n",
-    "                source,\n",
-    "                target[:, idx * n_points : (idx + 1) * n_points],\n",
-    "                img[..., idx * n_points : (idx + 1) * n_points],\n",
-    "                **kwargs,\n",
-    "            )\n",
+    "        for t, i in zip(\n",
+    "            target.chunk(self.n_patches, dim=1),\n",
+    "            img.chunk(self.n_patches, dim=-1),\n",
+    "        ):\n",
+    "            partial = self.renderer(density, source, t, i, **kwargs)\n",
     "            partials.append(partial)\n",
     "        img = torch.cat(partials, dim=-1)\n",
     "\n",