"""Lever-arm calibration from RTK positions and full IMU preintegration.""" from __future__ import annotations from dataclasses import dataclass import numpy as np from scipy.sparse import coo_matrix, csr_matrix, eye from scipy.sparse.linalg import lsqr, splu from scipy.spatial.transform import Rotation, Slerp from imu_lidar.geometry import make_transform, orthonormalize_rotation, so3_exp from .rtk_imu_rotation import RotationCalibrationResult, RotationSession, _attitude_rows @dataclass(frozen=True) class TranslationCalibrationResult: lever_IMU_to_RTK_in_IMU_m: np.ndarray t_RTK_IMU_m: np.ndarray T_RTK_IMU: np.ndarray translation_std_m: np.ndarray lever_information_singular_values: np.ndarray lever_precision_rank: int position_residual_rms_xyz_m: np.ndarray velocity_residual_rms_xyz_m_s: np.ndarray accel_bias_by_session_m_s2: dict[str, np.ndarray] knot_count_by_session: dict[str, int] loo_delta_m: dict[str, np.ndarray] ok: bool notes: tuple[str, ...] @dataclass(frozen=True) class _SessionFactors: session: RotationSession knot_t_s: np.ndarray position_enu_m: np.ndarray R_ENU_IMU: np.ndarray delta_p: tuple[np.ndarray, ...] delta_v: tuple[np.ndarray, ...] J_p_ba: tuple[np.ndarray, ...] J_v_ba: tuple[np.ndarray, ...] duration_s: np.ndarray def _preintegrate_translation_interval( times_s: np.ndarray, gyro_rad_s: np.ndarray, acc_m_s2: np.ndarray, t0: float, t1: float, gyro_bias_rad_s: np.ndarray, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, float]: """Fast nominal ``delta_p/delta_v`` and accel-bias Jacobians. Rotation covariance and gyro-bias Jacobians are deliberately omitted here: rotation and gyro bias have already been fixed by Phase R1, while the translation linear system only consumes the accelerometer-bias Jacobians. """ left = max(int(np.searchsorted(times_s, t0, side='left') - 1), 0) right = min(int(np.searchsorted(times_s, t1, side='right')), times_s.size - 1) delta_r = np.eye(3) delta_v = np.zeros(3) delta_p = np.zeros(3) j_v_ba = np.zeros((3, 3)) j_p_ba = np.zeros((3, 3)) for index in range(left, right): sample_t0 = float(times_s[index]) sample_t1 = float(times_s[index + 1]) if sample_t1 <= t0 or sample_t0 >= t1: continue segment_t0 = max(sample_t0, t0) segment_t1 = min(sample_t1, t1) dt = segment_t1 - segment_t0 if dt <= 0.0: continue sample_dt = max(sample_t1 - sample_t0, 1e-12) u0 = (segment_t0 - sample_t0) / sample_dt u1 = (segment_t1 - sample_t0) / sample_dt gyro0 = (1.0 - u0) * gyro_rad_s[index] + u0 * gyro_rad_s[index + 1] gyro1 = (1.0 - u1) * gyro_rad_s[index] + u1 * gyro_rad_s[index + 1] acc0 = (1.0 - u0) * acc_m_s2[index] + u0 * acc_m_s2[index + 1] acc1 = (1.0 - u1) * acc_m_s2[index] + u1 * acc_m_s2[index + 1] omega = 0.5 * (gyro0 + gyro1) - gyro_bias_rad_s acc = 0.5 * (acc0 + acc1) r_i = delta_r delta_p = delta_p + delta_v * dt + 0.5 * r_i @ acc * dt**2 delta_v = delta_v + r_i @ acc * dt j_p_ba = j_p_ba + j_v_ba * dt - 0.5 * r_i * dt**2 j_v_ba = j_v_ba - r_i * dt delta_r = orthonormalize_rotation(delta_r @ so3_exp(omega * dt)) return delta_p, delta_v, j_p_ba, j_v_ba, float(max(t1 - t0, 0.0)) def _make_session_factors( session: RotationSession, rotation: RotationCalibrationResult, *, knot_step_s: float, ) -> _SessionFactors: t_attitude, r_enu_rtk = _attitude_rows(session.rtk, rotation.convention) position_valid = session.rtk.position_valid t_position = session.rtk.t_s[position_valid] position = session.rtk.position_enu_m[position_valid] time_offset_s = rotation.applied_time_offset_s start = max(float(t_attitude[0]), float(t_position[0]), float(session.imu.t_s[0] - time_offset_s)) end = min(float(t_attitude[-1]), float(t_position[-1]), float(session.imu.t_s[-1] - time_offset_s)) if end - start < 5.0: raise ValueError(f"{session.session_id}: less than 5 s common RTK/IMU support") knot_t = np.arange(start + 0.25, end - 0.25, knot_step_s) if knot_t.size < 4: raise ValueError(f"{session.session_id}: not enough translation knots") position_knots = np.column_stack( [np.interp(knot_t, t_position, position[:, axis]) for axis in range(3)] ) r_enu_rtk_knots = Slerp(t_attitude, Rotation.from_matrix(r_enu_rtk))(knot_t).as_matrix() r_enu_imu = r_enu_rtk_knots @ rotation.R_RTK_IMU bg = rotation.gyro_bias_by_session_rad_s[session.session_id] delta_p: list[np.ndarray] = [] delta_v: list[np.ndarray] = [] j_p_ba: list[np.ndarray] = [] j_v_ba: list[np.ndarray] = [] durations = [] for t0, t1 in zip(knot_t[:-1], knot_t[1:]): dp, dv, jp, jv, duration = _preintegrate_translation_interval( session.imu.t_s, session.imu.gyro_rad_s, session.imu.acc_m_s2, float(t0 + time_offset_s), float(t1 + time_offset_s), bg, ) delta_p.append(dp) delta_v.append(dv) j_p_ba.append(jp) j_v_ba.append(jv) durations.append(duration) return _SessionFactors( session=session, knot_t_s=knot_t, position_enu_m=position_knots, R_ENU_IMU=r_enu_imu, delta_p=tuple(delta_p), delta_v=tuple(delta_v), J_p_ba=tuple(j_p_ba), J_v_ba=tuple(j_v_ba), duration_s=np.asarray(durations), ) def _append_block( rows: list[int], cols: list[int], values: list[float], rhs: list[float], groups: list[int], matrix_blocks: list[tuple[int, np.ndarray]], vector: np.ndarray, sigma: np.ndarray, group: int, ) -> None: row0 = len(rhs) for axis in range(3): rhs.append(float(vector[axis] / sigma[axis])) groups.append(group) for col0, block in matrix_blocks: for local_col in range(block.shape[1]): value = float(block[axis, local_col] / sigma[axis]) if value != 0.0: rows.append(row0 + axis) cols.append(col0 + local_col) values.append(value) def _build_system( factors: list[_SessionFactors], *, position_sigma_xyz_m: np.ndarray, velocity_sigma_xyz_m_s: np.ndarray, ) -> tuple[csr_matrix, np.ndarray, np.ndarray, dict[str, tuple[int, int]], list[tuple[str, int, str]]]: # x = [shared lever(3), per-session ba(3), per-knot velocities(3*K)] offsets: dict[str, tuple[int, int]] = {} variable_count = 3 for item in factors: ba_offset = variable_count velocity_offset = ba_offset + 3 offsets[item.session.session_id] = (ba_offset, velocity_offset) variable_count = velocity_offset + 3 * item.knot_t_s.size rows: list[int] = [] cols: list[int] = [] values: list[float] = [] rhs: list[float] = [] groups: list[int] = [] factor_labels: list[tuple[str, int, str]] = [] gravity = np.array([0.0, 0.0, -9.80665]) group = 0 for item in factors: ba_offset, velocity_offset = offsets[item.session.session_id] for index, dt in enumerate(item.duration_s): r_i = item.R_ENU_IMU[index] r_j = item.R_ENU_IMU[index + 1] dp_rtk = item.position_enu_m[index + 1] - item.position_enu_m[index] constant_p = dp_rtk - 0.5 * gravity * dt**2 - r_i @ item.delta_p[index] _append_block( rows, cols, values, rhs, groups, [ (0, r_i - r_j), (ba_offset, -r_i @ item.J_p_ba[index]), (velocity_offset + 3 * index, -dt * np.eye(3)), ], -constant_p, position_sigma_xyz_m, group, ) factor_labels.append((item.session.session_id, group, "position")) group += 1 constant_v = -gravity * dt - r_i @ item.delta_v[index] _append_block( rows, cols, values, rhs, groups, [ (ba_offset, -r_i @ item.J_v_ba[index]), (velocity_offset + 3 * index, -np.eye(3)), (velocity_offset + 3 * (index + 1), np.eye(3)), ], -constant_v, velocity_sigma_xyz_m_s, group, ) factor_labels.append((item.session.session_id, group, "velocity")) group += 1 # Loose physical bias prior. It prevents an unobservable constant # acceleration from masquerading as gravity while remaining data-led. _append_block( rows, cols, values, rhs, groups, [(ba_offset, np.eye(3))], np.zeros(3), np.full(3, 0.5), group, ) factor_labels.append((item.session.session_id, group, "bias_prior")) group += 1 matrix = coo_matrix((values, (rows, cols)), shape=(len(rhs), variable_count)).tocsr() return matrix, np.asarray(rhs), np.asarray(groups), offsets, factor_labels def _irls(matrix: csr_matrix, rhs: np.ndarray, groups: np.ndarray) -> tuple[np.ndarray, np.ndarray]: row_weights = np.ones(rhs.size) solution = np.zeros(matrix.shape[1]) for _ in range(5): weighted = matrix.multiply(row_weights[:, None]) solution = lsqr(weighted, rhs * row_weights, atol=1e-10, btol=1e-10, iter_lim=3000)[0] residual = matrix @ solution - rhs new_weights = np.ones_like(row_weights) for group in np.unique(groups): selection = groups == group norm = float(np.linalg.norm(residual[selection])) if norm > 3.0: new_weights[selection] = np.sqrt(3.0 / norm) if np.max(np.abs(new_weights - row_weights)) < 1e-3: row_weights = new_weights break row_weights = new_weights return solution, row_weights def _solve_factors( factors: list[_SessionFactors], position_sigma: np.ndarray, velocity_sigma: np.ndarray, ) -> tuple[np.ndarray, np.ndarray, csr_matrix, np.ndarray, dict[str, tuple[int, int]], np.ndarray, np.ndarray]: matrix, rhs, groups, offsets, labels = _build_system( factors, position_sigma_xyz_m=position_sigma, velocity_sigma_xyz_m_s=velocity_sigma, ) solution, row_weights = _irls(matrix, rhs, groups) weighted = matrix.multiply(row_weights[:, None]).tocsr() residual = matrix @ solution - rhs data_groups = {group for _, group, kind in labels if kind != "bias_prior"} data_rows = np.isin(groups, list(data_groups)) variance = float(np.sum((residual[data_rows] * row_weights[data_rows]) ** 2) / max(np.count_nonzero(data_rows) - solution.size, 1)) information = (weighted.T @ weighted).tocsc() + eye(weighted.shape[1], format="csc") * 1e-10 h_ll = information[:3, :3].toarray() h_ln = information[:3, 3:] h_nn = information[3:, 3:] nuisance_solve = splu(h_nn).solve(h_ln.T.toarray()) schur = h_ll - h_ln.toarray() @ nuisance_solve covariance_lever = np.linalg.pinv(schur, rcond=1e-10) * variance return solution, covariance_lever, matrix, rhs, offsets, groups, residual def solve_rtk_imu_translation( sessions: list[RotationSession] | tuple[RotationSession, ...], rotation: RotationCalibrationResult, *, knot_step_s: float = 2.0, compute_loo: bool = True, ) -> TranslationCalibrationResult: """Estimate the shared IMU-to-RTK lever arm and return ``T_RTK_IMU``.""" items = list(sessions) factors = [_make_session_factors(session, rotation, knot_step_s=knot_step_s) for session in items] position_sigma = np.array([0.025, 0.025, 0.060]) velocity_sigma = np.array([0.08, 0.08, 0.12]) solution, covariance_l, matrix, rhs, offsets, groups, residual = _solve_factors( factors, position_sigma, velocity_sigma ) lever = solution[:3] t_rtk_imu = -rotation.R_RTK_IMU @ lever covariance_t = rotation.R_RTK_IMU @ covariance_l @ rotation.R_RTK_IMU.T std_t = np.sqrt(np.maximum(np.diag(covariance_t), 0.0)) schur_information = np.linalg.pinv(covariance_l, rcond=1e-12) singular_values = np.linalg.svd(schur_information, compute_uv=False) threshold = max(float(singular_values[0]) * 1e-4, 1e-9) rank = int(np.count_nonzero(singular_values > threshold)) # Recover physical residuals: system rows are grouped in XYZ triples and # alternate position/velocity, followed by one bias prior per session. position_errors: list[np.ndarray] = [] velocity_errors: list[np.ndarray] = [] cursor = 0 for item in factors: for _ in range(item.knot_t_s.size - 1): position_errors.append(residual[cursor : cursor + 3] * position_sigma) cursor += 3 velocity_errors.append(residual[cursor : cursor + 3] * velocity_sigma) cursor += 3 cursor += 3 pos_rms = np.sqrt(np.mean(np.asarray(position_errors) ** 2, axis=0)) vel_rms = np.sqrt(np.mean(np.asarray(velocity_errors) ** 2, axis=0)) biases = { item.session.session_id: solution[offsets[item.session.session_id][0] : offsets[item.session.session_id][0] + 3].copy() for item in factors } loo: dict[str, np.ndarray] = {} if compute_loo and len(factors) >= 3: for omitted in factors: kept = [item for item in factors if item.session.session_id != omitted.session.session_id] loo_solution, *_ = _solve_factors(kept, position_sigma, velocity_sigma) loo[omitted.session.session_id] = (-rotation.R_RTK_IMU @ loo_solution[:3]) - t_rtk_imu max_loo_xy = max((float(np.linalg.norm(value[:2])) for value in loo.values()), default=0.0) max_loo_z = max((abs(float(value[2])) for value in loo.values()), default=0.0) ok = bool( rank == 3 and float(np.max(std_t[:2])) <= 0.05 and float(std_t[2]) <= 0.10 and float(np.max(pos_rms[:2])) <= 0.10 and float(pos_rms[2]) <= 0.20 and max_loo_xy <= 0.10 and max_loo_z <= 0.20 ) notes = [ "lever l is vector IMU-origin -> RTK-origin expressed in IMU", "transform translation uses t_RTK_IMU = -R_RTK_IMU @ l", "RTK position is never differentiated; position and velocity preintegration factors are solved jointly", ] if not rotation.ok: notes.append("upstream rotation is not accepted, so translation is diagnostic only") ok = False if not ok: notes.append("translation failed one or more strict acceptance gates") return TranslationCalibrationResult( lever_IMU_to_RTK_in_IMU_m=lever, t_RTK_IMU_m=t_rtk_imu, T_RTK_IMU=make_transform(t_rtk_imu, rotation.R_RTK_IMU), translation_std_m=std_t, lever_information_singular_values=singular_values, lever_precision_rank=rank, position_residual_rms_xyz_m=pos_rms, velocity_residual_rms_xyz_m_s=vel_rms, accel_bias_by_session_m_s2=biases, knot_count_by_session={item.session.session_id: int(item.knot_t_s.size) for item in factors}, loo_delta_m=loo, ok=ok, notes=tuple(notes), )