Structured MVN sampling dispatch — keeping the root cheap

from __future__ import annotations

import time
import warnings
warnings.filterwarnings("ignore", message=r".*IProgress.*")

import einx
import jax
import jax.numpy as jnp
import lineax as lx
import matplotlib.pyplot as plt
import numpy as np
import gaussx
from gaussx import (
    BlockDiag,
    Kronecker,
    LowRankUpdate,
    MultivariateNormal,
    cholesky,
    sqrt as op_sqrt,
)

jax.config.update("jax_enable_x64", True)
KEY = jax.random.PRNGKey(0)
plt.rcParams.update({
    "figure.dpi": 110,
    "axes.grid": True,
    "axes.grid.which": "both",
    "xtick.minor.visible": True,
    "ytick.minor.visible": True,
    "grid.alpha": 0.3,
})


def psd_op(M):
    return lx.MatrixLinearOperator(M, lx.positive_semidefinite_tag)

n = 6
key, k_M, k_eps = jax.random.split(KEY, 3)
M = jax.random.normal(k_M, (n, n))
Sigma_dense = M @ M.T + 0.5 * jnp.eye(n)
Sigma_op    = psd_op(Sigma_dense)

L_op = cholesky(Sigma_op)
L    = L_op.as_matrix()
print("L type :", type(L_op).__name__)
print("L L^T - Sigma  ||F  =",
      f"{float(jnp.linalg.norm(L @ L.T - Sigma_dense)):.2e}")

# Sample 50000 draws and compare empirical covariance.
S = 50_000
mu = jnp.zeros(n)
eps = jax.random.normal(k_eps, (S, n))
X = mu + einx.dot("s a, b a -> s b", eps, L)
Sigma_emp = einx.dot("s i, s j -> i j", X, X) / S

print("dense  ||Sigma_emp - Sigma||_F  =",
      f"{float(jnp.linalg.norm(Sigma_emp - Sigma_dense)):.2e}    "
      f"(MC noise ~ ||Sigma||_F * n / sqrt(S))")

key, k_eps = jax.random.split(KEY)
n_diag = 8
diag_var = jnp.array([4.0, 1.0, 9.0, 0.25, 2.5, 0.5, 16.0, 1.5])
D_op = lx.DiagonalLinearOperator(diag_var)

L_D = cholesky(D_op)
print("L type :", type(L_D).__name__)
print("L diag :", np.asarray(L_D.diagonal))
print("L diag = sqrt(var)?", bool(jnp.allclose(L_D.diagonal, jnp.sqrt(diag_var))))

# Sample 50000 draws -- one element-wise multiply per draw.
S = 50_000
eps = jax.random.normal(k_eps, (S, n_diag))
X = jax.vmap(L_D.mv)(eps)
diag_emp = einx.var("s d -> d", X)
print("\nempirical diag(Sigma):", np.asarray(diag_emp))
print("true diag(Sigma)     :", np.asarray(diag_var))
print(f"max relative error    = {float(jnp.max(jnp.abs(diag_emp - diag_var) / diag_var)):.2e}")

n_A, n_B = 4, 5
key, k_A, k_B, k_eps = jax.random.split(KEY, 4)

MA = jax.random.normal(k_A, (n_A, n_A))
A  = MA @ MA.T + 0.3 * jnp.eye(n_A)
MB = jax.random.normal(k_B, (n_B, n_B))
B  = MB @ MB.T + 0.3 * jnp.eye(n_B)

A_op = psd_op(A); B_op = psd_op(B)
K_op = Kronecker(A_op, B_op)

L_K  = cholesky(K_op)
print("L type            :", type(L_K).__name__)
print("L_K factor types  :", [type(o).__name__ for o in L_K.operators])

# Sanity: chol(A) ⊗ chol(B) really is the Cholesky of A ⊗ B.
LA = jnp.linalg.cholesky(A); LB = jnp.linalg.cholesky(B)
LK_dense_ref = jnp.kron(LA, LB)
print("\nL_K - chol(A) ⊗ chol(B)  ||F  =",
      f"{float(jnp.linalg.norm(L_K.as_matrix() - LK_dense_ref)):.2e}")

