numpynumpy-indexing

NumPy : How to determine the index of the first axis of an ndarray according to some condition?


Consider the following ndarray a -

In [117]: a                                                                                          
Out[117]: 
array([[[nan, nan],
        [nan, nan],
        [nan, nan]],

       [[ 3., 11.],
        [ 7., 13.],
        [12., 16.]],

       [[ 0.,  4.],
        [ 6.,  1.],
        [ 5.,  8.]],

       [[17., 10.],
        [15.,  9.],
        [ 2., 14.]]])

The minimum computed on the first axis is -

In [118]: np.nanmin(a, 0)                                                                            
Out[118]: 
array([[0., 4.],
       [6., 1.],
       [2., 8.]])

which is a[2] from visual inspection. What is the most efficient way to calculate this index 2


Solution

  • as suggested by @Divakar you can use np.nanargmin

    import numpy as np
    
    a = np.array([[[np.nan, np.nan],
            [np.nan, np.nan],
            [np.nan, np.nan]],
    
           [[ 3., 11.],
            [ 7., 13.],
            [12., 16.]],
    
           [[ 0.,  4.],
            [ 6.,  1.],
            [ 5.,  8.]],
    
           [[17., 10.],
            [15.,  9.],
            [ 2., 14.]]])
    minIdx = np.nanargmin(np.sum(a,(1,2)))
    minIdx
    2
    a[minIdx]
    array([[0., 4.],
           [6., 1.],
           [5., 8.]])