24def plot_one_cen(
 25    tracks: list[Track],
 26    outdir: str,
 27    chrom: str,
 28    settings: PlotSettings,
 29) -> tuple[Figure, np.ndarray, list[str]]:
 30    """
 31    Plot a single centromere figure from a list of `Track`s.
 32
 33    # Args
 34    * `tracks`
 35        * List of tracks to plot. The order in the list determines placement on the figure.
 36    * `outdir`
 37        * Output directory.
 38    * `chrom`
 39        * Chromosome name to filter for in `Track.data`
 40    * `settings`
 41        * Settings for output plots.
 42
 43    # Returns
 44    * Figure, its axes, and the output filename(s).
 45
 46    # Usage
 47    ```python
 48    import cenplot
 49
 50    chrom = "chm13_chr10:38568472-42561808"
 51    track_list, settings = cenplot.read_one_cen_tracks("tracks_example_api.toml", chrom=chrom)
 52    fig, axes, outfiles = cenplot.plot_one_cen(track_list.tracks, "plots", chrom, settings)
 53    ```
 54    """
 55    # Show chrom trimmed of spaces for logs and filenames.
 56    logging.info(f"Plotting {chrom}...")
 57
 58    if not settings.xlim:
 59        # Get min and max position of all tracks for this cen.
 60        _, min_st_pos = get_min_max_track(tracks, typ="min")
 61        _, max_end_pos = get_min_max_track(tracks, typ="max", default_col="chrom_end")
 62    else:
 63        min_st_pos = settings.xlim[0]
 64        max_end_pos = settings.xlim[1]
 65
 66    # # Scale height based on track length.
 67    # adj_height = height * (trk_max_end / max_end_pos)
 68    # height = height if adj_height == 0 else adj_height
 69
 70    fig, axes, track_indices = create_subplots(
 71        tracks,
 72        settings,
 73    )
 74    if settings.legend_pos == LegendPosition.Left:
 75        track_col, legend_col = 1, 0
 76    else:
 77        track_col, legend_col = 0, 1
 78
 79    track_labels: list[str] = []
 80
 81    def get_track_label(chrom: str, track: Track, all_track_labels: list[str]) -> str:
 82        if not track.title:
 83            return ""
 84        try:
 85            fmt_track_label = track.title.format(chrom=chrom)
 86        except KeyError:
 87            fmt_track_label = track.title
 88
 89        track_label = fmt_track_label.encode("ascii", "ignore").decode("unicode_escape")
 90
 91        # Update track label for each overlap.
 92        if track.pos == TrackPosition.Overlap:
 93            try:
 94                track_label = f"{all_track_labels[-1]}\n{track_label}"
 95            except IndexError:
 96                pass
 97
 98        return track_label
 99
