Add gene_labels and labels to plot_genes()

Author: leehart

malariagen_data/anoph/genome_features.py

@@ -332,6 +332,8 @@ def plot_genes(
         x_range: Optional[gplt_params.x_range] = None,
         title: Optional[gplt_params.title] = None,
         output_backend: gplt_params.output_backend = gplt_params.output_backend_default,
+        gene_labels: Optional[gplt_params.gene_labels] = None,
+        labels: Optional[gplt_params.labels] = None,
     ) -> gplt_params.figure:
         debug = self._log.debug
 
@@ -407,6 +409,101 @@ def plot_genes(
             line_width=0,
         )
 
+        if gene_labels:
+            debug("determine new figure height and range to accommodate gene labels")
+
+            # Increase the figure height by a certain factor, to accommodate labels.
+            height_increase_factor = 1.2
+            fig.height = int(fig.height * height_increase_factor)
+
+            # Get the original y_range.
+            # Note: fig.y_range is not subscriptable.
+            orig_y_range = fig.y_range.start, fig.y_range.end
+
+            # Determine the midpoint of the original range, to rescale outward from there.
+            orig_mid_y_range = (orig_y_range[0] + orig_y_range[1]) / 2
+            orig_y_range_extent = orig_y_range[1] - orig_y_range[0]
+
+            # Determine the new start and end points of the extended range.
+            new_y_range_extent = orig_y_range_extent * height_increase_factor
+            new_y_range_extent_half = new_y_range_extent / 2
+            new_y_start = orig_mid_y_range - new_y_range_extent_half
+            new_y_end = orig_mid_y_range + new_y_range_extent_half
+
+            # Set the new y_range.
+            fig.y_range = bokeh.models.Range1d(new_y_start, new_y_end)
+
+            debug("determine midpoint of each gene rectangle")
+            data["mid_x"] = (data["start"] + data["end"]) / 2
+
+            debug("make gene labels and pointers")
+
+            # Put gene_labels into a new column, where the gene_id matches.
+            # Fill unmapped genes with empty strings, otherwise "NaN" would be displayed.
+            data["gene_label"] = data["ID"].map(gene_labels).fillna("")
+
+            # Put gene pointers (▲ or ▼) in a new column, depending on the strand.
+            # Except if the gene_label is null or an empty string, which should not be shown.
+            data["gene_pointer"] = data.apply(
+                lambda row: ("▼" if row["strand"] == "+" else "▲")
+                if row["gene_label"]
+                else "",
+                axis=1,
+            )
+
+            # Put the pointer above or below the gene rectangle, depending on + or - strand.
+            neg_strand_pointer_y = orig_mid_y_range - 1.2
+            pos_strand_pointer_y = orig_mid_y_range + 1.2
+            data["pointer_y"] = data["strand"].apply(
+                lambda strand: pos_strand_pointer_y
+                if strand == "+"
+                else neg_strand_pointer_y
+            )
+
+            # Put the label above or below the gene rectangle, depending on + or - strand.
+            neg_strand_label_y = orig_mid_y_range - 1.15
+            pos_strand_label_y = orig_mid_y_range + 1.25
+            data["label_y"] = data["strand"].apply(
+                lambda strand: pos_strand_label_y
+                if strand == "+"
+                else neg_strand_label_y
+            )
+
+            # Get the data as a ColumnDataSource.
+            data_as_cds = bokeh.models.ColumnDataSource(data)
+
+            # Create a LabelSet for the gene pointers.
+            gene_pointers_ls = bokeh.models.LabelSet(
+                source=data_as_cds,
+                x="mid_x",
+                y="pointer_y",
+                text="gene_pointer",
+                text_align="center",
+                text_baseline="middle",
+                text_font_size="9pt",
+                text_color="#444444",
+            )
+
+            # Create a LabelSet for the gene labels.
+            gene_labels_ls = bokeh.models.LabelSet(
+                source=data_as_cds,
+                x="mid_x",
+                y="label_y",
+                text="gene_label",
+                text_align="left",
+                text_baseline="middle",
+                text_font_size="9pt",
+                text_color="#444444",
+                x_offset=8,
+            )
+
+            # Add the markers and labels to the figure.
+            fig.add_layout(gene_pointers_ls)
+            fig.add_layout(gene_labels_ls)
+
+        if labels:
+            fig.add_layout(labels)
+
         debug("tidy up the plot")
         fig.ygrid.visible = False
         yticks = [0.4, 1.4]

malariagen_data/anoph/gplt_params.py

@@ -113,3 +113,13 @@
 ]
 
 colors: TypeAlias = Annotated[Sequence[str], "List of colors."]
+
+gene_labels: TypeAlias = Annotated[
+    Mapping[str, str],
+    "A mapping of gene identifiers to custom labels, which will appear in the plot.",
+]
+
+labels: TypeAlias = Annotated[
+    bokeh.models.LabelSet,
+    "A LabelSet to use in the plot.",
+]

notebooks/plot_genes.ipynb

@@ -241,12 +241,67 @@
    "id": "c27741d9",
    "metadata": {},
    "outputs": [],
+   "source": [
+    "ag3.plot_genes(\n",
+    "  region=\"2R\",\n",
+    "  gene_labels={\n",
+    "    \"AGAP001096\": \"far left + gene\",\n",
+    "    \"AGAP001099\": \"far left - gene\",\n",
+    "    \"AGAP002942\": \"central + gene\",\n",
+    "    \"AGAP002949\": \"central - gene\",\n",
+    "    \"AGAP004676\": \"far right + gene\",\n",
+    "    \"AGAP004674\": \"far right - gene\",\n",
+    "  }\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "523197c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import bokeh\n",
+    "import pandas as pd\n",
+    "data = pd.DataFrame.from_dict(\n",
+    "  [\n",
+    "    {'pos': 10_000_000, 'y': 2, 'label': 'Custom Label A'},\n",
+    "    {'pos': 30_000_000, 'y': 0.9, 'label': 'Custom Label B'},\n",
+    "    {'pos': 50_000_000, 'y': -0.1, 'label': 'Custom Label C'},\n",
+    "  ]\n",
+    ")\n",
+    "data_as_cds = bokeh.models.ColumnDataSource(data)\n",
+    "labels = bokeh.models.LabelSet(\n",
+    "    source=data_as_cds,\n",
+    "    x='pos',\n",
+    "    y='y',\n",
+    "    text='label',\n",
+    "    text_align='center',\n",
+    "    text_baseline='middle',\n",
+    "    text_font_size='9pt',\n",
+    "    text_color='blue',\n",
+    ")\n",
+    "\n",
+    "ag3.plot_genes(\n",
+    "  region=\"2R\",\n",
+    "  labels=labels,\n",
+    "  height=200,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e2a3a927",
+   "metadata": {},
+   "outputs": [],
    "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "mgen_data_py3.11",
    "language": "python",
    "name": "python3"
   },
@@ -260,7 +315,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.12"
+   "version": "3.11.5"
   },
   "widgets": {
    "application/vnd.jupyter.widget-state+json": {