# Sample via the *structured* Kronecker root: each draw is a matvec on the
# small factors -- never instantiate the (n_A n_B) x (n_A n_B) matrix.
S = 80_000
n = n_A * n_B
eps = jax.random.normal(k_eps, (S, n))
X = jax.vmap(L_K.mv)(eps)
Sigma_emp = einx.dot("s i, s j -> i j", X, X) / S
Sigma_true = jnp.kron(A, B)
print(f"||Sigma_emp - kron(A,B)||_F  = "
      f"{float(jnp.linalg.norm(Sigma_emp - Sigma_true)):.2e}    "
      f"(MC noise ~ ||Sigma||_F * n / sqrt(S))")

key, k_C, k_D, k_E, k_eps = jax.random.split(KEY, 5)
def make_psd(key, k):
    M = jax.random.normal(key, (k, k))
    return psd_op(M @ M.T + 0.3 * jnp.eye(k))

C_op = make_psd(k_C, 3)
D_op = make_psd(k_D, 4)
E_op = make_psd(k_E, 2)
BD_op = BlockDiag(C_op, D_op, E_op)

L_BD = cholesky(BD_op)
print("L type         :", type(L_BD).__name__)
print("inner blocks   :", [type(b).__name__ for b in L_BD.operators])

# Per-block reference Choleskys
LC = jnp.linalg.cholesky(C_op.as_matrix())
LD = jnp.linalg.cholesky(D_op.as_matrix())
LE = jnp.linalg.cholesky(E_op.as_matrix())

block_errs = [
    float(jnp.linalg.norm(L_BD.operators[k].as_matrix() - L_ref))
    for k, L_ref in enumerate([LC, LD, LE])
]
print("per-block ||L_dispatch - L_ref||_F:", [f"{e:.2e}" for e in block_errs])

# Monte-Carlo sample-covariance check
S = 60_000
n_total = 3 + 4 + 2
eps = jax.random.normal(k_eps, (S, n_total))
X   = jax.vmap(L_BD.mv)(eps)
Sigma_emp = einx.dot("s i, s j -> i j", X, X) / S
Sigma_true = jax.scipy.linalg.block_diag(C_op.as_matrix(), D_op.as_matrix(), E_op.as_matrix())
print(f"\n||Sigma_emp - block-diag||_F  = "
      f"{float(jnp.linalg.norm(Sigma_emp - Sigma_true)):.2e}    "
      f"(MC noise ~ ||Sigma||_F * n / sqrt(S))")

from gaussx import solve as gx_solve

N_blk, d_blk = 12, 3
key, k_eps = jax.random.split(KEY)

# Build a random SPD block-tridiagonal *precision* Lambda: diag blocks SPD,
# sub-diagonal blocks small enough that the whole matrix is SPD.
key, k_diag, k_sub = jax.random.split(key, 3)
diag_raw = jax.random.normal(k_diag, (N_blk, d_blk, d_blk))
diag_blk = einx.dot("n i a, n j a -> n i j", diag_raw, diag_raw) \
         + 3.0 * jnp.broadcast_to(jnp.eye(d_blk), (N_blk, d_blk, d_blk))
sub_blk  = 0.2 * jax.random.normal(k_sub, (N_blk - 1, d_blk, d_blk))

Lambda_op = gaussx.BlockTriDiag(diag_blk, sub_blk, tags=lx.positive_semidefinite_tag)
L_BT      = cholesky(Lambda_op)                          # Lambda = L L^T, L lower block-bidiagonal
print("Lambda type :", type(Lambda_op).__name__)
print("L type      :", type(L_BT).__name__)

# Cholesky round-trip: L L^T = Lambda
L_dense      = L_BT.as_matrix()
Lambda_dense = Lambda_op.as_matrix()
print(f"||L L^T - Lambda||_F  =  {float(jnp.linalg.norm(L_dense @ L_dense.T - Lambda_dense)):.2e}")

