Utility functions to a) create a zoom of an existing line plot and b) copy lines from a source axis to a destination axis (used in (a) )

zoom_ax.py

import numpy as np
from matplotlib.axes import Axes, mlines

def create_zoom(ax_to_zoom, inset_size, lines=None, loc='lower left', loc1=1, loc2=2,
                xlim=None, ylim=None, xticks=2, yticks=2, ec='C7',
                inset_kwargs=None, box_kwargs=None, connector_kwargs=None, **kwargs):
    """Zoom into an axis with an inset plot

    The valid locations (for `loc`, `loc1`, and `loc2`) are: 'upper right' : 1, 'upper left' : 2, 'lower left' : 3,
    'lower right' : 4,
    'right' : 5, 'center left' : 6, 'center right' : 7, 'lower center' : 8, 'upper center' : 9, 'center' : 10

    :param ax_to_zoom: Source axis which data will be copied from and inset axis inserted.
    :type ax_to_zoom: Axes
    :param inset_size: Zoom factor for `zoomed_inset_axes` if argument is a float, else is width and height,
        respectively, for `inset_axes`.
    :type inset_size: float or tuple
    :param lines: Lines to plot in inset axis. If None, all lines are plotted.
    :type lines: List[Line2D]
    :param loc: Location to place the inset axes.
    :param loc1: Corner to use for connecting the inset axes and the area in the
        parent axes.
    :param loc2: Corner to use for connecting the inset axes and the area in the
        parent axes.
    :param xlim: Limits of x-axis. Also limits data **plotted** when copying to the inset axis.
    :type xlim: float or Tuple[float, float]
    :param ylim: Limits of y-axis.
    :type ylim: float or Tuple[float, float]
    :param xticks: Number of ticks (int) or location of ticks (list) or no x-axis (False).
    :type xticks: int or list
    :param yticks: Number of ticks (int) or location of ticks (list) or no y-axis (False).
    :type yticks: int or list
    :param ec: Edge color for border of zoom and connecting lines
    :param inset_kwargs: Keywords for `inset_axes` or `zoomed_inset_axes`.
            E.g. dict(bbox_to_anchor=(1,1), bbox_transform=ax_to_zoom.transAxes)
    :type inset_kwargs: dict
    :param box_kwargs: Keywords for `mpl_toolkits.axes_grid1.inset_locator.mark_inset`.
        To remove not mark the inset axis at all, set box_kwargs to `'None'`
    :type box_kwargs: dict or str
    :param connector_kwargs: Keywords for connecting lines between inset axis and source axis.
        To remove connecting lines set to '`None'`.
    :type connector_kwargs: dict or str
    :param kwargs: Additional keyword arguments for plotting. See `Axes` keywords.

    :return: inset axis
    :rtype: Axes
    """
    from mpl_toolkits.axes_grid1.inset_locator import inset_axes, zoomed_inset_axes, mark_inset
    if inset_kwargs is None:
        inset_kwargs = dict(bbox_to_anchor=None, bbox_transform=None)
    elif 'bbox_to_anchor' in inset_kwargs:
        if inset_kwargs['bbox_to_anchor'] is None:
            inset_kwargs['bbox_transform'] = None
        elif 'bbox_transform' not in inset_kwargs or\
                ('bbox_transform' in inset_kwargs and inset_kwargs['bbox_transform'] is None):
            inset_kwargs['bbox_transform'] = ax_to_zoom.transAxes
    if box_kwargs is None:
        box_kwargs = dict()
    if connector_kwargs is None:
        connector_kwargs = dict()

    axes_kwargs = dict(facecolor="white")

    if type(inset_size) is tuple:
        ax_inset: Axes = inset_axes(ax_to_zoom, width=inset_size[0], height=inset_size[1],
                                    loc=loc, axes_kwargs=axes_kwargs, **inset_kwargs)
    else:
        ax_inset: Axes = zoomed_inset_axes(ax_to_zoom, zoom=inset_size,
                                           loc=loc, axes_kwargs=axes_kwargs, **inset_kwargs)

    copy_lines(lines or ax_to_zoom, ax_inset, xlim=xlim, **kwargs)

    ax_inset.set_xlim(xlim)
    ax_inset.set_ylim(ylim)
    for _ticks, _ticks_axis in zip([xticks, yticks], ['x', 'y']):
        get_axis = ax_inset.get_xaxis if _ticks_axis == 'x' else ax_inset.get_yaxis
        if _ticks:
            if type(_ticks) is int:
                ax_inset.locator_params(nbins=_ticks, axis=_ticks_axis, tight=True)
            else:
                get_axis().set_ticks(_ticks)
        else:
            get_axis().set_visible(False)

    for spine in ax_inset.spines:
        ax_inset.spines[spine].set_visible(True)
        ax_inset.spines[spine].set_color(ec)
    if box_kwargs != 'None':
        box_patch, p1, p2 = mark_inset(ax_to_zoom, ax_inset, loc1=loc1, loc2=loc2, ec=ec,
                                       **box_kwargs)
        for loc, spine in ax_inset.spines.items():
            spine.set(**box_kwargs)  # consistency between inset border and marked box
        if type(connector_kwargs) is dict:
            p1.set(**connector_kwargs)
            p2.set(**connector_kwargs)
        elif connector_kwargs == 'None':
            p1.set(color='None')
            p2.set(color='None')
    return ax_inset


