def index_map(i, d):
    """
    The index map used mapping the from 2D index to 1D index

    Parameters
    ----------
    i : int, np.array
        the index in the 2D case.
    d : int
        the dimension to use for the 2D index.

    Returns
    -------
    int, np.array
        1D index.

    """
    return 2 * i + d