import json import time import numpy as np def joint_prune(z_u, z_v, eta, lam): """Exact proximal map for a shared paired 2-lambda column penalty.""" keep = (np.sum(z_u * z_u, axis=0) + np.sum(z_v * z_v, axis=0)) > 4.0 * eta * lam out_u, out_v = z_u.copy(), z_v.copy() out_u[:, ~keep] = 0.0 out_v[:, ~keep] = 0.0 return out_u, out_v, keep def balance(u, v, eps=0.0): """Reciprocal column rescaling; eps=0 is the exact mathematical map.""" nu = np.linalg.norm(u, axis=0) nv = np.linalg.norm(v, axis=0) a = np.sqrt(nv / (nu + eps)) return u * a[None, :], v / a[None, :] def factor_step(u, v, x, eta, lam, mu): residual = u @ v.T - x gu = residual @ v + mu * u z_u = u - eta * gu # alternating update as in the stated formula gv = (z_u @ v.T - x).T @ z_u + mu * v z_v = v - eta * gv u, v, keep = joint_prune(z_u, z_v, eta, lam) if np.any(keep): # balancing only retained pairs avoids 0/0 and preserves the product ub, vb = balance(u[:, keep], v[:, keep]) u[:, keep], v[:, keep] = ub, vb return u, v, keep def run_path(x, u0, v0, lambdas, eta=0.01, mu=1e-3, steps=250): u, v = u0.copy(), v0.copy() rows = [] for lam in lambdas: for _ in range(steps): u, v, _ = factor_step(u, v, x, eta, lam, mu) active = np.linalg.norm(u, axis=0) > 1e-10 rows.append({ "lambda": float(lam), "rank": int(active.sum()), "loss": float(0.5 * np.sum((x - u @ v.T) ** 2)), "objective": float(0.5 * np.sum((x - u @ v.T) ** 2) + 0.5 * mu * (np.sum(u*u) + np.sum(v*v)) + 2 * lam * active.sum()) }) return u, v, rows def main(): rng = np.random.default_rng(1352) # Low-rank positive target with deliberately excessive factor rank. n, m, true_rank, r = 16, 13, 3, 8 x = rng.normal(size=(n, true_rank)) @ rng.normal(size=(m, true_rank)).T x /= np.linalg.norm(x, 'fro') / np.sqrt(n*m) # SVD initialization gives useful columns plus weak redundant columns. p, s, qt = np.linalg.svd(x, full_matrices=False) u0 = np.zeros((n, r)); v0 = np.zeros((m, r)) for j in range(r): if j < len(s): u0[:, j] = p[:, j] * np.sqrt(max(s[j], 1e-12)) v0[:, j] = qt[j, :] * np.sqrt(max(s[j], 1e-12)) else: u0[:, j] = 0.03 * rng.normal(size=n) v0[:, j] = 0.03 * rng.normal(size=m) # Make the surplus columns genuinely weak but nonzero. u0[:, true_rank:] *= 0.12 v0[:, true_rank:] *= 0.12 # Prediction 1: exact joint threshold boundary, swept over scales/lambdas. eta = 0.01 threshold_errors = [] boundary_cases = [] for lam in np.geomspace(1e-4, 2.0, 24): threshold = 4 * eta * lam for scale in np.geomspace(0.05, 20.0, 17): zu = np.array([[scale, 1.0]]) zv = np.array([[1.0, scale]]) q = np.sum(zu*zu, axis=0) + np.sum(zv*zv, axis=0) _, _, keep = joint_prune(zu, zv, eta, lam) threshold_errors.append(abs(float(keep[0]) - float(q[0] > threshold))) if abs(q[0] - threshold) < max(threshold * 0.02, 1e-12): boundary_cases.append((q[0], threshold, bool(keep[0]))) prox_check = {"mismatches": int(sum(threshold_errors)), "tested": len(threshold_errors), "near_boundary_cases": boundary_cases[:4]} # Prediction 2: reciprocal scaling must preserve product and equalize norms. products, norm_gaps, scale_invariance = [], [], [] for _ in range(100): uu = rng.normal(size=(7, 1)); vv = rng.normal(size=(6, 1)) c = 10 ** rng.uniform(-6, 6) uu *= c; vv /= c ub, vb = balance(uu, vv, eps=0.0) products.append(np.linalg.norm(uu @ vv.T - ub @ vb.T, 'fro')) norm_gaps.append(abs(np.linalg.norm(ub) - np.linalg.norm(vb))) scale_invariance.append(abs(np.linalg.norm(ub @ vb.T, 'fro') - np.linalg.norm(uu @ vv.T, 'fro'))) balance_check = {"max_product_error": float(max(products)), "max_balanced_norm_gap": float(max(norm_gaps)), "max_product_norm_error": float(max(scale_invariance))} # Prediction 3: increasing lambda on a warm start removes columns monotonically. lambdas = [0.0, 0.02, 0.08, 0.2, 0.5, 1.0, 2.0, 4.0] t0 = time.perf_counter() _, _, path = run_path(x, u0, v0, lambdas, eta=eta, mu=1e-3, steps=300) warm_time = time.perf_counter() - t0 ranks = [row["rank"] for row in path] monotone = all(ranks[i+1] <= ranks[i] for i in range(len(ranks)-1)) # Standard comparison: independently optimized fixed-rank factorizations. baseline = [] for rr in [true_rank, r]: ub, vb = u0[:, :rr].copy(), v0[:, :rr].copy() t1 = time.perf_counter() for _ in range(sum([300] * len(lambdas))): # lambda=0 is ordinary factorized least squares with scale regularization ub, vb, _ = factor_step(ub, vb, x, eta, 0.0, 1e-3) baseline.append({"rank": rr, "loss": float(0.5*np.sum((x-ub@vb.T)**2)), "seconds": time.perf_counter()-t1}) selected = min((row for row in path if row["loss"] <= baseline[1]["loss"] * 1.05), key=lambda row: row["rank"], default=path[-1]) result = { "prox_prediction": prox_check, "balance_prediction": balance_check, "path_prediction": {"lambdas": lambdas, "ranks": ranks, "monotone_nonincreasing": monotone, "losses": [row["loss"] for row in path], "warm_seconds": warm_time}, "baseline": baseline, "prediction_summary": { "threshold_rule": {"predicted_mismatches": 0, "observed_mismatches": int(sum(threshold_errors)), "tested": len(threshold_errors)}, "balancing_rule": {"predicted_product_error": 0.0, "observed_max_product_error": float(max(products)), "predicted_norm_gap": 0.0, "observed_max_norm_gap": float(max(norm_gaps))}, "lambda_zero_no_pruning": {"predicted_rank": r, "observed_rank": ranks[0]}, "warm_path_rank": {"predicted_nonincreasing": True, "observed_nonincreasing": monotone, "observed_transition": f"{ranks[0]}->{ranks[-1]}"} }, "selected_path_point": selected, "setup": {"shape": [n,m], "true_rank": true_rank, "max_rank": r, "eta": eta, "steps_per_stage": 300} } print(json.dumps(result, indent=2)) if __name__ == '__main__': main()