matplotlib subplots helper

plt_subplots_helper.py

#
# helper function to matplotlib subplots make pretty and easy to handle.
#

from matplotlib import pyplot as plt 
plt.style.use('default')  # To avoid chart invisible, due to browser theme etc.


# Create raws x cols square subplots, fits to given width
# returned axes is 1D array, so that easily handle with simple loop
def create_figax(raws, cols, width=40):
  size = int(width/cols)
  fig, axes = plt.subplots(raws, cols, figsize=(size*cols, size*raws))
  axes = np.array(axes).reshape(-1) # make axes 1D
  return fig, axes

# Hide axis after specified index (for unused subplots in created axes)
def clear_remainig_axes(axes, index):
  for i in range(index-1, len(axes)):
    axes[i].axis("off")

# Just a helper for setting commonly used attributes
def set_axes_attributes(ax, 
                        title=None,
                        title_size=None, 
                        xlabel=None, 
                        ylabel=None,
                        label_size=None):
  if title:
    ax.set_title(title, size=title_size)
  if xlabel:
    ax.set_xlabel(xlabel, size=label_size)
  if ylabel:
    ax.set_ylabel(ylabel, size=label_size)


#
# Usage example
#

fig, axes = create_figax(5,6) # create 5 x 6 subplots

# plot something with simple 1D loop
for i, ax in enumerate(axes[:8]):
  ax.plot([1,i],[1,i], marker="o")
  set_axes_attributes(ax, title=i, title_size=20, xlabel=i+1, label_size=10)

# since subplots after axes[9], delete their axis
clear_remainig_axes(axes, 9)

# call tight_layout to adjust positions
fig.tight_layout()