# Sampling target: x ~ N(0, Sigma) with Sigma = Lambda^{-1} (the actual Markov
# chain marginal). The recipe from [](#eq:blocktridiag-root) is x = L^{-T} eps,
# which we compute via gaussx.solve(L.T, eps) -- a banded triangular solve in
# O(N d^3), never materialising the dense Sigma.
n = N_blk * d_blk
S = 40_000
eps = jax.random.normal(k_eps, (S, n))
X = jax.vmap(lambda e: gx_solve(L_BT.T, e))(eps)
Sigma_emp = einx.dot("s i, s j -> i j", X, X) / S

# Reference Sigma = Lambda^{-1} (dense, only built here for the cross-check).
Sigma_dense_ref = jnp.linalg.inv(Lambda_dense)
print(f"||Sigma_emp - Lambda^-1||_F  =  "
      f"{float(jnp.linalg.norm(Sigma_emp - Sigma_dense_ref)):.2e}    "
      f"(MC noise ~ ||Sigma||_F * n / sqrt(S))")

# Sanity: multiplying eps by L instead would have given draws from N(0, Lambda),
# not N(0, Lambda^{-1}). Show the difference.
X_wrong = jax.vmap(L_BT.mv)(eps)
Cov_wrong = einx.dot("s i, s j -> i j", X_wrong, X_wrong) / S
print(f"||Cov(L @ eps) - Lambda||_F  =  "
      f"{float(jnp.linalg.norm(Cov_wrong - Lambda_dense)):.2e}    "
      f"(this would be wrong target; matches Lambda, not Lambda^-1)")

key, k_U, k_eps0, k_eps1 = jax.random.split(KEY, 4)
n_lr = 6
r    = 2                                # rank of the update

base_M = jax.random.normal(jax.random.PRNGKey(7), (n_lr, n_lr))
Sigma0 = base_M @ base_M.T + 0.4 * jnp.eye(n_lr)
Sigma0_op = psd_op(Sigma0)

U = jax.random.normal(k_U, (n_lr, r))
D = jnp.diag(jnp.array([1.0, 0.5]))     # rank-2 PSD diag
Sigma_full = Sigma0 + einx.dot("i a, a b, j b -> i j", U, D, U)

# Structured sampling: no Cholesky of the sum, just two independent draws.
L0 = jnp.linalg.cholesky(Sigma0)
sqrtD = jnp.sqrt(jnp.diag(D))
S = 80_000
eps0 = jax.random.normal(k_eps0, (S, n_lr))
eps1 = jax.random.normal(k_eps1, (S, r))
X = einx.dot("s a, b a -> s b", eps0, L0) \
  + einx.multiply("s b, b -> s b", eps1, sqrtD) @ U.T

Sigma_emp = einx.dot("s i, s j -> i j", X, X) / S
print(f"||Sigma_emp - (Sigma0 + U D U^T)||_F  = "
      f"{float(jnp.linalg.norm(Sigma_emp - Sigma_full)):.2e}    "
      f"(MC noise ~ ||Sigma||_F * n / sqrt(S))")

S_op = op_sqrt(Sigma_op)
S_mat = S_op.as_matrix()
print("symmetric: S = S^T?         ",
      bool(jnp.allclose(S_mat, S_mat.T, atol=1e-10)))
print("symmetric: S @ S = Sigma?   ",
      f"{float(jnp.linalg.norm(S_mat @ S_mat - Sigma_dense)):.2e}")
print("Cholesky:  L L^T = Sigma?   ",
      f"{float(jnp.linalg.norm(L @ L.T - Sigma_dense)):.2e}")

# Same noise, two roots -> two correlated point clouds.
n_dense = Sigma_dense.shape[0]
S_n = 4_000
key_, kc = jax.random.split(KEY)
eps_dense = jax.random.normal(kc, (S_n, n_dense))
X_chol = einx.dot("s a, b a -> s b", eps_dense, L)
X_sqrt = einx.dot("s a, b a -> s b", eps_dense, S_mat)

# Their *covariances* match (same Gaussian); their *coordinates* differ.
print("\nSame eps, different roots:")
print(f"  ||cov(X_chol) - cov(X_sqrt)||_F = "
      f"{float(jnp.linalg.norm(jnp.cov(X_chol.T) - jnp.cov(X_sqrt.T))):.2e}")
