Clip painted regions to chromosome outlines

src/jcvi/graphics/chromosome.py

@@ -11,6 +11,7 @@
 import sys
 from typing import List, Optional, Tuple
 
+from matplotlib.path import Path
 from natsort import natsorted
 import numpy as np
 
@@ -60,7 +61,8 @@ def __init__(
         y1, y2 = sorted((y1, y2))
         super().__init__(ax)
         pts, r = self.get_pts(x, y1, y2, width)
-        self.append(Polygon(pts, fill=False, lw=lw, ec=ec, zorder=zorder))
+        self.outline = Polygon(pts, fill=False, lw=lw, ec=ec, zorder=zorder)
+        self.append(self.outline)
         if patch:
             rr = r * 0.9  # Shrink a bit for the patches
             # First patch is colored if there is an even number of patches, otherwise not colored
@@ -74,6 +76,22 @@ def __init__(
                 )
 
         self.add_patches()
+        self.clip_patches()
+
+    def clip(self, artist):
+        """Clip an artist (e.g. a painted region) to the chromosome outline"""
+        artist.set_clip_path(self.outline)
+
+    def clip_patches(self):
+        # Clip after add_patches(): set_clip_path() captures the outline's
+        # transform, which is only set once the outline is on the axes
+        for p in self[1:]:
+            if isinstance(p, Rectangle):
+                self.clip(p)
+
+    def set_transform(self, tr):
+        super().set_transform(tr)
+        self.clip_patches()  # re-clip against the new transform
 
     def get_pts(self, x, y1, y2, width):
         w = width / 2
@@ -112,7 +130,8 @@ def __init__(
         x1, x2 = sorted((x1, x2))
         super().__init__(ax)
         pts, r = self.get_pts(x1, x2, y, height, style=style)
-        self.append(Polygon(pts, fill=False, lw=lw, ec=ec, zorder=zorder + 1))
+        self.outline = Polygon(pts, fill=False, lw=lw, ec=ec, zorder=zorder + 1)
+        self.append(self.outline)
 
         if fc:
             pts, r = self.get_pts(x1, x2, y, height / 2, style=style)
@@ -130,6 +149,20 @@ def __init__(
                 )
 
         self.add_patches()
+        self.clip_patches()
+
+    def clip(self, artist):
+        """Clip an artist (e.g. a painted region) to the chromosome outline"""
+        artist.set_clip_path(self.outline)
+
+    def clip_patches(self):
+        for p in self[1:]:
+            if isinstance(p, Rectangle):
+                self.clip(p)
+
+    def set_transform(self, tr):
+        super().set_transform(tr)
+        self.clip_patches()  # re-clip against the new transform
 
     def get_pts(self, x1, x2, y, height, style="auto") -> Tuple[list, float]:
         h = height / 2
@@ -160,16 +193,26 @@ def __init__(self, ax, x, y1, y2, y3, width=0.015, fc="k", fill=False, zorder=2)
         pts += [[x - r, y1 - r], [x - r, y2 + r]]
         pts += plot_cap((x, y2 + r), np.radians(range(180, 360)), r)
         pts += [[x + r, y2 + r], [x + r, y1 - r]]
-        ax.add_patch(Polygon(pts, fc=fc, fill=fill, zorder=zorder))
+        upper_arm = Polygon(pts, fc=fc, fill=fill, zorder=zorder)
+        ax.add_patch(upper_arm)
         pts = []
         pts += plot_cap((x, y2 - r), np.radians(range(180)), r)
         pts += [[x - r, y2 - r], [x - r, y3 + r]]
         pts += plot_cap((x, y3 + r), np.radians(range(180, 360)), r)
         pts += [[x + r, y3 + r], [x + r, y2 - r]]
-        ax.add_patch(Polygon(pts, fc=fc, fill=fill, zorder=zorder))
+        lower_arm = Polygon(pts, fc=fc, fill=fill, zorder=zorder)
+        ax.add_patch(lower_arm)
         ax.add_patch(
             CirclePolygon((x, y2), radius=r * 0.5, fc="k", ec="k", zorder=zorder)
         )
+        self.ax = ax
+        self.outline_path = Path.make_compound_path(
+            upper_arm.get_path(), lower_arm.get_path()
+        )
+
+    def clip(self, artist):
+        """Clip an artist (e.g. a painted region) to both chromosome arms"""
+        artist.set_clip_path(self.outline_path, self.ax.transData)
 
 
 class ChromosomeMap(object):
@@ -582,17 +625,20 @@ def draw_chromosomes(
 
     # first the chromosomes
     chr_locations = {}
+    chr_glyphs = {}
     for a, (chr, clen) in enumerate(natsorted(chr_lens.items())):
         xx = xstart + a * xinterval + 0.5 * xwidth
         chr_locations[chr] = xx
         root.text(xx, ystart + 0.01, str(get_number(chr)), ha="center")
         if centromeres:
             yy = ystart - centromeres[chr] * ratio
-            ChromosomeWithCentromere(
+            chr_glyphs[chr] = ChromosomeWithCentromere(
                 root, xx, ystart, yy, ystart - clen * ratio, width=xwidth
             )
         else:
-            Chromosome(root, xx, ystart, ystart - clen * ratio, width=xwidth)
+            chr_glyphs[chr] = Chromosome(
+                root, xx, ystart, ystart - clen * ratio, width=xwidth
+            )
 
     alpha = 1
     # color the regions
@@ -611,16 +657,17 @@ def draw_chromosomes(
                 start = prev_end
             yystart = ystart - end * ratio
             yyend = ystart - start * ratio
-            root.add_patch(
-                Rectangle(
-                    (xx, yystart),
-                    xwidth,
-                    yyend - yystart,
-                    fc=class_colors.get(klass, "lightslategray"),
-                    lw=0,
-                    alpha=alpha,
-                )
+            painted = Rectangle(
+                (xx, yystart),
+                xwidth,
+                yyend - yystart,
+                fc=class_colors.get(klass, "lightslategray"),
+                lw=0,
+                alpha=alpha,
             )
+            # Keep painted regions inside the rounded chromosome caps
+            chr_glyphs[chr].clip(painted)
+            root.add_patch(painted)
             prev_end, prev_klass = b.end, klass
 
             if imagemap:

tests/graphics/test_chromosome.py

@@ -0,0 +1,95 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+
+import matplotlib
+
+matplotlib.use("Agg")
+from matplotlib.patches import Rectangle
+import matplotlib.pyplot as plt
+from matplotlib.transforms import Affine2D
+import numpy as np
+
+from jcvi.graphics.chromosome import (
+    Chromosome,
+    ChromosomeWithCentromere,
+    HorizontalChromosome,
+)
+
+
+def _canvas():
+    fig = plt.figure(figsize=(2, 2), dpi=100)
+    ax = fig.add_axes((0, 0, 1, 1))
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.set_axis_off()
+    return fig, ax
+
+
+def _render(fig):
+    fig.canvas.draw()
+    rgb = np.asarray(fig.canvas.buffer_rgba())[:, :, :3]
+    plt.close(fig)
+    return rgb
+
+
+def _is_red(rgb, x, y):
+    # Image rows run top-to-bottom; convert from axes fraction
+    h, w = rgb.shape[:2]
+    r, g, b = rgb[int((1 - y) * (h - 1)), int(x * (w - 1))]
+    return r > 200 and g < 80 and b < 80
+
+
+def test_chromosome_clip_keeps_painted_region_inside_caps():
+    fig, ax = _canvas()
+    chrom = Chromosome(ax, 0.5, 0.1, 0.9, width=0.2)
+    painted = Rectangle((0.4, 0.8), 0.2, 0.1, fc="red", lw=0)
+    chrom.clip(painted)
+    ax.add_patch(painted)
+    rgb = _render(fig)
+    assert _is_red(rgb, 0.5, 0.85)  # inside the cap
+    assert not _is_red(rgb, 0.41, 0.895)  # corner outside the rounded cap
+
+
+def test_chromosome_patch_option_is_clipped():
+    fig, ax = _canvas()
+    Chromosome(
+        ax, 0.5, 0.1, 0.9, width=0.2, patch=[0.1, 0.2, 0.8, 0.9], patchcolor="red"
+    )
+    rgb = _render(fig)
+    assert _is_red(rgb, 0.5, 0.85)
+    assert not _is_red(rgb, 0.42, 0.895)
+
+
+def test_horizontal_chromosome_patch_option_is_clipped():
+    fig, ax = _canvas()
+    HorizontalChromosome(
+        ax, 0.1, 0.9, 0.5, height=0.2, patch=[0.1, 0.2, 0.8, 0.9], patchcolor="red"
+    )
+    rgb = _render(fig)
+    assert _is_red(rgb, 0.85, 0.5)
+    assert not _is_red(rgb, 0.895, 0.42)
+
+
+def test_chromosome_with_centromere_clips_to_both_arms():
+    fig, ax = _canvas()
+    chrom = ChromosomeWithCentromere(ax, 0.5, 0.9, 0.5, 0.1, width=0.2)
+    painted = Rectangle((0.4, 0.1), 0.2, 0.8, fc="red", lw=0)
+    chrom.clip(painted)
+    ax.add_patch(painted)
+    rgb = _render(fig)
+    assert _is_red(rgb, 0.5, 0.7)  # upper arm
+    assert _is_red(rgb, 0.5, 0.3)  # lower arm
+    assert not _is_red(rgb, 0.41, 0.895)  # outside the top cap
+    assert not _is_red(rgb, 0.41, 0.105)  # outside the bottom cap
+
+
+def test_horizontal_chromosome_patch_follows_set_transform():
+    fig, ax = _canvas()
+    hc = HorizontalChromosome(
+        ax, 0.1, 0.9, 0.3, height=0.2, patch=[0.1, 0.2, 0.8, 0.9], patchcolor="red"
+    )
+    hc.set_transform(Affine2D().translate(0, 0.4) + ax.transData)
+    rgb = _render(fig)
+    assert _is_red(rgb, 0.85, 0.7)  # inside the shifted cap
+    assert not _is_red(rgb, 0.895, 0.62)  # corner outside the shifted cap
+    assert not _is_red(rgb, 0.85, 0.3)  # nothing left at the old position