Recover stereo1D example

_scripts/run_stereo_study.py

@@ -1,5 +1,7 @@
 import matplotlib.pylab as plt
 import numpy as np
+from cert_tools.linalg_tools import rank_project
+from cert_tools.sdp_solvers import solve_sdp_cvxpy
 
 from auto_template.learner import Learner
 from auto_template.sim_experiments import (
@@ -8,9 +10,10 @@
     run_scalability_new,
     run_scalability_plot,
 )
+from lifters.stereo1d_lifter import Stereo1DLifter
 from lifters.stereo2d_lifter import Stereo2DLifter
 from lifters.stereo3d_lifter import Stereo3DLifter
-from utils.plotting_tools import savefig
+from utils.plotting_tools import plot_matrix, savefig
 
 # matplotlib.use("TkAgg")
 # plt.ion()
@@ -145,5 +148,101 @@ def run_all(n_seeds, recompute, tightness=True, scalability=True):
             stereo_tightness(d=3)
 
 
+def run_stereo_1d():
+    np.random.seed(0)
+    lifter = Stereo1DLifter(n_landmarks=2)
+    # lifter.theta = np.array([3.0])
+    # lifter.landmarks = np.array([2.3, 4.6])
+
+    print("theta:", lifter.theta, "landmarks:", lifter.landmarks)
+    print("x:", lifter.get_x())
+    # AutoTight
+
+    A_known_all = lifter.get_A_known(add_known_redundant=True)
+    A_known = lifter.get_A_known()
+    # shortcut: A_learned = lifter.get_A_learned_simple(A_known=A_known)
+
+    Y = lifter.generate_Y(factor=1.0)
+
+    # with A_known
+    basis, S = lifter.get_basis(Y, A_known=A_known)
+    print("with known:", S)
+    A_red = lifter.generate_matrices_simple(basis=basis)
+    fig, axs = plt.subplots(1, len(A_red) + len(A_known), sharey=True)
+    for i, Ai in enumerate(A_known):
+        title = f"$A_{{k,{i}}}$"
+        fig, ax, im = plot_matrix(Ai, ax=axs[i], colorbar=False)
+        ax.set_title(title)
+        print(f"known {i}", Ai.toarray())
+    for j, Ai in enumerate(A_red):
+        title = f"$A_{{\ell,{j}}}$"
+        fig, ax, im = plot_matrix(Ai, ax=axs[i + 1 + j], colorbar=False)
+        ax.set_title(title)
+        print(f"learned {j}", Ai.toarray())
+
+    # without A_known:
+    basis, S = lifter.get_basis(Y, A_known=[])
+    print("without known:", S)
+    A_all = lifter.generate_matrices_simple(basis=basis)
+    fig_raw, axs_raw = plt.subplots(1, len(A_all), sharey=True)
+    for i, Ai in enumerate(A_all):
+        title = f"$A_{{\ell,{i}}}$"
+        plot_matrix(Ai, ax=axs_raw[i], colorbar=False)
+        axs_raw[i].set_title(title)
+
+    # hard coded for better comparability
+    fig, axs = plt.subplots(3, len(A_all), sharey=True)
+    plot_matrix(A_all[0] * np.sqrt(2), ax=axs[0, 0], colorbar=False)
+    plot_matrix(A_all[1] * np.sqrt(2), ax=axs[0, 1], colorbar=False)
+    plot_matrix(-A_all[2], ax=axs[0, 2], colorbar=False)
+    axs[0, 0].set_title(f"$A_{{\ell,{1}}}$")
+    axs[0, 1].set_title(f"$A_{{\ell,{2}}}$")
+    axs[0, 2].set_title(f"$A_{{\ell,{3}}}$")
+    print(np.round(A_all[0].toarray() * np.sqrt(2), 4))
+    print(np.round(A_all[1].toarray() * np.sqrt(2), 4))
+    print(np.round(-A_all[2].toarray(), 4))
+
+    plot_matrix(A_known[0], ax=axs[1, 0], colorbar=False)
+    plot_matrix(A_known[1], ax=axs[1, 1], colorbar=False)
+    plot_matrix(A_red[0] * np.sqrt(2) / 2, ax=axs[1, 2], colorbar=False)
+    axs[1, 0].set_title(f"$A_{{k,{1}}}$")
+    axs[1, 1].set_title(f"$A_{{k,{2}}}$")
+    axs[1, 2].set_title(f"$A_{{\ell,{1}}}$")
+
+    plot_matrix(A_known_all[0], ax=axs[2, 0], colorbar=False)
+    plot_matrix(A_known_all[1], ax=axs[2, 1], colorbar=False)
+    plot_matrix(A_known_all[2], ax=axs[2, 2], colorbar=False)
+    axs[2, 0].set_title(f"$A_{{k,{1}}}$")
+    axs[2, 1].set_title(f"$A_{{k,{2}}}$")
+    axs[2, 2].set_title(f"$A_{{k,{3}}}$")
+
+    Q, y = lifter.get_Q()
+    fig, ax, im = plot_matrix(Q)
+    x = lifter.get_x()
+
+    print("theta gt:", lifter.theta)
+
+    theta_hat, info_local = lifter.local_solver(t_init=lifter.theta, y=y)
+    print("theta global:", theta_hat, "cost:", info_local["cost"])
+
+    # sanity check
+    x_hat = lifter.get_x(theta=theta_hat)
+    assert abs(x_hat.T @ Q @ x_hat - info_local["cost"]) / info_local["cost"] < 1e-10
+    for A_list, label in zip([A_known, A_all], ["known", "learned"]):
+        lifter.test_constraints(A_list)
+        Constraints = [(lifter.get_A0(), 1.0)] + [(Ai, 0.0) for Ai in A_list]
+        X, info_sdp = solve_sdp_cvxpy(Q=Q, Constraints=Constraints, verbose=False)
+        x_round, info_rank = rank_project(X, 1)
+        error = abs(X[0, 1] - theta_hat) / theta_hat
+        RDG = abs(info_local["cost"] - info_sdp["cost"])
+        print(
+            f"{label}, theta sdp:{X[0,1]}, theta rounded:{x_round[1]} EVR:{info_rank['EVR']}, cost:{info_sdp['cost']}"
+        )
+        print("eigvals:", np.linalg.eigvalsh(X))
+        print("RDG:", RDG)
+        print("error", error)
+
+
 if __name__ == "__main__":
