alsgls 0.1.0__tar.gz

alsgls-0.1.0/PKG-INFO

@@ -0,0 +1,78 @@
+Metadata-Version: 2.4
+Name: alsgls
+Version: 0.1.0
+Summary: Lightweight low-rank+diag GLS/SUR via ALS with EM baseline
+Author-email: Gaurav Sood <contact@gsood.com>
+License: MIT
+Project-URL: Homepage, https://github.com/finite-sample/alsgls
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.21
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+
+## A Lightweight ALS Solver for Iterative GLS
+
+When a GLS problem involves hundreds of equations, the $K × K$ covariance matrix becomes the computational bottleneck.  A simple statistical remedy is to assume that most of the cross‑equation dependence can be captured by a *handful of latent factors* plus equation‑specific noise.  This “low‑rank + diagonal” assumption slashes the number of unknowns from roughly $K^²$ to about $K×k$ parameters, where **k** (the latent factor rank) is much smaller than $K$.  The model alone, however, does **not** guarantee speed: we still have to fit the parameters.
+
+### Installation
+
+Install the library from PyPI:
+
+```bash
+pip install alsgls
+```
+
+For local development, clone the repo and use an editable install:
+
+```bash
+pip install -e .
+```
+
+### Usage
+
+```python
+from alsgls import als_gls, simulate_sur, nll_per_row, XB_from_Blist
+
+Xs_tr, Y_tr, Xs_te, Y_te = simulate_sur(N_tr=240, N_te=120, K=60, p=3, k=4)
+B, F, D, mem, _ = als_gls(Xs_tr, Y_tr, k=4)
+Yhat_te = XB_from_Blist(Xs_te, B)
+nll = nll_per_row(Y_te - Yhat_te, F, D)
+```
+
+See `examples/compare_als_vs_em.py` for a complete ALS versus EM comparison.
+
+### Documentation and notebooks
+
+Background material and reproducible experiments are available in the notebooks under [`als_sim/`](als_sim/), such as [`als_sim/als_comparison.ipynb`](als_sim/als_comparison.ipynb) and [`als_sim/als_sur.ipynb`](als_sim/als_sur.ipynb).
+
+### Solving low‑rank GLS: EM versus ALS
+
+The classic EM algorithm alternates between updating the regression coefficients $\beta$ and updating the factor loadings $F$ and the diagonal noise $D$.  Even though $\hat{\Sigma}$ is low‑rank, EM’s M‑step recreates the **full** $K × K$ inverse, wiping out the memory win.
+
+An alternative is **Alternating‑Least‑Squares (ALS)**. The Woodbury identity reduces the expensive inverse to a tiny k × k system, and the β‑update can be written without explicitly forming the dense matrix at all.  In practice, ALS converges in 5–6 sweeps and never allocates more than $O(K k)$ memory, while EM allocates $O(K^²)$.
+
+**Rule of thumb:** if your GLS routine keeps looping between $\beta$ and a fresh $\hat{\Sigma}$, replacing the $\hat{\Sigma}$‑update by a factor‑ALS step yields the same statistical fit with an order‑of‑magnitude smaller memory footprint.
+
+### Beyond SUR: where the idea travels
+
+Random‑effects models, feasible GLS with estimated heteroskedastic weights, optimal‑weight GMM, and spatial autoregressive GLS all iterate β ↔ Σ̂.  Each can adopt the same ALS trick: treat the weight matrix as low‑rank + diagonal, invert only the k × k core, and avoid the dense K × K algebra.  Memory savings in published examples range from 5× to 20×, depending on k.
+
+### A concrete case‑study: Seemingly‑Unrelated Regressions
+
+To show the magnitude, we ran a Monte‑Carlo experiment with N = 300 observations, three regressors, rank‑3 factors, and K set to 50, 80, 120.  EM was given 45 iterations; ALS, six sweeps.  The largest array EM holds is the dense Σ⁻¹, whereas ALS’s largest is the skinny factor matrix F.  The table summarises six replications:
+
+|   K | β‑RMSE EM | β‑RMSE ALS | Peak MB EM | Peak MB ALS | Memory ratio |
+| --: | :-------: | :--------: | ---------: | ----------: | -----------: |
+|  50 |   0.021   |    0.021   |     0.020  |      0.002  |         10×  |
+|  80 |   0.020   |    0.020   |     0.051  |      0.003  |         17×  |
+| 120 |   0.020   |    0.020   |     0.115  |      0.004  |         29×  |
+
+Statistically, the two estimators are indistinguishable (paired‑test p ≥ 0.14).  Computationally, ALS needs only a few megabytes whereas EM needs tens to hundreds.
+
+### 5  Choosing a solver in practice
+
+For small systems ($K < 50$), dense GLS or even separate OLS is fine.  Between 50 and 300 equations, a low‑rank **factor‑ALS** solver gives the same estimates at roughly one‑tenth the memory and runs happily on a GPU.  Once K enters the hundreds, any dense inverse becomes prohibitive; structured approaches such as factor‑ALS or sparse/banded $\hat{\Sigma}$ are mandatory.

alsgls-0.1.0/README.md

