How to plot vertical lines on subplots that are generated from a function outside of the function

Question

I have the following function:

def test():
  f, axs = plt.subplots(2,2)
  x=np.array([1,2,3])
  y=np.array([1,2,3])
  for i, ax in enumerate(axs.flat):
    axes.append(ax.plot(x,y))
  return f, axs

Update: I would like to add vertical lines for each of the 4 graphs. For example, say I want to plot a vertical line for the first subplot:

Answer 1

I might have missed something, but it's not clear what axes.append(ax.plot(x, y)) is going to do, since you don't define axes . But since axs is already an array, and since you are mutating its contents, I don't think you need to do any appending.

Anyway, if you want to add a vertical line, you could do this outside your function like so:

import numpy as np
import matplotlib.pyplot as plt

def test():
    fig, axs = plt.subplots(2, 2)
    x = np.array([1, 2, 3])
    y = np.array([1, 2, 3])
    for i, ax in enumerate(axs.flat):
        ax.plot(x, y)
    return fig, axs

fig, axs = test()

_ = axs[0, 1].axvline(2, color='red')

Result:

Or you could equally do ax.axvline() or ax.axhline() inside the function.