3d numpy array perform operation along column i only

I have a 3d numpy array like this (as an example)

a = np.array([
        [[1, -2, 3, 4],[5, 6, 7, 8]],
        [[9, 10, 11, 12],[13, 14, 15, 16]]
])

I want to apply the following operations only to elements within the column with index 1 in the inner dimension. The elements are [-2,6,10,14] for the example above. The operations would be:

# variables used in the operations
v1, v2, v3 = 12, 4, 2

# the following two operations should only be applied to specified column across all the array

# 1st operation
a[a >= v1] = v1

# output
a = np.array([
        [[1, -2, 3, 4],[5, 6, 7, 8]],
        [[9, 10, 11, 12],[13, 12, 15, 16]]
])

# 2nd operation
f = lambda x: -2 if(x==-2) else (x-v3)/(v2-v3)
a = f(a)

# output

a = np.array([
        [[1, -2, 3, 4],[5, 2, 7, 8]],
        [[9, 4, 11, 12],[13, 5, 15, 16]]
])

Can someone help me? I have looked into several NumPy methods but can’t seem to adapt to my example.

Answer

You need to change you function to be vectorial (i.e accept an array and input and return an array as output), and slice to only apply it on the desired “column”:

f = lambda x: np.where(x==-2, -2, (x-v3)/(v2-v3))

a[...,[1]] = f(a[...,[1]])

output:

array([[[ 1, -2,  3,  4],
        [ 5,  2,  7,  8]],

       [[ 9,  4, 11, 12],
        [13,  5, 15, 16]]])