optional plot in plot function-CodePudding

This should be a pretty straightforward question. I have a function which plots a number of results, and I want certain data plotted everytime; however, I want the option to plot additional material (e.g. a polynomial fit). Example Data:

x1 = np.linspace(0, 10, 10)
y1 = x1**(3/2)-x1*2
x2 = np.linspace(0, 10, 10)
y2 = 0.5*x2   x2**0.5

fit1 = np.poly1d(np.polyfit(x1, y1, 2))
fit2 = np.poly1d(np.polyfit(x2, y2, 2))

Example Plotting Function:

def plot():
    fig, ax = plt.subplots(1, 2, sharex=True, sharey=True)
    
    ax[0].scatter(x1, y1)
    ax[1].scatter(x2, y2)
    
    plt.show()

So I'd like to add some parameters to plot() which allows me to plot fit1 and fit2 but only if I want to. I'm not sure if this is allowed. I'll be plotting my chart many times, and the want the option to compare different fit lines to the data.

CodePudding user response：

In your function plot you can add two if statements and the conditions (boolean values) are arguments of your function.

For example:

def plot(condition_1, condition_2):
    fig, ax = plt.subplots(1, 2, sharex=True, sharey=True)
    
    if condition_1:
        ax[0].scatter(x1, y1)

    if condition_2:    
        ax[1].scatter(x2, y2)
    
    plt.show()

Something like this would work.