Indexing single dimensional Axes array with two indices

Question

I want to write a formula for plotting charts. I have data like this

import numpy as np
import matplotlib.pyplot as plt

x = (['2020-08-01', '2020-09-01', '2020-08-01'])
y = ([1,3, 8])

And my formula is

def fire_m(x,y, i,j):
    ax[i,j].plot(x, y)

When i do

fig, ax = plt.subplots(3,2)
fire_m(x, y, 0, 0)
plt.show()

That's ok. But if i need just one column

fig, ax = plt.subplots(3,1)
    fire_m(x, y, 0, 0)
    plt.show()

I get

IndexError: too many indices for array: array is 1-dimensional, but 2 were indexed

I tryed to improve

def fire_m(x,y, i,j=np.nan):
        ax[i,j].plot(x, y)

But it's not correct. How to fix it?

What you are referring to as "formula" is typically called a "function" fyi — busybear, Jan 09 '21 at 22:25
You can add `squeeze=False` to `plt.subplots(...., squeeze=False)` to have always 2D axes. — JohanC, Jan 10 '21 at 01:41

score 1 · Answer 1 · answered Jan 09 '21 at 22:23

If you want to use nan as the workaround, you can try this:

def fire_m(x, y, i, j=np.nan):
    if j == np.nan:
        ax[i].plot(x, y)
    else:
        ax[i,j].plot(x, y)

Another option is to use numpy.ravel_multi_index. This will convert the i, j-coordinate to a single index in a flattened version of the array.

def fire_m(x, y, i, j):
    n = np.ravel_multi_index([i, j], ax.shape)
    ax.flatten()[n].plot(x, y)

1 Answers1