canbiplotme 0.3.0__py3-none-any.whl

canbiplotme/dynamic.py

@@ -0,0 +1,105 @@
+"""Dynamic extension (section 3.7): yearly group deviations, centred pairs
+bootstrap test and trajectories in the reference canonical plane."""
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+
+from .inference import holm
+from .stage1 import _design, _fml, _pf, _prepare
+
+
+def _year_design(d, year, frictions, years):
+    cols, names = {}, []
+    for t in years[1:]:
+        for f in frictions:
+            nm = f"d__{t}__{f}"
+            cols[nm] = np.where(d[year].values == t, d[f].values, 0.0)
+            names.append(nm)
+    return pd.concat([d.reset_index(drop=True), pd.DataFrame(cols)], axis=1), names
+
+
+def _dyn_fit(pf, d, sp, years, vcov):
+    ex = sorted(d[sp["origin"]].unique())
+    dd, names = _design(d, sp["origin"], sp["frictions"], ex)
+    dd, ynames = _year_design(dd, sp["year"], sp["frictions"], years)
+    m = pf.fepois(_fml(sp["flow"], names, sp["controls"], sp["origin"], sp["destination"], sp["year"], ynames),
+                  data=dd, vcov=vcov, iwls_maxiter=100)
+    coef = m.coef()
+    present = [n for n in ynames if n in coef.index and np.isfinite(coef[n])]
+    delta = np.zeros((len(years), len(sp["frictions"])))
+    for i, t in enumerate(years[1:], start=1):
+        for k, f in enumerate(sp["frictions"]):
+            nm = f"d__{t}__{f}"
+            delta[i, k] = coef[nm] if nm in present else np.nan
+    vc = pd.DataFrame(m._vcov, index=coef.index, columns=coef.index).loc[present, present].values
+    return delta, coef[present].values, vc, present
+
+
+def gravity_dynamic(data, flow, origin, destination, year, frictions, group, controls=None) -> dict:
+    """beta_{i,t} = beta_i + delta_{g,t}, delta_{g,first year} = 0, by group."""
+    pf = _pf()
+    d0 = _prepare(data, origin, destination, frictions)
+    years = sorted(d0[year].unique())
+    sp = dict(flow=flow, origin=origin, destination=destination, year=year, frictions=list(frictions),
+              group=group, controls=list(controls or []), years=years)
+    out = {}
+    for g in sorted(d0[group].dropna().astype(str).unique()):
+        d = d0[d0[group].astype(str) == g]
+        delta, b, V, names = _dyn_fit(pf, d, sp, years, {"CRV1": "_pair"})
+        out[g] = {"delta": pd.DataFrame(delta, index=years, columns=frictions), "coef": b, "V": V, "names": names}
+    return {"groups": out, "spec": sp}
+
+
+def _wald(b, V):
+    return float(b @ np.linalg.solve(V, b))
+
+
+def dynamic_test(dyn: dict, data, R=199, seed=20251025) -> pd.DataFrame:
+    """Wald test of delta_{g,t} = 0 calibrated by a centred pairs bootstrap, Holm across groups."""
+    pf = _pf()
+    sp = dyn["spec"]
+    d0 = _prepare(data, sp["origin"], sp["destination"], sp["frictions"])
+    rows = []
+    for k, (g, e) in enumerate(dyn["groups"].items()):
+        W_obs = _wald(e["coef"], e["V"])
+        d = d0[d0[sp["group"]].astype(str) == g]
+        by_pair = d.groupby("_pair").indices
+        pairs = np.array(sorted(by_pair))
+        Ws = []
+        for r in range(R):
+            rng = np.random.default_rng(seed + 10000 * (k + 1) + r + 1)
+            s = pairs[rng.integers(len(pairs), size=len(pairs))]
+            db = d.iloc[np.concatenate([by_pair[p] for p in s])].copy()
+            db["_pair"] = db["_pair"] + "_" + np.repeat(np.arange(len(s)), [len(by_pair[p]) for p in s]).astype(str)
+            try:
+                _, b, V, names = _dyn_fit(pf, db, sp, sp["years"], {"CRV1": "_pair"})
+            except Exception:
+                Ws.append(np.nan)
+                continue
+            common = [n for n in names if n in e["names"]]
+            ib = [names.index(n) for n in common]
+            ie = [e["names"].index(n) for n in common]
+            Ws.append(_wald(b[ib] - e["coef"][ie], V[np.ix_(ib, ib)]))
+        Ws = np.array(Ws)
+        rows.append({"group": g, "W": W_obs, "df": len(e["coef"]),
+                     "p": (1 + np.nansum(Ws >= W_obs)) / (1 + np.isfinite(Ws).sum()),
+                     "R_valid": int(np.isfinite(Ws).sum())})
+    out = pd.DataFrame(rows)
+    out["p_holm"] = holm(out["p"].values)
+    return out
+
+
+def cbme_trajectory(fit, dyn) -> dict:
+    """a_g(t) = V2'(mu_g + delta_{g,t} - mean_t delta_{g,t} - mbar) (Definition 3.7.1)."""
+    deltas = {g: e["delta"].values for g, e in dyn["groups"].items()} if "groups" in dyn else dyn
+    V2 = fit.V[:, :fit.A.shape[1]]
+    coords, disp, path = {}, {}, {}
+    for gi, g in enumerate(fit.data.levels):
+        dlt = np.nan_to_num(np.asarray(deltas[g], float))
+        mt = dlt - dlt.mean(0) + fit.mu[gi]
+        a = (mt - fit.mbar) @ V2
+        coords[g] = a
+        disp[g] = float(np.sqrt(((a[-1] - a[0]) ** 2).sum()))
+        path[g] = float(np.sqrt((np.diff(a, axis=0) ** 2).sum(1)).sum())
+    return {"coords": coords, "displacement": pd.Series(disp), "path_length": pd.Series(path)}

