updated njit in vec_utils

Author: alexlib

openptv_python/track.py

@@ -289,6 +289,7 @@ def angle_acc(
     return float(angle), float(acc)
 
 
+
 def candsearch_in_pix(
     next_frame: List[Target],
     num_targets: int,

openptv_python/vec_utils.py

@@ -5,87 +5,81 @@
 # system to decide whether to invest in loop peeling etc. Here we write
 # the logical structure, and allow optimizing for size as well.
 
-import math
-
 import numpy as np
-from numba import njit
-
-# Define the np.ndarray type as an numpy array of 3 floats
-# vec3d = np.empty(3, dtype=float)
+from numba import float64, int32, njit
 
-# and 2 floats
-#  = np.empty(2, dtype=float)
 
-@njit
-def norm(x: float, y: float, z: float) -> float:
-    """Return the norm of a 3D vector given by 3 float components."""
-    return vec_norm(vec_set(x, y, z))
+@njit(float64(float64[:]),fastmath=True, cache=True, nogil=True)
+def vec_norm(vec: np.ndarray) -> float:
+    """vec_norm() gives the norm of a vector."""
+    return np.sqrt(vec[0]**2 + vec[1]**2 + vec[2]**2)
 
-@njit
+@njit(float64[:](float64,float64,float64),fastmath=True, cache=True, nogil=True)
 def vec_set(x: float, y: float, z: float) -> np.ndarray:
     """Set the components of a  3D vector from separate doubles."""
     return np.array([x, y, z])
 
+@njit(float64(float64,float64,float64),fastmath=True, cache=True, nogil=True)
+def norm(x: float, y: float, z: float) -> float:
+    """Return the norm of a 3D vector given by 3 float components."""
+    return vec_norm(vec_set(x, y, z))
 
+
+@njit(float64[:](float64[:]),fastmath=True, cache=True, nogil=True)
 def vec_copy(src: np.ndarray) -> np.ndarray:
     """Copy one 3D vector into another."""
     return src.copy()
 
-@njit
+@njit(float64[:](float64[:],float64[:]),fastmath=True, cache=True, nogil=True)
 def vec_subt(from_: np.ndarray, sub: np.ndarray) -> np.ndarray:
     """Subtract two 3D vectors."""
     return from_ - sub
 
-@njit
+@njit(float64[:](float64[:],float64[:]),fastmath=True, cache=True, nogil=True)
 def vec_add(vec1: np.ndarray, vec2: np.ndarray) -> np.ndarray:
     """Add two 3D vectors."""
     return vec1 + vec2
 
-@njit
+@njit(float64[:](float64[:],float64),fastmath=True, cache=True, nogil=True)
 def vec_scalar_mul(vec: np.ndarray, scalar: float) -> np.ndarray:
     """vec_scalar_mul(np.ndarray, scalar) multiplies a vector by a scalar."""
     return vec * scalar
 
-@njit
+@njit(float64(float64[:],float64[:]),fastmath=True, cache=True, nogil=True)
 def vec_diff_norm(vec1: np.ndarray, vec2: np.ndarray) -> float:
     """vec_diff_norm() gives the norm of the difference between two vectors."""
-    # return np.linalg.norm(vec1 - vec2)
     vec = vec1 - vec2
-    return math.sqrt(vec[0]**2 + vec[1]**2 + vec[2]**2)
+    return np.sqrt(vec[0]**2 + vec[1]**2 + vec[2]**2)
 
-@njit
-def vec_norm(vec: np.ndarray) -> float:
-    """vec_norm() gives the norm of a vector."""
-    return math.sqrt(vec[0]**2 + vec[1]**2 + vec[2]**2)
 
-@njit
+@njit(float64(float64[:],float64[:]),fastmath=True, cache=True, nogil=True)
 def vec_dot(vec1: np.ndarray, vec2: np.ndarray) -> float:
     """vec_dot() gives the dot product of two vectors as lists of floats."""
     # return np.dot(vec1, vec2)
     return float(vec1[0]*vec2[0] + vec1[1]*vec2[1] + vec1[2]*vec2[2])
 
-@njit
+@njit(float64[:](float64[:],float64[:]),fastmath=True, cache=True, nogil=True)
 def vec_cross(vec1: np.ndarray, vec2: np.ndarray) -> np.ndarray:
     """Cross product of two vectors."""
     # return np.cross(vec1, vec2)
     return np.array([vec1[1]*vec2[2] - vec1[2]*vec2[1],
                  vec1[2]*vec2[0] - vec1[0]*vec2[2],
                  vec1[0]*vec2[1] - vec1[1]*vec2[0]])
 
-@njit
+@njit("boolean(float64[:],float64[:],float64)",fastmath=True, cache=True, nogil=True)
 def vec_cmp(vec1: np.ndarray, vec2: np.ndarray, tol: float = 1e-6) -> bool:
     """vec_cmp() checks whether two vectors are equal within a tolerance."""
     return np.allclose(vec1, vec2, atol=tol)
 
-@njit
+@njit(float64[:](float64[:]),fastmath=True, cache=True, nogil=True)
 def unit_vector(vec: np.ndarray) -> np.ndarray:
     """Normalize a vector to a unit vector."""
     magnitude = vec_norm(vec)
     if magnitude == 0:
         return vec  # Avoid division by zero for zero vectors
     return vec / magnitude
 
-@njit
-def vec_init(length=3) -> np.ndarray:
+@njit(float64[:](int32),fastmath=True, cache=True, nogil=True)
+def vec_init(length: int=3) -> np.ndarray:
     """Initialize a vector to zero."""
     return np.zeros(length, dtype=float)

tests/test_tracking_run.py

@@ -285,8 +285,8 @@ def test_trackback(self):
         trackcorr_c_finish(run, run.seq_par.last)
 
 
-        run.tpar.dvxmin = run.tpar.dvymin = run.tpar.dvzmin = -50
-        run.tpar.dvxmax = run.tpar.dvymax = run.tpar.dvzmax = 50
+        run.tpar.dvxmin = run.tpar.dvymin = run.tpar.dvzmin = -50.0
+        run.tpar.dvxmax = run.tpar.dvymax = run.tpar.dvzmax = 50.0
 
 
         run.lmax = vec_norm(

tests/test_vec_utils.py

@@ -1,3 +1,4 @@
+"""Tests for the vec_utils module."""
 import numpy as np
 
 from openptv_python.vec_utils import (