custom-index.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1f09f76e-a2af-47db-8a67-98020637bd31",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "\n",
    "xr.set_options(display_style=\"text\", display_expand_data=False)\n",
    "rng = np.random.default_rng(seed=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5a225768-3d01-420c-9b63-0b9aed23dbf3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre>&lt;xarray.DataArray (x: 3600, y: 1800)&gt;\n",
       "0.1257 -0.1321 0.6404 0.1049 -0.5357 ... 1.813 1.865 -0.1664 0.6568 -0.5354\n",
       "Coordinates:\n",
       "    lat      (x, y) float64 -90.0 -89.9 -89.8 -89.7 ... 89.6 89.7 89.8 89.9\n",
       "    lon      (x, y) float64 0.0 0.0 0.0 0.0 0.0 ... 359.9 359.9 359.9 359.9\n",
       "Dimensions without coordinates: x, y</pre>"
      ],
      "text/plain": [
       "<xarray.DataArray (x: 3600, y: 1800)>\n",
       "0.1257 -0.1321 0.6404 0.1049 -0.5357 ... 1.813 1.865 -0.1664 0.6568 -0.5354\n",
       "Coordinates:\n",
       "    lat      (x, y) float64 -90.0 -89.9 -89.8 -89.7 ... 89.6 89.7 89.8 89.9\n",
       "    lon      (x, y) float64 0.0 0.0 0.0 0.0 0.0 ... 359.9 359.9 359.9 359.9\n",
       "Dimensions without coordinates: x, y"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lat = np.arange(-90, 90, 0.1)\n",
    "lon = np.arange(0, 360, 0.1)\n",
    "\n",
    "lat2d, lon2d = np.meshgrid(lat, lon)\n",
    "\n",
    "data = rng.standard_normal(lat2d.shape)\n",
    "arr = xr.DataArray(\n",
    "    data,\n",
    "    dims=(\"x\", \"y\"),\n",
    "    coords={\"lat\": ((\"x\", \"y\"), lat2d), \"lon\": ((\"x\", \"y\"), lon2d)},\n",
    ")\n",
    "arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5f31aef1-93f8-4871-a888-2ee025b323ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Indexes:"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arr.xindexes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4fd48635-cecb-4623-b4d2-782dac4878e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.spatial import KDTree\n",
    "\n",
    "\n",
    "class KDTreeIndex(xr.core.indexes.Index):\n",
    "    def __init__(self, data, names, dim, **options):\n",
    "        self.names = names\n",
    "        self.dim = dim\n",
    "        self.shape = data.shape\n",
    "        self.index = KDTree(data.reshape(-1, self.shape[-1]), **options)\n",
    "\n",
    "    @classmethod\n",
    "    def from_variables(cls, variables, **options):\n",
    "        data = np.concatenate(\n",
    "            [var.data[..., None] for var in variables.values()], axis=-1\n",
    "        )\n",
    "        dims = {var.dims for var in variables.values()}\n",
    "        if len(dims) != 1:\n",
    "            raise ValueError(\"variables need to have the same dimensions\")\n",
    "        (dims,) = dims\n",
    "        names = list(variables.keys())\n",
    "        return cls(data, names, dims, **options)\n",
    "\n",
    "    def sel(self, indexers):\n",
    "        unknown_dimensions = set(indexers) - set(self.names)\n",
    "        if unknown_dimensions:\n",
    "            raise ValueError(\"unknown dimensions:\", list(unknown_dimensions))\n",
    "\n",
    "        points = np.concatenate(\n",
    "            [indexers[name][..., None] for name in self.names],\n",
    "            axis=-1,\n",
    "        )\n",
    "\n",
    "        distances, indices_ = self.index.query(points)\n",
    "        indices = np.unravel_index(indices_, self.shape[:-1])\n",
    "\n",
    "        isel_indexers = {\n",
    "            dim: xr.DataArray(\n",
    "                data, coords={\"distance\": (\"points\", distances)}, dims=\"points\"\n",
    "            )\n",
    "            for dim, data in zip(self.dim, indices)\n",
    "        }\n",
    "\n",
    "        return xr.core.indexes.IndexSelResult(isel_indexers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2e265e78-3fea-4476-8a57-67915b111e08",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<__main__.KDTreeIndex at 0x7f020877eee0>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tree = KDTreeIndex.from_variables(arr.coords)\n",
    "tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "87977264-3a5a-4378-a50d-49c05e377608",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre>&lt;xarray.DataArray (x: 3600, y: 1800)&gt;\n",
       "0.1257 -0.1321 0.6404 0.1049 -0.5357 ... 1.813 1.865 -0.1664 0.6568 -0.5354\n",
       "Coordinates:\n",
       "  * lat      (x, y) float64 -90.0 -89.9 -89.8 -89.7 ... 89.6 89.7 89.8 89.9\n",
       "  * lon      (x, y) float64 0.0 0.0 0.0 0.0 0.0 ... 359.9 359.9 359.9 359.9\n",
       "  * x        (x) int64 0 1 2 3 4 5 6 7 ... 3593 3594 3595 3596 3597 3598 3599\n",
       "  * y        (y) int64 0 1 2 3 4 5 6 7 ... 1793 1794 1795 1796 1797 1798 1799</pre>"
      ],
      "text/plain": [
       "<xarray.DataArray (x: 3600, y: 1800)>\n",
       "0.1257 -0.1321 0.6404 0.1049 -0.5357 ... 1.813 1.865 -0.1664 0.6568 -0.5354\n",
       "Coordinates:\n",
       "  * lat      (x, y) float64 -90.0 -89.9 -89.8 -89.7 ... 89.6 89.7 89.8 89.9\n",
       "  * lon      (x, y) float64 0.0 0.0 0.0 0.0 0.0 ... 359.9 359.9 359.9 359.9\n",
       "  * x        (x) int64 0 1 2 3 4 5 6 7 ... 3593 3594 3595 3596 3597 3598 3599\n",
       "  * y        (y) int64 0 1 2 3 4 5 6 7 ... 1793 1794 1795 1796 1797 1798 1799"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_arr = xr.DataArray(\n",
    "    arr.variable,\n",
    "    coords={\n",
    "        \"lat\": xr.Variable((\"x\", \"y\"), lat2d),\n",
    "        \"lon\": xr.Variable((\"x\", \"y\"), lon2d),\n",
    "        \"x\": xr.Variable(\"x\", np.arange(lat2d.shape[0])),\n",
    "        \"y\": xr.Variable(\"y\", np.arange(lat2d.shape[1])),\n",
    "    },\n",
    "    indexes={\n",
    "        \"lat\": tree,\n",
    "        \"lon\": tree,\n",
    "        \"x\": xr.core.indexes.PandasIndex.from_variables({\"x\": xr.Variable(\"x\", np.arange(lat2d.shape[0]))}),\n",
    "        \"y\": xr.core.indexes.PandasIndex.from_variables({\"y\": xr.Variable(\"y\", np.arange(lat2d.shape[1]))}),\n",
    "    },\n",
    "    fastpath=True,\n",
    ")\n",
    "new_arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c0da7bfb-da63-4f55-8dc7-7e1deb221b85",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Indexes:\n",
       "lat: <__main__.KDTreeIndex object at 0x7f020877eee0>\n",
       "lon: <__main__.KDTreeIndex object at 0x7f020877eee0>\n",
       "x: <xarray.core.indexes.PandasIndex object at 0x7f01bf793400>\n",
       "y: <xarray.core.indexes.PandasIndex object at 0x7f01b5f0d180>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_arr.xindexes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b064f7bb-2409-4c54-a280-a03f854512a5",
   "metadata": {},
   "source": [
    "# try the possible operations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "128e4763-1e93-41cc-820e-d6f8ce25dd1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "indexers = {\n",
    "    \"lat\": np.array([0.742, 10.213, 17.648]),\n",
    "    \"lon\": np.array([8.873, 3.12, 9.15]),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1e49cad-26fc-4027-88f4-cb362d810dce",
   "metadata": {},
   "source": [
    "`sel`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "80313cf5-f77c-4279-adb2-2ea9ad115064",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre>&lt;xarray.DataArray (points: 3)&gt;\n",
       "-1.736 1.034 -1.5\n",
       "Coordinates:\n",
       "    lat       (points) float64 0.7 10.2 17.6\n",
       "    lon       (points) float64 8.9 3.1 9.1\n",
       "    x         (points) int64 89 31 91\n",
       "    y         (points) int64 907 1002 1076\n",
       "    distance  (points) float64 0.04993 0.02385 0.06931\n",
       "Dimensions without coordinates: points</pre>"
      ],
      "text/plain": [
       "<xarray.DataArray (points: 3)>\n",
       "-1.736 1.034 -1.5\n",
       "Coordinates:\n",
       "    lat       (points) float64 0.7 10.2 17.6\n",
       "    lon       (points) float64 8.9 3.1 9.1\n",
       "    x         (points) int64 89 31 91\n",
       "    y         (points) int64 907 1002 1076\n",
       "    distance  (points) float64 0.04993 0.02385 0.06931\n",
       "Dimensions without coordinates: points"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_arr.sel(indexers)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "510f6306-f2ec-4ab1-9d3f-271c8dd75e99",
   "metadata": {},
   "source": [
    "`isel`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f61efc45-c290-4faf-b00f-18b4849b9b26",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Indexes:\n",
       "x: <xarray.core.indexes.PandasIndex object at 0x7f01bf793720>\n",
       "y: <xarray.core.indexes.PandasIndex object at 0x7f01bf793540>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sliced = new_arr.isel(x=slice(5, 10), y=slice(10, 20))\n",
    "sliced.xindexes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "103d002e-8ae5-40b2-ba3d-da6e7646bbe5",
   "metadata": {},
   "source": [
    "`roll`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "12d28c27-82c1-4de6-97a2-ca793750db40",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Indexes:\n",
       "x: <xarray.core.indexes.PandasIndex object at 0x7f020878bc20>\n",
       "y: <xarray.core.indexes.PandasIndex object at 0x7f01bf793c70>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rolled = new_arr.roll(x=5, y=10, roll_coords=True)\n",
    "rolled.xindexes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a74aadc2-2a1f-4dfb-a13d-5bc030d4f50c",
   "metadata": {},
   "source": [
    "`stack`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1325e936-6681-49ee-b49a-0f331cd9adad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Indexes:\n",
       "z: <xarray.core.indexes.PandasMultiIndex object at 0x7f01b6048c80>\n",
       "x: <xarray.core.indexes.PandasMultiIndex object at 0x7f01b6048c80>\n",
       "y: <xarray.core.indexes.PandasMultiIndex object at 0x7f01b6048c80>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stacked = new_arr.stack(z=(\"x\", \"y\"))\n",
    "stacked.xindexes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef69e3da-e8a8-4bb2-9c74-53ef47f58ec4",
   "metadata": {},
   "source": [
    "`copy`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "872e4bed-6097-4407-899a-60714d0204e8",
   "metadata": {},
   "outputs": [
    {
     "ename": "NotImplementedError",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNotImplementedError\u001b[0m                       Traceback (most recent call last)",
      "Input \u001b[0;32mIn [13]\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mnew_arr\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcopy\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/Documents/Programming/xarray/xarray/core/dataarray.py:1049\u001b[0m, in \u001b[0;36mDataArray.copy\u001b[0;34m(self, deep, data)\u001b[0m\n\u001b[1;32m   1047\u001b[0m     indexes \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_indexes\n\u001b[1;32m   1048\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1049\u001b[0m     indexes \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mxindexes\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcopy_indexes\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdeep\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdeep\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1050\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_replace(variable, coords, indexes\u001b[38;5;241m=\u001b[39mindexes)\n",
      "File \u001b[0;32m~/Documents/Programming/xarray/xarray/core/indexes.py:1166\u001b[0m, in \u001b[0;36mIndexes.copy_indexes\u001b[0;34m(self, deep)\u001b[0m\n\u001b[1;32m   1164\u001b[0m new_indexes \u001b[38;5;241m=\u001b[39m {}\n\u001b[1;32m   1165\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m idx, coords \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgroup_by_index():\n\u001b[0;32m-> 1166\u001b[0m     new_idx \u001b[38;5;241m=\u001b[39m \u001b[43midx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcopy\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdeep\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdeep\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1167\u001b[0m     new_indexes\u001b[38;5;241m.\u001b[39mupdate({k: new_idx \u001b[38;5;28;01mfor\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m coords})\n\u001b[1;32m   1169\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m new_indexes\n",
      "File \u001b[0;32m~/Documents/Programming/xarray/xarray/core/indexes.py:105\u001b[0m, in \u001b[0;36mIndex.copy\u001b[0;34m(self, deep)\u001b[0m\n\u001b[1;32m    104\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcopy\u001b[39m(\u001b[38;5;28mself\u001b[39m, deep: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m):  \u001b[38;5;66;03m# pragma: no cover\u001b[39;00m\n\u001b[0;32m--> 105\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mNotImplementedError\u001b[39;00m()\n",
      "\u001b[0;31mNotImplementedError\u001b[0m: "
     ]
    }
   ],
   "source": [
    "new_arr.copy()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}