Tag: Modules

A useful Python plotting module

A useful Python plotting module

First of all, let’s talk a little bit about online Python documentation. At the top of my list is the official Python Documentation. It is extremely well written, with plenty of examples. I always start from the Tutorial page and dig from there. For the specific “how-to-do-that” questions, do a Google search and then ONLY choose answers from stackoverflow or stackexchange. Any other sites will probably flash more ads in front of your eyes than you can handle. Not to mention, their answers are likely copied from the two sites I mentioned.

Some of my examples throughout this blog use a plotting function that I made using Plotly. As far as graphing packages in Python are concerned, Plotly and Matplotlib are my favorites. However, I chose the former for this because its graphs are more interactive, which can be helpful when looking at data. In my opinion, you will get the most out of the function if you’re using VSCode and have the Jupyter notebook extensions installed. The Plotly graphs will show right there in the interactive window.

The plotting function should be inside a Python module, which I named utils.py. It’s a starting place for someone to add other functions or classes for tasks that are done routinely in programs. You can copy the code below and save it in a utils.py module and import it as needed when running some of the examples. Check out how to use it here.

""" utils.py

Contains a useful plotting function that is used in the coding examples.
The function was built using Plotly instead of Matplotlib due to its
interactive graphs and because it runs better on Raspberry Pi Linux.

Author: Eduardo Nigro
    rev 0.0.6
    2022-01-24
"""
import numpy as np
import plotly.io as pio
import plotly.graph_objects as go
from plotly.subplots import make_subplots

# Setting plotting modified template as default
mytemplate = pio.templates["plotly_white"]
mytemplate.layout["paper_bgcolor"] = "rgb(250, 250, 250)"
pio.templates.default = mytemplate


def plot_line(
    x, y, xname="Time (s)", yname=None, axes="single",
    figsize=None, line=True, marker=False, legend=None):
    """
    Plot lines using plotly.

    :param x: x values for plotting.
        List or ndarray. List of lists or ndarrays is also supported.
    :type x: list(float), list(list), list(ndarray)

    :param y: y values for plotting.
        List or ndarray. List of lists or ndarrays is also supported.
    :type y: list(float), list(list), list(ndarray)

    :param xname: The x axis title. Default value is ``'Time (s)'``.
        If ``'Angle (deg.)'`` is used, axes ticks are configured in 360 degree
        increments.
    :type xname: str

    :param yname: The y axis title.
        A string or list of strings containing the names of the y axis titles.
        If ``None``, the y axis titles will be ``'y0'``, ``'y1'``, etc.
    :type yname: str, list(str)

    :param axes: The configuration of axis on the plot.
        If ``'single'``, multiple curves are plotted on the same axis.
        If ``'multi'``, each curve is plotted on its own axis.
    :type axes: str

    :param figsize: The figure size (``width``, ``height``) in pixels.
    :type figsize: tuple(int)

    :param line: Displays the curves if ``True``.
    :type line: bool, list(bool)

    :param marker: Displays markers on the curves if ``True``.
    :type marker: bool, list(bool)

    :param legend: List of legend names for multiple curves.
        Length of `legend` must be the same as length of `y`.
    :type legend: list(str)

    Example:
        >>> import numpy as np
        >>> from utils import plot_line
        >>> t = np.linspace(0,2,100)
        >>> y0 = np.sin(1*np.pi*t)
        >>> y1 = np.cos(1*np.pi*t)
        >>> plot_line(
                [t]*2, [y0, y1],
                yname=['sin & cos'],
                legend=['sin(pi x t)', 'cos(pi x t)']
                )

    """
    # Making sure x and y inputs are put in lists if needed
    if type(x) != list:
        x = [x]
    else:
        if type(x[0]) not in [list, np.ndarray]:
            x = [x]
    if type(y) != list:
        y = [y]
    else:
        if type(y[0]) not in [list, np.ndarray]:
            y = [y]
    # Doing a simple check for consistent x and y inputs
    if len(x) != len(y):
        raise Exception("'x' and 'y' inputs must have the same length.")
    # Adjusting y axis title based on input
    if not yname:
        yname = ["y" + str(i) for i in range(len(y))]
    elif type(yname) != list:
        yname = [yname]
    if (len(yname) == 1) and (len(y) > 1):
        yname = yname * len(y)
    # Setting legend display option
    if legend is not None:
        if len(legend) == len(y):
            showlegend = True
        else:
            raise Exception("'y' and 'legend' must have the same length.")
    else:
        showlegend = False
        legend = [None] * len(y)
    # Checking for single (with multiple curves)
    # or multiple axes with one curve per axes
    if axes == "single":
        naxes = 1
        iaxes = [0] * len(y)
        colors = [
            "#1F77B4",
            "#FF7F0E",
            "#2CA02C",
            "#D62728",
            "#9467BD",
            "#8C564B",
            "#E377C2",
            "#7F7F7F",
            "#BCBD22",
            "#17BECF",
        ]
    elif axes == "multi":
        naxes = len(y)
        iaxes = range(0, len(y))
        colors = ["rgb(50, 100, 150)"] * len(y)
    else:
        raise Exception("Valid axes options are: 'single' or 'multi'.")
    # Checking for line and marker options
    if type(line) != list:
        line = [line] * len(y)
    if type(marker) != list:
        marker = [marker] * len(y)
    mode = []
    markersize = []
    for linei, markeri in zip(line, marker):
        if linei and markeri:
            mode.append("lines+markers")
            markersize.append(2)
        elif linei and not markeri:
            mode.append("lines")
            markersize.append(2)
        elif not linei and markeri:
            mode.append("markers")
            markersize.append(8)
    # Setting figure parameters
    if figsize:
        wfig, hfig = figsize
    else:
        wfig = 650
        hfig = 100 + 150*naxes
    m0 = 10
    margin = dict(l=6 * m0, r=3 * m0, t=3 * m0, b=3 * m0)
    # Plotting results
    fig = make_subplots(rows=naxes, cols=1)
    for i, xi, yi, ynamei, legendi, colori, modei, markersizei in zip(
            iaxes, x, y, yname, legend, colors, mode, markersize):
        # Adding x, y traces to appropriate plot
        fig.add_trace(
            go.Scatter(
                x=xi,
                y=yi,
                name=legendi,
                mode=modei,
                line=dict(width=1, color=colori),
                marker=dict(size=markersizei, color=colori),
            ),
            row=i + 1,
            col=1,
        )
        # Adding x axes ticks
        if xname.lower().find("angle") < 0:
            # Regular x axes
            fig.update_xaxes(matches="x", row=i + 1, col=1)
        else:
            # Special case where x axes has angular values
            fig.update_xaxes(
                tickmode="array",
                tickvals=np.arange(0, np.round(xi[-1]) + 360, 360),
                matches="x",
                row=i + 1,
                col=1,
            )
        # Adding y axis title to all plots
        fig.update_yaxes(title_text=ynamei, row=i + 1, col=1)
    # Adding x axis title to bottom plot only
    fig.update_xaxes(title_text=xname, row=i + 1, col=1)
    # Applying figure size, margins, and legend
    fig.update_layout(
        margin=margin, width=wfig, height=hfig, showlegend=showlegend)
    fig.show()