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

111def merge_plots(
112    figures: list[tuple[Figure, np.ndarray, list[str]]], outfile: str
113) -> None:
114    """
115    Merge plots produced by `plot_one_cen`.
116
117    # Args
118    * `figures`
119        * List of figures, their axes, and the name of the output files. Only pngs are concatentated.
120    * `outfile`
121        * Output merged file.
122        * Either `png` or `pdf`
123
124    # Returns
125    * None
126    """
127    if outfile.endswith(".pdf"):
128        with PdfPages(outfile) as pdf:
129            for fig, _, _ in figures:
130                pdf.savefig(fig)
131    else:
132        merged_images = np.concatenate(
133            [
134                plt.imread(file)
135                for _, _, files in figures
136                for file in files
137                if file.endswith("png")
138            ]
139        )
140        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    ymin, ymax = (
 97        df_track["y"].min(),
 98        df_track["y"].max(),
 99    )
100
101    ax.set_ylim(ymin, ymax)
102
103    if legend_ax and legend:
104        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):  # type: ignore[arg-type]
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]}" for cs in track.options.colorscale.keys()
22        ]
23    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    label = track.options.label
24
25    if track.pos != TrackPosition.Overlap:
26        spines = ("right", "top")
27    else:
28        spines = None
29
30    format_ax(
31        ax,
32        xticks=hide_x,
33        xticklabel_fontsize=track.options.fontsize,
34        yticklabel_fontsize=track.options.fontsize,
35        spines=spines,
36    )
37
38    plot_options = {"zorder": zorder, "alpha": alpha}
39    if color:
40        plot_options["color"] = color
41    elif "color" in track.data.columns:
42        plot_options["color"] = track.data["color"]
43    else:
44        plot_options["color"] = track.options.DEF_COLOR
45
46    # Add bar
47    ax.bar(
48        track.data["chrom_st"],
49        track.data["name"],
50        track.data["chrom_end"] - track.data["chrom_st"],
51        label=label,
52        **plot_options,
53    )  # type: ignore[arg-type]
54    # Trim plot to margins
55    ax.margins(x=0, y=0)
56
57    set_ylim(ax, track)
58
59    if legend_ax and legend:
60        draw_uniq_entry_legend(
61            legend_ax,
62            track,
63            ref_ax=ax,
64            ncols=track.options.legend_ncols,
65            loc="center left",
66            alignment="left",
67        )

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

 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        )