A factor analysis of the sleep quality scores from Karpathy's post

karpathy_sleep_scores_factor_analysis.py

import arviz as az
import numpy as np
import pymc as pm
import pytensor.tensor as pt

# --- Data ---
S = np.array([
    [1.00, 0.42, 0.47, 0.14],
    [0.42, 1.00, 0.59, 0.38],
    [0.47, 0.59, 1.00, 0.65],
    [0.14, 0.38, 0.65, 1.00],
], dtype=float)

names = ["Autosleep", "8Sleep", "Oura", "Whoop"]
p = S.shape[0]

# Choose an "effective n" (sample size). If you know the true n, plug it in.
# This is what turns the ML objective into an actual likelihood strength.
n_eff = 70

with pm.Model() as factor_model:
    # Data container (optional, but nice if you want to swap matrices later)
    S_data = pm.Data("S", S)

    # One-factor loadings. (If you ever want negative loadings, use a signed prior.)
    loadings = pm.Uniform("loadings", lower=0.0, upper=1.0, shape=p)

    # Correlation-model uniqueness to keep diag(Sigma)=1:
    # Sigma = λλ' + diag(1 - λ^2)
    unique_vars = 1.0 - loadings**2
    Sigma = pt.outer(loadings, loadings) + pt.diag(unique_vars)

    # Stable log|Sigma|
    sign, logdet_Sigma = pt.linalg.slogdet(Sigma)
    # With your construction and bounds, Sigma should be PD, so sign should be +1.
    pm.Potential("pd_guard", pt.switch(pt.eq(sign, 1.0), 0.0, -np.inf))

    # trace(S Sigma^{-1}) computed via solve for stability
    # Solve Sigma * X = S  =>  X = Sigma^{-1} S  =>  tr(Sigma^{-1} S) = tr(X)
    X = pt.linalg.solve(Sigma, S_data)
    trace_term = pt.nlinalg.trace(X)

    # ML / Wishart-style objective (constants dropped):
    # log p(S | Sigma) ∝ -(n_eff/2) * ( log|Sigma| + tr(S Sigma^{-1}) )
    pm.Potential("ml_fit", -(n_eff / 2.0) * (logdet_Sigma + trace_term))

    trace = pm.sample(
        2000, tune=2000, target_accept=0.95,
        return_inferencedata=True, random_seed=42,
    )

summary = az.summary(trace, var_names=["loadings"])
print(summary)

loadings_mean = trace.posterior["loadings"].mean(dim=["chain", "draw"]).values
loadings_sd = trace.posterior["loadings"].std(dim=["chain", "draw"]).values

Sigma_mean = np.outer(loadings_mean, loadings_mean) + np.diag(1 - loadings_mean**2)

print("\nObserved Correlation Matrix (S):")
print(S)

print("\nImplied Correlation Matrix (Sigma at posterior mean loadings):")
print(np.round(Sigma_mean, 2))

print("\nFactor Loadings by Item (sorted by |mean|):")
for name, mean, sd in sorted(zip(names, loadings_mean, loadings_sd), key=lambda k: -abs(k[1])):
    print(f"- {name}: {mean:.2f} (± {sd:.3f})")