diff --git a/eks/multicam_smoother.py b/eks/multicam_smoother.py index 076b11c..7a5892b 100644 --- a/eks/multicam_smoother.py +++ b/eks/multicam_smoother.py @@ -835,16 +835,26 @@ def h_fn(x): def triangulate_3d_models(marker_array, camgroup) -> np.ndarray: - """Triangulate per-model, per-kpt, per-frame: (M,K,T,3).""" + """Triangulate per-model, per-kpt, per-frame: (M,K,T,3). + + M*K calls are independent so we parallelise over all available CPU cores. + """ + from joblib import Parallel, delayed + M, C, T, K, _ = marker_array.shape raw = marker_array.get_array() # (M,C,T,K,3) + + def _tri(m, k): + xy_views = raw[m, :, :, k, :2] # (C, T, 2) + return m, k, camgroup.triangulate(xy_views, fast=True) # (T, 3) + + results = Parallel(n_jobs=-1, prefer="threads")( + delayed(_tri)(m, k) for m in range(M) for k in range(K) + ) + tri = np.zeros((M, K, T, 3), dtype=float) - for m in range(M): - for k in range(K): - # Batch all T frames together: shape (C, T, 2) - xy_views = raw[m, :, :, k, :2] # (C, T, 2) - # triangulate expects (C, N, 2) and returns (N, 3) - tri[m, k, :, :] = camgroup.triangulate(xy_views, fast=True) + for m, k, arr in results: + tri[m, k] = arr return tri @@ -863,24 +873,13 @@ def project_3d_covariance_to_2d(ms_k, Vs_k, h_cam, inflated_vars_k): var_y: (T,) - y-direction posterior variances """ - # Compute Jacobian of projection function at each 3D point - def project_single_point(x_3d): - return h_cam(x_3d) - - # Compute Jacobian for each time point - jacobians = [] - for t in range(ms_k.shape[0]): - jac = jax.jacfwd(project_single_point)(ms_k[t]) - jacobians.append(jac) - - jacobians = np.array(jacobians) # (T, 2, 3) + # Jacobians via vmap over T — one vectorized call instead of T individual dispatches + jacobians = np.array(vmap(jax.jacfwd(h_cam))(jnp.asarray(ms_k))) # (T, 2, 3) - # Project 3D covariance to 2D: Cov_2D = J * Cov_3D * J^T - cov2d_proj = np.zeros((ms_k.shape[0], 2, 2)) - for t in range(ms_k.shape[0]): - J = jacobians[t] # (2, 3) - V_3d = Vs_k[t] # (3, 3) - cov2d_proj[t] = J @ V_3d @ J.T # (2, 2) + # Project 3D covariance to 2D: Cov_2D = J @ Cov_3D @ J^T, vectorized over T + J = jacobians # (T, 2, 3) + V = np.array(Vs_k) # (T, 3, 3) + cov2d_proj = J @ V @ J.transpose(0, 2, 1) # (T, 2, 2) # Extract x and y variances and add ensemble variance var_x = cov2d_proj[:, 0, 0] + inflated_vars_k[:, 0] diff --git a/tests/test_multicam_smoother.py b/tests/test_multicam_smoother.py index 5c3f03c..4054e85 100644 --- a/tests/test_multicam_smoother.py +++ b/tests/test_multicam_smoother.py @@ -611,8 +611,11 @@ def test_project_3d_covariance_to_2d_matches_fd_linearization(): var_x, var_y = project_3d_covariance_to_2d(ms_k, Vs_k, h_cam, inflated) - # check a few timesteps - for t in [0, 3, 5, 7]: + assert var_x.shape == (T,) + assert var_y.shape == (T,) + + # check all timesteps + for t in range(T): J_fd = _finite_diff_jacobian(lambda x: np.array(h_cam(jnp.asarray(x))), ms_k[t]) cov2d_fd = J_fd @ Vs_k[t] @ J_fd.T np.testing.assert_allclose(var_x[t] - inflated[t, 0], cov2d_fd[0, 0], rtol=1e-4, atol=1e-6)