Certifiable Landmark SLAM on SE(2)

import numpy as np
import plotly.graph_objects as go

import gtsam
from gtsam.symbol_shorthand import L, R, T

np.set_printoptions(precision=4, suppress=True)
rng = np.random.default_rng(6)

rows, cols, spacing = 4, 4, 2.0

# Boustrophedon ("lawnmower") visit order through the grid.
order, cell_index = [], {}
for row in range(rows):
    columns = range(cols) if row % 2 == 0 else reversed(range(cols))
    for col in columns:
        cell_index[(row, col)] = len(order)
        order.append((row, col))

positions = np.array([[col * spacing, row * spacing] for row, col in order])

# Each pose faces along its direction of travel.
truth = []
for k in range(len(order)):
    step = positions[min(k + 1, len(order) - 1)] - positions[max(k - 1, 0)]
    truth.append(
        gtsam.Pose2(positions[k][0], positions[k][1], np.arctan2(step[1], step[0]))
    )

# Landmarks sit off the grid, each seen from the poses within range.
landmarks = np.array([[-2.0, 2.0], [8.0, 1.0], [3.0, 8.0]])
sight_range = 6.0

odometry = [(k, k + 1) for k in range(len(order) - 1)]
sightings = [
    (k, b)
    for k in range(len(order))
    for b in range(len(landmarks))
    if np.linalg.norm(landmarks[b] - positions[k]) < sight_range
]
print(f"{len(order)} poses, {len(landmarks)} landmarks")
print(f"{len(odometry)} odometry edges, {len(sightings)} sightings")

kappa, tau, nu = 100.0, 25.0, 40.0  # rotation, odometry, sighting precisions

# Isotropic Langevin on rotation has sigma = 1 / sqrt(2 kappa) in the
# small-dispersion limit; the translation blocks are plain Gaussians.
sigma_R = 1.0 / np.sqrt(2.0 * kappa)
sigma_t = 1.0 / np.sqrt(tau)
sigma_v = 1.0 / np.sqrt(nu)

graph = gtsam.NonlinearFactorGraph()
rot_noise = gtsam.noiseModel.Isotropic.Variance(1, 1.0 / kappa)

for i, j in odometry:
    relative = truth[i].between(truth[j])
    measured = gtsam.Pose2(
        relative.x() + rng.normal(scale=sigma_t),
        relative.y() + rng.normal(scale=sigma_t),
        relative.theta() + rng.normal(scale=sigma_R),
    )
    graph.add(
        gtsam.FrobeniusBetweenFactorRot2(R(i), R(j), measured.rotation(), rot_noise)
    )
    graph.add(
        gtsam.RelativeTranslationFactor2(R(i), T(i), T(j), measured.translation(), tau)
    )

for k, b in sightings:
    # The sighting is expressed in the body frame of pose k.
    body = truth[k].transformTo(landmarks[b]) + rng.normal(scale=sigma_v, size=2)
    graph.add(gtsam.RelativeTranslationFactor2(R(k), T(k), L(b), body, nu))

print(f"rotation:  kappa = {kappa:g}, sigma = {np.rad2deg(sigma_R):.2f} deg")
print(f"odometry:  tau   = {tau:g}, sigma = {sigma_t:.2f} m")
print(f"sighting:  nu    = {nu:g}, sigma = {sigma_v:.2f} m")
print(f"{graph.size()} factors")

rank = 2  # p = d for the first staircase level

initial = gtsam.Values()
for index in range(len(truth)):
    perturbed = gtsam.Rot2.fromAngle(truth[index].theta() + rng.normal(scale=0.3))
    initial.insert(R(index), perturbed.matrix().T)                   # d x p
    initial.insert(T(index), rng.normal(scale=0.5, size=(1, rank)))  # 1 x p

# A landmark is one unconstrained row, the same shape as a translation.
for b in range(len(landmarks)):
    initial.insert(L(b), rng.normal(scale=0.5, size=(1, rank)))

print("rotation slice: ", initial.atMatrix(R(0)).shape)
print("landmark slice: ", initial.atMatrix(L(0)).shape)

alm_params = gtsam.AugmentedLagrangianParams()
alm_params.maxIterations = 100

params = gtsam.RiemannianStaircaseParams()
params.pMin = rank
params.pMax = 6
params.eta = 1e-3
params.setAlmParams(alm_params)

result = gtsam.RiemannianStaircaseOptimizer(graph, initial, params).optimize()

print(f"Certified:    {result.certified}")
print(f"Final rank:   {result.finalRank}")
print(f"Ranks tried:  {np.asarray(result.getRanksVisited()).astype(int)}")
print(f"Solver time:  {result.totalTime:.4f} s")

if not result.hasRoundedSolution():
    raise RuntimeError("The staircase did not return a rounded solution.")

rounded = result.roundedValues()
blocks = {i: rounded.atMatrix(R(i)).copy() for i in range(len(truth))}
rows = {i: rounded.atMatrix(T(i)).copy() for i in range(len(truth))}
marks = {b: rounded.atMatrix(L(b)).copy() for b in range(len(landmarks))}

# Certifiable methods optimize over O(d); undo a global reflection if present.
if sum(np.linalg.det(block) < 0.0 for block in blocks.values()) > len(truth) // 2:
    for block in blocks.values():
        block[:, -1] *= -1.0
    for row in list(rows.values()) + list(marks.values()):
        row[:, -1] *= -1.0