@@ -0,0 +1,62 @@
+## A Lightweight ALS Solver for Iterative GLS
+
+When a GLS problem involves hundreds of equations, the $K × K$ covariance matrix becomes the computational bottleneck.  A simple statistical remedy is to assume that most of the cross‑equation dependence can be captured by a *handful of latent factors* plus equation‑specific noise.  This “low‑rank + diagonal” assumption slashes the number of unknowns from roughly $K^²$ to about $K×k$ parameters, where **k** (the latent factor rank) is much smaller than $K$.  The model alone, however, does **not** guarantee speed: we still have to fit the parameters.
+
+### Installation
+
+Install the library from PyPI:
+
+```bash
+pip install alsgls
+```
+
+For local development, clone the repo and use an editable install:
+
+```bash
+pip install -e .
+```
+
+### Usage
+
+```python
+from alsgls import als_gls, simulate_sur, nll_per_row, XB_from_Blist
+
+Xs_tr, Y_tr, Xs_te, Y_te = simulate_sur(N_tr=240, N_te=120, K=60, p=3, k=4)
+B, F, D, mem, _ = als_gls(Xs_tr, Y_tr, k=4)
+Yhat_te = XB_from_Blist(Xs_te, B)
+nll = nll_per_row(Y_te - Yhat_te, F, D)
+```
+
+See `examples/compare_als_vs_em.py` for a complete ALS versus EM comparison.
+
+### Documentation and notebooks
+
+Background material and reproducible experiments are available in the notebooks under [`als_sim/`](als_sim/), such as [`als_sim/als_comparison.ipynb`](als_sim/als_comparison.ipynb) and [`als_sim/als_sur.ipynb`](als_sim/als_sur.ipynb).
+
+### Solving low‑rank GLS: EM versus ALS
+
+The classic EM algorithm alternates between updating the regression coefficients $\beta$ and updating the factor loadings $F$ and the diagonal noise $D$.  Even though $\hat{\Sigma}$ is low‑rank, EM’s M‑step recreates the **full** $K × K$ inverse, wiping out the memory win.
+
+An alternative is **Alternating‑Least‑Squares (ALS)**. The Woodbury identity reduces the expensive inverse to a tiny k × k system, and the β‑update can be written without explicitly forming the dense matrix at all.  In practice, ALS converges in 5–6 sweeps and never allocates more than $O(K k)$ memory, while EM allocates $O(K^²)$.
+
+**Rule of thumb:** if your GLS routine keeps looping between $\beta$ and a fresh $\hat{\Sigma}$, replacing the $\hat{\Sigma}$‑update by a factor‑ALS step yields the same statistical fit with an order‑of‑magnitude smaller memory footprint.
+
+### Beyond SUR: where the idea travels
+
+Random‑effects models, feasible GLS with estimated heteroskedastic weights, optimal‑weight GMM, and spatial autoregressive GLS all iterate β ↔ Σ̂.  Each can adopt the same ALS trick: treat the weight matrix as low‑rank + diagonal, invert only the k × k core, and avoid the dense K × K algebra.  Memory savings in published examples range from 5× to 20×, depending on k.
+
+### A concrete case‑study: Seemingly‑Unrelated Regressions
+
+To show the magnitude, we ran a Monte‑Carlo experiment with N = 300 observations, three regressors, rank‑3 factors, and K set to 50, 80, 120.  EM was given 45 iterations; ALS, six sweeps.  The largest array EM holds is the dense Σ⁻¹, whereas ALS’s largest is the skinny factor matrix F.  The table summarises six replications:
+
+|   K | β‑RMSE EM | β‑RMSE ALS | Peak MB EM | Peak MB ALS | Memory ratio |
+| --: | :-------: | :--------: | ---------: | ----------: | -----------: |
+|  50 |   0.021   |    0.021   |     0.020  |      0.002  |         10×  |
+|  80 |   0.020   |    0.020   |     0.051  |      0.003  |         17×  |
+| 120 |   0.020   |    0.020   |     0.115  |      0.004  |         29×  |
+
+Statistically, the two estimators are indistinguishable (paired‑test p ≥ 0.14).  Computationally, ALS needs only a few megabytes whereas EM needs tens to hundreds.
+
+### 5  Choosing a solver in practice
+
+For small systems ($K < 50$), dense GLS or even separate OLS is fine.  Between 50 and 300 equations, a low‑rank **factor‑ALS** solver gives the same estimates at roughly one‑tenth the memory and runs happily on a GPU.  Once K enters the hundreds, any dense inverse becomes prohibitive; structured approaches such as factor‑ALS or sparse/banded $\hat{\Sigma}$ are mandatory.

alsgls-0.1.0/alsgls/__init__.py

@@ -0,0 +1,7 @@
+from .als import als_gls
+from .em import em_gls
+from .metrics import mse, nll_per_row
+from .sim import simulate_sur, simulate_gls
+from .ops import XB_from_Blist
+
+__all__ = ["als_gls", "em_gls", "mse", "nll_per_row", "simulate_sur", "simulate_gls", "XB_from_Blist"]

alsgls-0.1.0/alsgls/als.py