100    num_hor_split = 0
101    for idx, track in enumerate(tracks):
102        track_row = track_indices[idx]
103        track_label = get_track_label(chrom, track, track_labels)
104
105        try:
106            track_ax: Axes = axes[track_row, track_col]
107        except IndexError:
108            print(f"Cannot get track ({track_row, track_col}) for {track}.")
109            continue
110        try:
111            legend_ax: Axes = axes[track_row, legend_col]
112        except IndexError:
113            legend_ax = None
114
115        # Set xaxis limits
116        track_ax.set_xlim(min_st_pos, max_end_pos)
117
118        if track.opt == TrackType.Legend:
119            draw_legend(track_ax, axes, track, tracks, track_row, track_col)
120        elif track.opt == TrackType.Position:
121            # Hide everything but x-axis
122            format_ax(
123                track_ax,
124                grid=True,
125                xticklabel_fontsize=track.options.legend_fontsize,
126                yticks=True,
127                yticklabel_fontsize=track.options.legend_fontsize,
128                spines=("right", "left", "top"),
129            )
130        elif track.opt == TrackType.Spacer:
131            # Hide everything.
132            format_ax(
133                track_ax,
134                grid=True,
135                xticks=True,
136                yticks=True,
137                spines=("right", "left", "top", "bottom"),
138            )
139        else:
140            # Switch track option. {bar, label, ident, hor}
141            # Add legend.
142            if track.opt == TrackType.HOR or track.opt == TrackType.HORSplit:
143                draw_fn = draw_hor
144            elif track.opt == TrackType.HOROrt:
145                draw_fn = draw_hor_ort
146            elif track.opt == TrackType.Label:
147                draw_fn = draw_label
148            elif track.opt == TrackType.SelfIdent:
149                draw_fn = draw_self_ident
150            elif track.opt == TrackType.LocalSelfIdent:
151                draw_fn = draw_local_self_ident
152            elif track.opt == TrackType.Bar:
153                draw_fn = draw_bar
154            elif track.opt == TrackType.Line:
155                draw_fn = draw_line
156            elif track.opt == TrackType.Strand:
157                draw_fn = draw_strand
158            else:
159                raise ValueError("Invalid TrackType. Unreachable.")
160
161            draw_fn(
162                ax=track_ax,
163                legend_ax=legend_ax,
164                track=track,
165                zorder=idx,
166            )
167
168        # Store label if more overlaps.
169        track_labels.append(track_label)
170
171        # Set labels for both x and y axis.
172        set_both_labels(track_label, track_ax, track)
173
174        if not legend_ax:
175            continue
176
177        # Make legend title invisible for HORs split after 1.
178        if track.opt == TrackType.HORSplit:
179            legend_ax_legend = legend_ax.get_legend()
180            if legend_ax_legend and num_hor_split != 0:
181                legend_title = legend_ax_legend.get_title()
182                legend_title.set_alpha(0.0)
183            num_hor_split += 1
184
185        # Minimalize all legend cols except self-ident
186        if track.opt != TrackType.SelfIdent or (
187            track.opt == TrackType.SelfIdent and not track.options.legend
188        ):
189            format_ax(
190                legend_ax,
191                grid=True,
192                xticks=True,
193                xticklabel_fontsize=track.options.legend_fontsize,
194                yticks=True,
195                yticklabel_fontsize=track.options.legend_fontsize,
196                spines=("right", "left", "top", "bottom"),
197            )
198        else:
199            format_ax(
200                legend_ax,
201                grid=True,
202                xticklabel_fontsize=track.options.legend_fontsize,
203                yticklabel_fontsize=track.options.legend_fontsize,
204                spines=("right", "top"),
205            )
206
207    # Add title
208    if settings.title:
209        title = settings.title.format(chrom=chrom)
210        fig.suptitle(
211            title,
212            x=settings.title_x,
213            y=settings.title_y,
214            horizontalalignment=settings.title_horizontalalignment,
215            fontsize=settings.title_fontsize,
216        )
217
218    os.makedirs(outdir, exist_ok=True)
219    if isinstance(settings.format, str):
220        output_format = [settings.format]
221    else:
222        output_format = settings.format
223    # Pad between axes.
224    fig.get_layout_engine().set(h_pad=settings.axis_h_pad)
225
226    # PNG must always be plotted last.
227    # Matplotlib modifies figure settings causing formatting errors in vectorized image formats (svg, pdf)
228    png_output = "png" in output_format
229    if png_output:
230        output_format.remove("png")
231
232    outfiles = []
233    for fmt in output_format:
234        outfile = os.path.join(outdir, f"{chrom}.{fmt}")
235        fig.savefig(outfile, dpi=settings.dpi, transparent=settings.transparent)
236        outfiles.append(outfile)
237
238    if png_output:
239        outfile = os.path.join(outdir, f"{chrom}.png")
240        fig.savefig(
241            outfile,
242            dpi=settings.dpi,
243            transparent=settings.transparent,
244        )
245        outfiles.append(outfile)
246
247    return fig, axes, outfiles

 90def merge_plots(
 91    figures: list[tuple[Figure, np.ndarray, list[str]]], outfile: str
 92) -> None:
 93    """
 94    Merge plots produced by `plot_one_cen`.
 95
 96    # Args
 97    * `figures`
 98        * List of figures, their axes, and the name of the output files. Only pngs are concatentated.
 99    * `outfile`
100        * Output merged file.
101        * Either `png` or `pdf`
102
103    # Returns
104    * None
105    """
106    if outfile.endswith(".pdf"):
107        with PdfPages(outfile) as pdf:
108            for fig, _, _ in figures:
109                pdf.savefig(fig)
110    else:
111        merged_images = np.concatenate(
112            [
113                plt.imread(file)
114                for _, _, files in figures
115                for file in files
116                if file.endswith("png")
117            ]
118        )
119        plt.imsave(outfile, merged_images)

