quinncasey22
quinncasey22

Reputation: 39

optional plot in plot function

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.

Upvotes: 0

Views: 225

Answers (1)

Pedro Rodrigues
Pedro Rodrigues

Reputation: 36

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.

Upvotes: 1

Related Questions