Merge pull request #65 from theislab/fix/gwp_prepare

fixed attr arguments
theislab · Mar 8, 2024 · 47989fb · 47989fb
2 parents 5856363 + e061b3c
commit 47989fb
Show file tree

Hide file tree

Showing 4 changed files with 200 additions and 51 deletions.
diff --git a/examples/problems/100_tagged_arrays.ipynb b/examples/problems/100_tagged_arrays.ipynb
@@ -36,6 +36,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import warnings\n",
+    "\n",
+    "warnings.simplefilter(\"ignore\", FutureWarning)\n",
+    "\n",
     "from moscot import datasets\n",
     "from moscot.problems.generic import GWProblem\n",
     "\n",
@@ -58,7 +62,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "id": "dec86f59",
    "metadata": {},
    "outputs": [
@@ -72,7 +76,7 @@
        "    varm: 'PCs'"
       ]
      },
-     "execution_count": 3,
+     "execution_count": 2,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -96,13 +100,26 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "id": "16f0c3a9",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34mINFO    \u001b[0m Ordering \u001b[1;35mIndex\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[32m'0-0'\u001b[0m, \u001b[32m'1-0'\u001b[0m, \u001b[32m'2-0'\u001b[0m, \u001b[32m'3-0'\u001b[0m, \u001b[32m'4-0'\u001b[0m, \u001b[32m'5-0'\u001b[0m, \u001b[32m'6-0'\u001b[0m, \u001b[32m'7-0'\u001b[0m, \u001b[32m'8-0'\u001b[0m, \u001b[32m'9-0'\u001b[0m,                     \n",
+      "                \u001b[32m'10-0'\u001b[0m, \u001b[32m'11-0'\u001b[0m, \u001b[32m'12-0'\u001b[0m, \u001b[32m'13-0'\u001b[0m, \u001b[32m'14-0'\u001b[0m, \u001b[32m'15-0'\u001b[0m, \u001b[32m'16-0'\u001b[0m, \u001b[32m'17-0'\u001b[0m, \u001b[32m'18-0'\u001b[0m,                            \n",
+      "                \u001b[32m'19-0'\u001b[0m, \u001b[32m'0-1'\u001b[0m, \u001b[32m'1-1'\u001b[0m, \u001b[32m'2-1'\u001b[0m, \u001b[32m'3-1'\u001b[0m, \u001b[32m'4-1'\u001b[0m, \u001b[32m'5-1'\u001b[0m, \u001b[32m'6-1'\u001b[0m, \u001b[32m'7-1'\u001b[0m, \u001b[32m'8-1'\u001b[0m,                             \n",
+      "                \u001b[32m'9-1'\u001b[0m, \u001b[32m'10-1'\u001b[0m, \u001b[32m'11-1'\u001b[0m, \u001b[32m'12-1'\u001b[0m, \u001b[32m'13-1'\u001b[0m, \u001b[32m'14-1'\u001b[0m, \u001b[32m'15-1'\u001b[0m, \u001b[32m'16-1'\u001b[0m, \u001b[32m'17-1'\u001b[0m,                             \n",
+      "                \u001b[32m'18-1'\u001b[0m, \u001b[32m'19-1'\u001b[0m\u001b[1m]\u001b[0m,                                                                                   \n",
+      "               \u001b[33mdtype\u001b[0m=\u001b[32m'object'\u001b[0m\u001b[1m)\u001b[0m in ascending order.                                                                 \n"
+     ]
+    }
+   ],
    "source": [
     "gwp = GWProblem(adata)\n",
-    "gwp = gwp.prepare(key=\"batch\", joint_attr=\"X_pca\", GW_x=\"X_pca\", GW_y=\"X_pca\")"
+    "gwp = gwp.prepare(key=\"batch\", x_attr=\"X_pca\", y_attr=\"X_pca\", joint_attr=\"X_pca\")"
    ]
   },
   {
@@ -116,7 +133,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "id": "e4ca2aea",
    "metadata": {},
    "outputs": [
@@ -137,10 +154,10 @@
        "           [ 2.024, -1.597, -0.591, ..., -0.114,  0.29 , -0.332],\n",
        "           [-0.938,  2.426, -0.128, ...,  0.194,  0.829, -0.438],\n",
        "           [ 2.709, -2.885, -0.925, ...,  0.315,  0.334,  0.078]],\n",
-       "          dtype=float32), tag='point_cloud', cost=<ott.geometry.costs.SqEuclidean object at 0x29e7bb310>)"
+       "          dtype=float32), tag=<Tag.POINT_CLOUD: 'point_cloud'>, cost=<ott.geometry.costs.SqEuclidean object at 0x000002837922CC10>)"
       ]
      },
-     "execution_count": 5,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -162,17 +179,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "id": "7b70c639",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "('point_cloud', <ott.geometry.costs.SqEuclidean at 0x29e7bb310>)"
+       "(<Tag.POINT_CLOUD: 'point_cloud'>,\n",
+       " <ott.geometry.costs.SqEuclidean at 0x2837922cc10>)"
       ]
      },
-     "execution_count": 6,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -198,7 +216,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 6,
    "id": "e50ce79b",
    "metadata": {},
    "outputs": [
@@ -223,7 +241,7 @@
        "           dtype=float32))"
       ]
      },
-     "execution_count": 7,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -245,7 +263,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 7,
    "id": "e3d67c03",
    "metadata": {},
    "outputs": [
@@ -263,7 +281,7 @@
        " None)"
       ]
      },
-     "execution_count": 8,
+     "execution_count": 7,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -293,14 +311,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 8,
    "id": "9b1a71ba",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "('point_cloud',\n",
+       "(<Tag.POINT_CLOUD: 'point_cloud'>,\n",
        " ArrayView([[-1.076, -3.572, -0.315, ...,  0.492, -0.226,  0.194],\n",
        "            [ 0.852, -1.033,  3.834, ...,  0.252, -0.115,  0.243],\n",
        "            [-0.411, -2.689, -1.863, ..., -0.347,  0.601,  0.005],\n",
@@ -311,7 +329,7 @@
        "           dtype=float32))"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -331,17 +349,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 9,
    "id": "a723c855",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "('cost_matrix', None)"
+       "(<Tag.COST_MATRIX: 'cost_matrix'>, None)"
       ]
      },
-     "execution_count": 10,
+     "execution_count": 9,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -369,7 +387,7 @@
   "kernelspec": {
    "display_name": "moscot",
    "language": "python",
-   "name": "moscot"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -381,7 +399,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.9"
+   "version": "3.11.5"
   }
  },
  "nbformat": 4,

diff --git a/examples/problems/200_custom_cost_matrices.ipynb b/examples/problems/200_custom_cost_matrices.ipynb
@@ -39,6 +39,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import warnings\n",
+    "\n",
+    "warnings.simplefilter(\"ignore\", FutureWarning)\n",
+    "\n",
     "from moscot import datasets\n",
     "from moscot.problems.generic import GWProblem\n",
     "\n",
@@ -102,6 +106,22 @@
    "id": "36ef5120",
    "metadata": {},
    "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34mINFO    \u001b[0m Ordering \u001b[1;35mIndex\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[32m'0-0'\u001b[0m, \u001b[32m'1-0'\u001b[0m, \u001b[32m'2-0'\u001b[0m, \u001b[32m'3-0'\u001b[0m, \u001b[32m'4-0'\u001b[0m, \u001b[32m'5-0'\u001b[0m, \u001b[32m'6-0'\u001b[0m, \u001b[32m'7-0'\u001b[0m, \u001b[32m'8-0'\u001b[0m, \u001b[32m'9-0'\u001b[0m,                     \n",
+      "                \u001b[32m'10-0'\u001b[0m, \u001b[32m'11-0'\u001b[0m, \u001b[32m'12-0'\u001b[0m, \u001b[32m'13-0'\u001b[0m, \u001b[32m'14-0'\u001b[0m, \u001b[32m'15-0'\u001b[0m, \u001b[32m'16-0'\u001b[0m, \u001b[32m'17-0'\u001b[0m, \u001b[32m'18-0'\u001b[0m,                            \n",
+      "                \u001b[32m'19-0'\u001b[0m, \u001b[32m'0-1'\u001b[0m, \u001b[32m'1-1'\u001b[0m, \u001b[32m'2-1'\u001b[0m, \u001b[32m'3-1'\u001b[0m, \u001b[32m'4-1'\u001b[0m, \u001b[32m'5-1'\u001b[0m, \u001b[32m'6-1'\u001b[0m, \u001b[32m'7-1'\u001b[0m, \u001b[32m'8-1'\u001b[0m,                             \n",
+      "                \u001b[32m'9-1'\u001b[0m, \u001b[32m'10-1'\u001b[0m, \u001b[32m'11-1'\u001b[0m, \u001b[32m'12-1'\u001b[0m, \u001b[32m'13-1'\u001b[0m, \u001b[32m'14-1'\u001b[0m, \u001b[32m'15-1'\u001b[0m, \u001b[32m'16-1'\u001b[0m, \u001b[32m'17-1'\u001b[0m,                             \n",
+      "                \u001b[32m'18-1'\u001b[0m, \u001b[32m'19-1'\u001b[0m, \u001b[32m'0-2'\u001b[0m, \u001b[32m'1-2'\u001b[0m, \u001b[32m'2-2'\u001b[0m, \u001b[32m'3-2'\u001b[0m, \u001b[32m'4-2'\u001b[0m, \u001b[32m'5-2'\u001b[0m, \u001b[32m'6-2'\u001b[0m, \u001b[32m'7-2'\u001b[0m,                            \n",
+      "                \u001b[32m'8-2'\u001b[0m, \u001b[32m'9-2'\u001b[0m, \u001b[32m'10-2'\u001b[0m, \u001b[32m'11-2'\u001b[0m, \u001b[32m'12-2'\u001b[0m, \u001b[32m'13-2'\u001b[0m, \u001b[32m'14-2'\u001b[0m, \u001b[32m'15-2'\u001b[0m, \u001b[32m'16-2'\u001b[0m,                              \n",
+      "                \u001b[32m'17-2'\u001b[0m, \u001b[32m'18-2'\u001b[0m, \u001b[32m'19-2'\u001b[0m\u001b[1m]\u001b[0m,                                                                           \n",
+      "               \u001b[33mdtype\u001b[0m=\u001b[32m'object'\u001b[0m\u001b[1m)\u001b[0m in ascending order.                                                                 \n",
+      "\u001b[34mINFO    \u001b[0m Computing pca with `\u001b[33mn_comps\u001b[0m=\u001b[1;36m30\u001b[0m` for `xy` using `adata.X`                                                  \n",
+      "\u001b[34mINFO    \u001b[0m Computing pca with `\u001b[33mn_comps\u001b[0m=\u001b[1;36m30\u001b[0m` for `xy` using `adata.X`                                                  \n"
+     ]
+    },
     {
      "data": {
       "text/plain": [
@@ -115,7 +135,7 @@
    ],
    "source": [
     "gwp = GWProblem(adata)\n",
-    "gwp = gwp.prepare(key=\"batch\", joint_attr=\"X_pca\", GW_x=\"X_pca\", GW_y=\"X_pca\")\n",
+    "gwp = gwp.prepare(key=\"batch\", x_attr=\"X_pca\", y_attr=\"X_pca\")\n",
     "gwp"
    ]
   },
@@ -275,10 +295,25 @@
    "execution_count": 7,
    "id": "4b99031c",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34mINFO    \u001b[0m Ordering \u001b[1;35mIndex\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[32m'0-0'\u001b[0m, \u001b[32m'1-0'\u001b[0m, \u001b[32m'2-0'\u001b[0m, \u001b[32m'3-0'\u001b[0m, \u001b[32m'4-0'\u001b[0m, \u001b[32m'5-0'\u001b[0m, \u001b[32m'6-0'\u001b[0m, \u001b[32m'7-0'\u001b[0m, \u001b[32m'8-0'\u001b[0m, \u001b[32m'9-0'\u001b[0m,                     \n",
+      "                \u001b[32m'10-0'\u001b[0m, \u001b[32m'11-0'\u001b[0m, \u001b[32m'12-0'\u001b[0m, \u001b[32m'13-0'\u001b[0m, \u001b[32m'14-0'\u001b[0m, \u001b[32m'15-0'\u001b[0m, \u001b[32m'16-0'\u001b[0m, \u001b[32m'17-0'\u001b[0m, \u001b[32m'18-0'\u001b[0m,                            \n",
+      "                \u001b[32m'19-0'\u001b[0m, \u001b[32m'0-1'\u001b[0m, \u001b[32m'1-1'\u001b[0m, \u001b[32m'2-1'\u001b[0m, \u001b[32m'3-1'\u001b[0m, \u001b[32m'4-1'\u001b[0m, \u001b[32m'5-1'\u001b[0m, \u001b[32m'6-1'\u001b[0m, \u001b[32m'7-1'\u001b[0m, \u001b[32m'8-1'\u001b[0m,                             \n",
+      "                \u001b[32m'9-1'\u001b[0m, \u001b[32m'10-1'\u001b[0m, \u001b[32m'11-1'\u001b[0m, \u001b[32m'12-1'\u001b[0m, \u001b[32m'13-1'\u001b[0m, \u001b[32m'14-1'\u001b[0m, \u001b[32m'15-1'\u001b[0m, \u001b[32m'16-1'\u001b[0m, \u001b[32m'17-1'\u001b[0m,                             \n",
+      "                \u001b[32m'18-1'\u001b[0m, \u001b[32m'19-1'\u001b[0m, \u001b[32m'0-2'\u001b[0m, \u001b[32m'1-2'\u001b[0m, \u001b[32m'2-2'\u001b[0m, \u001b[32m'3-2'\u001b[0m, \u001b[32m'4-2'\u001b[0m, \u001b[32m'5-2'\u001b[0m, \u001b[32m'6-2'\u001b[0m, \u001b[32m'7-2'\u001b[0m,                            \n",
+      "                \u001b[32m'8-2'\u001b[0m, \u001b[32m'9-2'\u001b[0m, \u001b[32m'10-2'\u001b[0m, \u001b[32m'11-2'\u001b[0m, \u001b[32m'12-2'\u001b[0m, \u001b[32m'13-2'\u001b[0m, \u001b[32m'14-2'\u001b[0m, \u001b[32m'15-2'\u001b[0m, \u001b[32m'16-2'\u001b[0m,                              \n",
+      "                \u001b[32m'17-2'\u001b[0m, \u001b[32m'18-2'\u001b[0m, \u001b[32m'19-2'\u001b[0m\u001b[1m]\u001b[0m,                                                                           \n",
+      "               \u001b[33mdtype\u001b[0m=\u001b[32m'object'\u001b[0m\u001b[1m)\u001b[0m in ascending order.                                                                 \n"
+     ]
+    }
+   ],
    "source": [
     "joint_attr = {\"key\": \"cost_matrices\", \"tag\": \"cost_matrix\"}\n",
-    "gwp = gwp.prepare(key=\"batch\", joint_attr=joint_attr, GW_x=\"X_pca\", GW_y=\"X_pca\")"
+    "gwp = gwp.prepare(key=\"batch\", joint_attr=joint_attr, x_attr=\"X_pca\", y_attr=\"X_pca\")"
    ]
   },
   {
@@ -287,19 +322,44 @@
    "id": "b3ed4f02",
    "metadata": {},
    "source": [
-    "If we want to use only quadratic custom cost matrices, we need to modify `GW_x` and `GW_y`."
+    "If we want to use only quadratic custom cost matrices, we need to modify `x_attr` and `y_attr`."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 8,
    "id": "7c6f8eb6",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34mINFO    \u001b[0m Ordering \u001b[1;35mIndex\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[32m'0-0'\u001b[0m, \u001b[32m'1-0'\u001b[0m, \u001b[32m'2-0'\u001b[0m, \u001b[32m'3-0'\u001b[0m, \u001b[32m'4-0'\u001b[0m, \u001b[32m'5-0'\u001b[0m, \u001b[32m'6-0'\u001b[0m, \u001b[32m'7-0'\u001b[0m, \u001b[32m'8-0'\u001b[0m, \u001b[32m'9-0'\u001b[0m,                     \n",
+      "                \u001b[32m'10-0'\u001b[0m, \u001b[32m'11-0'\u001b[0m, \u001b[32m'12-0'\u001b[0m, \u001b[32m'13-0'\u001b[0m, \u001b[32m'14-0'\u001b[0m, \u001b[32m'15-0'\u001b[0m, \u001b[32m'16-0'\u001b[0m, \u001b[32m'17-0'\u001b[0m, \u001b[32m'18-0'\u001b[0m,                            \n",
+      "                \u001b[32m'19-0'\u001b[0m, \u001b[32m'0-1'\u001b[0m, \u001b[32m'1-1'\u001b[0m, \u001b[32m'2-1'\u001b[0m, \u001b[32m'3-1'\u001b[0m, \u001b[32m'4-1'\u001b[0m, \u001b[32m'5-1'\u001b[0m, \u001b[32m'6-1'\u001b[0m, \u001b[32m'7-1'\u001b[0m, \u001b[32m'8-1'\u001b[0m,                             \n",
+      "                \u001b[32m'9-1'\u001b[0m, \u001b[32m'10-1'\u001b[0m, \u001b[32m'11-1'\u001b[0m, \u001b[32m'12-1'\u001b[0m, \u001b[32m'13-1'\u001b[0m, \u001b[32m'14-1'\u001b[0m, \u001b[32m'15-1'\u001b[0m, \u001b[32m'16-1'\u001b[0m, \u001b[32m'17-1'\u001b[0m,                             \n",
+      "                \u001b[32m'18-1'\u001b[0m, \u001b[32m'19-1'\u001b[0m, \u001b[32m'0-2'\u001b[0m, \u001b[32m'1-2'\u001b[0m, \u001b[32m'2-2'\u001b[0m, \u001b[32m'3-2'\u001b[0m, \u001b[32m'4-2'\u001b[0m, \u001b[32m'5-2'\u001b[0m, \u001b[32m'6-2'\u001b[0m, \u001b[32m'7-2'\u001b[0m,                            \n",
+      "                \u001b[32m'8-2'\u001b[0m, \u001b[32m'9-2'\u001b[0m, \u001b[32m'10-2'\u001b[0m, \u001b[32m'11-2'\u001b[0m, \u001b[32m'12-2'\u001b[0m, \u001b[32m'13-2'\u001b[0m, \u001b[32m'14-2'\u001b[0m, \u001b[32m'15-2'\u001b[0m, \u001b[32m'16-2'\u001b[0m,                              \n",
+      "                \u001b[32m'17-2'\u001b[0m, \u001b[32m'18-2'\u001b[0m, \u001b[32m'19-2'\u001b[0m\u001b[1m]\u001b[0m,                                                                           \n",
+      "               \u001b[33mdtype\u001b[0m=\u001b[32m'object'\u001b[0m\u001b[1m)\u001b[0m in ascending order.                                                                 \n"
+     ]
+    }
+   ],
    "source": [
-    "GW_x = {\"attr\": \"obsp\", \"key\": \"cost_matrices\", \"tag\": \"cost_matrix\", \"cost\": \"custom\"}\n",
-    "GW_y = {\"attr\": \"obsp\", \"key\": \"cost_matrices\", \"tag\": \"cost_matrix\", \"cost\": \"custom\"}\n",
-    "gwp = gwp.prepare(key=\"batch\", joint_attr=\"X_pca\", GW_x=GW_x, GW_y=GW_y)"
+    "x_attr = {\n",
+    "    \"attr\": \"obsp\",\n",
+    "    \"key\": \"cost_matrices\",\n",
+    "    \"tag\": \"cost_matrix\",\n",
+    "    \"cost\": \"custom\",\n",
+    "}\n",
+    "y_attr = {\n",
+    "    \"attr\": \"obsp\",\n",
+    "    \"key\": \"cost_matrices\",\n",
+    "    \"tag\": \"cost_matrix\",\n",
+    "    \"cost\": \"custom\",\n",
+    "}\n",
+    "gwp = gwp.prepare(key=\"batch\", joint_attr=\"X_pca\", x_attr=x_attr, y_attr=y_attr)"
    ]
   },
   {
@@ -316,9 +376,24 @@
    "execution_count": 9,
    "id": "cbb8b363",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34mINFO    \u001b[0m Ordering \u001b[1;35mIndex\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[32m'0-0'\u001b[0m, \u001b[32m'1-0'\u001b[0m, \u001b[32m'2-0'\u001b[0m, \u001b[32m'3-0'\u001b[0m, \u001b[32m'4-0'\u001b[0m, \u001b[32m'5-0'\u001b[0m, \u001b[32m'6-0'\u001b[0m, \u001b[32m'7-0'\u001b[0m, \u001b[32m'8-0'\u001b[0m, \u001b[32m'9-0'\u001b[0m,                     \n",
+      "                \u001b[32m'10-0'\u001b[0m, \u001b[32m'11-0'\u001b[0m, \u001b[32m'12-0'\u001b[0m, \u001b[32m'13-0'\u001b[0m, \u001b[32m'14-0'\u001b[0m, \u001b[32m'15-0'\u001b[0m, \u001b[32m'16-0'\u001b[0m, \u001b[32m'17-0'\u001b[0m, \u001b[32m'18-0'\u001b[0m,                            \n",
+      "                \u001b[32m'19-0'\u001b[0m, \u001b[32m'0-1'\u001b[0m, \u001b[32m'1-1'\u001b[0m, \u001b[32m'2-1'\u001b[0m, \u001b[32m'3-1'\u001b[0m, \u001b[32m'4-1'\u001b[0m, \u001b[32m'5-1'\u001b[0m, \u001b[32m'6-1'\u001b[0m, \u001b[32m'7-1'\u001b[0m, \u001b[32m'8-1'\u001b[0m,                             \n",
+      "                \u001b[32m'9-1'\u001b[0m, \u001b[32m'10-1'\u001b[0m, \u001b[32m'11-1'\u001b[0m, \u001b[32m'12-1'\u001b[0m, \u001b[32m'13-1'\u001b[0m, \u001b[32m'14-1'\u001b[0m, \u001b[32m'15-1'\u001b[0m, \u001b[32m'16-1'\u001b[0m, \u001b[32m'17-1'\u001b[0m,                             \n",
+      "                \u001b[32m'18-1'\u001b[0m, \u001b[32m'19-1'\u001b[0m, \u001b[32m'0-2'\u001b[0m, \u001b[32m'1-2'\u001b[0m, \u001b[32m'2-2'\u001b[0m, \u001b[32m'3-2'\u001b[0m, \u001b[32m'4-2'\u001b[0m, \u001b[32m'5-2'\u001b[0m, \u001b[32m'6-2'\u001b[0m, \u001b[32m'7-2'\u001b[0m,                            \n",
+      "                \u001b[32m'8-2'\u001b[0m, \u001b[32m'9-2'\u001b[0m, \u001b[32m'10-2'\u001b[0m, \u001b[32m'11-2'\u001b[0m, \u001b[32m'12-2'\u001b[0m, \u001b[32m'13-2'\u001b[0m, \u001b[32m'14-2'\u001b[0m, \u001b[32m'15-2'\u001b[0m, \u001b[32m'16-2'\u001b[0m,                              \n",
+      "                \u001b[32m'17-2'\u001b[0m, \u001b[32m'18-2'\u001b[0m, \u001b[32m'19-2'\u001b[0m\u001b[1m]\u001b[0m,                                                                           \n",
+      "               \u001b[33mdtype\u001b[0m=\u001b[32m'object'\u001b[0m\u001b[1m)\u001b[0m in ascending order.                                                                 \n"
+     ]
+    }
+   ],
    "source": [
-    "gwp = gwp.prepare(key=\"batch\", joint_attr=joint_attr, GW_x=GW_x, GW_y=GW_y)"
+    "gwp = gwp.prepare(key=\"batch\", joint_attr=joint_attr, x_attr=x_attr, y_attr=y_attr)"
    ]
   }
  ],
@@ -344,7 +419,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.13"
+   "version": "3.11.5"
   },
   "vscode": {
    "interpreter": {