24def draw_hor(
25    ax: Axes,
26    track: Track,
27    *,
28    zorder: float = 1.0,
29    legend_ax: Axes | None = None,
30):
31    """
32    Draw HOR plot on axis with the given `Track`.
33    """
34    hide_x = track.options.hide_x
35    legend = track.options.legend
36    border = track.options.bg_border
37    bg_color = track.options.bg_color
38
39    if track.pos != TrackPosition.Overlap:
40        spines = (
41            ("right", "left", "top", "bottom") if hide_x else ("right", "left", "top")
42        )
43    else:
44        spines = None
45
46    format_ax(
47        ax,
48        xticks=hide_x,
49        xticklabel_fontsize=track.options.fontsize,
50        yticks=True,
51        yticklabel_fontsize=track.options.fontsize,
52        spines=spines,
53    )
54
55    ylim = ax.get_ylim()
56    height = ylim[1] - ylim[0]
57
58    if track.options.mode == "hor":
59        colname = "name"
60    else:
61        colname = "mer"
62
63    # Add HOR track.
64    for row in track.data.iter_rows(named=True):
65        start = row["chrom_st"]
66        end = row["chrom_end"]
67        color = row["color"]
68        rect = Rectangle(
69            (start, 0),
70            end + 1 - start,
71            height,
72            color=color,
73            lw=0,
74            label=row[colname],
75            zorder=zorder,
76        )
77        ax.add_patch(rect)
78
79    if border:
80        # Ensure border is always on top.
81        add_rect(ax, height, zorder + 1.0)
82
83    if bg_color:
84        # Ensure bg is below everything.
85        add_rect(ax, height, zorder - 1.0, fill=True, color=bg_color)
86
87    if legend_ax and legend:
88        draw_uniq_entry_legend(
89            legend_ax,
90            track,
91            ref_ax=ax,
92            ncols=track.options.legend_ncols,
93            loc="center left",
94            alignment="left",
95        )

