Add ADMM results

_scripts/generate_cliques.py

@@ -4,7 +4,7 @@
 import matplotlib.pylab as plt
 import numpy as np
 import scipy.sparse as sp
-from cert_tools.admm_clique import ADMMClique
+from cert_tools.admm_clique import ADMMClique, initialize_overlap
 from cert_tools.base_clique import BaseClique
 
 from lifters.matweight_lifter import MatWeightLocLifter
@@ -318,6 +318,7 @@ def create_clique_list_slam(
                 err = abs(x.T @ A @ x - b)
                 assert err < 1e-6, f"constraint {i}: {err}"
             cost_total += x.T @ clique.Q @ x
+    initialize_overlap(clique_list)
 
     if DEBUG:
         Q, y = lifter.get_Q()
@@ -328,4 +329,11 @@ def create_clique_list_slam(
         np.testing.assert_allclose(Q_mat.toarray(), Q.toarray())
         assert abs(cost_test - cost_total) / cost_total < 1e-5, (cost_test, cost_total)
 
+        for i, clique in enumerate(clique_list):
+            left = clique_list[i - 1].X if i > 0 else None
+            right = clique_list[i + 1].X if i < len(clique_list) - 1 else None
+            clique.generate_g(left=left, right=right)
+            error = clique.evaluate_F(clique.X)
+            np.testing.assert_allclose(error, 0)
+
     return clique_list

_scripts/run_time_study.py

@@ -77,18 +77,34 @@ def generate_results(lifter: MatWeightLocLifter, n_params_list=[10]):
         time_dict["t dSDP"] = time.time() - t1
         time_dict["cost dSDP"] = info["cost"]
 
+        x_dSDP, evr_mean = extract_solution(new_lifter, X_list)
+        time_dict["evr dSDP"] = evr_mean
+
         print("running ADMM...", end="")
         t1 = time.time()
+        # X0 = []
+        # for c in clique_list:
+        #     x_clique = new_lifter.get_x(theta=theta_est, var_subset=c.var_dict)
+        #     X0.append(np.outer(x_clique, x_clique))
+        X0 = None
+        print("target", info["cost"])
         X_list, info = solve_alternating(
             clique_list,
-            use_fusion=True,
-            verbose=VERBOSE,
+            X0=X0,
+            use_fusion=False,
+            rho_start=1e4,
+            verbose=True,
+            early_stop=False,
+            max_iter=20,
+            mu_rho=2.0,
+            tau_rho=2.0,
         )
+        print(info["msg"], end="...")
         time_dict["t ADMM"] = time.time() - t1
         time_dict["cost ADMM"] = info["cost"]
 
-        x_dSDP, evr_mean = extract_solution(new_lifter, X_list)
-        time_dict["evr dSDP"] = evr_mean
+        x_ADMM, evr_mean = extract_solution(new_lifter, X_list)
+        time_dict["evr ADMM"] = evr_mean
 
         print("creating constraints...", end="")
         t1 = time.time()
@@ -137,8 +153,8 @@ def generate_results(lifter: MatWeightLocLifter, n_params_list=[10]):
     n_params_list = np.logspace(1, 2, 10).astype(int)
     # n_params_list = np.arange(10, 101, step=10).astype(int)
     # n_params_list = np.arange(10, 19)  # runs out of memory at 19!
-    fname = f"_results/{lifter}_time_admm.pkl"
-    overwrite = False
+    fname = f"_results/{lifter}_time_test.pkl"
+    overwrite = True
 
     try:
         assert overwrite is False
@@ -150,19 +166,26 @@ def generate_results(lifter: MatWeightLocLifter, n_params_list=[10]):
     for label, plot in zip(["t", "cost"], [sns.scatterplot, sns.barplot]):
         df_long = df.melt(
             id_vars=["n params"],
-            value_vars=[f"{label} local", f"{label} dSDP", f"{label} SDP"],
+            value_vars=[
+                f"{label} local",
+                f"{label} dSDP",
+                f"{label} SDP",
+                f"{label} ADMM",
+            ],
             value_name=label,
             var_name="solver type",
         )
         fig, ax = plt.subplots()
         plot(df_long, x="n params", y=label, hue="solver type", ax=ax)
         ax.set_yscale("log")
-        ax.set_xscale("log")
+        if label != "cost":
+            ax.set_xscale("log")
         ax.grid("on")
 
     fig, ax = plt.subplots()
     ax.loglog(df["n params"], df["evr SDP"], label="SDP")
     ax.loglog(df["n params"], df["evr dSDP"], label="dSDP")
+    ax.loglog(df["n params"], df["evr ADMM"], label="ADMM")
     ax.legend()
     ax.grid("on")
     ax.set_ylabel("EVR")

lifters/matweight_lifter.py

@@ -145,7 +145,7 @@ def get_x(self, theta=None, parameters=None, var_subset=None):
                 t_ki_i = C_i0 @ t_k0_0 - t_i0_i
                 vec += [t_ki_i.flatten()]
             else:  # Other variables
-                vec += [theta[var].flatten()]
+                vec += [theta[var].flatten("F")]
         return np.hstack(vec)
 
     def get_dim_x(self, var_subset=None):