2
\$\begingroup\$

For plotting I often use a dictionary that contains my data of interest, regarding my numerical simulations. For example I have a numpy array let's call it U the contains a discretization of a equation in every simulation timestep. Simply put, my U is a 2D np.array where each row is a given timestep (first row first time step) and each column is a spatial discretization. For example:

U = np.array([ [ 0, 0.25, 0.5, 0.25, 0 ],
               [ 0.125, 0.225, 0.3, 0.225, 0.125 ],
               [ 0.2, 0.2, 0.2, 0.2, 0.2] ])

Hence, U[0] is the initial positions, U1 is the first iteration, U2 the last one, of course, is easy to get to thousands timesteps, lets move from here, since isn't important how U (or any simulation) is calculated.

Each simulation is packed in a "data structure" dictionary to guide the plotting, that contains some key information as: x_data spatial discretization, y_data functions evaluated in each point in spatial discretization in all time steps (The U above), data limits x_lim, y_lim, if I want to set them, label that is what y_data is related to, title a title to my graph, linestyle matplotlib linestyle, color color to plot this data. For example:

# These were pre-calculated somewhere else in code
V_x_support = np.linspace(...)
V_t_support = np.linspace(...)
V = np.array(...)

# Here starts the data to plot
ts_range = (0,3000,300)      # I just want to pick some data to graph

V_dataplots = []
for pos in np.arange(*ts_range):
data = {
        'x_data': V_x_support,
        'y_data': V[pos],
        'x_lim': [-1, 1],
        'y_lim': [-1, 1],
        'label': f'{V_t_support[pos]:.3g} s',
        'title': f'{V_t_support[pos]:.3g} s',
        'linestyle': '--',
        'color': '.4',
       }

Once I calculated some arrays, I would love to compare them using matplotlib, but I hate to repeat myself, hence I started to write functions to make the drawings, where I can just pass to my plot function

fig = panelize(
    data = [V_dataplots],
    plot_function = single_plot, 
    title = 'V Mimetic 4th order',
    labely='Velocity',
    labelx='Position',
    sharex=True,
    sharey=True
)

This will pick my manually chosen positions in time and plot it in a panel. The function to do the panels and the plots are below.

How can I reduce the number o lines, keeping the code more cleaner and more readable while avoiding to repeat myself. There are any good MatPlotLib techniques to handle this?

First two examples of the image output

single plot

compare plot

# TODO:
#   - If a panelized axes do not have associated data to plot it will
#     not configure the axes and therefor will not be "prettified"
#     need to fix this (or remove at all the not drawn plots).
#
#   - Draw legend only if some plot has set a label, how to do this? 
#
#   - Make the font scale with less plots are drawn, less axes, bigger
#     fonts.

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import warnings

warnings.filterwarnings("ignore", category=mpl.MatplotlibDeprecationWarning)


def get_axes_diagonal(ax):
    axes_limits = list(map(ax.transData.transform, zip(ax.get_xlim(), ax.get_ylim())))
    return np.linalg.norm(axes_limits[1] - axes_limits[0])


def countour_every(ax, every, x_data, y_data,
        color='black', linestyle='-', marker='o', **kwargs):
    """Draw a line with countour marks at each every points"""

    ax_diag_size = get_axes_diagonal(ax)
    diag_ratio = np.round(ax_diag_size/150, 1)
    lw = .5*diag_ratio                  # line width
    ms = 2*diag_ratio                   # Marker size

    # Some variable parameters
    every=int(len(x_data)/every)

    line, = ax.plot(x_data, y_data, linestyle)
    line.set_linewidth(lw)
    line.set_color(color)

    mark, = ax.plot(x_data, y_data, marker)
    mark.set_markevery(every)
    mark.set_markersize(ms)
    mark.set_color('white')

    contour, = ax.plot(x_data, y_data, marker)
    contour.set_markevery(every)
    contour.set_markersize(ms)
    contour.set_color(color)
    contour.set_fillstyle('none')
    contour.set_markeredgewidth(lw)

    return line


