Beschleunigung der Rotationsmatrixberechnung mit numba

Question

Ich versuche, eine Reihe von Rotationsmatrixberechnungen zu beschleunigen, die zu einer 3D-Matrix führen (Dimensionen = 3x3xnumv, wobei numv die Anzahl der Eckpunkte ist). Bisher führt meine jit-Funktion zu einer deutlich langsameren Berechnung.

from numpy import sin, cos, ones, sqrt, array, float64, zeros, isnan, shape 
from numpy.linalg import norm 
from numba import jit 
from numba import float64 as _float64 

def calculate_rot_matrix(rot_edges, kb, k): 
''' 
Calculates rotation matrices for set of input 2 edges 
Returns rot matrix with shape (3, 3, max_edges) 
edges are different for vertices vs. edges (but only vertices are kept) 
''' 
    b   = kb/k # global kb 
    b[isnan(b)] = 0.0 
    sin_theta = norm(rot_edges, axis=1).reshape(-1, 1) * k/2.0 
    cos_theta = sqrt(ones(shape(sin_theta)) - sin_theta ** 2.0) 
    n1, n2, n3 = b[:, 0], b[:, 1], b[:, 2] 
    s, c  = sin_theta.reshape(-1), cos_theta.reshape(-1) 
    # get rotation matrices 
    R = array([[c + n1**(2.0) * (1.0 - c), n1*n2*(1.0 - c) - s*n3, n3*n1 * (1.0 - c) + s*n2], 
      [n1*n2*(1.0 - c) + s*n3, c + n2**(2.0) * (1.0 - c), n3*n2 * (1.0 - c) - s*n1], 
      [n1*n3*(1.0 - c) - s*n2, n2*n3*(1.0 - c) + s*n1, c + n3**(2.0) * (1.0 - c)]]) 
    # fix empty rotations 
    R[isnan(R)] = 0.0 
    return R 

@jit((_float64[:,:], _float64[:,:], _float64[:])) 
def jit_calculate_rot_matrix(rot_edges, kb, k): 
''' 
Calculates rotation matrices for set of input 2 edges 
Returns rot matrix with shape (3, 3, max_edges) 
edges are different for vertices vs. edges (but only vertices are kept) 
''' 
    b   = kb/k # global kb 
    b[isnan(b)] = 0.0 
    sin_theta = norm(rot_edges, axis=1).reshape(-1, 1) * k/2.0 
    cos_theta = sqrt(ones(shape(sin_theta)) - sin_theta ** 2.0) 
    n1, n2, n3 = b[:, 0], b[:, 1], b[:, 2] 
    s, c  = sin_theta.reshape(-1), cos_theta.reshape(-1) 
    # get rotation matrices 
    R = array([[c + n1**(2.0) * (1.0 - c), n1*n2*(1.0 - c) - s*n3, n3*n1 * (1.0 - c) + s*n2], 
      [n1*n2*(1.0 - c) + s*n3, c + n2**(2.0) * (1.0 - c), n3*n2 * (1.0 - c) - s*n1], 
      [n1*n3*(1.0 - c) - s*n2, n2*n3*(1.0 - c) + s*n1, c + n3**(2.0) * (1.0 - c)]]) 
    # fix empty rotations 
    R[isnan(R)] = 0.0 
    return R 

if __name__ == '__main__': 
    import cProfile 
    import pstats 
    import cStringIO 
    import traceback 

    numv = 100 
    rot_edges = zeros((numv, 3), dtype=float64) 
    rot_edges[:, 1] = 1.0 
    kb = zeros((numv, 3), dtype=float64) 
    # k = norm(kb, axis=1).reshape(-1, 1) 
    k = ones((numv, 1), dtype=float64) 

    profile = cProfile.Profile() 
    profile.enable() 
    # ======================================================================= 
    # profile enabled 
    # ======================================================================= 
    for i in range(10000): 
     R = calculate_rot_matrix(rot_edges, kb, k) 
    for i in range(10000): 
     R_jit = jit_calculate_rot_matrix(rot_edges, kb, k) 
    # ======================================================================= 
    # profile disabled 
    # ======================================================================= 
    profile.disable() 
    stream = cStringIO.StringIO() 
    sortby = 'cumulative' 
    ps = pstats.Stats(profile, stream=stream).sort_stats(sortby) 
    ps.strip_dirs() 
    ps.sort_stats(1) 
    ps.print_stats(20) 
    print stream.getvalue() 

auf der Dokumentations-Based, ich denke, die Geschwindigkeit gewinnt ich vom Laufen der jitted Funktion mit nopython = True als Parameter wäre zu bekommen. Während einige Operationen auf Arrays funktionieren (sin, cos), möchte ich jedoch wissen, ob es irgendeine "Norm" -Funktion gibt (die auf einer numv × 3-Matrix von Vektoren arbeitet, die einen numv × 1-Vektor erzeugt). Ich rufe auch mehrmals um, um in die richtige Form übertragen zu können, und ich denke, da dies eine "Python" -Funktion ist, kann es nicht in jit nopython übersetzt werden.

Answer 1

1. Umformung ist keine teure Operation, da normalerweise nur Schritte manipuliert werden;

2. „Ich möchte wissen, ob es eine ist‚Norm‘Funktionstyp (auf einer numv x 3 Matrix von Vektoren arbeitet, ein numv x 1 Vektor-Herstellung)“ Ich denke numpy.linalg.norm() schon tut, was Sie wollen - nur verwenden seine axis Parameter:

   np.linalg.norm(some_array, axis=0) 
   

3. Die meisten Ihrer Operationen bereits vektorisiert sind und wahrscheinlich intern (auf numpy) in C geschrieben, und ich sehe nicht, wie viel Sie davon ab, diesen Code mit numba beschleunigt gewinnen würde.