feat: Enhance playground.py with new 3D tracking and affinity calculations

- Added functions for calculating perpendicular distances between predicted 3D tracking points and camera rays, improving 3D tracking accuracy.
- Introduced a new function for calculating 3D affinity scores based on distances and time differences, enhancing the integration of 3D tracking with existing systems.
- Updated existing functions to support new data types and improved documentation for clarity on parameters and return values.
- Refactored affinity calculation logic to utilize JAX for performance optimization in distance computations.

app/camera/__init__.py

@@ -103,8 +103,10 @@ def unproject_points_onto_plane(
 
     (i.e. back-project points onto a plane)
 
+    `intersect_image_rays_with_plane`/`compute_ray_plane_intersections`
+
     Args:
-        points_2d: [..., 2] image pixel coordinates
+        points_2d: [..., 2] image pixel coordinates (with camera distortion)
         plane_normal: (3,) normal vector of the plane in world coords
         plane_point: (3,) a known point on the plane in world coords
         K: Camera intrinsic matrix
@@ -118,7 +120,7 @@ def unproject_points_onto_plane(
     Returns:
         [..., 3] world-space intersection points
     """
-    # Step 1: undistort (no-op here)
+    # Step 1: undistort
     pts = undistort_points(
         np.asarray(points_2d), np.asarray(K), np.asarray(dist_coeffs)
     )
@@ -313,6 +315,13 @@ class CameraParams:
             object.__setattr__(self, "_proj", pm)
         return pm
 
+    @property
+    def location(self) -> Num[Array, "3"]:
+        """
+        The 3D location of camera (relative to world coordinate system)
+        """
+        return self.pose_matrix[:3, -1].reshape((3,))
+
 
 @jaxtyped(typechecker=beartype)
 @dataclass(frozen=True)
@@ -390,7 +399,7 @@ class Camera:
         Un-project 2D points to 3D points on a plane at z = constant.
 
         Args:
-            points_2d: 2D points in pixel coordinates
+            points_2d: 2D points in pixel coordinates (with camera distortion)
             z: z-coordinate of the plane (default: 0.0,  i.e. ground/horizon/floor plane)
 
         Returns: