tredplot.py 文件源码

def ax_imshow(ax, P_h1h2, cmap, label, h1_hat, h2_hat, samplekey,
              r=4, draw_circle=True, ticks=True):
    im = ax.imshow(P_h1h2, cmap=cmap, origin="lower")

    from mpl_toolkits.axes_grid1 import make_axes_locatable
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=.05)
    cb = plt.colorbar(im, cax)
    cb.set_label(label)
    if not ticks:
        cb.set_ticks([])

    if draw_circle:
        circle = plt.Circle((h1_hat, h2_hat), r, ec='w', fill=False)
        ax.add_artist(circle)

    annotation = "$\hat{h_1}=%d, \hat{h_2}=%d$" % (h1_hat, h2_hat)
    ax.text(200, 100, annotation, color=lsg, ha="center", va="center")

    ax.set_xlabel(r"$h_1$")
    ax.set_ylabel(r"$h_2$")
    title = "Joint probability density for $h_1$ and $h_2$ ({})"\
            .format(samplekey)
    ax.set_title(title)