canbiplotme/inference.py

@@ -0,0 +1,281 @@
+"""Parametric bootstrap under the null (G0), pairwise tests and influence."""
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+from scipy import stats as sps
+
+from .stage2 import CBMEFit, fit_core
+
+
+def _rng(seed):
+    from .rng import RRNG
+    return seed if isinstance(seed, (np.random.Generator, RRNG)) else np.random.default_rng(seed)
+
+
+def sqrt_psd(S: np.ndarray) -> np.ndarray:
+    S = np.asarray(S, float)
+    ev, evec = np.linalg.eigh((S + S.T) / 2)
+    return evec @ np.diag(np.sqrt(np.maximum(ev, 0.0)))
+
+
+def _prepare(fit: CBMEFit):
+    x = fit.data
+    if x.D is None and x.Omega is not None and "Omega_sqrt" not in fit.cache:
+        fit.cache["Omega_sqrt"] = {lev: sqrt_psd(x.Omega[lev]) for lev in x.block_levels}
+    return fit
+
+
+def simulate_coefs(fit: CBMEFit, means="common", errors: str = "joint", deviations: str = "normal",
+                   scale: float = 1.0, rng=None) -> np.ndarray:
+    """Simulate coefficient vectors from the fitted two-level model (Definition 3.6.1).
+
+    ``means='common'`` gives the generator G0 (null hypothesis); ``'group'`` uses
+    the estimated group means (deviations from the common mean multiplied by
+    ``scale``); a (G, K) array gives arbitrary means. ``errors`` is ``'joint'``,
+    ``'independent'`` (G0-ind) or ``'replicate'`` (one random bootstrap replicate
+    per group); ``deviations`` is ``'normal'`` or ``'t4'``.
+    """
+    rng = _rng(rng)
+    _prepare(fit)
+    x = fit.data
+    K, N = x.K, x.N
+    if not x.has_joint and errors != "independent":
+        errors = "independent"
+    if errors == "replicate" and x.D is None:
+        raise ValueError("errors='replicate' needs bootstrap replicates.")
+    gi = np.array([x.levels.index(g) for g in x.group])
+    if isinstance(means, np.ndarray):
+        M = means[gi]
+    elif means == "common":
+        M = np.tile(fit.mu0, (N, 1))
+    else:
+        M = (fit.mu0 + scale * (fit.mu - fit.mu0))[gi]
+    # deviations u* = L z with T + 1e-10 I = L L' (Cholesky), drawn individual by
+    # individual, then the errors group by group in the order of the levels:
+    # the order of the code of the thesis (rev_funciones.R::generar), so that
+    # with canbiplotme.rng.RRNG the replicates are those of R.
+    Lt = np.linalg.cholesky(fit.T + 1e-10 * np.eye(K))
+    if deviations == "normal":
+        Z = rng.standard_normal((N, K))
+    else:
+        Z = rng.standard_t(4, (N, K)) / np.sqrt(2)
+    U = Z @ Lt.T
+    E = np.zeros((N, K))
+    ok = np.isfinite(x.beta)
+    blk = x.block_of
+    for lev in x.block_levels:          # error blocks (the groups by default)
+        idx = np.where(blk == lev)[0]
+        if not len(idx):
+            continue
+        if errors == "independent":
+            for i in idx:
+                k = np.where(ok[i])[0]
+                if not len(k):
+                    continue
+                ev, evec = np.linalg.eigh(x.V[i][np.ix_(k, k)])
+                ev, evec = ev[::-1], evec[:, ::-1]
+                E[i, k] = evec @ (np.sqrt(np.maximum(ev, 0.0)) * rng.standard_normal(len(k)))
+            continue
+        if not ok[idx].any():
+            continue
+        if x.D is not None:
+            D = x.D[lev]
+            R = D.shape[0]
+            e = (D[rng.integers(R)] * np.sqrt(R / (R - 1)) if errors == "replicate"
+                 else D.T @ rng.standard_normal(R) / np.sqrt(R - 1))
+        else:
+            Sq = fit.cache["Omega_sqrt"][lev]
+            e = Sq @ rng.standard_normal(Sq.shape[1])
+        E[idx] = e.reshape(len(idx), K)
+    B = M + U + E
+    B[~np.isfinite(x.beta)] = np.nan
+    return B
+
+
+def refit(fit: CBMEFit, beta, rows=None, fill_missing: bool = True):
+    rows = np.arange(fit.data.N) if rows is None else np.asarray(rows)
+    try:
+        return fit_core(fit.data, beta, rows, fit.settings, fit.cache, fill_missing)
+    except np.linalg.LinAlgError:
+        return None
+
+
+def _with_core(fit: CBMEFit, core: dict, beta=None) -> CBMEFit:
+    import copy
+    f = copy.copy(fit)
+    f.__dict__ = dict(fit.__dict__)
+    f._update(core)
+    if beta is not None:
+        d = copy.copy(fit.data)
+        d.__dict__ = dict(fit.data.__dict__)
+        d.beta = beta
+        f.data = d
+    return f
+
+
+class TestResult:
+    def __init__(self, **kw):
+        self.__dict__.update(kw)
+
+    def __repr__(self) -> str:
+        s = [f"Null-hypothesis parametric bootstrap (generator {'G0' if self.errors == 'joint' else 'G0-ind'}), "
+             f"{self.B_valid} valid replicates" + (f" ({self.failed} failed)" if self.failed else ""), "",
+             "Equality of group means",
+             pd.DataFrame({"observed": [self.observed["trace"], self.observed["lambda1"]],
+                           "p": [self.global_["trace"], self.global_["lambda1"]]},
+                          index=["trace", "largest root"]).round(4).to_string(), "",
+             "Pairwise comparisons (Holm)", self.pairs.round(4).to_string(index=False), "",
+             f"{int(self.pairs['different'].sum())} of {len(self.pairs)} pairs differ"]
+        if self.composition is not None:
+            c = self.composition
+            s.append(f"\nComposition of the first dimension: S = {c['S']:.4f} (95th percentile under H0 "
+                     f"{c['q_S']:.4f}, p = {c['p_S']:.4f}); S_plane = {c['S_plane']:.4f}; rule holds: {c['holds']}")
+        return "\n".join(s)
+
+
+def holm(p) -> np.ndarray:
+    """Holm adjustment (same as R ``p.adjust(method = 'holm')``)."""
+    p = np.asarray(p, float)
+    m = len(p)
+    o = np.argsort(p, kind="stable")
+    adj = np.minimum(1.0, np.maximum.accumulate((m - np.arange(m)) * p[o]))
+    out = np.empty(m)
+    out[o] = adj
+    return out
+
+
+def boot_pvalue(obs, ref) -> float:
+    """(1 + #{r: t*_r >= t}) / (1 + R), with R the number of finite reference values."""
+    ref = np.asarray(ref, float)
+    fin = np.isfinite(ref)
+    return float((1 + np.sum(ref[fin] >= obs)) / (1 + fin.sum()))
+
+
+def cbme_test(fit: CBMEFit, B: int = 1999, errors: str = "joint", S_min: float = 0.5,
+              level: float = 0.95, seed=None, keep: bool = True, r_compatible: bool = True) -> TestResult:
+    """Tests by null-hypothesis parametric bootstrap (sections 3.6.2-3.6.5).
+
+    Equality of the group means (trace and largest root), every pair of groups
+    (Mahalanobis D2 with Holm correction) and the composition of the first
+    dimension (statistic S). The second stage is re-estimated in every replicate.
+
+    With an integer ``seed`` and ``r_compatible=True`` (default), replicate ``b``
+    is generated with ``canbiplotme.rng.RRNG(seed + b)``, the generator of R
+    after ``set.seed(seed + b)``: the replicates are those of
+    ``gcbiplot::gcb_test`` and of the code of the thesis. Otherwise a
+    ``numpy.random.Generator`` is used.
+    """
+    from .rng import RRNG
+    _prepare(fit)
+    if r_compatible and seed is not None and not isinstance(seed, np.random.Generator):
+        gen = lambda b: RRNG(int(seed) + b)
+    else:
+        rng = _rng(seed)
+        gen = lambda b: rng
+    reps = []
+    for b in range(1, B + 1):
+        core = refit(fit, simulate_coefs(fit, "common", errors, rng=gen(b)))
+        if core is not None:
+            reps.append(core)
+    Bv = len(reps)
+    keys = list(fit.stats.keys())
+    S_star = pd.DataFrame([[c["stats"][k] for k in keys] for c in reps], columns=keys)
+    D_star = np.array([c["D2"].values for c in reps])
+    glob = {"trace": boot_pvalue(fit.stats["trace"], S_star["trace"].values),
+            "lambda1": boot_pvalue(fit.stats["lambda1"], S_star["lambda1"].values)}
+    p_pair = np.array([boot_pvalue(fit.D2.values[j], D_star[:, j]) for j in range(len(fit.D2))])
+    lo, hi = [], []
+    for j in range(len(fit.D2)):
+        r = D_star[:, j]; n = int(np.isfinite(r).sum()); c = int(np.sum(r[np.isfinite(r)] >= fit.D2.values[j]))
+        ci = sps.binomtest(c, n).proportion_ci(confidence_level=0.95, method="exact")
+        lo.append((1 + ci.low * n) / (1 + n)); hi.append((1 + ci.high * n) / (1 + n))
+    dfk = int(np.isfinite(np.diag(fit.mu_cov[0])).sum()) if hasattr(fit, "mu_cov") else fit.W.shape[0]
+    pairs = pd.DataFrame({"pair": fit.D2.index, "D2": fit.D2.values, "p": p_pair, "p_mc_low": lo,
+                          "p_mc_high": hi, "p_holm": holm(p_pair),
+                          "p_chisq": sps.chi2.sf(fit.D2.values, dfk)})
+    pairs["different"] = pairs["p_holm"] < 1 - level
+    comp = None
+    if "S" in fit.stats:
+        qS = float(np.nanquantile(S_star["S"].values, level))
+        S = fit.stats["S"]
+        comp = {"S": S, "S_plane": fit.stats["S_plane"], "q_S": qS, "p_S": boot_pvalue(S, S_star["S"].values),
+                "holds": bool(S > qS and (S_min is None or np.isnan(S_min) or S >= S_min))}
+    return TestResult(global_=glob, pairs=pairs, composition=comp, B=B, B_valid=Bv, failed=B - Bv,
+                      errors=errors, level=level, observed=fit.stats, seed=seed,
+                      reference={"stats": S_star, "D2": D_star} if keep else None)
+
+
+def procrustes_align(X: np.ndarray, Y: np.ndarray) -> np.ndarray:
+    """Rotate X to best match Y (orthogonal Procrustes, no translation or scaling)."""
+    u, _, vt = np.linalg.svd(X.T @ Y)
+    return X @ (u @ vt)
+
+
+def jack_rho(fit: CBMEFit, beta: np.ndarray) -> np.ndarray:
+    N = beta.shape[0]
+    A = fit.A
+    out = np.full(N, np.nan)
+    allrows = np.arange(N)
+    for i in range(N):
+        keep = np.delete(allrows, i)
+        f = refit(fit, beta[keep], rows=keep)
+        if f is None or np.any(f["n"] == 0):
+            continue
+        out[i] = np.sqrt(np.sum((f["A"] - procrustes_align(A, f["A"])) ** 2)) / np.sqrt(np.sum(A ** 2))
+    return out
+
+
+def cbme_influence(fit: CBMEFit, B_ref: int = 200, level: float = 0.95, seed=None) -> dict:
+    """Jackknife influence rho_(i) and reference of rho_max (section 3.6.6)."""
+    rng = _rng(seed)
+    _prepare(fit)
+    rho = jack_rho(fit, fit.data.beta)
+    ref = None
+    if B_ref > 0:
+        ref = []
+        for _ in range(B_ref):
+            bs = simulate_coefs(fit, "group", rng=rng)
+            core = refit(fit, bs)
+            if core is None:
+                ref.append(np.nan)
+                continue
+            f2 = _with_core(fit, core)
+            ref.append(np.nanmax(jack_rho(f2, bs)))
+        ref = np.array(ref)
+    rmax = np.nanmax(rho)
+    q = float(np.nanquantile(ref, level)) if ref is not None else np.nan
+    p = (1 + np.nansum(ref >= rmax)) / (1 + np.isfinite(ref).sum()) if ref is not None else np.nan
+    tab = pd.DataFrame({"id": fit.data.id, "group": fit.data.group, "rho": rho}).sort_values("rho", ascending=False)
+    return {"rho": tab, "rho_max": rmax, "reference": ref, "q": q, "p": p, "excessive": bool(rmax > q)}
+
+
+def cbme_permutation_test(fit: CBMEFit, B: int = 999, seed=None) -> dict:
+    """Permutation test of the equality of group means (section 3.6.1).
+
+    Permutes the group labels and re-estimates the second stage. Exact only
+    under exchangeability; liberal when the measurement error differs between
+    groups. Comparator of :func:`cbme_test`."""
+    import copy
+    from .stage2 import Settings
+    rng = _rng(seed)
+    x = fit.data
+    st = copy.copy(fit.settings)
+    st.T_errors = st.mean_errors = "independent"
+    xp = copy.copy(x)
+    xp.__dict__ = dict(x.__dict__)
+    xp.D = None
+    xp.Omega = None
+    cache = {"pos": fit.cache["pos"], "Omega": None}
+    ref = []
+    for _ in range(B):
+        xp.group = rng.permutation(x.group)
+        try:
+            f = fit_core(xp, x.beta, np.arange(x.N), st, cache)
+            ref.append((f["stats"]["trace"], f["stats"]["lambda1"]))
+        except np.linalg.LinAlgError:
+            pass
+    ref = np.array(ref)
+    obs = (fit.stats["trace"], fit.stats["lambda1"])
+    return {"trace": (1 + np.sum(ref[:, 0] >= obs[0] - 1e-12)) / (1 + len(ref)),
+            "lambda1": (1 + np.sum(ref[:, 1] >= obs[1] - 1e-12)) / (1 + len(ref))}

