Add the "atldld dataset preview" command (#66)

Also:
* Add the "constants" module
* Add the "plot" module

Author: Stannislav

src/atldld/cli.py

@@ -15,6 +15,8 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 """The command line interface (CLI) for Atlas-Download-Tools."""
+from typing import Any, Dict, Optional, Sequence
+
 import click
 
 import atldld
@@ -60,47 +62,73 @@ def cache_folder():
 
 
 # ============================= Dataset subcommand =============================
-
-
-@click.group(help="Commands related to atlas datasets")
-def dataset():
-    """Run dataset subcommands."""
-
-
-@click.command("info", help="Get information for a given dataset ID")
-@click.argument("dataset_id", type=int)
-def dataset_info(dataset_id):
-    """Get information for a given dataset ID."""
-    import textwrap
-
-    import atldld.dataset
+def get_dataset_meta_or_abort(
+    dataset_id: int, include: Optional[Sequence[str]] = None
+) -> Dict[str, Any]:
+    """Download the dataset metadata.
+
+    Parameters
+    ----------
+    dataset_id
+        The dataset ID.
+    include
+        The include keys to use in the RMA query.
+
+    Returns
+    -------
+    meta : dict
+        The dataset metadata.
+
+    Raises
+    ------
+    click.Abort
+        Whenever the metadata download fails or yields unexpected results.
+    """
     from atldld import requests
 
     # Send request
     rma_parameters = requests.RMAParameters(
         "SectionDataSet",
         criteria={"id": dataset_id},
-        include=("genes", "section_images"),
+        include=include,
     )
     try:
         msg = requests.rma_all(rma_parameters)
     except requests.RMAError as exc:
-        click.secho(
-            f"An error occurred while querying the AIBS servers: {str(exc)}",
-            fg="red",
+        raise click.ClickException(
+            f"An error occurred while querying the AIBS servers: {str(exc)}"
         )
-        raise click.Abort
 
     # Check response
     if len(msg) == 0:
-        click.secho(f"Dataset with ID {dataset_id} does not exist", fg="red")
-        raise click.Abort
+        raise click.ClickException(f"Dataset with ID {dataset_id} does not exist")
     elif len(msg) > 1:
-        click.secho("Something went wrong: got more than one dataset", fg="red")
-        raise click.Abort
+        raise click.ClickException("Something went wrong: got more than one dataset")
 
-    # Print response
     meta = msg[0]
+    if not isinstance(meta, dict) or not all(isinstance(key, str) for key in meta):
+        raise click.ClickException("Got an unexpected dataset information format")
+
+    return meta
+
+
+@click.group(help="Commands related to atlas datasets")
+def dataset():
+    """Run dataset subcommands."""
+
+
+@click.command("info", help="Get information for a given dataset ID")
+@click.argument("dataset_id", type=int)
+def dataset_info(dataset_id):
+    """Get information for a given dataset ID."""
+    import textwrap
+
+    import atldld.dataset
+
+    # Download the dataset metadata
+    meta = get_dataset_meta_or_abort(dataset_id, include=["genes", "section_images"])
+
+    # Print response
     section_images = meta.pop("section_images")
     r_str = meta["red_channel"] or "-"
     g_str = meta["green_channel"] or "-"
@@ -124,5 +152,57 @@ def dataset_info(dataset_id):
     click.secho(textwrap.dedent(output).strip(), fg="green")
 
 
+@click.command("preview", help="Plot a preview of dataset slices")
+@click.argument("dataset_id", type=int)
+@click.option(
+    "-o",
+    "--output-dir",
+    type=click.Path(file_okay=False, writable=True, resolve_path=True),
+    help="""
+    The output directory for the plot figure. If not provided the current
+    working directory will be used.
+    """,
+)
+def dataset_preview(dataset_id, output_dir):
+    """Plot a sketch of section images mapped into the reference space."""
+    import pathlib
+
+    import atldld.dataset
+    import atldld.utils
+    from atldld import plot
+
+    # Download the dataset metadata
+    meta = get_dataset_meta_or_abort(dataset_id, include=["section_images"])
+    plane_of_section = atldld.dataset.PlaneOfSection(meta["plane_of_section_id"])
+    section_image_metas = meta.pop("section_images")
+    section_image_metas.sort(key=lambda image_meta_: image_meta_["section_number"])
+
+    click.secho("Fetching the corner coordinates of the section images...", fg="green")
+    all_corners = []
+    with click.progressbar(section_image_metas) as progress:
+        for image_meta in progress:
+            corners = atldld.utils.get_corners_in_ref_space(
+                image_meta["id"],
+                image_meta["image_width"],
+                image_meta["image_height"],
+            )
+            all_corners.append(corners)
+
+    click.secho("Plotting...", fg="green")
+    img_file_name = f"dataset-id-{dataset_id}-preview.png"
+    if output_dir is None:
+        img_path = pathlib.Path.cwd() / img_file_name
+    else:
+        img_path = pathlib.Path(output_dir) / img_file_name
+        img_path.parent.mkdir(exist_ok=True, parents=True)
+    fig = plot.dataset_preview(all_corners, plane_of_section)
+    fig.suptitle(f"Dataset ID {dataset_id}", fontsize=32)
+    fig.set_dpi(200)
+    fig.savefig(img_path)
+    click.secho("Figure was saved in ", fg="green", nl=False)
+    click.secho(f"{img_path.resolve().as_uri()}", fg="yellow", bold=True)
+
+
 root.add_command(dataset)
 dataset.add_command(dataset_info)
+dataset.add_command(dataset_preview)

src/atldld/constants.py

@@ -19,3 +19,7 @@
 
 GLOBAL_CACHE_FOLDER = pathlib.Path.home() / ".atldld"
 GLOBAL_CACHE_FOLDER.mkdir(exist_ok=True)
+
+# Volume dimensions
+REF_DIM_1UM = (13200, 8000, 11400)
+REF_DIM_25UM = (528, 320, 456)

src/atldld/plot.py

@@ -0,0 +1,146 @@
+# The package atldld is a tool to download atlas data.
+#
+# Copyright (C) 2021 EPFL/Blue Brain Project
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""Different plotting routines."""
+from typing import Iterable
+
+import numpy as np
+from matplotlib.figure import Figure
+from matplotlib.lines import Line2D
+
+from atldld.constants import REF_DIM_1UM
+from atldld.dataset import PlaneOfSection
+
+
+def dataset_preview(
+    all_corners: Iterable[np.ndarray],
+    plane_of_section: PlaneOfSection,
+) -> Figure:
+    """Plot a preview of how section images fit into the reference space.
+
+    Parameters
+    ----------
+    all_corners
+        The corners of all section images. Each element in this iterable should
+        be a NumPy array of shape (4, 3). The format of this array corresponds
+        to that returned by the `atldld.utils.get_corners_in_ref_space`
+        function.
+
+        The first axis refers to the four corners of a section image in the
+        following order:
+
+        1. Lower left (0, 0)
+        2. Lower right (0, 1)
+        3. Upper right (1, 1)
+        4. Upper left (1, 0)
+
+        This corresponds to following the corners counterclockwise starting with
+        the corner in the axes origin. The second array axis contains the 3D
+        coordinates of the corners in the standard PIR references space.
+
+    plane_of_section
+        The plane of section of the dataset. Can be either coronal or sagittal.
+
+    Returns
+    -------
+    fig
+        The figure with the plot.
+    """
+    # A semi-arbitrary choice of the reference space scale. This choice only
+    # changes the ticks on the axes, but not the overall plot. The 25µm scale
+    # is one of the common scales used for volumes.
+    ref_space_scale = 25
+    ref_space_size = np.array(REF_DIM_1UM) / ref_space_scale
+    p, i, r = 0, 1, 2
+    labels = {
+        p: "p (coronal)",
+        i: "i (transverse)",
+        r: "r (sagittal)",
+    }
+    # We'll plot the views of all four edges in counterclockwise order starting
+    # with the bottom edge. The last two x-axes are inverted so that the edge
+    # vertices on the right of a plot appear on the left of the following plot.
+    edges = [[0, 1], [1, 2], [2, 3], [3, 0]]
+    titles = ["Bottom Edges", "Right Edges", "Top Edges", "Left Edges"]
+    inverts = [False, False, True, True]
+
+    # Depending on the plane of section the views are different
+    if plane_of_section == PlaneOfSection.CORONAL:
+        y_axis = p
+        x_axes = [r, i, r, i]
+    elif plane_of_section == PlaneOfSection.SAGITTAL:
+        y_axis = r
+        x_axes = [p, i, p, i]
+    else:  # pragma: no cover
+        raise NotImplementedError(f"Unknown plane of section: {plane_of_section}")
+
+    # The figure width is fixed and arbitrary. (Is there a more clever choice?)
+    # The figure height is computed to match the ratio between the total width
+    # and height of all subplots. This way the ratios of the x and y axes are
+    # roughly the same (but not quite since the in-between spaces, titles, etc.
+    # are not taken into account...)
+    plot_width = sum(ref_space_size[x_axis] for x_axis in x_axes)
+    plot_height = ref_space_size[y_axis]
+    fig_width_inches = 14
+    fig_height_inches = fig_width_inches * plot_height / plot_width
+
+    fig = Figure(figsize=(fig_width_inches, fig_height_inches))
+    fig.set_tight_layout(True)
+    # The width ratios of subplots are based on the reference volume dimensions,
+    # this way the scales of the x-axes roughly match.
+    axs = fig.subplots(
+        ncols=4,
+        sharey=True,
+        gridspec_kw={"width_ratios": [ref_space_size[x_axis] for x_axis in x_axes]},
+    )
+    # Y-label only on the left-most plot because it's the same for all plots
+    axs[0].set_ylabel(labels[y_axis], fontsize=16)
+
+    # Add the legend for the reference space lines
+    ref_space_line_style = {"color": "blue", "linestyle": ":"}
+    line = Line2D([], [], **ref_space_line_style)
+    axs[0].legend(
+        [line],
+        ["Reference space boundary"],
+        loc="upper left",
+        bbox_to_anchor=(0, -0.2),
+        borderaxespad=0,
+        frameon=False,
+    )
+
+    # The actual plotting
+    for ax, edge, x_axis, invert, title in zip(axs, edges, x_axes, inverts, titles):
+        # Axes setup
+        ax.grid(True, linestyle=":", color="gray")
+        ax.set_title(title)
+        ax.set_xlabel(labels[x_axis], fontsize=16)
+
+        # Reference space boundary lines
+        ax.axvline(0, **ref_space_line_style)
+        ax.axvline(ref_space_size[x_axis], **ref_space_line_style)
+        ax.axhline(0, **ref_space_line_style)
+        ax.axhline(ref_space_size[y_axis], **ref_space_line_style)
+
+        # Plot the section image edges
+        for corners in all_corners:
+            points = corners[np.ix_(edge, [x_axis, y_axis])]
+            coords = points.T / ref_space_scale
+            ax.plot(*coords, color="green")
+            ax.scatter(*coords, color="red")
+        if invert:
+            ax.invert_xaxis()
+
+    return fig

src/atldld/utils.py

@@ -142,6 +142,67 @@ def get_image(
     return img
 
 
+def get_corners_in_ref_space(
+    image_id: int,
+    image_width: int,
+    image_height: int,
+) -> np.ndarray:
+    """Get the corner coordinates of a section image in the reference space.
+
+    Parameters
+    ----------
+    image_id
+        The section image ID.
+    image_width
+        The width of the section image.
+    image_height
+        The height of the section image.
+
+    Returns
+    -------
+    ref_corners : np.ndarray
+
+    Notes
+    -----
+    The x and y coordinates in the API requests refer to the mathematical
+    axes with the origin in the lower left corner of the plotted image. This
+    is not the same as the array indices of ``image`` since the element
+    ``image[0, 0]`` is mapped to the upper left corner, ``image[i_max, 0]`` to
+    the lower left corner, etc.
+
+    Mathematical coordinates:
+
+    .. code-block:: text
+
+        ^ (0, 1)            (1, 1)
+        |
+        |
+        +------------------------->
+          (0, 0)            (1, 0)
+
+        Corresponding elements of the image array:
+        ^ image[0, 0]       image[0, j_max]
+        |
+        |
+        +------------------------->
+          image[i_max, 0]   image[i_max, j_max]
+
+    """
+    # Can we do this with affine transforms without sending a separate query
+    # for each corner???
+    ref_corners = []
+    for x, y in (
+        (0, 0),
+        (image_width, 0),
+        (image_width, image_height),
+        (0, image_height),
+    ):
+        pir = xy_to_pir_API_single(x, y, image_id)
+        ref_corners.append(pir)
+
+    return np.array(ref_corners)
+
+
 def get_experiment_list_from_gene(gene_name, axis="sagittal"):
     """Get Allen's experiments IDs for a given gene expression name.