rism3d_pressure.py 文件源码

def plot_and_save(x, func, xvv_inst, plot=False, fname=False):
    mat = x.T
    for m in range(xvv_inst.nsites):
        for n in range(m+1):
            if fname:
                mat = np.c_[mat, func[:,m,n].T]
            if plot:
                plt.plot(x, func[:,m,n], 
                         label='{}-{}'.format(xvv_inst.atom_names[m], 
                                              xvv_inst.atom_names[n]))
    if fname:
        np.savetxt(fname, mat)
    if plot:
        plt.legend()
        plt.savefig('graph.png', dpi=300)
        plt.show()