# noinspection SpellCheckingInspection
line_props = {
    'agg_filter':         'get_agg_filter',
    'alpha':              'get_alpha',
    'antialiased':        'get_antialiased',
    'color':              'get_color',
    'dash_capstyle':      'get_dash_capstyle',
    'dash_joinstyle':     'get_dash_joinstyle',
    'drawstyle':          'get_drawstyle',
    'fillstyle':          'get_fillstyle',
    'label':              'get_label',
    'linestyle':          'get_linestyle',
    'linewidth':          'get_linewidth',
    'marker':             'get_marker',
    'markeredgecolor':    'get_markeredgecolor',
    'markeredgewidth':    'get_markeredgewidth',
    'markerfacecolor':    'get_markerfacecolor',
    'markerfacecoloralt': 'get_markerfacecoloralt',
    'markersize':         'get_markersize',
    'markevery':          'get_markevery',
    'rasterized':         'get_rasterized',
    'solid_capstyle':     'get_solid_capstyle',
    'solid_joinstyle':    'get_solid_joinstyle',
    'zorder':             'get_zorder',
    }


# noinspection PyProtectedMember
def copy_lines(src, ax_dest, xlim=None, xunit=None, rel_lw=2., cmap=None, **kwargs):
    """ Copies the lines from a source Axes `ax_src` to a destination Axes `ax_dest`.
    By default, linewidth is doubled. To disable, set rel_lw to 1 or provide 'lw' keyword.

    :param src: Source Axes or list of Line2D objects.
    :type src: Axes or list
    :param ax_dest: Destination Axes.
    :type ax_dest: Axes
    :param lines: Provide the list of Line2D objects to plot. If None, all lines from ax_src will be re-plotted.
    :type lines: List[Axes]
    :param xlim: Set the domain of the axis. This will restrict what is plotted (hence should be accompanied by
        xunit), not just a call to `set_xlim`
    :type xlim: tuple or list or int
    :param xunit: The units used with xlim to restrict the domain.
    :type xunit: Quantity
    :param rel_lw: Relative linewidth change for each line in source. Default doubles the linewidth.
    :type rel_lw: float
    :param kwargs: Keywords to provide to `Axes.plot` that will overwrite source properties. The aliases to use are
        defined in `line_props`.
    """
    from matplotlib import cbook, colors
    kwargs = cbook.normalize_kwargs(kwargs, mlines.Line2D._alias_map)
    if isinstance(src, Axes):
        lines = src.get_lines()
    elif np.iterable(src):
        lines = src
    else:
        lines = [src]
    if xlim is not None and xunit is not None:
        if type(xlim) is tuple or type(xlim) is list:
            xlim = [int(xlim[0]/xunit), int(xlim[1]/xunit)]
        else:
            xlim = [int(xlim/xunit), -1]
    if cmap is not None and type(cmap) is str:
        cmap = plt.get_cmap(cmap)
    for i, line in enumerate(lines):
        props = {}
        for prop, getter in line_props.items():
            props[prop] = getattr(line, getter)()
        if cmap is not None:
            props['color'] = colors.to_hex(cmap(i), keep_alpha=False)
        props = {**props, **kwargs}
        if 'linewidth' not in kwargs:
            # change relative thickness of line
            props['linewidth'] *= rel_lw
        x, y = line.get_data()
        if xlim is not None and xunit is not None:
            x = x[xlim[0]:xlim[1]]
            y = y[xlim[0]:xlim[1]]
        ax_dest.plot(x, y, **props)