def _get_fig_and_ax(fig):
    """Get figure and current ax. Input can be either a figure or an ax."""
    import matplotlib.pyplot as plt
    if isinstance(fig, plt.Axes):
        ax = fig
        fig = ax.figure
    elif isinstance(fig, plt.Figure):
        ax = fig.gca()
    else:
        raise ValueError('fig should be a matplotlib Figure or Axes instance.')
    return fig, ax
