A collection of helpers for plotting using sympy

sympyplot_tools.py

from typing import Iterator, Tuple, Dict, Any, Iterable, Union, cast
from sympy import plot as splot
from sympy.plotting.plot import Plot
from matplotlib import pyplot
from itertools import filterfalse
from more_itertools import unzip
import random
import colorsys

Curve = Tuple[Any, Dict[str, Any]]
Range = Tuple[Any, None]
CurveOrRange = Union[Curve, Range]

def curves_iter(curves: Iterable[CurveOrRange]
) -> Iterator[Curve]:
    return cast(Iterator[Curve], 
        filterfalse(lambda expr: expr[1] is None, curves))

def splot_multiple(
    *exprs: CurveOrRange,
    plotf: splot, random_colors=False,
    **options):
    "plots multiple curves with extra kwargs for each one"
    show = options.pop("show", True)
    curves, curves_args = unzip(curves_iter(exprs))
    plot = plotf(*curves, show=False, **options)
    for i, curve_args in enumerate(curves_args):
        if random_colors:
            plot[i].line_color = [random_bright_rgb_color()]
        for key, value in curve_args.items():
            setattr(plot[i], key, value)
    plot.show() if show else None
    return plot

# Stolen from https://stackoverflow.com/a/43437435/5538719
def random_bright_rgb_color():
    "Picks a random \"bright\" RGB color"
    h = random.random()
    s = 0.5 + random.random()/2.0
    l = 0.4 + random.random()/5.0
    return colorsys.hls_to_rgb(h, l, s)

# Stolen from https://stackoverflow.com/a/60325901/5538719
def move_sympyplot_to_axes(p: Plot, ax) -> None:
    backend = p.backend(p)
    backend.ax = ax
    # Fix for > sympy v1.5
    backend._process_series(backend.parent._series, ax, backend.parent)
    backend.ax.spines['right'].set_color('none')
    backend.ax.spines['bottom'].set_position('zero')
    backend.ax.spines['top'].set_color('none')
    pyplot.close(backend.fig)