improved code, ray_tracing should be faster

Author: alexlib

openptv_python/lsqadj.py

@@ -1,8 +1,10 @@
 """Least squares adjustment of the camera parameters."""
 import numpy as np
+from numba import njit
 
 
-def ata(a: np.ndarray, ata: np.ndarray, m: int, n: int, n_large: int) -> None:
+@njit
+def ata(a: np.ndarray, at: np.ndarray, m: int, n: int, n_large: int) -> None:
     """
     Multiply transpose of a matrix A by matrix A itself, creating a symmetric matrix inplace.
 
@@ -16,9 +18,9 @@ def ata(a: np.ndarray, ata: np.ndarray, m: int, n: int, n_large: int) -> None:
     """
     for i in range(n):
         for j in range(n):
-            ata.flat[i * n + j] = 0.0
+            at.flat[i * n + j] = 0.0
             for k in range(m):
-                ata.flat[i * n + j] += a.flat[k * n_large + i] * a.flat[k * n_large + j]
+                at.flat[i * n + j] += a.flat[k * n_large + i] * a.flat[k * n_large + j]
 
 
 # def ata(a: np.ndarray, ata: np.ndarray, m: int, n: int, n_large: int) -> None:
@@ -45,7 +47,7 @@ def ata(a: np.ndarray, ata: np.ndarray, m: int, n: int, n_large: int) -> None:
 #                     a.flat[k * n_large + i] * a.flat[k * n_large + j]
 #                 )
 
-
+@njit
 def atl(u: np.ndarray, a: np.ndarray, b: np.ndarray, m: int, n: int, n_large: int):
     """Multiply transpose of a matrix A by vector b, creating vector u.
 
@@ -229,7 +231,7 @@ def atl(u: np.ndarray, a: np.ndarray, b: np.ndarray, m: int, n: int, n_large: in
 #             pa += k
 #         c += 1
 
-
+@njit
 def matmul(
     a: np.ndarray,
     b: np.ndarray,

openptv_python/ray_tracing.py

@@ -2,6 +2,7 @@
 from typing import Tuple
 
 import numpy as np
+from numba import njit
 
 from .calibration import Calibration
 from .lsqadj import matmul
@@ -48,70 +49,202 @@ def ray_tracing(
     -------
             _type_: _description_
     """
-    # Initial ray direction in global coordinate system
-    tmp1 = np.r_[x, y, -1 * cal.int_par.cc]
-    tmp1 = unit_vector(tmp1)
-    start_dir = np.empty(3, dtype=float)
-    matmul(start_dir, cal.ext_par.dm, tmp1, 3, 3, 1, 3, 3)
-
     primary_point = np.r_[cal.ext_par.x0, cal.ext_par.y0, cal.ext_par.z0]
