plotter.py 文件源码

def saturation_index_countour(lab, elem1, elem2, Ks, labels=False):
    plt.figure()
    plt.title('Saturation index %s%s' % (elem1, elem2))
    resoluion = 100
    n = math.ceil(lab.time.size / resoluion)
    plt.xlabel('Time')
    z = np.log10((lab.species[elem1]['concentration'][:, ::n] + 1e-8) * (
        lab.species[elem2]['concentration'][:, ::n] + 1e-8) / lab.constants[Ks])
    lim = np.max(abs(z))
    lim = np.linspace(-lim - 0.1, +lim + 0.1, 51)
    X, Y = np.meshgrid(lab.time[::n], -lab.x)
    plt.xlabel('Time')
    CS = plt.contourf(X, Y, z, 20, cmap=ListedColormap(sns.color_palette(
        "RdBu_r", 101)), origin='lower', levels=lim, extend='both')
    if labels:
        plt.clabel(CS, inline=1, fontsize=10, colors='w')
    # cbar = plt.colorbar(CS)
    if labels:
        plt.clabel(CS, inline=1, fontsize=10, colors='w')
    cbar = plt.colorbar(CS)
    plt.ylabel('Depth')
    ax = plt.gca()
    ax.ticklabel_format(useOffset=False)
    cbar.ax.set_ylabel('Saturation index %s%s' % (elem1, elem2))
    return ax