11def draw_hor_ort(
12    ax: Axes,
13    track: Track,
14    *,
15    zorder: float = 1.0,
16    legend_ax: Axes | None = None,
17):
18    """
19    Draw HOR ort plot on axis with the given `Track`.
20    """
21    draw_strand(ax, track, zorder=zorder, legend_ax=legend_ax)

 10def draw_label(
 11    ax: Axes,
 12    track: Track,
 13    *,
 14    zorder: float = 1.0,
 15    legend_ax: Axes | None = None,
 16) -> None:
 17    """
 18    Draw label plot on axis with the given `Track`.
 19    """
 20    hide_x = track.options.hide_x
 21    color = track.options.color
 22    alpha = track.options.alpha
 23    legend = track.options.legend
 24    border = track.options.bg_border
 25    edgecolor = track.options.edgecolor
 26
 27    patch_options: dict[str, Any] = {"zorder": zorder}
 28    patch_options["alpha"] = alpha
 29
 30    # Overlapping tracks should not cause the overlapped track to have their spines/ticks/ticklabels removed.
 31    if track.pos != TrackPosition.Overlap:
 32        spines = (
 33            ("right", "left", "top", "bottom") if hide_x else ("right", "left", "top")
 34        )
 35        yticks = True
 36    else:
 37        yticks = False
 38        spines = None
 39    format_ax(
 40        ax,
 41        xticks=hide_x,
 42        xticklabel_fontsize=track.options.fontsize,
 43        yticks=yticks,
 44        yticklabel_fontsize=track.options.fontsize,
 45        spines=spines,
 46    )
 47
 48    ylim = ax.get_ylim()
 49    height = ylim[1] - ylim[0]
 50
 51    patch_options["edgecolor"] = edgecolor
 52
 53    for row in track.data.iter_rows(named=True):
 54        start = row["chrom_st"]
 55        end = row["chrom_end"]
 56
 57        if row["name"] == "-" or not row["name"]:
 58            labels = {}
 59        else:
 60            labels = {"label": row["name"]}
 61
 62        # Allow override.
 63        if color:
 64            patch_options["facecolor"] = color
 65        elif "color" in row:
 66            patch_options["facecolor"] = row["color"]
 67
 68        if track.options.shape == "rect":
 69            rect = Rectangle(
 70                (start, 0),
 71                end + 1 - start,
 72                height,
 73                **labels,
 74                **patch_options,
 75            )
 76            ax.add_patch(rect)
 77        elif track.options.shape == "tri":
 78            midpt = ((end - start) / 2) + start
 79            vertices = [
 80                (start, height),
 81                (end, height),
 82                # tip
 83                (midpt, 0),
 84            ]
 85            ptch = Polygon(
 86                vertices,
 87                closed=True,
 88                **labels,
 89                **patch_options,
 90            )
 91            ax.add_patch(ptch)
 92
 93    if border:
 94        # Ensure border on top with larger zorder.
 95        add_rect(ax, height, fill=False, zorder=zorder + 1.0)
 96
 97    # Draw legend.
 98    if legend_ax and legend:
 99        draw_uniq_entry_legend(
100            legend_ax,
101            track,
102            ref_ax=ax,
103            ncols=track.options.legend_ncols,
104            label_order=track.options.legend_label_order,
105            loc="center left",
106            alignment="left",
107        )

55def draw_self_ident(
56    ax: Axes,
57    track: Track,
58    *,
59    zorder: float = 1.0,
60    legend_ax: Axes | None = None,
61) -> None:
62    """
63    Draw self identity plot on axis with the given `Track`.
64    """
65    hide_x = track.options.hide_x
66    invert = track.options.invert
67    legend = track.options.legend
68
69    colors, verts = [], []
70    spines = ("right", "left", "top", "bottom") if hide_x else ("right", "left", "top")
71    format_ax(
72        ax,
73        xticks=hide_x,
74        xticklabel_fontsize=track.options.fontsize,
75        yticks=True,
76        yticklabel_fontsize=track.options.fontsize,
77        spines=spines,
78    )
79
80    if invert:
81        df_track = track.data.with_columns(y=-pl.col("y"))
82    else:
83        df_track = track.data
84
85    for _, df_diam in df_track.group_by(["group"]):
86        df_points = df_diam.select("x", "y")
87        color = df_diam["color"].first()
88        colors.append(color)
89        verts.append(df_points)
90
91    # https://stackoverflow.com/a/29000246
92    polys = PolyCollection(verts, zorder=zorder)
93    polys.set(array=None, facecolors=colors)
94    ax.add_collection(polys)
95
96    ax.set_ylim(df_track["y"].min(), df_track["y"].max())
97
98    if legend_ax and legend:
99        draw_self_ident_hist(legend_ax, track, zorder=zorder)