def project_to_so(matrix):
    """Return the closest proper rotation to a square matrix."""
    left, _, right_transpose = np.linalg.svd(matrix)
    correction = np.eye(matrix.shape[0])
    correction[-1, -1] = np.linalg.det(left @ right_transpose)
    return left @ correction @ right_transpose

def rotation_from_matrix(matrix):
    """Rebuild a Rot2 from a proper 2-by-2 rotation matrix."""
    return gtsam.Rot2.fromAngle(np.arctan2(matrix[1, 0], matrix[0, 0]))

estimate = gtsam.Values()
for index in range(len(truth)):
    estimate.insert(R(index), rotation_from_matrix(project_to_so(blocks[index].T)))
    estimate.insert(T(index), gtsam.Point2(rows[index][0]))
for b in range(len(landmarks)):
    estimate.insert(L(b), gtsam.Point2(marks[b][0]))

truth_values = gtsam.Values()
for index in range(len(truth)):
    truth_values.insert(R(index), truth[index].rotation())
    truth_values.insert(T(index), gtsam.Point2(truth[index].translation()))
for b in range(len(landmarks)):
    truth_values.insert(L(b), gtsam.Point2(landmarks[b]))

relaxation_bound = float(np.asarray(result.getCostPerLevel())[-1])
rounded_objective = graph.error(estimate)
gap = (rounded_objective - relaxation_bound) / max(relaxation_bound, 1e-12)

print(f"Relaxation lower bound: {relaxation_bound:.6f}")
print(f"Rounded objective:      {rounded_objective:.6f}")
print(f"Relative gap:           {gap:.3e}")
print(f"Objective at truth:     {graph.error(truth_values):.6f}")

# Each odometry edge contributes d + 1 residuals and each sighting d; the
# gauge removes one pose.
residuals = len(odometry) * 3 + len(sightings) * 2
parameters = len(truth) * 3 + len(landmarks) * 2 - 3
dof = residuals - parameters
print(f"residuals {residuals} - parameters {parameters} = {dof} dof")
print(f"Expected optimal cost (0.5 * chi^2): {0.5 * dof:.1f}")
print(f"Observed optimal cost:               {rounded_objective:.1f}")
print(f"Ratio:                               {rounded_objective / (0.5 * dof):.2f}")

def procrustes(source, target):
    """Best rigid alignment of source onto target."""
    source_centre, target_centre = source.mean(0), target.mean(0)
    A, B = source - source_centre, target - target_centre
    left, _, right_transpose = np.linalg.svd(A.T @ B)
    return (A @ (left @ right_transpose)) + target_centre

positions_est = np.array([estimate.atPoint2(T(i)) for i in range(len(truth))])
landmarks_est = np.array([estimate.atPoint2(L(b)) for b in range(len(landmarks))])

# The relaxation fixes no gauge, so align pose and landmark estimates together.
aligned = procrustes(np.vstack([positions_est, landmarks_est]),
                     np.vstack([positions, landmarks]))
aligned_poses = aligned[: len(positions)]
aligned_landmarks = aligned[len(positions) :]

figure = go.Figure()
figure.add_trace(go.Scatter(x=positions[:, 0], y=positions[:, 1], mode="lines+markers",
                            name="ground truth", line=dict(color="#2a78d6", width=2)))
figure.add_trace(go.Scatter(x=aligned_poses[:, 0], y=aligned_poses[:, 1], mode="lines+markers",
                            name="estimate", line=dict(color="#eb6834", width=2)))
figure.add_trace(go.Scatter(x=landmarks[:, 0], y=landmarks[:, 1], mode="markers",
                            name="landmarks (truth)",
                            marker=dict(color="#2a78d6", size=13, symbol="square-open",
                                        line=dict(width=2))))
figure.add_trace(go.Scatter(x=aligned_landmarks[:, 0], y=aligned_landmarks[:, 1],
                            mode="markers", name="landmarks (estimate)",
                            marker=dict(color="#eb6834", size=13, symbol="x")))
figure.update_layout(width=680, height=560, xaxis_title="x [m]", yaxis_title="y [m]",
                     yaxis=dict(scaleanchor="x", scaleratio=1))
figure.show()

error = np.linalg.norm(aligned_poses - positions, axis=1)
print(f"RMS position error after alignment: {np.sqrt((error ** 2).mean()):.4f} m")

ranks = np.asarray(result.getRanksVisited()).astype(int)
phases = {
    "QCQP build": np.asarray(result.getQcqpBuildTimePerLevel()),
    "local solve": np.asarray(result.getNlpTimePerLevel()),
    "verify": np.asarray(result.getVerifyTimePerLevel()),
}

timing = go.Figure()
for name, seconds in phases.items():
    timing.add_trace(go.Bar(x=[f"p = {level}" for level in ranks], y=seconds, name=name))
timing.update_layout(
    barmode="stack", title="Time per staircase level",
    xaxis_title="staircase rank", yaxis_title="seconds",
    height=380, margin=dict(l=0, r=0, t=40, b=0))
timing.show()

for name, seconds in phases.items():
    print(f"{name:>12}: {seconds.sum():.5f} s")