保存来自函数的matplotlib绘图 python

Question

我写了一个函数，可以从数据集中取出一系列数值，然后输出一个图表。比如说：

my_plot(location_dataset, min_temperature, max_temperature) 这个函数会根据你指定的温度范围，返回一个降水量的图表。

假设我想保存加利福尼亚州在60到70华氏度之间的图表。那么，我可以调用这个函数 my_plot(California, 60, 70)，这样就能得到加利福尼亚在60到70华氏度时的降水量图表。

我想问的是：怎么把这个函数生成的图表保存为jpeg格式呢？

我知道 plt.savefig() 是用来保存图表的，但那是在没有调用函数的情况下。在我的情况下，我该怎么做呢？

谢谢！

更多细节：这是我的代码（简化版）：

import matplotlib.pyplot as plt

def my_plot(location_dataset, min_temperature, max_temperature):
    condition = (location_dataset['temperature'] > min_temperature) & (dataset['temperature'] <= max_temperature)
    subset = location_dataset[condition] # subset the data based on the temperature range

    x = subset['precipitation'] # takes the precipitation column only
    plt.figure(figsize=(8, 6))
    plt.plot(x)
    plt.show()

然后我这样调用这个函数： my_plot(California, 60, 70)，就得到了60到70华氏度范围内的图表。可是我该怎么保存这个图表，而不在函数定义里加 savefig 呢？这是因为我需要更改最小和最大温度的参数。

Answer 1

把figure的引用存到一个变量里，然后从你的函数中返回这个变量：

import matplotlib.pyplot as plt

def my_plot(location_dataset, min_temperature, max_temperature):
    condition = (location_dataset['temperature'] > min_temperature) & (dataset['temperature'] <= max_temperature)
    subset = location_dataset[condition] # subset the data based on the temperature range

    x = subset['precipitation'] # takes the precipitation column only
    # N.B. referenca taken to fig
    fig = plt.figure(figsize=(8, 6))
    plt.plot(x)
    plt.show()

    return fig

当你调用这个函数的时候，可以用这个引用来保存图形：

fig = my_plot(...)
fig.savefig("somefile.png")