def prettify_axes(ax, data):
    """Makes my plot pretty"""

    ax_diag_size = get_axes_diagonal(ax)

    # Some variable parameters
    diag_ratio = np.round(ax_diag_size/150, 1)
    lw = .5*diag_ratio                  # line width
    slw = 2*lw                          # spine line width
    ln = 3*diag_ratio                   # Tick length

    # Set spines
    ax.spines['left'].set_linewidth(slw)
    ax.spines['right'].set_linewidth(slw)
    ax.spines['bottom'].set_linewidth(slw)
    ax.spines['top'].set_linewidth(slw)

    if 'x_label' in data:
        ax.set_xlabel(data['x_label'])

    if 'y_label' in data:
        ax.set_ylabel(data['y_label'])

    if 'title' in data:
        ax.set_title(data['title'])

    if 'y_lim' in data:
        ax.set_ylim(data['y_lim'])

    if 'x_lim' in data:
        ax.set_xlim(data['x_lim'])

    # Draw legend only if labels were set (HOW TO DO IT?)
    # ax.get_legend_handles_labels() <-- maybe this can help
    # if ax("has_some_label_set"):
    ax.legend(loc='upper right', prop={'size': 6})

    ax.title.set_fontsize(7)
    ax.xaxis.set_tick_params(labelsize=6)
    ax.xaxis.set_tick_params(width=lw)
    ax.xaxis.set_tick_params(direction='in')
    ax.xaxis.set_tick_params(length=ln)
    ax.xaxis.set_tick_params(zorder=0)
    ax.xaxis.label.set_size(7)

    ax.yaxis.set_tick_params(labelsize=6)
    ax.yaxis.set_tick_params(width=lw)
    ax.yaxis.set_tick_params(direction='in')
    ax.yaxis.set_tick_params(length=ln)
    ax.yaxis.set_tick_params(zorder=0)
    ax.yaxis.label.set_size(7)


def scale_loglog(ax, data):
    """Set a plot to loglog scale"""


def prettify_second_axes(ax):

    ax_diag_size = get_axes_diagonal(ax)

    # Some variable parameters
    diag_ratio = np.round(ax_diag_size/150, 1)
    lw = .5*diag_ratio                  # line width
    slw = 2*lw                          # spine line width
    ln = 3*diag_ratio                   # Tick length

    ax.yaxis.set_tick_params(labelsize=6)
    ax.yaxis.set_tick_params(width=lw)
    ax.yaxis.set_tick_params(direction='in')
    ax.yaxis.set_tick_params(length=ln)
    ax.yaxis.set_tick_params(zorder=0)
    ax.yaxis.set_tick_params(labelcolor='red')
    ax.yaxis.label.set_size(7)


def compare_plot(ax, data):
    """Compare two plots and also gives the difference

    optional:
        show difference or not
        difference axis independent or not
    """
    ax_diag_size = get_axes_diagonal(ax)

    line1 = countour_every(ax, 10, **data[0])
    line2 = countour_every(ax, 10, **data[1])

    if 'label' in data[0]:
        line1.set_label(data[0]['label'])

    if 'label' in data[1]:
        line2.set_label(data[1]['label'])

    ax2 = ax.twinx()
    line3 = ax2.plot(
            data[0]['x_data'], 
            data[0]['y_data']-data[1]['y_data'], '-',
            color='red', alpha=.2, zorder=1, label='Diff')

    prettify_axes(ax, data[0])
    prettify_second_axes(ax2)

def single_plot(ax, data):
    """Plot a line data"""
    if isinstance(data, (list,tuple)):
        data = data[0]

    ax_diag_size = get_axes_diagonal(ax)

    line = countour_every(ax, 10, **data)

    if 'label' in data:
        line.set_label(data['label'])

    prettify_axes(ax, data)


def panelize(data, plot_function, 
        title=None, labelx=None, labely=None,
        sharex=False, sharey=False):
    """Put a group of data into a square panel"""


    if isinstance(data[0], (list, tuple)):
        ndata = len(data)
        nplots = len(data[0])
    else:
        ndata = 1
        nplots = len(data)

    # Wider than taller
    ncols = int(np.ceil(np.sqrt(nplots)))
    nrows = int(np.ceil(nplots/ncols))

    # Use grid instead subplots
    fig, axes = plt.subplots(nrows, ncols, sharex=sharex, sharey=sharey)

    # Calculate the panel distribution and size
    max_width = 6
    width_per_plot = max_width / ncols
    height_ratio = 1 
    height_per_plot = height_ratio*width_per_plot
    max_height = nrows*height_per_plot
    fig.set_size_inches(max_width,max_height)

    for n, (ax, d) in enumerate(zip(np.atleast_1d(axes).flatten(),zip(*data))): 
        plot_function(ax, d)

    if title is not None:
        fig.suptitle(title)
        fig.set_tight_layout({'rect': [0, 0.03, 1, 0.95]})
    else:
        fig.set_tight_layout(True) 

    if sharex and labelx:
        fig.text(0.5, 0.04, labelx, ha='center')

    if sharey and labely:
        fig.text(0.01, 0.5, labely, va='center', rotation='vertical')


    return fig
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.