What is the preferred way to import pylab at a function level in Python 2.7?

Question

I have written a relatively simple function in python that can be used to plot the time domain history of a data set as well as the frequency domain response of a data set after a fast fourier transform. In this function I use the command from pylab import * to bring in all the necessary functionality. However, despite successfully creating the plot, I get a warning stating

import * only allowed at the module level.

So if using the command from pylab import * is not the preferred methodology, how do I properly load all the necessary functionality from pylab. The code is attached below. Also, is there a way to close the figure after the function is exited, I have tried plt.close() which is not recognized for subplots?

def Time_Domain_Plot(Directory,Title,X_Label,Y_Label,X_Data,Y_Data):
    # Directory: The path length to the directory where the output file is 
    #            to be stored
    # Title:     The name of the output plot, which should end with .eps or .png
    # X_Label:   The X axis label
    # Y_Label:   The Y axis label
    # X_Data:    X axis data points (usually time at which Yaxis data was acquired
    # Y_Data:    Y axis data points, usually amplitude

    from pylab import *
    from matplotlib import rcParams
    rcParams.update({'figure.autolayout': True})
    Output_Location = Directory.rstrip() + Title.rstrip()
    fig,plt = plt.subplots()
    matplotlib.rc('xtick',labelsize=18)
    matplotlib.rc('ytick',labelsize=18)
    plt.set_xlabel(X_Label,fontsize=18)
    plt.set_ylabel(Y_Label,fontsize=18)
    plt.plot(X_Data,Y_Data,color='red')
    fig.savefig(Output_Location)
    plt.clear()

Solution

From the matplotlib documentation:

pylab is a convenience module that bulk imports matplotlib.pyplot (for plotting) and numpy (for mathematics and working with arrays) in a single name space. Although many examples use pylab, it is no longer recommended.

I would recommend not importing pylab at all, and instead try using

import matplotlib
import matplotlib.pyplot as plt

And then prefixing all of your pyplot functions with plt.

I also noticed that you assign the second return value from plt.subplots() to plt. You should rename that variable to something like fft_plot (for fast fourier transform) to avoid naming conflicts with pyplot.

With regards to your other question (about fig, save fig()) you're going to need to drop that first fig because it's not necessary, and you'll call savefig() with plt.savefig() because it is a function in the pyplot module. So that line will look like

plt.savefig(Output_Location)

Try something like this:

def Time_Domain_Plot(Directory,Title,X_Label,Y_Label,X_Data,Y_Data):
    # Directory: The path length to the directory where the output file is 
    #            to be stored
    # Title:     The name of the output plot, which should end with .eps or .png
    # X_Label:   The X axis label
    # Y_Label:   The Y axis label
    # X_Data:    X axis data points (usually time at which Yaxis data was acquired
    # Y_Data:    Y axis data points, usually amplitude

    import matplotlib
    from matplotlib import rcParams, pyplot as plt

    rcParams.update({'figure.autolayout': True})
    Output_Location = Directory.rstrip() + Title.rstrip()
    fig,fft_plot = plt.subplots()
    matplotlib.rc('xtick',labelsize=18)
    matplotlib.rc('ytick',labelsize=18)
    fft_plot.set_xlabel(X_Label,fontsize=18)
    fft_plot.set_ylabel(Y_Label,fontsize=18)
    plt.plot(X_Data,Y_Data,color='red')
    plt.savefig(Output_Location)
    plt.close()

Comments

[Jon]

Thank you for the suggestion. Unfortunately I have already tried that. Once I import matplotlib.pyplot as plt, the code no longer recognizes the command matplotlib.rc, nor fig, save fig(). If I import matplotlib.pyplot as plt, then how do I get the code to recognize the other commands?

[wpercy]

You also need to import maplotlib apparently. I'm updating my answer now

[Jon]

I will wait for your updated answer. I also tried to import matplotlib and that the code still does not recognize the command 'savefig'

[wpercy]

@Jon That's because savefig is in the pyplot module and you need to prefix it with plt.

[Jon]

Unfortunately that does not work wither. if I use plt.savefig I get a message stating that 'AxesSubplot' object has no attribute savefig. Also, one would think that if I import marplotlib, then there would be no more need for from matplotlib import rcParams, but if I remove that line, then it no longer recognizes rcParams

[wpercy]

@Jon You are clearly struggling with the structure of the matplotlib module and probably python modules in general. The big problem is that you override the value for plt with the actual plot, leaving yourself with no access to the pyplot module. You need to either have from matplotlib import rcParams OR if you don't want that and just want to import matplotlib, you need to prefix rcParams with matplotlib like matplotlib.rcParams

[Jon]

Never mind my previous comment, i just deleted it, I saw a minot error that changed everything. Your posted code did the trick and I appreciate your patience with me.