12def draw_self_ident_hist(ax: Axes, track: Track, *, zorder: float = 1.0):
13    """
14    Draw self identity histogram plot on axis with the given `Track`.
15    """
16    legend_bins = track.options.legend_bins
17    legend_xmin = track.options.legend_xmin
18    legend_asp_ratio = track.options.legend_asp_ratio
19    colorscale = track.options.colorscale
20    assert isinstance(colorscale, dict), (
21        f"Colorscale not a identity interval mapping for {track.title}"
22    )
23
24    cmap = IntervalTree(
25        Interval(rng[0], rng[1], color) for rng, color in colorscale.items()
26    )
27    cnts, values, bars = ax.hist(
28        track.data["percent_identity_by_events"], bins=legend_bins, zorder=zorder
29    )
30    ax.set_xlim(legend_xmin, 100.0)
31    ax.minorticks_on()
32    ax.set_xlabel(
33        "Mean nucleotide identity\nbetween pairwise intervals",
34        fontsize=track.options.legend_title_fontsize,
35    )
36    ax.set_ylabel(
37        "# of Intervals (thousands)", fontsize=track.options.legend_title_fontsize
38    )
39
40    # Ensure that legend is only a portion of the total height.
41    # Otherwise, take up entire axis dim.
42    ax.set_box_aspect(legend_asp_ratio)
43
44    for _, value, bar in zip(cnts, values, bars):
45        # Make value a non-null interval
46        # ex. (1,1) -> (1, 1.000001)
47        color = cmap.overlap(value, value + 0.00001)
48        try:
49            color = next(iter(color)).data
50        except Exception:
51            color = None
52        bar.set_facecolor(color)

 9def draw_local_self_ident(
10    ax: Axes,
11    track: Track,
12    *,
13    zorder: float = 1.0,
14    legend_ax: Axes | None = None,
15) -> None:
16    """
17    Draw local, self identity plot on axis with the given `Track`.
18    """
19    if not track.options.legend_label_order:
20        track.options.legend_label_order = [
21            f"{cs[0]}-{cs[1]}"
22            for cs in track.options.colorscale.keys()
23        ]
24    draw_label(ax, track, zorder=zorder, legend_ax=legend_ax)

 9def draw_bar(
10    ax: Axes,
11    track: Track,
12    *,
13    zorder: float = 1.0,
14    legend_ax: Axes | None = None,
15) -> None:
16    """
17    Draw bar plot on axis with the given `Track`.
18    """
19    hide_x = track.options.hide_x
20    color = track.options.color
21    alpha = track.options.alpha
22    legend = track.options.legend
23    ymin = track.options.ymin
24    ymax = track.options.ymax
25    label = track.options.label
26
27    if track.pos != TrackPosition.Overlap:
28        spines = ("right", "top")
29    else:
30        spines = None
31
32    format_ax(
33        ax,
34        xticks=hide_x,
35        xticklabel_fontsize=track.options.fontsize,
36        yticklabel_fontsize=track.options.fontsize,
37        spines=spines,
38    )
39
40    plot_options = {"zorder": zorder, "alpha": alpha}
41    if color:
42        plot_options["color"] = color
43    elif "color" in track.data.columns:
44        plot_options["color"] = track.data["color"]
45    else:
46        plot_options["color"] = track.options.DEF_COLOR
47
48    # Add bar
49    ax.bar(
50        track.data["chrom_st"],
51        track.data["name"],
52        track.data["chrom_end"] - track.data["chrom_st"],
53        label=label,
54        **plot_options,
55    )
56    # Trim plot to margins
57    ax.margins(x=0, y=0)
58    ax.set_ylim(ymin=ymin, ymax=ymax)
59
60    if legend_ax and legend:
61        draw_uniq_entry_legend(
62            legend_ax,
63            track,
64            ref_ax=ax,
65            ncols=track.options.legend_ncols,
66            loc="center left",
67            alignment="left",
68        )

 10def draw_line(
 11    ax: Axes,
 12    track: Track,
 13    *,
 14    zorder: float = 1.0,
 15    legend_ax: Axes | None = None,
 16) -> None:
 17    """
 18    Draw line plot on axis with the given `Track`.
 19    """
 20    hide_x = track.options.hide_x
 21    color = track.options.color
 22    alpha = track.options.alpha
 23    legend = track.options.legend
 24    ymin = track.options.ymin
 25    ymax = track.options.ymax
 26    label = track.options.label
 27    linestyle = track.options.linestyle
 28    linewidth = track.options.linewidth
 29    marker = track.options.marker
 30    markersize = track.options.markersize
 31
 32    if track.pos != TrackPosition.Overlap:
 33        spines = ("right", "top")
 34    else:
 35        spines = None
 36
 37    format_ax(
 38        ax,
 39        xticks=hide_x,
 40        xticklabel_fontsize=track.options.fontsize,
 41        yticklabel_fontsize=track.options.fontsize,
 42        spines=spines,
 43    )
 44
 45    plot_options = {"zorder": zorder, "alpha": alpha}
 46    if color:
 47        plot_options["color"] = color
 48    elif "color" in track.data.columns:
 49        plot_options["color"] = track.data["color"]
 50    else:
 51        plot_options["color"] = track.options.DEF_COLOR
 52
 53    if linestyle:
 54        plot_options["linestyle"] = linestyle
 55    if linewidth:
 56        plot_options["linewidth"] = linewidth
 57
 58    # Fill between cannot add markers
 59    if not track.options.fill:
 60        plot_options["marker"] = marker
 61        if markersize:
 62            plot_options["markersize"] = markersize
 63
 64    if track.options.position == "midpoint":
 65        df = track.data.with_columns(
 66            chrom_st=pl.col("chrom_st") + (pl.col("chrom_end") - pl.col("chrom_st")) / 2
 67        )
 68    else:
 69        df = track.data
 70
 71    if track.options.log_scale:
 72        ax.set_yscale("log")
 73
 74    # Add bar
 75    if track.options.fill:
 76        ax.fill_between(
 77            df["chrom_st"],
 78            df["name"],
 79            0,
 80            label=label,
 81            **plot_options,
 82        )
 83    else:
 84        ax.plot(
 85            df["chrom_st"],
 86            df["name"],
 87            label=label,
 88            **plot_options,
 89        )
 90
 91    # Trim plot to margins
 92    ax.margins(x=0, y=0)
 93    ax.set_ylim(ymin=ymin, ymax=ymax)
 94
 95    if legend_ax and legend:
 96        draw_uniq_entry_legend(
 97            legend_ax,
 98            track,
 99            ref_ax=ax,
100            ncols=track.options.legend_ncols,
101            loc="center left",
102            alignment="left",
103        )

 8def draw_strand(
 9    ax: Axes,
10    track: Track,
11    *,
12    zorder: float = 1.0,
13    legend_ax: Axes | None = None,
14):
15    """
16    Draw strand plot on axis with the given `Track`.
17    """
18    hide_x = track.options.hide_x
19    fwd_color = (
20        track.options.fwd_color if track.options.fwd_color else track.options.DEF_COLOR
21    )
22    rev_color = (
23        track.options.rev_color if track.options.rev_color else track.options.DEF_COLOR
24    )
25    scale = track.options.scale
26    legend = track.options.legend
27
28    if track.pos != TrackPosition.Overlap:
29        spines = (
30            ("right", "left", "top", "bottom") if hide_x else ("right", "left", "top")
31        )
32    else:
33        spines = None
34
35    format_ax(
36        ax,
37        xticks=hide_x,
38        xticklabel_fontsize=track.options.fontsize,
39        yticks=True,
40        yticklabel_fontsize=track.options.fontsize,
41        spines=spines,
42    )
43
44    ylim = ax.get_ylim()
45    height = ylim[1] - ylim[0]
46
47    for row in track.data.iter_rows(named=True):
48        # sample arrow
49        start = row["chrom_st"]
50        end = row["chrom_end"]
51        strand = row["strand"]
52        if strand == "-":
53            tmp_start = start
54            start = end
55            end = tmp_start
56            color = rev_color
57        else:
58            color = fwd_color
59
60        if track.options.use_item_rgb:
61            color = row["color"]
62
63        arrow = FancyArrowPatch(
64            (start, height * 0.5),
65            (end, height * 0.5),
66            mutation_scale=scale,
67            color=color,
68            clip_on=False,
69            zorder=zorder,
70            label=row["name"],
71        )
72        ax.add_patch(arrow)
73
74    if legend_ax and legend:
75        draw_uniq_entry_legend(
76            legend_ax,
77            track,
78            ref_ax=ax,
79            ncols=track.options.legend_ncols,
80            loc="center",
81        )