def _plotIDStrip(ax: matplotlib.axes, id_array: np.ndarray, label: str,
                 x_extents: List[float]) -> None:
    '''
    Plot a strip above the data sequence image that represents class labels at each time.
    '''

    id_array = np.reshape(id_array, (1, len(id_array)))
    ax.imshow(id_array, aspect="auto", interpolation="none", extent=x_extents + [0, 1])
    ax.set_ylabel(label)
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_ticks([])
Example #2
0
def plot_chromatograph(
    seq: SeqRecord, region: Tuple[int, int] = None, ax: mpl.axes = None
) -> plt.axes:
    """Plot Sanger chromatograph.

    region: include both start and end (1-based)
    """
    if seq is None:
        return ax

    if region is None:
        # turn into 0 based for better indexing
        region_start, region_end = 0, len(seq)
    else:
        region_start = max(region[0], 0)
        region_end = min(region[1], len(seq) - 1)

    if ax is None:
        _, ax = plt.subplots(1, 1, figsize=(16, 6))

    _colors = defaultdict(
        lambda: "purple", {"A": "g", "C": "b", "G": "k", "T": "r"}
    )

    # Get signals
    peaks = seq.annotations["peak positions"]
    trace_x = seq.annotations["trace_x"]
    traces_y = [seq.annotations["channel " + str(i)] for i in range(1, 5)]
    bases = seq.annotations["channels"]

    xlim_left, xlim_right = peaks[region_start] - 1, peaks[region_end] + 0.5

    # subset peak and sequence
    # TODO: this might fix the bug
    peak_start = peaks[0]
    peak_zip = [
        (p, s)
        for i, (p, s) in enumerate(zip(peaks, seq))
        if region_start <= i <= region_end
    ]
    peaks, seq = list(zip(*peak_zip))

    # subset trace_x and traces_y together
    trace_zip = [
        (x + peak_start, *ys)
        for x, *ys in zip(trace_x, *traces_y)
        if xlim_left <= x <= xlim_right
    ]
    if not trace_zip:
        return ax
    trace_x, *traces_y = list(zip(*trace_zip))

    # Plot traces
    trmax = max(map(max, traces_y))
    for base in bases:
        trace_y = [1.0 * ci / trmax for ci in traces_y[bases.index(base)]]
        ax.plot(trace_x, trace_y, color=_colors[base], lw=2, label=base)
        ax.fill_between(
            trace_x, 0, trace_y, facecolor=_colors[base], alpha=0.125
        )

    # Plot bases at peak positions
    for i, peak in enumerate(peaks):
        #  LOGGER.debug(f"{i}, {peak}, {seq[i]}, {xlim_left + i}")
        ax.text(
            peak,
            -0.11,
            seq[i],
            color=_colors[seq[i]],
            va="center",
            ha="center",
            alpha=0.66,
            fontsize="x-large",
            fontweight="bold",
        )

    ax.set_ylim(bottom=-0.15, top=1.05)
    #  peaks[0] - max(2, 0.02 * (peaks[-1] - peaks[0])),
    #  right=peaks[-1] + max(2, 0.02 * (peaks[-1] - peaks[0])),
    ax.set_xlim(xlim_left + 0.5, xlim_right)
    ax.set_xticks(peaks)
    ax.set_xticklabels(list(range(region_start + 1, region_end + 2)))
    # hide y axis
    ax.set_yticklabels([])
    ax.get_yaxis().set_visible(False)
    # hide border
    ax.spines["left"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["top"].set_visible(False)
    # hide grid
    ax.grid(False)
    # set legend
    ax.legend(loc="upper left", bbox_to_anchor=(0.95, 0.99))
    return ax