def downsample_grid(a, b, samples, ret_idx=False):
    """Content-based downsampling for faster visualization

    The arrays `a` and `b` make up a 2D scatter plot with high
    and low density values. This method takes out points at
    indices with high density.

    Parameters
    ----------
    a, b: 1d ndarrays
        The input arrays to downsample
    samples: int
        The desired number of samples
    remove_invalid: bool
        Remove nan and inf values before downsampling
    ret_idx: bool
        Also return a boolean array that corresponds to the
        downsampled indices in `a` and `b`.

    Returns
    -------
    dsa, dsb: 1d ndarrays of shape (samples,)
        The arrays `a` and `b` downsampled by evenly selecting
        points and pseudo-randomly adding or removing points
        to match `samples`.
    idx: 1d boolean array with same shape as `a`
        Only returned if `ret_idx` is True.
        A boolean array such that `a[idx] == dsa`
    """
    # fixed random state for this method
    rs = np.random.RandomState(seed=47).get_state()

    samples = int(samples)

    if samples and samples < a.size:
        # The events to keep
        keep = np.zeros_like(a, dtype=bool)

        # 1. Produce evenly distributed samples
        # Choosing grid-size:
        # - large numbers tend to show actual structures of the sample,
        #   which is not desired for plotting
        # - small numbers tend will not result in too few samples and,
        #   in order to reach the desired samples, the data must be
        #   upsampled again.
        # 300 is about the size of the plot in marker sizes and yields
        # good results.
        grid_size = 300
        xpx = norm(a, a, b) * grid_size
        ypx = norm(b, b, a) * grid_size
        # The events on the grid to process
        toproc = np.ones((grid_size, grid_size), dtype=bool)

        for ii in range(xpx.size):
            xi = xpx[ii]
            yi = ypx[ii]
            # filter for overlapping events
            if valid(xi, yi) and toproc[int(xi-1), int(yi-1)]:
                toproc[int(xi-1), int(yi-1)] = False
                # include event
                keep[ii] = True

        # 2. Make sure that we reach `samples` by adding or
        # removing events.
        diff = np.sum(keep) - samples
        if diff > 0:
            # Too many samples
            rem_indices = np.where(keep)[0]
            np.random.set_state(rs)
            rem = np.random.choice(rem_indices,
                                   size=diff,
                                   replace=False)
            keep[rem] = False
        elif diff < 0:
            # Not enough samples
            add_indices = np.where(~keep)[0]
            np.random.set_state(rs)
            add = np.random.choice(add_indices,
                                   size=abs(diff),
                                   replace=False)
            keep[add] = True

        assert np.sum(keep) == samples, "sanity check"
        asd = a[keep]
        bsd = b[keep]
        assert np.allclose(a[keep], asd, equal_nan=True), "sanity check"
        assert np.allclose(b[keep], bsd, equal_nan=True), "sanity check"
    else:
        keep = np.ones_like(a, dtype=bool)
        asd = a
        bsd = b

    if ret_idx:
        return asd, bsd, keep
    else:
        return asd, bsd
Content-based downsampling for faster visualization

    The arrays `a` and `b` make up a 2D scatter plot with high
    and low density values. This method takes out points at
    indices with high density.

    Parameters
    ----------
    a, b: 1d ndarrays
        The input arrays to downsample
    samples: int
        The desired number of samples
    remove_invalid: bool
        Remove nan and inf values before downsampling
    ret_idx: bool
        Also return a boolean array that corresponds to the
        downsampled indices in `a` and `b`.

    Returns
    -------
    dsa, dsb: 1d ndarrays of shape (samples,)
        The arrays `a` and `b` downsampled by evenly selecting
        points and pseudo-randomly adding or removing points
        to match `samples`.
    idx: 1d boolean array with same shape as `a`
        Only returned if `ret_idx` is True.
        A boolean array such that `a[idx] == dsa`