@@ -0,0 +1,112 @@
+import numpy as np
+from .ops import (
+    apply_siginv_to_matrix, woodbury_pieces,
+    stack_B_list, unstack_B_vec, XB_from_Blist, cg_solve
+)
+
+def als_gls(
+    Xs, Y, k,
+    lam_F=1e-3, lam_B=1e-3, sweeps=8, d_floor=1e-8,
+    cg_maxit=800, cg_tol=3e-7
+):
+    """
+    ALS for low-rank+diag GLS.
+    Returns (B_list, F, D, mem_MB_est, info)
+    """
+    # Input validation
+    if not isinstance(Xs, list) or len(Xs) == 0:
+        raise ValueError("Xs must be a non-empty list of arrays")
+    if Y.ndim != 2:
+        raise ValueError("Y must be a 2D array")
+    N, K = Y.shape
+    if len(Xs) != K:
+        raise ValueError(f"Number of X matrices ({len(Xs)}) must match Y columns ({K})")
+    for j, X in enumerate(Xs):
+        if X.ndim != 2 or X.shape[0] != N:
+            raise ValueError(f"X[{j}] must be 2D with {N} rows")
+    if not (1 <= k <= min(K, N)):
+        raise ValueError(f"k must be between 1 and min(K={K}, N={N})")
+    if lam_F < 0 or lam_B < 0:
+        raise ValueError("Regularization parameters must be non-negative")
+    
+    p_list = [X.shape[1] for X in Xs]
+
+    # init B (OLS per equation)
+    B = []
+    for j, X in enumerate(Xs):
+        XtX = X.T @ X + lam_B * np.eye(X.shape[1])
+        Xty = X.T @ Y[:, [j]]
+        B.append(np.linalg.solve(XtX, Xty))
+
+    R = Y - XB_from_Blist(Xs, B)
+    # PCA-like init for F with relative threshold
+    U, s, Vt = np.linalg.svd(R, full_matrices=False)
+    s_thresh = max(s[0] * 1e-10, 1e-8) if len(s) > 0 else 1e-8
+    r = min(k, (s > s_thresh).sum() or 1)
+    F = Vt.T[:, :r] * np.sqrt(np.maximum(s[:r], 1e-12))
+    if r < k:
+        F = np.pad(F, ((0, 0), (0, k - r)))
+    D = np.maximum(np.var(R, axis=0), d_floor)
+
+    # main ALS loop
+    prev = None
+    cg_info = None
+    for _ in range(sweeps):
+        # Precompute Woodbury pieces once per sweep
+        Dinv, Cf = woodbury_pieces(F, D)
+
+        def A_mv(bvec):
+            """Matrix-free normal operator H(B) = X^T Σ^{-1} X · b + lam_B b"""
+            B_dir = unstack_B_vec(bvec, p_list)
+            M = XB_from_Blist(Xs, B_dir)                 # N x K
+            S = apply_siginv_to_matrix(M, F, D, Dinv=Dinv, Cf=Cf)  # N x K
+            out_blocks = []
+            for j, X in enumerate(Xs):
+                out_blocks.append(X.T @ S[:, [j]])
+            out = np.concatenate(out_blocks, axis=0).ravel()
+            return out + lam_B * bvec
+
+        # simple diagonal preconditioner: approx diag of H
+        def M_pre(v):
+            diag_entries = []
+            # Rough diag: X_j^T (Σ^{-1} e_j e_j^T) X_j ≈ X_j^T (Dinv_j) X_j
+            for j, X in enumerate(Xs):
+                w = float(Dinv[j])
+                diag_entries.extend([w] * X.shape[1])
+            d = np.array(diag_entries) + lam_B
+            return v / np.maximum(d, 1e-8)
+
+        # β-step via CG
+        rhs_blocks = []
+        S_y = apply_siginv_to_matrix(Y, F, D, Dinv=Dinv, Cf=Cf)
+        for j, X in enumerate(Xs):
+            rhs_blocks.append(X.T @ S_y[:, [j]])
+        b = np.concatenate(rhs_blocks, axis=0).ravel()
+        bvec0 = stack_B_list(B)
+        bvec, cg_info = cg_solve(A_mv, b, x0=bvec0, maxit=cg_maxit, tol=cg_tol, M_pre=M_pre)
+        B = unstack_B_vec(bvec, p_list)
+
+        # factor step
+        R = Y - XB_from_Blist(Xs, B)
+        # update U (scores) and F (loadings) by two ridge solves
+        FtF = F.T @ F + lam_F * np.eye(F.shape[1])
+        U = R @ F @ np.linalg.inv(FtF)
+        UtU = U.T @ U + lam_F * np.eye(F.shape[1])
+        F = R.T @ U @ np.linalg.inv(UtU)
+
+        # diagonal noise
+        D = np.maximum(np.mean((R - U @ F.T) ** 2, axis=0), d_floor)
+
+        # cheap objective proxy: per-row NLL (using current F,D and R)
+        obj = 0.5 * np.mean((R * (1.0 / np.sqrt(np.maximum(D, 1e-12)))) ** 2)  # weighted residuals
+        if np.isfinite(obj):
+            if prev is not None:
+                rel = (prev - obj) / max(1.0, abs(prev))
+                if rel < 1e-6:
+                    break
+            prev = obj
+
+    # Memory estimate: F (K×k) + D (K) + U (N×k) + intermediate matrices
+    mem_mb_est = (K * F.shape[1] + K + N * F.shape[1] + K * F.shape[1]) * 8 / 1e6
+    info = {"p_list": p_list, "cg": cg_info}
+    return B, F, D, mem_mb_est, info

alsgls-0.1.0/alsgls/em.py

@@ -0,0 +1,98 @@
+import numpy as np
+from .ops import XB_from_Blist
+
+def em_gls(Xs, Y, k, lam_F=1e-3, lam_B=1e-3, iters=30, d_floor=1e-8):
+    """
+    Dense-ish EM baseline for low-rank+diag GLS.
+    Builds Σ^{-1} explicitly (KxK) in the β-step to mimic O(K^2) memory.
+    Returns (B_list, F, D, mem_MB_est, info)
+    """
+    # Input validation
+    if not isinstance(Xs, list) or len(Xs) == 0:
+        raise ValueError("Xs must be a non-empty list of arrays")
+    if Y.ndim != 2:
+        raise ValueError("Y must be a 2D array")
+    N, K = Y.shape
+    if len(Xs) != K:
+        raise ValueError(f"Number of X matrices ({len(Xs)}) must match Y columns ({K})")
+    for j, X in enumerate(Xs):
+        if X.ndim != 2 or X.shape[0] != N:
+            raise ValueError(f"X[{j}] must be 2D with {N} rows")
+    if not (1 <= k <= min(K, N)):
+        raise ValueError(f"k must be between 1 and min(K={K}, N={N})")
+    if lam_F < 0 or lam_B < 0:
+        raise ValueError("Regularization parameters must be non-negative")
+    
+    p_list = [X.shape[1] for X in Xs]
+
+    # init B (OLS per equation)
+    B = []
+    for j, X in enumerate(Xs):
+        XtX = X.T @ X + lam_B * np.eye(X.shape[1])
+        Xty = X.T @ Y[:, [j]]
+        B.append(np.linalg.solve(XtX, Xty))
+
+    R = Y - XB_from_Blist(Xs, B)
+    U, s, Vt = np.linalg.svd(R, full_matrices=False)
+    s_thresh = max(s[0] * 1e-10, 1e-8) if len(s) > 0 else 1e-8
+    r = min(k, (s > s_thresh).sum() or 1)
+    F = Vt.T[:, :r] * np.sqrt(np.maximum(s[:r], 1e-12))
+    if r < k:
+        F = np.pad(F, ((0, 0), (0, k - r)))
+    D = np.maximum(np.var(R, axis=0), d_floor)
+
+    # Precompute Gram blocks X_j^T X_l.
+    # Only upper-triangular blocks are formed and the lower triangle is
+    # recovered via symmetry when assembling the normal-equation matrix.
+    G = [[None] * K for _ in range(K)]
+    for j in range(K):
+        for l in range(j, K):
+            G[j][l] = Xs[j].T @ Xs[l]
+
+    for _ in range(iters):
+        # E-step-like: nothing explicit (we directly update F,D after β)
+
+        # β-step (dense normal equations using Σ^{-1})
+        Dinv = 1.0 / np.clip(D, 1e-12, None)
+        M = F.T @ (F * Dinv[:, None])             # k x k
+        Cf = np.linalg.inv(np.eye(k) + M)
+        # Build Σ^{-1} explicitly (KxK)
+        Sigma_inv = np.diag(Dinv) - (F * Dinv[:, None]) @ Cf @ (F.T * Dinv[None, :])
+
+        A = np.zeros((sum(p_list), sum(p_list)))
+        rhs = np.zeros((sum(p_list), 1))
+        p_offsets = np.cumsum([0] + p_list)
+        # Blocks A_{j,l} = Σ^{-1}_{l,j} * X_j^T X_l (symmetric in j,l)
+        for j in range(K):
+            Sj = Sigma_inv[:, j]
+            r0, r1 = p_offsets[j], p_offsets[j + 1]
+            rhs[r0:r1, :] = Xs[j].T @ (Y @ Sj.reshape(-1, 1))
+            for l in range(j, K):
+                c0, c1 = p_offsets[l], p_offsets[l + 1]
+                block = G[j][l]
+                scalar = Sj[l]
+                A[r0:r1, c0:c1] = scalar * block
+                if l != j:
+                    # Mirror to the symmetric block to maintain A symmetric
+                    A[c0:c1, r0:r1] = scalar * block.T
+        A += lam_B * np.eye(A.shape[0])
+        A = (A + A.T) * 0.5  # enforce symmetry
+        bvec = np.linalg.solve(A, rhs).ravel()
+        B = []
+        i = 0
+        for p in p_list:
+            B.append(bvec[i:i+p].reshape(p, 1))
+            i += p
+
+        # Update residuals and then F, D
+        R = Y - XB_from_Blist(Xs, B)
+        # Update scores/loadings by two ridge solves
+        FtF = F.T @ F + lam_F * np.eye(F.shape[1])
+        Uhat = R @ F @ np.linalg.inv(FtF)
+        UtU = Uhat.T @ Uhat + lam_F * np.eye(F.shape[1])
+        F = R.T @ Uhat @ np.linalg.inv(UtU)
+        D = np.maximum(np.mean((R - Uhat @ F.T) ** 2, axis=0), d_floor)
+
+    mem_mb_est = (K * K) * 8 / 1e6  # explicit Σ^{-1}
+    info = {"p_list": p_list}
+    return B, F, D, mem_mb_est, info