+    glass = np.r_[cal.glass_par.vec_x, cal.glass_par.vec_y, cal.glass_par.vec_z]
+    return fast_ray_tracing(
+        x,
+        y,
+        cal.int_par.cc,
+        cal.ext_par.dm,
+        primary_point,
+        glass,
+        mm.d[0],
+        mm.n1,
+        mm.n2[0],
+        mm.n3,
+    )
+
+@njit
+def fast_ray_tracing(
+    camera_x: float,
+    camera_y: float,
+    camera_cc: float,
+    distortion_matrix: np.ndarray,
+    primary_point: np.ndarray,
+    glass_vector: np.ndarray,
+    distance_param: float,
+    refractive_index_medium1: float,
+    refractive_index_medium2: float,
+    refractive_index_medium3: float,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Fast ray tracing.
 
-    tmp1 = np.r_[cal.glass_par.vec_x, cal.glass_par.vec_y, cal.glass_par.vec_z]
-    glass_dir = unit_vector(tmp1)
-    c = vec_norm(tmp1) + mm.d[0]
+    Parameters
+    ----------
+    - camera_x, camera_y, camera_cc: Camera parameters
+    - distortion_matrix: Distortion matrix
+    - primary_point: Primary point coordinates
+    - glass_vector: Glass vector
+    - distance_param: Distance parameter
+    - refractive_index_medium1, refractive_index_medium2, refractive_index_medium3: Refractive indices
+
+    Returns
+    -------
+    - Tuple containing the resulting point X and output direction vector
+    """
+    # Initial ray direction in global coordinate system
+    initial_ray_direction = np.array([camera_x, camera_y, -1 * camera_cc])
+    initial_ray_direction = unit_vector(initial_ray_direction)
+    transformed_direction = np.empty(3, dtype=float)
+    matmul(transformed_direction, distortion_matrix, initial_ray_direction, 3, 3, 1, 3, 3)
+
+    glass_direction = unit_vector(glass_vector)
+    c_param = vec_norm(glass_vector) + distance_param
 
     # Project start ray on glass vector to find n1/n2 interface.
-    dist_cam_glass = vec_dot(glass_dir, primary_point) - c
-    tmp1 = vec_dot(glass_dir, start_dir)
+    dist_cam_glass = vec_dot(glass_direction, primary_point) - c_param
+    dot_product_start_dir = float(vec_dot(glass_direction, transformed_direction))
 
     # avoid division by zero
-    if tmp1 == 0:
-        tmp1 = 1.0
-    d1 = -dist_cam_glass / tmp1
+    if dot_product_start_dir == 0.0:
+        dot_product_start_dir = 1.0
+    d1 = -dist_cam_glass / dot_product_start_dir
 
-    tmp1 = vec_scalar_mul(start_dir, d1)
-    Xb = vec_add(primary_point, tmp1)
+    transformed_direction_scaled = vec_scalar_mul(transformed_direction, d1)
+    Xb = vec_add(primary_point, transformed_direction_scaled)
 
     # Break down ray into glass-normal and glass-parallel components. */
-    n = vec_dot(start_dir, glass_dir)
-    tmp1 = vec_scalar_mul(glass_dir, n)
+    n = vec_dot(transformed_direction, glass_direction)
+    transformed_direction_parallel = vec_scalar_mul(glass_direction, n)
 
-    tmp2 = vec_subt(start_dir, tmp1)
-    bp = unit_vector(tmp2)
+    transformed_direction_perpendicular = vec_subt(transformed_direction, transformed_direction_parallel)
+    bp = unit_vector(transformed_direction_perpendicular)
 
     # Transform to direction inside glass, using Snell's law
-    p = np.sqrt(1 - n * n) * mm.n1 / mm.n2[0]
+    p = np.sqrt(1 - n * n) * refractive_index_medium1 / refractive_index_medium2
     # glass parallel
     n = -np.sqrt(1 - p * p)
     # glass normal
 
     # Propagation length in glass parallel to glass vector */
-    tmp1 = vec_scalar_mul(bp, p)
-    tmp2 = vec_scalar_mul(glass_dir, n)
-    a2 = vec_add(tmp1, tmp2)
+    transformed_direction_parallel_scaled = vec_scalar_mul(bp, p)
+    transformed_direction_perpendicular_scaled = vec_scalar_mul(glass_direction, n)
+    a2 = vec_add(transformed_direction_parallel_scaled, transformed_direction_perpendicular_scaled)
 
-    tmp1 = np.abs(vec_dot(glass_dir, a2))
+    abs_dot_product = np.abs(vec_dot(glass_direction, a2))
 
     # avoid division by zero
-    if tmp1 == 0:
-        tmp1 = 1.0
-    d2 = mm.d[0] / tmp1
+    if abs_dot_product == 0:
+        abs_dot_product = 1.0
+    d2 = distance_param / abs_dot_product
 
-    #   point on the horizontal plane between n2,n3  */
-    tmp1 = vec_scalar_mul(a2, d2)
-    X = vec_add(Xb, tmp1)
+    # point on the horizontal plane between n2,n3
+    a2_scaled = vec_scalar_mul(a2, d2)
+    X = vec_add(Xb, a2_scaled)
 
-    # Again, direction in next medium */
-    n = vec_dot(a2, glass_dir)
-    tmp2 = vec_subt(a2, tmp2)
-    bp = unit_vector(tmp2)
+    # Again, direction in next medium
+    n = vec_dot(a2, glass_direction)
+    transformed_direction_perpendicular_scaled = vec_subt(a2, transformed_direction_perpendicular_scaled)
+    bp = unit_vector(transformed_direction_perpendicular_scaled)
 
     p = np.sqrt(1 - n * n)
-    p = p * mm.n2[0] / mm.n3
+    p = p * refractive_index_medium2 / refractive_index_medium3
     n = -np.sqrt(1 - p * p)
 
-    tmp1 = vec_scalar_mul(bp, p)
-    tmp2 = vec_scalar_mul(glass_dir, n)
-    out = vec_add(tmp1, tmp2)
+    transformed_direction_parallel_scaled = vec_scalar_mul(bp, p)
+    transformed_direction_perpendicular_scaled = vec_scalar_mul(glass_direction, n)
+    out = vec_add(transformed_direction_parallel_scaled, transformed_direction_perpendicular_scaled)
 
     return X, out
+
+# def fast_ray_tracing(
+#     x,
+#     y,
+#     cc,
+#     dm,
+#     primary_point,
+#     glass,
+#     d,
+#     n1,
+#     n2,
+#     n3,
+#     ) -> Tuple[np.ndarray, np.ndarray]:
+#     """ Fast ray tracing.
+
+#     Parameters:
+#     - x, y, cc: Camera parameters
+#     - dm: Distortion matrix
+#     - primary_point: Primary point coordinates
+#     - glass: Glass vector
+#     - d: Distance parameter
+#     - n1, n2, n3: Refractive indices
+
+#     Returns:
+#     - Tuple containing the resulting point X and output direction vector
+#     """
+#     # Initial ray direction in global coordinate system
+#     tmp1 = np.array([x, y, -1 * cc])
+#     tmp1 = unit_vector(tmp1)
+#     start_dir = np.empty(3, dtype=float)
+#     matmul(start_dir, dm, tmp1, 3, 3, 1, 3, 3)
+
+
+#     glass_dir = unit_vector(glass)
+#     c = vec_norm(glass) + d
+
+#     # Project start ray on glass vector to find n1/n2 interface.
+#     dist_cam_glass = vec_dot(glass_dir, primary_point) - c
+#     tmp1 = float(vec_dot(glass_dir, start_dir))
+
+#     # avoid division by zero
+#     if tmp1 == 0.0:
+#         tmp1 = 1.0
+#     d1 = -dist_cam_glass / tmp1
+
+#     tmp1 = vec_scalar_mul(start_dir, d1)
+#     Xb = vec_add(primary_point, tmp1)
+
+#     # Break down ray into glass-normal and glass-parallel components. */
+#     n = vec_dot(start_dir, glass_dir)
+#     tmp1 = vec_scalar_mul(glass_dir, n)
+
+#     tmp2 = vec_subt(start_dir, tmp1)
+#     bp = unit_vector(tmp2)
+
+#     # Transform to direction inside glass, using Snell's law
+#     p = np.sqrt(1 - n * n) * n1 / n2
+#     # glass parallel
+#     n = -np.sqrt(1 - p * p)
+#     # glass normal
+
+#     # Propagation length in glass parallel to glass vector */
+#     tmp1 = vec_scalar_mul(bp, p)
+#     tmp2 = vec_scalar_mul(glass_dir, n)
+#     a2 = vec_add(tmp1, tmp2)
+
+#     tmp1 = np.abs(vec_dot(glass_dir, a2))
+
+#     # avoid division by zero
+#     if tmp1 == 0:
+#         tmp1 = 1.0
+#     d2 = d / tmp1
+
+#     #   point on the horizontal plane between n2,n3  */
+#     tmp1 = vec_scalar_mul(a2, d2)
+#     X = vec_add(Xb, tmp1)
+
+#     # Again, direction in next medium */
+#     n = vec_dot(a2, glass_dir)
+#     tmp2 = vec_subt(a2, tmp2)
+#     bp = unit_vector(tmp2)
+
+#     p = np.sqrt(1 - n * n)
+#     p = p * n2 / n3
+#     n = -np.sqrt(1 - p * p)
+
+#     tmp1 = vec_scalar_mul(bp, p)
+#     tmp2 = vec_scalar_mul(glass_dir, n)
+#     out = vec_add(tmp1, tmp2)
+
+#     return X, out

openptv_python/vec_utils.py

@@ -8,81 +8,84 @@
 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)
 
 # 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 float(np.linalg.norm(vec_set(x, y, z)))
