def update_plot_limits(ax, white_space):
    """Sets the limit options of a matplotlib plot.

    Args:
        ax: matplotlib axes
        white_space(float): whitespace added to surround the tight limit of the data

    Note: This relies on ax.dataLim (in 2d) and ax.[xy, zz]_dataLim being set in 3d
    """

    if hasattr(ax, 'zz_dataLim'):
        bounds = ax.xy_dataLim.bounds
        ax.set_xlim(bounds[0] - white_space, bounds[0] + bounds[2] + white_space)
        ax.set_ylim(bounds[1] - white_space, bounds[1] + bounds[3] + white_space)

        bounds = ax.zz_dataLim.bounds
        ax.set_zlim(bounds[0] - white_space, bounds[0] + bounds[2] + white_space)
    else:
        bounds = ax.dataLim.bounds
        assert not any(map(np.isinf, bounds)), 'Cannot set bounds if dataLim has infinite elements'
        ax.set_xlim(bounds[0] - white_space, bounds[0] + bounds[2] + white_space)
        ax.set_ylim(bounds[1] - white_space, bounds[1] + bounds[3] + white_space)
Sets the limit options of a matplotlib plot.

    Args:
        ax: matplotlib axes
        white_space(float): whitespace added to surround the tight limit of the data

    Note: This relies on ax.dataLim (in 2d) and ax.[xy, zz]_dataLim being set in 3d