alsgls-0.1.0/alsgls/metrics.py

@@ -0,0 +1,39 @@
+import numpy as np
+from .ops import woodbury_pieces
+
+def mse(Y, Yhat):
+    return float(np.mean((Y - Yhat) ** 2))
+
+def nll_per_row(R, F, D):
+    """Negative log-likelihood per row for residual matrix ``R``.
+
+    Parameters
+    ----------
+    R : ndarray (N × K)
+        Residual matrix.
+    F : ndarray (K × k)
+        Factor loadings.
+    D : ndarray (K,)
+        Diagonal noise variances.
+
+    Returns
+    -------
+    float
+        ``0.5 * [tr(R Σ^{-1} R^T) + logdet(Σ) + K log(2π)]``
+        averaged over rows.
+    """
+    K = R.shape[1]
+    Dinv, Cf = woodbury_pieces(F, D)
+    # tr(R Σ^{-1} R^T) = sum over rows of r Σ^{-1} r^T
+    # Efficiently: R Σ^{-1} = apply_siginv_to_matrix(R, F, D), but avoid circular import.
+    # Inline Woodbury:
+    RDinv = R * Dinv[None, :]
+    T1 = RDinv @ F          # N x k
+    T2 = T1 @ Cf            # N x k
+    RSinv = RDinv - T2 @ (F.T * Dinv)  # N x K
+    quad = float(np.sum(RSinv * R))
+    # logdet via matrix determinant lemma:
+    # det(FF^T + D) = det(D) det(I + F^T D^{-1} F)
+    logdet = float(np.sum(np.log(np.clip(D, 1e-12, None)))) \
+             + float(np.linalg.slogdet(np.eye(F.shape[1]) + F.T @ (F * Dinv[:, None]))[1])
+    return 0.5 * (quad / R.shape[0] + logdet + K * np.log(2 * np.pi))

alsgls-0.1.0/alsgls/ops.py

@@ -0,0 +1,87 @@
+import numpy as np
+
+def woodbury_pieces(F: np.ndarray, D: np.ndarray):
+    """
+    Return Dinv, Cf used in Woodbury:
+    Σ = F F^T + diag(D)
+    Σ^{-1} = D^{-1} - D^{-1} F (I + F^T D^{-1} F)^{-1} F^T D^{-1}
+    """
+    D = np.asarray(D)
+    Dinv = 1.0 / np.clip(D, 1e-12, None)
+    FtDinv = (F.T * Dinv)           # k x K (row-scale F.T by Dinv)
+    M = FtDinv @ F                  # k x k == F^T D^{-1} F (reuse FtDinv to avoid re-scaling F)
+    # solve small kxk: (I + M)^{-1}
+    Cf = np.linalg.inv(np.eye(F.shape[1]) + M)
+    return Dinv, Cf
+
+def apply_siginv_to_matrix(M: np.ndarray, F: np.ndarray, D: np.ndarray, *, Dinv=None, Cf=None):
+    """Right-multiply an N×K matrix ``M`` by ``Σ^{-1}`` using Woodbury.
+
+    Parameters
+    ----------
+    M : np.ndarray
+        Matrix to be multiplied on the right by ``Σ^{-1}``.
+    F : np.ndarray
+        Low-rank factor matrix.
+    D : np.ndarray
+        Diagonal entries of the noise covariance.
+    Dinv : np.ndarray, optional
+        Precomputed ``1/D`` vector.  If ``None`` (default), it will be
+        computed internally via :func:`woodbury_pieces`.
+    Cf : np.ndarray, optional
+        Precomputed ``(I + F^T D^{-1} F)^{-1}``.  If ``None`` (default), it
+        will be computed internally via :func:`woodbury_pieces`.
+    """
+    if Dinv is None or Cf is None:
+        Dinv, Cf = woodbury_pieces(F, D)
+    # M * Dinv - M*(Dinv F) Cf (F^T Dinv)
+    MDinv = M * Dinv[None, :]
+    T1 = MDinv @ F              # N x k
+    T2 = T1 @ Cf                # N x k
+    T3 = T2 @ (F.T * Dinv)      # N x K
+    return MDinv - T3
+
+def stack_B_list(B_list):
+    """Stack list of (p_j,1) into flat vector."""
+    return np.concatenate([b.ravel() for b in B_list], axis=0)
+
+def unstack_B_vec(bvec, p_list):
+    """Inverse of stack: vector -> list of (p_j,1)."""
+    out, i = [], 0
+    for p in p_list:
+        out.append(bvec[i:i+p].reshape(p, 1))
+        i += p
+    return out
+
+def XB_from_Blist(Xs, B_list):
+    """Return N x K matrix of predictions."""
+    return np.column_stack([Xs[j] @ B_list[j] for j in range(len(Xs))])
+
+def cg_solve(operator_mv, b, x0=None, maxit=500, tol=1e-7, M_pre=None):
+    """
+    Conjugate gradient for SPD operator A (matrix-free).
+    operator_mv(x) -> A x
+    M_pre(x) -> apply preconditioner M^{-1} x (optional)
+    """
+    x = np.zeros_like(b) if x0 is None else x0.copy()
+    r = b - operator_mv(x)
+    z = M_pre(r) if M_pre is not None else r
+    p = z.copy()
+    rz_old = float(r @ z)
+    iterations = 0
+    for _ in range(maxit):
+        iterations += 1
+        Ap = operator_mv(p)
+        alpha = rz_old / max(1e-30, float(p @ Ap))
+        x += alpha * p
+        r -= alpha * Ap
+        res_norm = np.linalg.norm(r)
+        if res_norm <= tol * (np.linalg.norm(b) + 1e-30):
+            break
+        z = M_pre(r) if M_pre is not None else r
+        rz_new = float(r @ z)
+        beta = rz_new / max(1e-30, rz_old)
+        p = z + beta * p
+        rz_old = rz_new
+    info = {"iterations": iterations, "residual": float(np.linalg.norm(r))}
+    return x, info

