Vertical arrangement in matplotlib subplots

Question

I am trying to plot a figure with 4 subplots. The following code is similar to what I have written.

import numpy as np
import matplotlib.pyplot as plt
def trial(x, r, ax=None):
    y = r*np.sin(x)
    ax[0].plot(x, r*x)
    ax[0].set_title('x')
    ax[1].plot(x, y)
    ax[1].set_title('sinx')

x = np.linspace(0, 4*np.pi)    
fig, (ax1, ax2) = plt.subplots(2,2)
trial(x, 1, ax1)
trial(x, 2, ax2)

This is the output I get.

If I only plot ax1(after removing last index from subplots), x and sin(x) are vertically arranged. But as soon as I add ax2, they become horizontally arranged. I want x and sin(x) to be vertically arranged with both ax1 and ax2. What correction should I make?

Answer 1

Let's try:

fig, axes = plt.subplots(2,2)
trial(x, 1, axes[:,0])
trial(x, 2, axes[:,1])

Output: