secret
Created

cython vs. numpy performance

  • Download Gist
rotT_cython.pyx
Cython
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23
#cython: boundscheck=False, wraparound=False
import numpy as np
cimport numpy as np
 
def rotT(np.ndarray[np.float64_t, ndim=4] T,
np.ndarray[np.float64_t, ndim=2] g):
cdef np.ndarray[np.float64_t, ndim=4] Tprime
cdef Py_ssize_t i, j, k, l, ii, jj, kk, ll
cdef np.float64_t gg
 
Tprime = np.zeros((3,3,3,3), dtype=T.dtype)
for i in range(3):
for j in range(3):
for k in range(3):
for l in range(3):
for ii in range(3):
for jj in range(3):
for kk in range(3):
for ll in range(3):
gg = g[ii,i]*g[jj,j]*g[kk,k]*g[ll,l]
Tprime[i,j,k,l] = Tprime[i,j,k,l] + \
gg*T[ii,jj,kk,ll]
return Tprime
test_rotT.py
Python
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49
#!/usr/bin/env python
# http://stackoverflow.com/questions/4962606/fast-tensor-rotation-with-numpy
import pyximport; pyximport.install() # pip install cython
from rotT_cython import rotT
 
import numpy as np
 
T = np.array([[[[ 4.66533067e+01, 5.84985000e-02, -5.37671310e-01],
[ 5.84985000e-02, 1.56722231e+01, 2.32831900e-02],
[ -5.37671310e-01, 2.32831900e-02, 1.33399259e+01]],
[[ 4.60051700e-02, 1.54658176e+01, 2.19568200e-02],
[ 1.54658176e+01, -5.18223500e-02, -1.52814920e-01],
[ 2.19568200e-02, -1.52814920e-01, -2.43874100e-02]],
[[ -5.35577630e-01, 1.95558600e-02, 1.31108757e+01],
[ 1.95558600e-02, -1.51342210e-01, -6.67615000e-03],
[ 1.31108757e+01, -6.67615000e-03, 6.90486240e-01]]],
[[[ 4.60051700e-02, 1.54658176e+01, 2.19568200e-02],
[ 1.54658176e+01, -5.18223500e-02, -1.52814920e-01],
[ 2.19568200e-02, -1.52814920e-01, -2.43874100e-02]],
[[ 1.57414726e+01, -3.86167500e-02, -1.55971950e-01],
[ -3.86167500e-02, 4.65601977e+01, -3.57741000e-02],
[ -1.55971950e-01, -3.57741000e-02, 1.34215636e+01]],
[[ 2.58256300e-02, -1.49072770e-01, -7.38843000e-03],
[ -1.49072770e-01, -3.63410500e-02, 1.32039847e+01],
[ -7.38843000e-03, 1.32039847e+01, 1.38172700e-02]]],
[[[ -5.35577630e-01, 1.95558600e-02, 1.31108757e+01],
[ 1.95558600e-02, -1.51342210e-01, -6.67615000e-03],
[ 1.31108757e+01, -6.67615000e-03, 6.90486240e-01]],
[[ 2.58256300e-02, -1.49072770e-01, -7.38843000e-03],
[ -1.49072770e-01, -3.63410500e-02, 1.32039847e+01],
[ -7.38843000e-03, 1.32039847e+01, 1.38172700e-02]],
[[ 1.33639532e+01, -1.26331100e-02, 6.84650400e-01],
[ -1.26331100e-02, 1.34222177e+01, 1.67851800e-02],
[ 6.84650400e-01, 1.67851800e-02, 4.89151396e+01]]]])
 
g = np.array([[ 0.79389393, 0.54184237, 0.27593346],
[-0.59925749, 0.62028664, 0.50609776],
[ 0.10306737, -0.56714313, 0.8171449 ]])
 
def rotT_philipp(T, g):
gg = np.outer(g, g)
gggg = np.outer(gg, gg).reshape(4 * g.shape)
axes = ((0, 2, 4, 6), (0, 1, 2, 3))
return np.tensordot(gggg, T, axes)
 
if __name__ == '__main__':
assert np.allclose(rotT_philipp(T, g), rotT(T, g))
# | rotT | 10000 loops, best of 3: 28 usec per loop |
# | rotT_philipp | 10000 loops, best of 3: 113 usec per loop |

Please sign in to comment on this gist.

Something went wrong with that request. Please try again.