alsgls-0.1.0/alsgls/sim.py

@@ -0,0 +1,109 @@
+import numpy as np
+from .ops import XB_from_Blist
+
+
+def simulate_sur(N_tr, N_te, K, p, k, seed=0):
+    """Simulate a Seemingly Unrelated Regression (SUR) dataset.
+
+    Parameters
+    ----------
+    N_tr : int
+        Number of training samples.
+    N_te : int
+        Number of test samples.
+    K : int
+        Number of response equations.
+    p : int
+        Number of features per equation.
+    k : int
+        Latent factor dimension controlling correlated noise.
+    seed : int, optional
+        Seed for the NumPy random number generator. Defaults to ``0``.
+
+    Returns
+    -------
+    X_tr : list of ndarray
+        Feature matrices for the training set. Each element has shape ``(N_tr, p)``.
+    Y_tr : ndarray
+        Training responses of shape ``(N_tr, K)``.
+    X_te : list of ndarray
+        Feature matrices for the test set. Each element has shape ``(N_te, p)``.
+    Y_te : ndarray
+        Test responses of shape ``(N_te, K)``.
+
+    Notes
+    -----
+    Randomness is controlled via ``numpy.random.default_rng(seed)``; pass a
+    different ``seed`` for different simulations.
+
+    Examples
+    --------
+    >>> Xtr, Ytr, Xte, Yte = simulate_sur(100, 20, K=3, p=5, k=2, seed=42)
+    """
+    rng = np.random.default_rng(seed)
+    N = N_tr + N_te
+    base = rng.standard_normal((N, p))
+    Xs = [base + 0.5 * rng.standard_normal((N, p)) for _ in range(K)]
+    B = [rng.standard_normal((p, 1)) for _ in range(K)]
+    F0 = 1.0 * rng.standard_normal((K, k))
+    D0 = 0.05 + 0.20 * rng.random(K)
+    U = rng.standard_normal((N, k))
+    Y = XB_from_Blist(Xs, B) + U @ F0.T + rng.standard_normal((N, K)) * np.sqrt(D0)[None, :]
+    return [X[:N_tr] for X in Xs], Y[:N_tr], [X[N_tr:] for X in Xs], Y[N_tr:]
+
+
+def simulate_gls(N_tr, N_te, p_list, k, seed=0):
+    """Simulate a generalized least squares (GLS) dataset.
+
+    This variant allows each response equation to have its own number of
+    features as specified by ``p_list``.
+
+    Parameters
+    ----------
+    N_tr : int
+        Number of training samples.
+    N_te : int
+        Number of test samples.
+    p_list : sequence of int
+        Number of features for each equation.
+    k : int
+        Latent factor dimension controlling correlated noise.
+    seed : int, optional
+        Seed for the NumPy random number generator. Defaults to ``0``.
+
+    Returns
+    -------
+    X_tr : list of ndarray
+        Feature matrices for the training set. ``X_tr[j]`` has shape
+        ``(N_tr, p_list[j])``.
+    Y_tr : ndarray
+        Training responses of shape ``(N_tr, K)`` where ``K = len(p_list)``.
+    X_te : list of ndarray
+        Feature matrices for the test set. ``X_te[j]`` has shape
+        ``(N_te, p_list[j])``.
+    Y_te : ndarray
+        Test responses of shape ``(N_te, K)``.
+
+    Notes
+    -----
+    Randomness is controlled via ``numpy.random.default_rng(seed)``; pass a
+    different ``seed`` for different simulations.
+
+    Examples
+    --------
+    >>> p_list = [3, 5, 2]
+    >>> Xtr, Ytr, Xte, Yte = simulate_gls(100, 20, p_list, k=2, seed=0)
+    """
+    rng = np.random.default_rng(seed)
+    K = len(p_list)
+    N = N_tr + N_te
+    Xs = []
+    for p in p_list:
+        base = rng.standard_normal((N, p))
+        Xs.append(base + 0.5 * rng.standard_normal((N, p)))
+    B = [rng.standard_normal((p, 1)) for p in p_list]
+    F0 = 1.0 * rng.standard_normal((K, k))
+    D0 = 0.05 + 0.20 * rng.random(K)
+    U = rng.standard_normal((N, k))
+    Y = XB_from_Blist(Xs, B) + U @ F0.T + rng.standard_normal((N, K)) * np.sqrt(D0)[None, :]
+    return [X[:N_tr] for X in Xs], Y[:N_tr], [X[N_tr:] for X in Xs], Y[N_tr:]

alsgls-0.1.0/alsgls.egg-info/PKG-INFO

