如何为变量分配绘图并将该变量用作Python函数中的返回值

Question

我正在创建两个Python脚本来为技术报告生成一些图.在第一个脚本中,我定义了从硬盘上的原始数据生成绘图的函数.每个函数都会生成一个我需要的特定类型的图.第二个脚本更像是一个批处理文件,它应该绕过这些函数并将生成的图存储在我的硬盘上.

我需要的是一种在Python中返回绘图的方法.所以基本上我想这样做:

fig = some_function_that_returns_a_plot(args)
fig.savefig('plot_name')

但我不知道的是如何使情节成为我可以返回的变量.这可能吗？是这样,怎么样？

Answer 1

您可以定义绘图功能

import numpy as np
import matplotlib.pyplot as plt

# an example graph type
def fig_barh(ylabels, xvalues, title=''):
    # create a new figure
    fig = plt.figure()

    # plot to it
    yvalues = 0.1 + np.arange(len(ylabels))
    plt.barh(yvalues, xvalues, figure=fig)
    yvalues += 0.4
    plt.yticks(yvalues, ylabels, figure=fig)
    if title:
        plt.title(title, figure=fig)

    # return it
    return fig

然后使用它们

from matplotlib.backends.backend_pdf import PdfPages

def write_pdf(fname, figures):
    doc = PdfPages(fname)
    for fig in figures:
        fig.savefig(doc, format='pdf')
    doc.close()

def main():
    a = fig_barh(['a','b','c'], [1, 2, 3], 'Test #1')
    b = fig_barh(['x','y','z'], [5, 3, 1], 'Test #2')
    write_pdf('test.pdf', [a, b])

if __name__=="__main__":
    main()