-
+    return vec_norm(vec_set(x, y, z))
 
+@njit
 def vec_set(x: float, y: float, z: float) -> np.ndarray:
     """Set the components of a  3D vector from separate doubles."""
-    return np.r_[x, y, z]
+    return np.array([x, y, z])
 
 
 def vec_copy(src: np.ndarray) -> np.ndarray:
     """Copy one 3D vector into another."""
     return src.copy()
 
-
+@njit
 def vec_subt(from_: np.ndarray, sub: np.ndarray) -> np.ndarray:
     """Subtract two 3D vectors."""
     return from_ - sub
 
-
+@njit
 def vec_add(vec1: np.ndarray, vec2: np.ndarray) -> np.ndarray:
     """Add two 3D vectors."""
     return vec1 + vec2
 
-
+@njit
 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
 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)
 
-
+@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
 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
 def vec_cross(vec1: np.ndarray, vec2: np.ndarray) -> np.ndarray:
     """Cross product of two vectors."""
     # return np.cross(vec1, vec2)
-    return np.r_[vec1[1]*vec2[2] - vec1[2]*vec2[1],
+    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]]
-
+                 vec1[0]*vec2[1] - vec1[1]*vec2[0]])
 
+@njit
 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
 def unit_vector(vec: np.ndarray) -> np.ndarray:
     """Normalize a vector to a unit vector."""
-    magnitude = np.linalg.norm(vec)
+    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:
     """Initialize a vector to zero."""
     return np.zeros(length, dtype=float)