@@ -0,0 +1,78 @@
+Metadata-Version: 2.4
+Name: alsgls
+Version: 0.1.0
+Summary: Lightweight low-rank+diag GLS/SUR via ALS with EM baseline
+Author-email: Gaurav Sood <contact@gsood.com>
+License: MIT
+Project-URL: Homepage, https://github.com/finite-sample/alsgls
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.21
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+
+## A Lightweight ALS Solver for Iterative GLS
+
+When a GLS problem involves hundreds of equations, the $K × K$ covariance matrix becomes the computational bottleneck.  A simple statistical remedy is to assume that most of the cross‑equation dependence can be captured by a *handful of latent factors* plus equation‑specific noise.  This “low‑rank + diagonal” assumption slashes the number of unknowns from roughly $K^²$ to about $K×k$ parameters, where **k** (the latent factor rank) is much smaller than $K$.  The model alone, however, does **not** guarantee speed: we still have to fit the parameters.
+
+### Installation
+
+Install the library from PyPI:
+
+```bash
+pip install alsgls
+```
+
+For local development, clone the repo and use an editable install:
+
+```bash
+pip install -e .
+```
+
+### Usage
+
+```python
+from alsgls import als_gls, simulate_sur, nll_per_row, XB_from_Blist
+
+Xs_tr, Y_tr, Xs_te, Y_te = simulate_sur(N_tr=240, N_te=120, K=60, p=3, k=4)
+B, F, D, mem, _ = als_gls(Xs_tr, Y_tr, k=4)
+Yhat_te = XB_from_Blist(Xs_te, B)
+nll = nll_per_row(Y_te - Yhat_te, F, D)
+```
+
+See `examples/compare_als_vs_em.py` for a complete ALS versus EM comparison.
+
+### Documentation and notebooks
+
+Background material and reproducible experiments are available in the notebooks under [`als_sim/`](als_sim/), such as [`als_sim/als_comparison.ipynb`](als_sim/als_comparison.ipynb) and [`als_sim/als_sur.ipynb`](als_sim/als_sur.ipynb).
+
+### Solving low‑rank GLS: EM versus ALS
+
+The classic EM algorithm alternates between updating the regression coefficients $\beta$ and updating the factor loadings $F$ and the diagonal noise $D$.  Even though $\hat{\Sigma}$ is low‑rank, EM’s M‑step recreates the **full** $K × K$ inverse, wiping out the memory win.
+
+An alternative is **Alternating‑Least‑Squares (ALS)**. The Woodbury identity reduces the expensive inverse to a tiny k × k system, and the β‑update can be written without explicitly forming the dense matrix at all.  In practice, ALS converges in 5–6 sweeps and never allocates more than $O(K k)$ memory, while EM allocates $O(K^²)$.
+
+**Rule of thumb:** if your GLS routine keeps looping between $\beta$ and a fresh $\hat{\Sigma}$, replacing the $\hat{\Sigma}$‑update by a factor‑ALS step yields the same statistical fit with an order‑of‑magnitude smaller memory footprint.
+
+### Beyond SUR: where the idea travels
+
+Random‑effects models, feasible GLS with estimated heteroskedastic weights, optimal‑weight GMM, and spatial autoregressive GLS all iterate β ↔ Σ̂.  Each can adopt the same ALS trick: treat the weight matrix as low‑rank + diagonal, invert only the k × k core, and avoid the dense K × K algebra.  Memory savings in published examples range from 5× to 20×, depending on k.
+
+### A concrete case‑study: Seemingly‑Unrelated Regressions
+
+To show the magnitude, we ran a Monte‑Carlo experiment with N = 300 observations, three regressors, rank‑3 factors, and K set to 50, 80, 120.  EM was given 45 iterations; ALS, six sweeps.  The largest array EM holds is the dense Σ⁻¹, whereas ALS’s largest is the skinny factor matrix F.  The table summarises six replications:
+
+|   K | β‑RMSE EM | β‑RMSE ALS | Peak MB EM | Peak MB ALS | Memory ratio |
+| --: | :-------: | :--------: | ---------: | ----------: | -----------: |
+|  50 |   0.021   |    0.021   |     0.020  |      0.002  |         10×  |
+|  80 |   0.020   |    0.020   |     0.051  |      0.003  |         17×  |
+| 120 |   0.020   |    0.020   |     0.115  |      0.004  |         29×  |
+
+Statistically, the two estimators are indistinguishable (paired‑test p ≥ 0.14).  Computationally, ALS needs only a few megabytes whereas EM needs tens to hundreds.
+
+### 5  Choosing a solver in practice
+
+For small systems ($K < 50$), dense GLS or even separate OLS is fine.  Between 50 and 300 equations, a low‑rank **factor‑ALS** solver gives the same estimates at roughly one‑tenth the memory and runs happily on a GPU.  Once K enters the hundreds, any dense inverse becomes prohibitive; structured approaches such as factor‑ALS or sparse/banded $\hat{\Sigma}$ are mandatory.

alsgls-0.1.0/alsgls.egg-info/SOURCES.txt

@@ -0,0 +1,18 @@
+README.md
+pyproject.toml
+alsgls/__init__.py
+alsgls/als.py
+alsgls/em.py
+alsgls/metrics.py
+alsgls/ops.py
+alsgls/sim.py
+alsgls.egg-info/PKG-INFO
+alsgls.egg-info/SOURCES.txt
+alsgls.egg-info/dependency_links.txt
+alsgls.egg-info/requires.txt
+alsgls.egg-info/top_level.txt
+tests/test_als.py
+tests/test_lowrank_gls.py
+tests/test_metrics.py
+tests/test_ops.py
+tests/test_simulation.py

alsgls-0.1.0/alsgls.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

alsgls-0.1.0/alsgls.egg-info/requires.txt

@@ -0,0 +1,4 @@
+numpy>=1.21
+
+[dev]
+pytest>=7.0

alsgls-0.1.0/alsgls.egg-info/top_level.txt

@@ -0,0 +1 @@
+alsgls

alsgls-0.1.0/pyproject.toml

