fix bad variable names in _get_output_array_size

Author: s-kganz

notebooks/functions.ipynb

@@ -11,7 +11,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "405f62a3-f187-436b-a94c-1658014c9aae",
+   "id": "0d976e10-f285-4682-a532-ecfe9106a9f8",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -36,16 +36,23 @@
     "        else:\n",
     "            # This is a resampled axis, determine the new size\n",
     "            # by the ratio of the batchgen window to the tensor size.\n",
-    "            window_size = ds.X_generator.input_dims[key]\n",
+    "            window_size = bgen.input_dims[key]\n",
     "            tensor_size = output_tensor_dim[key]\n",
     "            resample_ratio = tensor_size / window_size\n",
     "    \n",
-    "            temp_output_size = ds.X_generator.ds.sizes[key] * resample_ratio\n",
+    "            temp_output_size = bgen.ds.sizes[key] * resample_ratio\n",
     "            assert temp_output_size.is_integer()\n",
     "            output_size[key] = int(temp_output_size)\n",
-    "    return output_size\n",
-    "    \n",
-    "\n",
+    "    return output_size"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "eb717cb1-1715-4441-8b67-5593d1dbd6b8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "def predict_on_array(\n",
     "    dataset: MapDataset,\n",
     "    model: torch.nn.Module,\n",
@@ -62,24 +69,29 @@
     "        dims=tuple(output_size.keys())\n",
     "    )\n",
     "    output_n = xr.full_like(output_da, 0)\n",
-    "\n",
-    "    '''\n",
+    "    \n",
     "    # Prepare data laoder\n",
     "    loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)\n",
+    "\n",
+    "    '''\n",
     "    for i, batch in enumerate(loader):\n",
     "        out_batch = model(batch).detach().numpy()\n",
-    "        # TODO write each example to the output array\n",
+    "\n",
+    "        # Iterate over each example in the batch\n",
     "        for ib in range(out_batch.shape[0]):\n",
     "            # Get the slice object associated with this example\n",
     "            old_indexer = dataset.X_generator._batch_selectors.selectors[(i*batch_size)+ib][0]\n",
     "            # Only index into axes that are resampled, rescaling the bounds\n",
-    "            new_indexer = dict(\n",
-    "                dim: slice()\n",
-    "                for dim in old_indexer if dim in resample_dim\n",
-    "            )\n",
-    "            output_da.loc[indexer] += model_output\n",
-    "            output_n.loc[indexer] += 1\n",
+    "            new_indexer = dict()\n",
+    "            for key in old_indexer:\n",
+    "                if key in resample_dim:\n",
+    "                    resample_ratio = output_tensor_dim[key] / dataset.X_generator.input_dims[key]\n",
+    "\n",
+    "            \n",
+    "            output_da.loc[new_indexer] += out_batch[ib, ...]\n",
+    "            output_n.loc[new_indexer] += 1\n",
     "    '''\n",
+    "    \n",
     "\n",
     "    # TODO aggregate output\n",
     "    return output_da"