How to make matplotlib figure that adjusts its size w.r.t. number of columns and rows automatically?

Question

The code below plots this figure.

from matplotlib import pyplot as plt
from matplotlib.lines import Line2D

import re
import numpy as np
import os

def na