Fixed orthogonal & PCA warm starts

import warnings

warnings.filterwarnings("ignore")

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
import numpy as np
import optax

import gauss_flows as gf
from _style import SCATTER_KW, style_ax

jax.config.update("jax_enable_x64", True)
rng = np.random.default_rng(0)

d = 3
n = 4000
M = jr.normal(jr.key(0), (d, d))
X = jr.normal(jr.key(1), (n, d)) @ M.T          # correlated, anisotropic

rot = gf.FixedRotation.from_data(X)
Q = np.asarray(rot.matrix)
Y = jax.vmap(rot.transform)(X)

x0 = jr.normal(jr.key(2), (d,))
z0, ld = rot.transform_and_log_det(x0)
xi, _ = rot.inverse_and_log_det(z0)
print(f"orthogonal? ||Q^T Q - I|| = {np.abs(Q.T @ Q - np.eye(d)).max():.1e}")
print(f"log_det = {float(ld):+.1e}   round-trip = {float(jnp.abs(x0 - xi).max()):.1e}")

cov_in = np.cov(np.asarray(X).T)
cov_out = np.cov(np.asarray(Y).T)
print(f"max |off-diagonal| covariance:  before = {np.abs(cov_in - np.diag(np.diag(cov_in))).max():.3f}"
      f"   after = {np.abs(cov_out - np.diag(np.diag(cov_out))).max():.2e}")

fig, axes = plt.subplots(1, 3, figsize=(14, 4.1))
axes[0].scatter(X[:, 0], X[:, 1], color="tab:blue", **SCATTER_KW)
axes[0].set(title="input (correlated)", xlabel="$x_0$", ylabel="$x_1$")
axes[1].scatter(Y[:, 0], Y[:, 1], color="tab:orange", **SCATTER_KW)
axes[1].set(title="after PCA frame $y=Qx$ (decorrelated)", xlabel="$y_0$", ylabel="$y_1$")
for ax in axes[:2]:
    ax.axhline(0, color="k", lw=0.6); ax.axvline(0, color="k", lw=0.6)
    ax.set_aspect("equal"); style_ax(ax)
im = axes[2].imshow(np.abs(cov_out), cmap="viridis")
axes[2].set(title="|Cov| after (diagonal)", xticks=range(d), yticks=range(d))
fig.colorbar(im, ax=axes[2], fraction=0.046)
fig.tight_layout()

Xc = X - X.mean(0)

def offdiag_loss_matrix(W):
    """Sum of squared off-diagonal covariance of W x — a decorrelation proxy."""
    C = (Xc @ W.T).T @ (Xc @ W.T) / n
    return jnp.sum((C - jnp.diag(jnp.diag(C))) ** 2)

W = jnp.asarray(Q)                                   # start at the (orthogonal) PCA frame
opt = optax.adam(2e-2)
state = opt.init(W)

@jax.jit
def step(W, state):
    loss, g = jax.value_and_grad(offdiag_loss_matrix)(W)
    upd, state = opt.update(g, state)
    return optax.apply_updates(W, upd), state, loss

ortho_err, logdets = [], []
for _ in range(400):
    W, state, _ = step(W, state)
    ortho_err.append(float(jnp.abs(W.T @ W - jnp.eye(d)).max()))
    logdets.append(float(jnp.linalg.slogdet(W)[1]))

print(f"unconstrained matrix after 400 steps:")
print(f"  ||W^T W - I|| = {ortho_err[-1]:.3f}   (left the orthogonal manifold)")
print(f"  log|det W|    = {logdets[-1]:+.3f}   (a flow would wrongly book this as 0)")

fig, ax = plt.subplots(figsize=(7.4, 4.2))
ax.plot(ortho_err, color="tab:red", lw=2, label=r"$\|W^\top W - I\|_\infty$ (orthogonality error)")
ax.plot(np.abs(logdets), color="tab:purple", lw=2, label=r"$|\log|\det W||$ (mis-booked log-det)")
ax.axhline(0, color="k", lw=0.8, ls="--")
ax.set(title="An unconstrained matrix drifts off $O(d)$ under training",
       xlabel="step", ylabel="error")