@@ -0,0 +1,27 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "alsgls"
+version = "0.1.0"
+description = "Lightweight low-rank+diag GLS/SUR via ALS with EM baseline"
+readme = "README.md"
+requires-python = ">=3.9"
+authors = [{ name = "Gaurav Sood", email = "contact@gsood.com" }]
+license = { text = "MIT" }
+urls = { "Homepage" = "https://github.com/finite-sample/alsgls" }
+classifiers = [
+  "Programming Language :: Python :: 3",
+  "License :: OSI Approved :: MIT License",
+  "Operating System :: OS Independent",
+]
+dependencies = [
+  "numpy>=1.21",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=7.0"]
+
+[tool.setuptools]
+packages = ["alsgls"]

alsgls-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

alsgls-0.1.0/tests/test_als.py

@@ -0,0 +1,41 @@
+import sys
+from pathlib import Path
+
+import numpy as np
+
+# Ensure the package root is on the path when tests run from within the
+# ``tests`` directory.
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+
+from alsgls.sim import simulate_gls
+from alsgls.als import als_gls
+from alsgls.ops import XB_from_Blist
+
+
+def test_als_shapes_and_mse_improvement():
+    # Synthesize a small GLS problem
+    p_list = [3, 5, 4]
+    X_tr, Y_tr, X_te, Y_te = simulate_gls(40, 40, p_list, k=2, seed=0)
+
+    # Baseline per-equation ridge OLS
+    lam_B = 1e-3
+    B_ols = []
+    for j, X in enumerate(X_tr):
+        XtX = X.T @ X + lam_B * np.eye(X.shape[1])
+        Xty = X.T @ Y_tr[:, [j]]
+        B_ols.append(np.linalg.solve(XtX, Xty))
+    baseline_mse = np.mean((Y_te - XB_from_Blist(X_te, B_ols)) ** 2)
+
+    # Run ALS GLS solver
+    B_list, F, D, _, _ = als_gls(X_tr, Y_tr, k=2, lam_B=lam_B, sweeps=8)
+
+    # Assert shapes
+    assert len(B_list) == len(p_list)
+    for j, Bj in enumerate(B_list):
+        assert Bj.shape == (p_list[j], 1)
+    assert F.shape == (len(p_list), 2)
+    assert D.shape == (len(p_list),)
+
+    # Compare test MSE to baseline
+    final_mse = np.mean((Y_te - XB_from_Blist(X_te, B_list)) ** 2)
+    assert final_mse < baseline_mse

alsgls-0.1.0/tests/test_lowrank_gls.py

@@ -0,0 +1,121 @@
+import numpy as np
+import sys
+from pathlib import Path
+
+# Ensure package root is on sys.path
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+
+from alsgls import als_gls, em_gls, simulate_sur, mse, nll_per_row
+
+
+def test_als_vs_em_basic():
+    """Test that ALS and EM produce similar results on a small problem."""
+    rng = np.random.default_rng(1)
+    N_tr, N_te, K, p, k = 30, 10, 5, 3, 2
+
+    # Generate test data
+    Xs_tr, Y_tr, Xs_te, Y_te = simulate_sur(N_tr, N_te, K, p, k, seed=1)
+
+    # Run both algorithms
+    B_als, F_als, D_als, mem_als, info_als = als_gls(
+        Xs_tr, Y_tr, k=k, lam_F=1e-3, lam_B=1e-3, sweeps=6
+    )
+    B_em, F_em, D_em, mem_em, info_em = em_gls(
+        Xs_tr, Y_tr, k=k, lam_F=1e-3, lam_B=1e-3, iters=15
+    )
+
+    # Check output shapes and finiteness
+    assert F_als.shape == F_em.shape == (K, k)
+    assert D_als.shape == D_em.shape == (K,)
+    assert len(B_als) == len(B_em) == K
+    
+    for arr in [F_als, D_als, F_em, D_em]:
+        assert np.isfinite(arr).all()
+    for b_list in [B_als, B_em]:
+        for b in b_list:
+            assert np.isfinite(b).all()
+    for v in [mem_als, mem_em]:
+        assert np.isfinite(v) and v > 0
+
+    # Check that both achieve similar test MSE
+    from alsgls import XB_from_Blist
+    Y_pred_als = XB_from_Blist(Xs_te, B_als)
+    Y_pred_em = XB_from_Blist(Xs_te, B_em)
+    
+    mse_als = mse(Y_te, Y_pred_als)
+    mse_em = mse(Y_te, Y_pred_em)
+    
+    assert np.isfinite(mse_als) and np.isfinite(mse_em)
+    # They should be reasonably close (within 10% relative difference)
+    rel_diff = abs(mse_als - mse_em) / max(mse_als, mse_em)
+    assert rel_diff < 0.1
+
+    # Check that NLL computation works
+    R_als = Y_tr - XB_from_Blist(Xs_tr, B_als)
+    R_em = Y_tr - XB_from_Blist(Xs_tr, B_em)
+
+    nll_als = nll_per_row(R_als, F_als, D_als)
+    nll_em = nll_per_row(R_em, F_em, D_em)
+    
+    assert np.isfinite(nll_als) and np.isfinite(nll_em)
+
+
+def test_input_validation():
+    """Test that input validation works correctly."""
+    rng = np.random.default_rng(42)
+    N, K, p = 10, 3, 2
+    
+    # Valid inputs
+    Xs = [rng.normal(size=(N, p)) for _ in range(K)]
+    Y = rng.normal(size=(N, K))
+    
+    # Should work fine
+    B, F, D, mem, info = als_gls(Xs, Y, k=2)
+    assert len(B) == K
+    
+    # Test various invalid inputs
+    try:
+        als_gls([], Y, k=2)  # Empty Xs
+        assert False, "Should have raised ValueError"
+    except ValueError:
+        pass
+    
+    try:
+        als_gls(Xs, Y.ravel(), k=2)  # 1D Y
+        assert False, "Should have raised ValueError"
+    except ValueError:
+        pass
+    
+    try:
+        als_gls(Xs, Y, k=0)  # Invalid k
+        assert False, "Should have raised ValueError"
+    except ValueError:
+        pass
+    
+    try:
+        als_gls(Xs, Y, k=2, lam_F=-1)  # Negative regularization
+        assert False, "Should have raised ValueError"
+    except ValueError:
+        pass
+
+
+def test_memory_estimate_scaling():
+    """Test that memory estimates scale appropriately with problem size."""
+    rng = np.random.default_rng(123)
+    k = 3
+    
+    # Small problem
+    N1, K1, p1 = 20, 10, 2
+    Xs1 = [rng.normal(size=(N1, p1)) for _ in range(K1)]
+    Y1 = rng.normal(size=(N1, K1))
+    _, _, _, mem1, _ = als_gls(Xs1, Y1, k=k, sweeps=2)
+    
+    # Larger problem (double dimensions)
+    N2, K2, p2 = 40, 20, 4
+    Xs2 = [rng.normal(size=(N2, p2)) for _ in range(K2)]
+    Y2 = rng.normal(size=(N2, K2))
+    _, _, _, mem2, _ = als_gls(Xs2, Y2, k=k, sweeps=2)
+    
+    # Memory should scale up (roughly 4x since K doubled and N doubled)
+    assert mem2 > mem1
+    assert mem2 < 10 * mem1  # Reasonable upper bound

alsgls-0.1.0/tests/test_metrics.py

@@ -0,0 +1,25 @@
+import numpy as np
+import os
+import sys
+
+# Ensure package root on path
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+from alsgls.metrics import nll_per_row
+
+
+def test_nll_matches_explicit():
+    rng = np.random.default_rng(0)
+    N, K, k = 5, 4, 2
+    R = rng.normal(size=(N, K))
+    F = rng.normal(size=(K, k))
+    D = rng.uniform(0.5, 1.5, size=K)
+
+    nll_func = nll_per_row(R, F, D)
+
+    Sigma = F @ F.T + np.diag(D)
+    Sigma_inv = np.linalg.inv(Sigma)
+    logdet = np.linalg.slogdet(Sigma)[1]
+    quad = np.sum(R @ Sigma_inv * R)
+    nll_explicit = 0.5 * (quad / N + logdet + K * np.log(2 * np.pi))
+
+    assert np.allclose(nll_func, nll_explicit)

alsgls-0.1.0/tests/test_ops.py

@@ -0,0 +1,54 @@
+import numpy as np
+import sys
+from pathlib import Path
+
+# Ensure package root is on sys.path
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+
+from alsgls.ops import (
+    woodbury_pieces,
+    apply_siginv_to_matrix,
+    stack_B_list,
+    unstack_B_vec,
+)
+
+
+def test_woodbury_pieces_and_apply_siginv_to_matrix():
+    rng = np.random.default_rng(0)
+    K, k = 5, 2
+    F = rng.standard_normal((K, k))
+    D = rng.uniform(0.5, 2.0, size=K)
+
+    Sigma = F @ F.T + np.diag(D)
+    Sigma_inv = np.linalg.inv(Sigma)
+
+    # Validate woodbury_pieces
+    Dinv, Cf = woodbury_pieces(F, D)
+    Sigma_inv_wb = np.diag(Dinv) - (F * Dinv[:, None]) @ Cf @ (F.T * Dinv)
+    assert np.allclose(Sigma_inv_wb, Sigma_inv, atol=1e-12, rtol=1e-12)
+
+    # Validate apply_siginv_to_matrix against explicit inverse
+    M = rng.standard_normal((3, K))
+    expected = M @ Sigma_inv
+    # without cached pieces
+    got = apply_siginv_to_matrix(M, F, D)
+    assert np.allclose(got, expected, atol=1e-12, rtol=1e-12)
+    # with cached pieces
+    got_cached = apply_siginv_to_matrix(M, F, D, Dinv=Dinv, Cf=Cf)
+    assert np.allclose(got_cached, expected, atol=1e-12, rtol=1e-12)
+
+
+def test_stack_and_unstack_B_list():
+    """Stack heterogeneous B_j blocks and recover them."""
+    rng = np.random.default_rng(0)
+    # heterogeneous shapes for each B_j
+    p_list = [1, 3, 2]
+    B_list = [rng.standard_normal((p, 1)) for p in p_list]
+
+    # Stack into vector then unstack back to list
+    b_vec = stack_B_list(B_list)
+    recovered = unstack_B_vec(b_vec, p_list)
+
+    # Each recovered block should match the original exactly
+    for orig, rec in zip(B_list, recovered):
+        assert np.allclose(orig, rec)

alsgls-0.1.0/tests/test_simulation.py

@@ -0,0 +1,51 @@
+import sys
+from pathlib import Path
+
+import numpy as np
+
+# Ensure the package root is on the import path when tests are executed
+sys.path.append(str(Path(__file__).resolve().parents[1]))
+from alsgls import simulate_sur, simulate_gls
+
+
+def test_simulate_sur_shapes_and_reproducibility():
+    N_tr, N_te, K, p, k, seed = 20, 10, 3, 5, 2, 42
+    X_tr, Y_tr, X_te, Y_te = simulate_sur(N_tr, N_te, K, p, k, seed=seed)
+
+    assert len(X_tr) == K == len(X_te)
+    for X in X_tr:
+        assert X.shape == (N_tr, p)
+    for X in X_te:
+        assert X.shape == (N_te, p)
+    assert Y_tr.shape == (N_tr, K)
+    assert Y_te.shape == (N_te, K)
+
+    X_tr2, Y_tr2, X_te2, Y_te2 = simulate_sur(N_tr, N_te, K, p, k, seed=seed)
+    for X1, X2 in zip(X_tr, X_tr2):
+        assert np.array_equal(X1, X2)
+    for X1, X2 in zip(X_te, X_te2):
+        assert np.array_equal(X1, X2)
+    assert np.array_equal(Y_tr, Y_tr2)
+    assert np.array_equal(Y_te, Y_te2)
+
+
+def test_simulate_gls_shapes_and_reproducibility():
+    N_tr, N_te, p_list, k, seed = 15, 7, [4, 3, 5], 2, 123
+    X_tr, Y_tr, X_te, Y_te = simulate_gls(N_tr, N_te, p_list, k, seed=seed)
+    K = len(p_list)
+
+    assert len(X_tr) == K == len(X_te)
+    for X, p in zip(X_tr, p_list):
+        assert X.shape == (N_tr, p)
+    for X, p in zip(X_te, p_list):
+        assert X.shape == (N_te, p)
+    assert Y_tr.shape == (N_tr, K)
+    assert Y_te.shape == (N_te, K)
+
+    X_tr2, Y_tr2, X_te2, Y_te2 = simulate_gls(N_tr, N_te, p_list, k, seed=seed)
+    for X1, X2 in zip(X_tr, X_tr2):
+        assert np.array_equal(X1, X2)
+    for X1, X2 in zip(X_te, X_te2):
+        assert np.array_equal(X1, X2)
+    assert np.array_equal(Y_tr, Y_tr2)
+    assert np.array_equal(Y_te, Y_te2)