Utilities for sorting spatialdata points along z-order curve

.coveragerc_omit

@@ -12,5 +12,6 @@ omit =
     src/vitessce/data_utils/ome.py
     src/vitessce/data_utils/entities.py
     src/vitessce/data_utils/multivec.py
+    src/vitessce/data_utils/spatialdata_points_zorder.py
     src/vitessce/widget_plugins/demo_plugin.py
     src/vitessce/widget_plugins/spatial_query.py

docs/notebooks/spatial_data_xenium_morton.ipynb

@@ -0,0 +1,387 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "nbsphinx": "hidden"
+   },
+   "source": [
+    "# Vitessce Widget Tutorial"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Visualization of a SpatialData object"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Import dependencies\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from os.path import join, isfile, isdir\n",
+    "from urllib.request import urlretrieve\n",
+    "import zipfile\n",
+    "import shutil\n",
+    "\n",
+    "from vitessce import (\n",
+    "    VitessceConfig,\n",
+    "    ViewType as vt,\n",
+    "    CoordinationType as ct,\n",
+    "    CoordinationLevel as CL,\n",
+    "    SpatialDataWrapper,\n",
+    "    get_initial_coordination_scope_prefix\n",
+    ")\n",
+    "\n",
+    "from vitessce.data_utils import (\n",
+    "    sdata_morton_sort_points,\n",
+    "    sdata_points_process_columns,\n",
+    "    sdata_points_write_bounding_box_attrs,\n",
+    "    sdata_points_modify_row_group_size,\n",
+    "    sdata_morton_query_rect,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from spatialdata import read_zarr"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data_dir = \"data\"\n",
+    "zip_filepath = join(data_dir, \"xenium_rep1_io.spatialdata.zarr.zip\")\n",
+    "spatialdata_filepath = join(data_dir, \"xenium_rep1_io.spatialdata.zarr\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if not isdir(spatialdata_filepath):\n",
+    "    if not isfile(zip_filepath):\n",
+    "        os.makedirs(data_dir, exist_ok=True)\n",
+    "        urlretrieve('https://s3.embl.de/spatialdata/spatialdata-sandbox/xenium_rep1_io.zip', zip_filepath)\n",
+    "    with zipfile.ZipFile(zip_filepath,\"r\") as zip_ref:\n",
+    "        zip_ref.extractall(data_dir)\n",
+    "        os.rename(join(data_dir, \"data.zarr\"), spatialdata_filepath)\n",
+    "        \n",
+    "        # This Xenium dataset has an AnnData \"raw\" element.\n",
+    "        # Reference: https://github.com/giovp/spatialdata-sandbox/issues/55\n",
+    "        raw_dir = join(spatialdata_filepath, \"tables\", \"table\", \"raw\")\n",
+    "        if isdir(raw_dir):\n",
+    "            shutil.rmtree(raw_dir)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sdata = read_zarr(spatialdata_filepath)\n",
+    "sdata"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sdata[\"transcripts\"].shape[0].compute()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sdata.tables[\"table\"].X = sdata.tables[\"table\"].X.toarray()\n",
+    "sdata.tables[\"dense_table\"] = sdata.tables[\"table\"]\n",
+    "sdata.write_element(\"dense_table\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# TODO: store the two separate images as a single image with two channels.\n",
+    "# Similar to https://github.com/EricMoerthVis/tissue-map-tools/pull/12"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sdata.tables['table'].obs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sdata"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sdata.points['transcripts'].head()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Sorting Points and creating a new Points element in the SpatialData object"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Step 1. Sort rows with `sdata_morton_sort_points`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sdata = sdata_morton_sort_points(sdata, \"transcripts\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Step 2. Clean up columns with `sdata_points_process_columns`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Add feature_index column to dataframe, and reorder columns so that feature_name (dict column) is the rightmost column.\n",
+    "ddf = sdata_points_process_columns(sdata, \"transcripts\", var_name_col=\"feature_name\", table_name=\"table\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ddf.head()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Step 3. Save sorted dataframe to new Points element"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sdata[\"transcripts_with_morton_codes\"] = ddf\n",
+    "sdata.write_element(\"transcripts_with_morton_codes\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Step 4. Write bounding box metadata with `sdata_points_write_bounding_box_attrs`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sdata_points_write_bounding_box_attrs(sdata, \"transcripts_with_morton_codes\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Step 5. Modify the row group sizes of the Parquet files with `sdata_points_modify_row_group_size`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sdata_points_modify_row_group_size(sdata, \"transcripts_with_morton_codes\", row_group_size=25_000)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Done"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Optionally, check the number of row groups in one of the parquet file parts.\n",
+    "import pyarrow.parquet as pq\n",
+    "from os.path import join\n",
+    "\n",
+    "parquet_file = pq.ParquetFile(join(sdata.path, \"points\", \"transcripts_with_morton_codes\", \"points.parquet\", \"part.0.parquet\"))\n",
+    "\n",
+    "# Get the number of row groups in this part-0 file.\n",
+    "num_groups = parquet_file.num_row_groups\n",
+    "num_groups"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Configure Vitessce\n",
+    "\n",
+    "Vitessce needs to know which pieces of data we are interested in visualizing, the visualization types we would like to use, and how we want to coordinate (or link) the views."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vc = VitessceConfig(\n",
+    "    schema_version=\"1.0.18\",\n",
+    "    name='MERFISH SpatialData Demo',\n",
+    ")\n",
+    "# Add data to the configuration:\n",
+    "wrapper = SpatialDataWrapper(\n",
+    "    sdata_path=spatialdata_filepath,\n",
+    "    # The following paths are relative to the root of the SpatialData zarr store on-disk.\n",
+    "    image_path=\"images/rasterized\",\n",
+    "    table_path=\"tables/table\",\n",
+    "    obs_feature_matrix_path=\"tables/table/X\",\n",
+    "    obs_spots_path=\"shapes/cells\",\n",
+    "    coordinate_system=\"global\",\n",
+    "    coordination_values={\n",
+    "        # The following tells Vitessce to consider each observation as a \"spot\"\n",
+    "        \"obsType\": \"cell\",\n",
+    "    }\n",
+    ")\n",
+    "dataset = vc.add_dataset(name='MERFISH').add_object(wrapper)\n",
+    "\n",
+    "# Add views (visualizations) to the configuration:\n",
+    "spatial = vc.add_view(\"spatialBeta\", dataset=dataset)\n",
+    "feature_list = vc.add_view(\"featureList\", dataset=dataset)\n",
+    "layer_controller = vc.add_view(\"layerControllerBeta\", dataset=dataset)\n",
+    "obs_sets = vc.add_view(\"obsSets\", dataset=dataset)\n",
+    "\n",
+    "vc.link_views_by_dict([spatial, layer_controller], {\n",
+    "    'spotLayer': CL([{\n",
+    "        'obsType': 'cell',\n",
+    "    }]),\n",
+    "}, scope_prefix=get_initial_coordination_scope_prefix(\"A\", \"obsSpots\"))\n",
+    "\n",
+    "vc.link_views([spatial, layer_controller, feature_list, obs_sets], ['obsType'], [wrapper.obs_type_label])\n",
+    "\n",
+    "# Layout the views\n",
+    "vc.layout(spatial | (feature_list / layer_controller / obs_sets));"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Render the widget"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vw = vc.widget()\n",
+    "vw"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}