ax.legend(fontsize=8); style_ax(ax)
fig.tight_layout()

def decompose_to_reflectors(Q):
    """Reflection vectors P such that gauss_flows' product H(P[-1])...H(P[0]) = Q.

    Column-by-column Householder zeroing reduces Q to a sign diagonal; each
    remaining -1 is one more axis reflection. Returns P ordered for gf's
    convention (gf applies reflections in array order, composing right-to-left).
    """
    Q = np.asarray(Q, float)
    d = Q.shape[0]
    R = Q.copy()
    us = []
    for k in range(d - 1):
        x = R[k:, k].copy()
        e = np.zeros_like(x); e[0] = 1.0
        alpha = -np.sign(x[0] or 1.0) * np.linalg.norm(x)   # stable sign choice
        v = x - alpha * e
        v = v / np.linalg.norm(v)
        u = np.zeros(d); u[k:] = v
        us.append(u)
        R[k:, k:] -= 2.0 * np.outer(v, v @ R[k:, k:])       # R <- H(v) R
    for j in np.where(np.sign(np.diag(R)) < 0)[0]:          # absorb the sign diagonal
        ej = np.zeros(d); ej[j] = 1.0
        us.append(ej)
    # H(us[-1])...H(us[0]) Q = I  =>  Q = H(us[0])...H(us[-1]); gf wants reversed order
    return np.array(us[::-1])

P = decompose_to_reflectors(Q)
warm = eqx.tree_at(lambda m: m.params,
                   gf.HouseholderRotation(n_reflections=P.shape[0], shape=(d,)),
                   jnp.asarray(P))
Q_warm = jax.vmap(warm.transform)(jnp.eye(d)).T
print(f"decomposed PCA frame into {P.shape[0]} reflections")
print(f"warm-started HouseholderRotation reproduces Q? ||Q - Q_warm|| = "
      f"{float(jnp.abs(jnp.asarray(Q) - Q_warm).max()):.1e}")
print(f"matches FixedRotation output on data? max|Δ| = "
      f"{float(jnp.abs(jax.vmap(warm.transform)(X) - Y).max()):.1e}")

def offdiag_loss_layer(layer):
    Yl = jax.vmap(layer.transform)(Xc)
    C = Yl.T @ Yl / n
    return jnp.sum((C - jnp.diag(jnp.diag(C))) ** 2)

def train_layer(layer, steps=400, lr=2e-2):
    params, static = eqx.partition(layer, eqx.is_inexact_array)
    opt = optax.adam(lr); state = opt.init(params)

    @eqx.filter_jit
    def step(params, state):
        loss, g = eqx.filter_value_and_grad(
            lambda p: offdiag_loss_layer(eqx.combine(p, static)))(params)
        upd, state = opt.update(g, state)
        return eqx.apply_updates(params, upd), state, loss

    hist = []
    for _ in range(steps):
        params, state, loss = step(params, state)
        hist.append(float(loss))
    return hist

cold_hist = train_layer(gf.HouseholderRotation(n_reflections=d, shape=(d,)))
warm_hist = train_layer(warm)
print(f"cold start: off-diag loss {cold_hist[0]:.2e} -> {cold_hist[-1]:.2e}")
print(f"warm start: off-diag loss {warm_hist[0]:.2e} -> {warm_hist[-1]:.2e}  (begins at the answer)")

fig, ax = plt.subplots(figsize=(7.6, 4.3))
ax.semilogy(np.maximum(cold_hist, 1e-32), color="tab:blue", lw=2, label="cold start (random reflectors)")
ax.semilogy(np.maximum(warm_hist, 1e-32), color="tab:green", lw=2, label="warm start (PCA decomposition)")
ax.set(title="Warm-starting the rotation at the PCA frame",
       xlabel="step", ylabel="off-diagonal covariance loss")
ax.legend(fontsize=9); style_ax(ax)
fig.tight_layout()