print(f"  ||X_chol - X_sqrt||_F           = "
      f"{float(jnp.linalg.norm(X_chol - X_sqrt)):.2e}    (large -- different points)")

def time_one(fn, n_runs=10):
    fn().block_until_ready()
    t0 = time.perf_counter()
    for _ in range(n_runs):
        fn().block_until_ready()
    return (time.perf_counter() - t0) / n_runs

sizes = [4, 8, 16, 24, 32]
S = 256
records = {"dense": [], "structured": [], "ref": []}

# Each timed path runs *factorisation + sampling* end-to-end so the speedups
# reflect what you'd actually feel as a user, not just the matvec. The dense
# path's Cholesky is recomputed inside f_dense, so the structured path's
# cholesky is also recomputed inside f_struct (and the explicit Kronecker
# product is rebuilt inside f_ref). Apples to apples.
for n_AB in sizes:
    key, ka, kb, ke = jax.random.split(jax.random.PRNGKey(123 + n_AB), 4)
    MA = jax.random.normal(ka, (n_AB, n_AB))
    A_  = MA @ MA.T + 0.3 * jnp.eye(n_AB)
    MB = jax.random.normal(kb, (n_AB, n_AB))
    B_  = MB @ MB.T + 0.3 * jnp.eye(n_AB)
    A_op_ = psd_op(A_); B_op_ = psd_op(B_)
    K_dense = jnp.kron(A_, B_)
    K_op_   = Kronecker(A_op_, B_op_)
    eps     = jax.random.normal(ke, (S, n_AB * n_AB))

    # 1) Dense path — Cholesky of the full Kronecker product, then sample.
    @jax.jit
    def f_dense():
        L = jnp.linalg.cholesky(K_dense)
        return einx.dot("s a, b a -> s b", eps, L)

    # 2) Structured: gaussx.cholesky returns Kronecker, sample via mv.
    #    Includes the cholesky call inside the timed function so it's
    #    measured the same way as f_dense above.
    @jax.jit
    def f_struct():
        L = cholesky(K_op_)
        return jax.vmap(L.mv)(eps)

    # 3) Reference: explicit kron of small Choleskys (factor + sample).
    @jax.jit
    def f_ref():
        LA_ = jnp.linalg.cholesky(A_); LB_ = jnp.linalg.cholesky(B_)
        L_kron = jnp.kron(LA_, LB_)
        return einx.dot("s a, b a -> s b", eps, L_kron)

    records["dense"     ].append(time_one(f_dense))
    records["structured"].append(time_one(f_struct))
    records["ref"       ].append(time_one(f_ref))

records = {k: np.asarray(v) for k, v in records.items()}
ns = np.array(sizes)
print(f"{'n_A=n_B':>8s} {'dense (s)':>12s} {'struct (s)':>12s} {'ref (s)':>10s}  speedup")
for i, n_AB in enumerate(sizes):
    sp = records["dense"][i] / records["structured"][i]
    print(f"{n_AB:>8d} {records['dense'][i]:>12.4f} {records['structured'][i]:>12.4f} "
          f"{records['ref'][i]:>10.4f}   {sp:>5.1f}x")

fig, ax = plt.subplots(figsize=(7.0, 4.0))
ax.plot(ns, records["dense"],      "o-", color="crimson",   lw=1.8, label=r"dense $(n_A n_B)^3$")
ax.plot(ns, records["structured"], "s-", color="steelblue", lw=1.8, label=r"gaussx.cholesky(Kronecker)")
ax.plot(ns, records["ref"],        "^-", color="forestgreen", lw=1.6, label=r"explicit chol(A) $\otimes$ chol(B)")
ax.set_xscale("log"); ax.set_yscale("log")
ax.set_xlabel(r"$n_A = n_B$  (so total size $= n_A n_B$)")
ax.set_ylabel("seconds per sample batch (S=256)")
ax.set_title("Kronecker dispatch keeps sampling cheap")
ax.legend(frameon=False, fontsize=8)
plt.tight_layout(); plt.show()