kalman.py

# ====================== Embedded helpers ======================
class KalmanCV3D:
    """Constant-velocity Kalman filter for a 3D point (pos+vel)."""
    def __init__(self, x0, v0=None, meas_sigma=0.01, accel_sigma=1.0,
                 p0_pos=0.10, p0_vel=1.0):
        import numpy as np
        self.np = np
        self.x = np.zeros(6, dtype=np.float64)
        self.x[:3] = np.asarray(x0, dtype=np.float64)
        if v0 is not None:
            self.x[3:] = np.asarray(v0, dtype=np.float64)
        self.P = np.diag([p0_pos, p0_pos, p0_pos, p0_vel, p0_vel, p0_vel]).astype(np.float64)
        self.R = np.eye(3) * (meas_sigma ** 2)
        self.accel_sigma = float(accel_sigma)

    def _F_Q(self, dt):
        np = self.np
        I = np.eye(3)
        F = np.block([[I, dt*I],
                      [np.zeros((3,3)), I]])
        dt2, dt3, dt4 = dt*dt, dt**3, dt**4
        q = self.accel_sigma**2
        Q = q * np.block([
            [ (dt4/4.0)*I, (dt3/2.0)*I ],
            [ (dt3/2.0)*I,    dt2*I    ],
        ])
        return F, Q

    def predict(self, dt):
        F, Q = self._F_Q(dt)
        self.x = F @ self.x
        self.P = F @ self.P @ F.T + Q
        return self.x[:3].copy()

    def update(self, z):
        np = self.np
        z = np.asarray(z, dtype=np.float64)
        H = np.block([np.eye(3), np.zeros((3,3))])
        S = H @ self.P @ H.T + self.R
        K = self.P @ H.T @ np.linalg.inv(S)
        y = z - (H @ self.x)
        self.x = self.x + K @ y
        self.P = (np.eye(6) - K @ H) @ self.P
        return self.x[:3].copy()

@staticmethod
def registration_icp_z_only(
    source, target, max_correspondence_distance=0.05,
    init=None, axis=(0.0, 0.0, 1.0), max_iters=50, tolerance=1e-7,
):
    """
    One-DoF ICP: translation along `axis` only. Returns an object with `.transformation`,
    `.inlier_rmse`, `.fitness` (Open3D-like).
    """
    import numpy as np, open3d as o3d
    from types import SimpleNamespace

    if init is None:
        init = np.eye(4, dtype=np.float64)
    if len(source.points) == 0 or len(target.points) == 0:
        return SimpleNamespace(transformation=init, inlier_rmse=np.inf, fitness=0.0)

    u = np.asarray(axis, dtype=np.float64)
    nu = np.linalg.norm(u)
    if nu == 0:
        raise ValueError("axis must be non-zero")
    u /= nu

    S = np.asarray(source.points, dtype=np.float64)
    R0, t0 = init[:3, :3], init[:3, 3]
    S0 = (S @ R0.T) + t0

    Tpts = np.asarray(target.points, dtype=np.float64)
    kdt = o3d.geometry.KDTreeFlann(target)
    d2_th = float(max_correspondence_distance ** 2)

    t_scalar = 0.0
    rmse = np.inf
    inlier_count = 0

    for _ in range(max_iters):
        Sc = S0 + t_scalar * u
        idx_src, idx_tgt, d2s = [], [], []

        for i, p in enumerate(Sc):
            k, idxs, d2 = kdt.search_knn_vector_3d(p, 1)
            if k and d2[0] <= d2_th:
                idx_src.append(i); idx_tgt.append(idxs[0]); d2s.append(d2[0])

        m = len(idx_src)
        if m < 3:
            break

        S_sel = S0[np.asarray(idx_src)]
        T_sel = Tpts[np.asarray(idx_tgt)]

        t_new = float(np.mean((T_sel - S_sel) @ u))
        if abs(t_new - t_scalar) < tolerance:
            t_scalar = t_new
            rmse = float(np.sqrt(np.mean(d2s))) if d2s else np.inf
            inlier_count = m
            break

        t_scalar = t_new
        rmse = float(np.sqrt(np.mean(d2s))) if d2s else np.inf
        inlier_count = m

    T_final = np.eye(4, dtype=np.float64)
    T_final[:3, :3] = R0
    T_final[:3, 3] = t0 + t_scalar * u

    fitness = (inlier_count / max(1, len(S))) if len(S) else 0.0
    return SimpleNamespace(transformation=T_final, inlier_rmse=rmse, fitness=float(fitness))
# ==================== end embedded helpers ====================