Index n dimensional array with (n-1) d array

by

Make use of advanced-indexing –

m,n = a.shape[1:]
I,J = np.ogrid[:m,:n]
a_max_values = a[idx, I, J]
b_max_values = b[idx, I, J]

For the general case:

def argmax_to_max(arr, argmax, axis):
    """argmax_to_max(arr, arr.argmax(axis), axis) == arr.max(axis)"""
    new_shape = list(arr.shape)
    del new_shape[axis]

    grid = np.ogrid[tuple(map(slice, new_shape))]
    grid.insert(axis, argmax)

    return arr[tuple(grid)]

Quite a bit more awkward than such a natural operation should be, unfortunately.

For indexing a n dim array with a (n-1) dim array, we could simplify it a bit to give us the grid of indices for all axes, like so –

def all_idx(idx, axis):
    grid = np.ogrid[tuple(map(slice, idx.shape))]
    grid.insert(axis, idx)
    return tuple(grid)

Hence, use it to index into input arrays –

axis = 0
a_max_values = a[all_idx(idx, axis=axis)]
b_max_values = b[all_idx(idx, axis=axis)]

More Related Contents:

  1. Is there a NumPy function to return the first index of something in an array?
  2. Understanding NumPy’s einsum
  3. How to calculate rolling / moving average using python + NumPy / SciPy?
  4. Saving a Numpy array as an image
  5. Python 32-bit memory limits on 64bit windows
  6. AttributeError: ‘Tensor’ object has no attribute ‘numpy’
  7. best way to preserve numpy arrays on disk
  8. Share Large, Read-Only Numpy Array Between Multiprocessing Processes
  9. How do I catch a numpy warning like it’s an exception (not just for testing)?
  10. Compute a confidence interval from sample data
  11. Interweaving two numpy arrays
  12. How can I attach a pyplot function to a figure instance?
  13. How does __contains__ work for ndarrays?
  14. numpy max vs amax vs maximum
  15. Replace negative values in an numpy array
  16. Python RuntimeWarning: overflow encountered in long scalars
  17. Matplotlib figure to image as a numpy array
  18. Selecting Random Windows from Multidimensional Numpy Array Rows
  19. Find elements of array one nearest to elements of array two
  20. How to copy data from a numpy array to another
  21. find and delete from more-dimensional numpy array
  22. Initialise numpy array of unknown length
  23. Saving Matplotlib graphs to image as full screen
  24. Read a base 64 encoded image from memory using OpenCv python library
  25. Pandas groupby with categories with redundant nan
  26. Easy way to test if each element in an numpy array lies between two values?
  27. Python numpy: cannot convert datetime64[ns] to datetime64[D] (to use with Numba)
  28. ImportError: No module named numpy on spark workers
  29. numpy arange: how to make “precise” array of floats?
  30. Sort array’s rows by another array in Python

Leave a Comment