Plotting quadratic equations with varying parameters

Question

On Wikipedia, there is a graph shown:

It shows the quadratic \$y = ax2 + bx + c\$, varying each coefficient separately while the other coefficients are fixed (at values a = 1, b = 0, c = 0)

As a learning experience I decided to replicate these plots in Matplotlib as follows:

import numpy as np
import matplotlib.pyplot as plt
import math

#Plot the quadratic function y = ax2 + bx + c
#Varying each coefficient [a, b, c] separately while the other coefficients are fixed (at values a = 1, b = 0, c = 0)

#Iterate these 5 coeficients and plot each line
coefs = [-2, -1, 0, 1, 2]

#set up the plot and 3 subplots (to show the effect of varying each coefficient)
f, (ax1, ax2, ax3) = plt.subplots(1, 3, sharey=True, figsize=(18, 6))

#some x values to plot
x = np.linspace(-2, 2, 30)

for idx, val in enumerate([ax1, ax2, ax3]):

    for v in coefs:        
        a, b, c = 1, 0, 0

        if idx == 0:
            a = v
        elif idx == 1:
            b = v
        else:
            c = v

        y = a * (x**2) + (b * x) + c

        val.plot(x, y, label="Coeficient is " + str(coefs[i]))

        val.axhline(y=0, color='k')
        val.axvline(x=0, color='k')    
        val.grid()
        val.legend(loc='lower center')

plt.show()

It works fine:

but I am new to programming and I have an uneasy feeling that using if statements isn't optimal. It feels like I should be iterating over something rather than using ifs.

What other ways could I iterate over the coefficients a, b & c to produce the 3 different subplots?

Answer 1

tl;dr: Use numpy's broadcasting features:

coefs = np.array([-2, -1, 0, 1, 2])

#set up the plot and 3 subplots (to show the effect of varying each coefficient)
f, axes = plt.subplots(1, 3, sharey=True, figsize=(18, 6))

#some x values to plot
x = np.linspace(-2, 2, 30)

# calculate the y-values, varying each coefficient separately
y_values = [coefs[:, None] * x[None, :] ** 2,                 # vary a
            x[None, :] ** 2 + coefs[:, None] * x[None, :],    # vary b
            x[None, :] ** 2 + coefs[:, None]]                 # vary c

for ax, ys in zip(axes, y_values):
    for c, y in zip(coefs, ys):
        ax.plot(x, y, label="Coefficient is {}".format(c))
    ax.axhline(y=0, color='k')
    ax.axvline(x=0, color='k')
    ax.grid()
    ax.legend()
plt.show()

Calculating all necessary y-values would make this a 3D array (one dimension being x, one being the different plots and one being the different values of the coefficients).

To make it slightly easier than that we can handle each plot as a 2D array and make a list of the plots. This way we don't have to deal with the fact that the quadratic term should always appear.

If you really wanted to you could probably make that into a vectorized computation as well, but three dimensions was a bit too much to get straight in my head right now.

---

Explanation of broadcasting:

You can use the fact that numpy has something called broadcasting. It automatically adapts the shapes of vectors you do some operations on, if they are consistent. It can make things look very confusing but allows you to vectorize operations along multiple dimensions (using the fact that doing so loops at C speed, instead of at Python speed).

A small example: You want to multiply two arrays, both of shape (2, 2). numpy then simply does an element-wise multiplication (so \$c_{ij} = a_{ij} \cdot b_{ij}\$).

But if you multiply two arrays of shape (1, 2) and (2, 2), this is not directly possible. So instead it repeats the first array along its first dimension (since it is only of lenth one) and then does an element-wise multiplication. In other words \$c_{ij} = a_{0j} \cdot b_{ij}\$.

To make this abstract example concrete:

>>> a = np.arange(2).reshape(1,2)
array([[0, 1]])
>>> b = np.arange(4).reshape(2,2)
array([[0, 1],
       [2, 3]])
>>> a * b
array([[0, 1],
       [0, 3]])

Now, the final piece is the ability to cast arrays to different shapes (as I already did in the above example using numpy.reshape. Apart from that command you can also use fancy slicing to achieve the same effect:

>>> np.arange(2)[None, :].shape
(1, 2)
>>> np.arange(2)[:, None].shape
(2, 1)

So it repeats it along the axis where I put None (and puts all values in the axis where I put :).