canbiplotme/meta.py

@@ -0,0 +1,118 @@
+"""Multivariate random-effects meta-analysis and meta-regression (REML / ML),
+the Python counterpart of ``mixmeta``/``mvmeta`` and of ``metafor::rma.mv``
+with an unstructured between-study covariance.
+
+Model for study i with outcome vector y_i (missing outcomes allowed):
+    y_i = (I_K kron x_i') beta + b_i + e_i,  b_i ~ N(0, Psi),  e_i ~ N(0, S_i)
+with S_i known.
+"""
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+from scipy import optimize, stats as sps
+
+
+def _theta_to_psi(th, K):
+    L = np.zeros((K, K)); L[np.tril_indices(K)] = th
+    L[np.diag_indices(K)] = np.exp(L[np.diag_indices(K)])
+    return L @ L.T
+
+
+def _psi_to_theta(P, K):
+    L = np.linalg.cholesky(P + 1e-8 * np.eye(K))
+    L[np.diag_indices(K)] = np.log(np.diag(L))
+    return L[np.tril_indices(K)]
+
+
+def _prep(y, S, X):
+    n, K = y.shape
+    rows = []
+    for i in range(n):
+        ok = np.where(np.isfinite(y[i]))[0]
+        if not len(ok):
+            continue
+        Xi = np.kron(np.eye(K), X[i][None, :])[ok]
+        rows.append((ok, y[i, ok], S[i][np.ix_(ok, ok)], Xi))
+    return rows
+
+
+def _ll(th, rows, K, p, reml, logdet_xx=0.0):
+    Psi = _theta_to_psi(th, K)
+    XtVX = np.zeros((p, p)); XtVy = np.zeros(p); ld = 0.0; nobs = 0
+    chol = []
+    for ok, yi, Si, Xi in rows:
+        try:
+            c = np.linalg.cholesky(Si + Psi[np.ix_(ok, ok)])
+        except np.linalg.LinAlgError:
+            return -np.inf, None, None
+        ld += 2 * np.log(np.diag(c)).sum(); nobs += len(ok)
+        Xc = np.linalg.solve(c, Xi); yc = np.linalg.solve(c, yi)
+        XtVX += Xc.T @ Xc; XtVy += Xc.T @ yc
+        chol.append((Xc, yc))
+    b = np.linalg.solve(XtVX, XtVy)
+    q = sum(np.sum((yc - Xc @ b) ** 2) for Xc, yc in chol)
+    if reml:
+        # + 0.5 log|X'X| as in mixmeta and metafor
+        ll = -0.5 * ((nobs - p) * np.log(2 * np.pi) + ld + np.linalg.slogdet(XtVX)[1] + q) + 0.5 * logdet_xx
+    else:
+        ll = -0.5 * (nobs * np.log(2 * np.pi) + ld + q)
+    return ll, b, XtVX
+
+
+def mvmeta(y, S, mods=None, method: str = "reml", outcome_names=None, Psi0=None) -> dict:
+    """Fit the multivariate meta-analysis / meta-regression.
+
+    ``y``: n x K estimates (NaN for missing outcomes); ``S``: n x K x K
+    within-study covariances (or n x K variances, taken as diagonal);
+    ``mods``: n x q moderators (an intercept is added). ``method``: 'reml',
+    'ml' or 'fixed'. Returns coefficients (q+1 x K), their covariance, Psi,
+    log-likelihood, AIC/BIC, Cochran's multivariate Q test of heterogeneity and
+    the I^2 statistic."""
+    y = np.asarray(y, float)
+    n, K = y.shape
+    S = np.asarray(S, float)
+    if S.ndim == 2:
+        S = np.stack([np.diag(s) for s in S])
+    X = np.ones((n, 1)) if mods is None else np.column_stack([np.ones(n), np.asarray(mods, float).reshape(n, -1)])
+    q1 = X.shape[1]; p = q1 * K
+    rows = _prep(y, S, X)
+    nobs = sum(len(r[0]) for r in rows)
+    # fixed-effects fit and Q statistic
+    Xlong = np.vstack([r[3] for r in rows])
+    th_fixed = np.zeros((K, K)); th_fixed[np.diag_indices(K)] = -40.0
+    th_fixed = th_fixed[np.tril_indices(K)]
+    ll_f, b_f, XtVX_f = _ll(th_fixed, rows, K, p, False)
+    Q = 0.0
+    for ok, yi, Si, Xi in rows:
+        r = yi - Xi @ b_f; Q += r @ np.linalg.solve(Si, r)
+    dfQ = nobs - p
+    if method == "fixed":
+        Psi, b, XtVX, ll, conv = np.zeros((K, K)), b_f, XtVX_f, ll_f, True
+    else:
+        reml = method == "reml"
+        if Psi0 is None:
+            yc = np.where(np.isfinite(y), y, np.nan)
+            v = np.nanvar(yc, axis=0, ddof=1) - np.nanmean(np.stack([np.diag(s) for s in S]), 0)
+            Psi0 = np.diag(np.maximum(v, 0.1 * np.nanvar(yc, axis=0, ddof=1)))
+        lxx = np.linalg.slogdet(Xlong.T @ Xlong)[1] if reml else 0.0
+        f = lambda th: -_ll(th, rows, K, p, reml, lxx)[0]
+        th0 = _psi_to_theta(np.asarray(Psi0, float), K)
+        r1 = optimize.minimize(f, th0, method="BFGS", options={"gtol": 1e-9, "maxiter": 5000})
+        r2 = optimize.minimize(f, r1.x, method="Nelder-Mead", options={"xatol": 1e-10, "fatol": 1e-13, "maxiter": 40000})
+        best = r2 if r2.fun < r1.fun else r1
+        ll, b, XtVX = _ll(best.x, rows, K, p, reml, lxx)
+        Psi, conv = _theta_to_psi(best.x, K), bool(r1.success or r2.success)
+    names = outcome_names or [f"y{k + 1}" for k in range(K)]
+    terms = ["(Intercept)"] + [f"x{j}" for j in range(1, q1)]
+    vc = np.linalg.inv(XtVX)
+    coef = pd.DataFrame(b.reshape(K, q1).T, index=terms, columns=names)
+    se = pd.DataFrame(np.sqrt(np.diag(vc)).reshape(K, q1).T, index=terms, columns=names)
+    o = [k * q1 + j for j in range(q1) for k in range(K)]          # term-major, as mixmeta
+    lab = [f"{t}.{nm}" for t in terms for nm in names]
+    vc = pd.DataFrame(vc[np.ix_(o, o)], index=lab, columns=lab)
+    npar = p + (K * (K + 1) // 2 if method != "fixed" else 0)
+    return {"coef": coef, "se": se, "vcov": vc, "Psi": pd.DataFrame(Psi, index=names, columns=names),
+            "logLik": ll, "AIC": -2 * ll + 2 * npar, "BIC": -2 * ll + np.log(nobs - (p if method == "reml" else 0)) * npar,
+            "Q": Q, "df_Q": dfQ, "p_Q": sps.chi2.sf(Q, dfQ), "I2": max(0.0, (Q - dfQ) / Q) * 100,
+            "method": method, "nobs": nobs, "converged": conv}

canbiplotme/plot.py

@@ -0,0 +1,48 @@
+"""Plot of the canonical biplot (matplotlib)."""
+from __future__ import annotations
+
+import numpy as np
+
+
+def plot_cbme(fit, regions=None, region="R3", individuals=True, arrows=True, colors=None, ax=None,
+              arrow_scale=None, labels_components=None):
+    """Group markers, component arrows, individuals and (optionally) one region per group."""
+    import matplotlib.pyplot as plt
+    if ax is None:
+        _, ax = plt.subplots(figsize=(7, 6.5))
+    A = fit.A
+    levs = fit.data.levels
+    if colors is None:
+        cmap = plt.get_cmap("tab10")
+        colors = {g: cmap(i % 10) for i, g in enumerate(levs)}
+    pols = {}
+    if regions is not None:
+        rr = {g: region for g in levs} if isinstance(region, str) else region
+        pols = {g: regions["regions"][rr[g]][g]["polygon"] for g in levs}
+    pts = [A] + ([fit.ind] if individuals and len(fit.ind) else []) + [p for p in pols.values() if p is not None]
+    lim = np.abs(np.vstack(pts)).max() * 1.05
+    if individuals and len(fit.ind):
+        gi = [fit.data.group[fit.data.id.index(i)] for i in fit.ind_ids]
+        ax.scatter(fit.ind[:, 0], fit.ind[:, 1], s=12, facecolors="none",
+                   edgecolors=[colors[g] for g in gi], alpha=0.6)
+    for g_i, g in enumerate(levs):
+        if pols.get(g) is not None:
+            ax.fill(pols[g][:, 0], pols[g][:, 1], color=colors[g], alpha=0.12)
+            ax.plot(pols[g][:, 0], pols[g][:, 1], color=colors[g], lw=1.5)
+        ax.scatter(A[g_i, 0], A[g_i, 1], color=colors[g], s=50, zorder=3)
+        ax.annotate(g, (A[g_i, 0], A[g_i, 1]), textcoords="offset points", xytext=(0, 6), ha="center",
+                    color=colors[g], fontweight="bold", fontsize=9)
+    if arrows:
+        C = fit.C
+        esc = arrow_scale if arrow_scale is not None else 0.8 * lim / np.sqrt((C ** 2).sum(1)).max()
+        names = labels_components or fit.data.components
+        for k in range(C.shape[0]):
+            ax.annotate("", xy=(esc * C[k, 0], esc * C[k, 1]), xytext=(0, 0),
+                        arrowprops=dict(arrowstyle="->", color="0.3"))
+            ax.text(esc * C[k, 0] * 1.12, esc * C[k, 1] * 1.12, names[k], color="0.3", fontsize=8, ha="center")
+    ax.axhline(0, ls=":", color="0.6"); ax.axvline(0, ls=":", color="0.6")
+    ax.set_xlim(-lim, lim); ax.set_ylim(-lim, lim); ax.set_aspect("equal")
+    lam = fit.lambda_
+    ax.set_xlabel(f"Dimension 1 ({100 * lam[0] / lam.sum():.0f}%)")
+    ax.set_ylabel(f"Dimension 2 ({100 * lam[1] / lam.sum():.0f}%)")
+    return ax

canbiplotme/regions.py

@@ -0,0 +1,94 @@
+"""Confidence regions R1-R4 for the group markers (section 3.5)."""
+from __future__ import annotations
+
+import numpy as np
+from scipy import stats as sps
+
+from .inference import _rng, procrustes_align, refit
+from .stage2 import CBMEFit
+
+
+def ellipse_polygon(center, S, q, n=100):
+    th = np.linspace(0, 2 * np.pi, n)
+    ev, evec = np.linalg.eigh(S)
+    ev, evec = ev[::-1], evec[:, ::-1]
+    P = np.column_stack([np.cos(th), np.sin(th)]) @ (evec @ np.diag(np.sqrt(np.maximum(ev, 0) * q))).T
+    return P + np.asarray(center)
+
+
+def polygon_area(P):
+    x, y = P[:, 0], P[:, 1]
+    return abs(np.sum(x * np.roll(y, -1) - np.roll(x, -1) * y)) / 2
+
+
+def cbme_regions(fit: CBMEFit, types=("R1", "R2", "R3", "R4"), level=0.95, B=999, npoints=100, seed=None) -> dict:
+    """Regions R1 (circle), R2 (Krzanowski ellipse), R3 (Hotelling type with
+    measurement error) and R4 (bootstrap of individuals aligned by Procrustes).
+
+    Returns ``{"regions": {type: {group: {"S", "q", "polygon", "area"}}}, "level", "A"}``.
+    """
+    if fit.A.shape[1] != 2:
+        raise ValueError("Regions are defined for the plane (dim = 2).")
+    rng = _rng(seed)
+    G, K, N = fit.A.shape[0], fit.W.shape[0], fit.n.sum()
+    levs = fit.data.levels
+    q2 = sps.chi2.ppf(level, 2)
+    out = {}
+    if "R1" in types:
+        out["R1"] = [{"S": np.eye(2) / fit.n[g], "q": q2} for g in range(G)]
+    if "R2" in types:
+        kk = G - 1
+        s = min(kk, K)
+        gam = kk / (N - G)
+        th = 1 + fit.lambda_[:s]
+        xi = fit.Afull[:, :s]
+        R2 = []
+        for g in range(G):
+            v = []
+            for l in range(2):
+                ex = sum(th[l] * (th[j] + gam * th[l]) / (gam * (th[j] - th[l]) ** 2) * xi[g, j] ** 2
+                         for j in range(s) if j != l)
+                v.append(1 / fit.n[g] + (0.5 * xi[g, l] ** 2 + ex) / (N - G))
+            cc = -th[0] * th[1] * (1 + gam) / ((N - G) * gam * (th[0] - th[1]) ** 2) * xi[g, 0] * xi[g, 1]
+            R2.append({"S": np.array([[v[0], cc], [cc, v[1]]]), "q": q2})
+        out["R2"] = R2
+    if "R3" in types:
+        V2 = fit.V[:, :2]
+        out["R3"] = [{"S": V2.T @ fit.mu_cov[g] @ V2,
+                      "q": 2 * (fit.n[g] - 1) / (fit.n[g] - 2) * sps.f.ppf(level, 2, fit.n[g] - 2)
+                      if fit.n[g] > 2 else np.nan} for g in range(G)]
+    if "R4" in types:
+        x = fit.data
+        reps = []
+        for _ in range(B):
+            idx = np.concatenate([rng.choice(np.where(x.group == lev)[0], size=np.sum(x.group == lev), replace=True)
+                                  for lev in levs])
+            core = refit(fit, x.beta[idx], rows=idx)
+            if core is not None:
+                reps.append(procrustes_align(core["A"], fit.A))
+        reps = np.array(reps)
+        out["R4"] = [{"S": np.cov(reps[:, g, :].T), "q": q2} for g in range(G)]
+    res = {}
+    for t, lst in out.items():
+        res[t] = {}
+        for g, r in enumerate(lst):
+            r["polygon"] = ellipse_polygon(fit.A[g], r["S"], r["q"], npoints) if np.isfinite(r["q"]) else None
+            r["area"] = polygon_area(r["polygon"]) if r["polygon"] is not None else np.nan
+            res[t][levs[g]] = r
+    return {"regions": res, "level": level, "A": fit.A, "levels": levs}
+
+
+def choose_region(regions: dict, coverage, min_cov: float = 0.90) -> dict:
+    """Rule of section 3.5.5: coverage closest to the nominal level without
+    falling below ``min_cov``; otherwise the region with the largest area.
+    ``coverage`` is a DataFrame (groups x region types)."""
+    out = {}
+    for g in regions["levels"]:
+        row = coverage.loc[g]
+        ok = row[row >= min_cov]
+        if len(ok):
+            out[g] = (ok - regions["level"]).abs().idxmin()
+        else:
+            areas = {t: regions["regions"][t][g]["area"] for t in regions["regions"]}
+            out[g] = max(areas, key=areas.get)
+    return out