""" lta.py -- LATENT TRANSITION ANALYSIS / latent Markov model (from scratch). Backs the notebooks in "Latent Transition Analysis". Latent transition analysis (LTA) is LCA put in motion: each subject occupies a latent class at every time point, and the class EVOLVES over time as a Markov chain. It is a per-subject Hidden Markov Model with parameters SHARED across subjects: S_{i,1} ~ Categorical(pi) (initial-state distribution) S_{i,t} | S_{i,t-1}=a ~ Categorical(tau_{a,.}) (C x C transition matrix) x_{i,t,j} | S_{i,t}=c ~ Categorical(delta_{c,j,.}) (measurement model, time-invariant) pi is the class mix at wave 1, tau the transition probabilities (tau_{ab} = P(move from class a to class b between consecutive waves), delta_{c,j,.} class c's response-probability vector for item j -- exactly the LCA measurement model, assumed invariant across time so that a "class" means the same thing at every wave. Ordinary LCA is the special case T=1. The from-scratch sampler is a DATA-AUGMENTATION Gibbs whose key step is FORWARD-FILTER BACKWARD-SAMPLE (FFBS): draw each subject's whole latent state sequence S_{i,1:T} jointly from its exact conditional (the same sampler used for Markov-switching / state-space models). Given the sequences, pi, tau and delta all have conjugate Dirichlet full conditionals: pi ~ Dirichlet(1 + counts of first-wave states) tau_{a,.} ~ Dirichlet(1 + counts of a->b transitions, pooled over subjects and time) delta_{c,j,.} ~ Dirichlet(1 + response-level counts among all (i,t) in class c) Latent states are ordered post-hoc by their mean emitted level so the transition matrix is interpretable (state 0 = lowest), which also resolves label switching. """ import numpy as np # --------------------------------------------------------------------------- # # emission log-likelihood log P(x_{i,t} | class) # # --------------------------------------------------------------------------- # def _emis_loglik(X, delta): """X:(N,T,J) int in 0..L-1 ; delta:(C,J,L) -> logB:(N,T,C) = sum_j log delta[c,j,x].""" N, T, J = X.shape; C = delta.shape[0] logB = np.zeros((N, T, C)) for c in range(C): acc = np.zeros((N, T)) for j in range(J): acc += np.log(delta[c, j])[X[:, :, j]] # gather (N,T) logB[:, :, c] = acc return logB def _cat_rows(P, rng): """vectorised categorical draw, one per row of the (n,C) probability matrix P.""" u = rng.random(P.shape[0]) return (np.cumsum(P, axis=1) < u[:, None]).sum(1).clip(0, P.shape[1] - 1) # --------------------------------------------------------------------------- # # forward filter + backward sample (FFBS), vectorised over subjects # # --------------------------------------------------------------------------- # def _ffbs(logB, pi, tau, rng): """Sample latent state sequences S:(N,T) given emissions and Markov parameters. Also returns the observed-data log-likelihood (sum of log forward scales).""" N, T, C = logB.shape B = np.exp(logB - logB.max(axis=2, keepdims=True)) # scale for stability (cancels in normalisation) alpha = np.empty((N, T, C)); scale = np.empty((N, T)) a = pi[None, :] * B[:, 0, :] scale[:, 0] = a.sum(1); alpha[:, 0, :] = a / scale[:, 0][:, None] for t in range(1, T): pred = alpha[:, t - 1, :] @ tau # (N,C) a = pred * B[:, t, :] scale[:, t] = a.sum(1); alpha[:, t, :] = a / scale[:, t][:, None] # backward sample S = np.empty((N, T), int) S[:, T - 1] = _cat_rows(alpha[:, T - 1, :], rng) for t in range(T - 2, -1, -1): w = alpha[:, t, :] * tau[:, S[:, t + 1]].T # (N,C) w = w / w.sum(1, keepdims=True) S[:, t] = _cat_rows(w, rng) loglik = float((np.log(scale) + logB.max(axis=2)).sum()) return S, loglik # --------------------------------------------------------------------------- # # Gibbs sampler for latent transition analysis # # --------------------------------------------------------------------------- # def lta_gibbs(X, C, rng, L=None, draws=4000, burn=2000, order_states=True): """FFBS data-augmentation Gibbs for the latent Markov / LTA model. X:(N,T,J) integer responses coded 0..L-1. Returns posterior draws of the initial distribution pi:(draws,C), transition matrix tau:(draws,C,C) and emission profiles delta:(draws,C,J,L). States are ordered by mean emitted level for interpretability.""" X = np.asarray(X); N, T, J = X.shape if L is None: L = int(X.max()) + 1 # init: random states S = rng.integers(0, C, (N, T)) PI = np.empty((draws, C)); TAU = np.empty((draws, C, C)); DELTA = np.empty((draws, C, J, L)) lev = np.arange(L) for it in range(draws + burn): # ---- emission delta | S (Dirichlet, conjugate) ---- delta = np.empty((C, J, L)) for c in range(C): mask = S == c for j in range(J): vals = X[:, :, j][mask] cnt = np.bincount(vals, minlength=L) if vals.size else np.zeros(L) delta[c, j] = rng.dirichlet(1.0 + cnt) # ---- pi | S and tau | S (Dirichlet, conjugate) ---- pi = rng.dirichlet(1.0 + np.bincount(S[:, 0], minlength=C)) trans = np.zeros((C, C)) for t in range(1, T): np.add.at(trans, (S[:, t - 1], S[:, t]), 1) tau = np.array([rng.dirichlet(1.0 + trans[a]) for a in range(C)]) # ---- S | delta, pi, tau (FFBS) ---- logB = _emis_loglik(X, delta) S, _ = _ffbs(logB, pi, tau, rng) # ---- order states by mean emitted level (interpretable + resolves switching) ---- if order_states: score = (delta.mean(1) * lev[None, :]).sum(1) # mean level per class o = np.argsort(score) inv = np.argsort(o) delta = delta[o]; pi = pi[o]; tau = tau[o][:, o] S = inv[S] if it >= burn: i = it - burn; PI[i] = pi; TAU[i] = tau; DELTA[i] = delta return dict(pi=PI, tau=TAU, delta=DELTA, L=L) def lta_loglik(X, pi, tau, delta): """observed-data log-likelihood at a point estimate (forward algorithm).""" logB = _emis_loglik(np.asarray(X), delta) N, T, C = logB.shape B = np.exp(logB - logB.max(axis=2, keepdims=True)) a = pi[None, :] * B[:, 0, :]; ll = np.log(a.sum(1)) + logB[:, 0, :].max(1) at = a / a.sum(1, keepdims=True) for t in range(1, T): a = (at @ tau) * B[:, t, :] ll += np.log(a.sum(1)) + logB[:, t, :].max(1) at = a / a.sum(1, keepdims=True) return float(ll.sum()) # --------------------------------------------------------------------------- # # simulation with known LTA parameters # # --------------------------------------------------------------------------- # def simulate_lta(N, T, pi, tau, delta, rng): """Simulate an LTA cohort. pi:(C,), tau:(C,C), delta:(C,J,L). Returns X:(N,T,J), S:(N,T).""" pi = np.asarray(pi); tau = np.asarray(tau); delta = np.asarray(delta) C, J, L = delta.shape S = np.empty((N, T), int) S[:, 0] = [rng.choice(C, p=pi) for _ in range(N)] for t in range(1, T): for i in range(N): S[i, t] = rng.choice(C, p=tau[S[i, t - 1]]) X = np.empty((N, T, J), int) for i in range(N): for t in range(T): for j in range(J): X[i, t, j] = rng.choice(L, p=delta[S[i, t], j]) return X, S def state_prevalence(S, C): """proportion in each class at each wave (the marginal class trajectory).""" T = S.shape[1] return np.array([[ (S[:, t] == c).mean() for c in range(C)] for t in range(T)])