-    run_all(n_seeds=1, recompute=True)
+    run_stereo_1d()
+    # run_all(n_seeds=1, recompute=True)

lifters/stereo1d_lifter.py

@@ -5,6 +5,7 @@
 
 class Stereo1DLifter(StateLifter):
     PARAM_LEVELS = ["no", "p"]
+    NOISE = 0.1
 
     def __init__(self, n_landmarks, param_level="no"):
         self.n_landmarks = n_landmarks
@@ -119,9 +120,12 @@ def get_param_idx_dict(self, var_subset=None):
             param_dict_[f"p_{n}"] = n + 1
         return param_dict_
 
-    def get_Q(self, noise: float = 1e-3) -> tuple:
+    def get_Q(self, noise: float = None) -> tuple:
         from poly_matrix.least_squares_problem import LeastSquaresProblem
 
+        if noise is None:
+            noise = self.NOISE
+
         y = 1 / (self.theta - self.landmarks) + np.random.normal(
             scale=noise, loc=0, size=self.n_landmarks
         )
@@ -168,6 +172,7 @@ def get_cost(self, t, y):
     def local_solver(
         self, t_init, y, num_iters=100, eps=1e-5, W=None, verbose=False, **kwargs
     ):
+        info = {}
         a = self.landmarks
         x_op = t_init
         for i in range(num_iters):
@@ -180,11 +185,13 @@ def local_solver(
                 x_op = x_op + dx
                 if np.abs(dx) < eps:
                     msg = f"converged dx after {i} it"
-                    return x_op, msg, self.get_cost(x_op, y)
+                    info = {"msg": msg, "cost": self.get_cost(x_op, y)}
+                    return x_op, info
             else:
                 msg = f"converged in du after {i} it"
-                return x_op, msg, self.get_cost(x_op, y)
-        return None, "didn't converge", None
+                info = {"msg": msg, "cost": self.get_cost(x_op, y)}
+                return x_op, info
+        return None, {"msg": "didn't converge", "cost": None}
 
     def __repr__(self):
         return f"stereo1d_{self.param_level}"