CODEWITHSUNDEEP
pythonnumpyarray-broadcasting
mzzx
mzzx

Reputation: 2164

numpy - vectorize functions: apply_over_axes / apply_along_axis

I want to calculate the determinant of mm subarrays of a mm*n dimensional arrays, and would like to do this in a fast/more elegant way. The brute-force approach works:

import numpy as n

array=n.array([[[0.,1.,2.,3.],[2,1,1,0]],[[0.5, 0.5,2,2],[0.5,1,0,2]]])
detarray=n.zeros(4)
for i in range(4):
    detarray[i]= n.linalg.det(array[:,:,i])

I would have tried doing this with apply_along_axis, but I know this is only for 1D arguments to the function, and so I presume I can't get this to work.

However, I thought apply_over_axes should also work:

n.apply_over_axes(n.linalg.det, array, [0,1])

but this gives me an error: "det() takes exactly 1 argument (2 given)"

Does anyone know why this doesn't work? If this type of calculation is really not possible with apply_over_axes, is there a better way to do this rather than the for-loop?

Upvotes: 4

Views: 1917

Answers (2)

kmario23
kmario23

Reputation: 61365

Utilizing the transpose semantics of NumPy for 3D arrays, you can simply pass the transposed array to numpy.linalg.det() as in:

In [13]: arr = np.array([[[0.,1.,2.,3.], 
                          [2, 1, 1, 0]], 

                         [[0.5, 0.5,2,2],
                          [0.5, 1, 0, 2]]])

In [14]: np.linalg.det(arr.T)
Out[14]: array([-1. ,  0.5, -2. ,  6. ])

Performance wise this approach seems to be twice as fast as the other approach of manually moving the axes using numpy.moveaxis

In [29]: %timeit np.linalg.det(np.moveaxis(arr, 2, 0))
12.9 µs ± 192 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [30]: %timeit np.linalg.det(arr.T)
6.2 µs ± 136 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Upvotes: 3

Paul Panzer
Paul Panzer

Reputation: 53029

numpy.linalg.det is vectorized out of the box. All you need to do is move the outer axes to the left:

>>> np.linalg.det(np.moveaxis(array, 2, 0))
array([-1. ,  0.5, -2. ,  6. ])

Upvotes: 1

Related Questions