canbiplotme 0.3.0__py3-none-any.whl

canbiplotme/reml.py

@@ -0,0 +1,124 @@
+"""REML / ML estimation of the two-level model (comparator, section 4.8).
+
+Python counterpart of ``cbme_reml`` (which uses ``metafor::rma.mv`` in R):
+each identified coefficient is an observation, the sampling covariance is
+block diagonal by individual (``errors='independent'``) or by group with the
+joint covariance (``errors='joint'``), the fixed effects are the means of every
+component in every group and the random effect of the individual has an
+unstructured covariance (the REML estimate of T).
+"""
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+from scipy import optimize, stats as sps
+
+from .data import MEData
+from .stage2 import make_cache
+
+
+def _long(x: MEData, errors: str):
+    ok = np.isfinite(x.beta)
+    ind, comp = np.where(ok)                       # individual-major order
+    y = x.beta[ind, comp]
+    n = len(y)
+    Vm = np.zeros((n, n))
+    if errors == "independent":
+        for i in np.unique(ind):
+            r = np.where(ind == i)[0]
+            Vm[np.ix_(r, r)] = x.V[i][np.ix_(comp[r], comp[r])]
+    else:
+        cache = make_cache(x)
+        for lev in x.levels:
+            r = np.where(x.group[ind] == lev)[0]
+            cols = cache["pos"][ind[r]] * x.K + comp[r]
+            Vm[np.ix_(r, r)] = cache["Omega"][lev][np.ix_(cols, cols)]
+    return y, Vm, ind, comp
+
+
+def _blocks(x: MEData, ind, errors):
+    """Index sets of independent blocks of the marginal covariance."""
+    if errors == "independent":
+        return [np.where(ind == i)[0] for i in np.unique(ind)]
+    return [np.where(x.group[ind] == lev)[0] for lev in x.levels]
+
+
+def _theta_to_T(th, K):
+    Lc = np.zeros((K, K))
+    Lc[np.tril_indices(K)] = th
+    Lc[np.diag_indices(K)] = np.exp(Lc[np.diag_indices(K)])
+    return Lc @ Lc.T
+
+
+def _T_to_theta(T, K):
+    Lc = np.linalg.cholesky(T + 1e-10 * np.eye(K))
+    Lc[np.diag_indices(K)] = np.log(np.diag(Lc))
+    return Lc[np.tril_indices(K)]
+
+
+def _loglik(th, y, X, Vm, ind, comp, blocks, K, reml):
+    T = _theta_to_T(th, K)
+    n, p = X.shape
+    logdet = 0.0
+    XtSiX = np.zeros((p, p)); XtSiy = np.zeros(p); ytSiy = 0.0
+    SiX_blocks = []
+    for r in blocks:
+        S = Vm[np.ix_(r, r)] + np.where(ind[r][:, None] == ind[r][None, :], T[np.ix_(comp[r], comp[r])], 0.0)
+        try:
+            c = np.linalg.cholesky(S)
+        except np.linalg.LinAlgError:
+            return -np.inf, None, None
+        logdet += 2 * np.log(np.diag(c)).sum()
+        Xi = np.linalg.solve(c, X[r]); yi = np.linalg.solve(c, y[r])
+        XtSiX += Xi.T @ Xi; XtSiy += Xi.T @ yi; ytSiy += yi @ yi
+    b = np.linalg.solve(XtSiX, XtSiy)
+    q = ytSiy - 2 * b @ XtSiy + b @ XtSiX @ b
+    if reml:
+        # same constant as metafor: + 0.5 log|X'X|
+        ll = -0.5 * ((n - p) * np.log(2 * np.pi) + logdet + np.linalg.slogdet(XtSiX)[1] + q) \
+            + 0.5 * np.linalg.slogdet(X.T @ X)[1]
+    else:
+        ll = -0.5 * (n * np.log(2 * np.pi) + logdet + q)
+    return ll, b, XtSiX
+
+
+def _fit(y, X, Vm, ind, comp, blocks, K, reml, T0):
+    f = lambda th: -_loglik(th, y, X, Vm, ind, comp, blocks, K, reml)[0]
+    th0 = _T_to_theta(T0, K)
+    best = optimize.minimize(f, th0, method="BFGS", options={"gtol": 1e-8, "maxiter": 2000})
+    res2 = optimize.minimize(f, best.x, method="Nelder-Mead", options={"xatol": 1e-10, "fatol": 1e-12, "maxiter": 20000})
+    if res2.fun < best.fun:
+        best = res2
+    ll, b, XtSiX = _loglik(best.x, y, X, Vm, ind, comp, blocks, K, reml)
+    return {"T": _theta_to_T(best.x, K), "b": b, "vcov": np.linalg.inv(XtSiX), "logLik": ll,
+            "converged": bool(best.success or res2.success)}
+
+
+def cbme_reml(x: MEData, errors: str = "independent", T0=None) -> dict:
+    """REML estimate of T and of the group means, and likelihood-ratio test
+    (ML) of group-specific versus common means. The size of the test is not
+    validated; it is a reference, as in the thesis."""
+    if errors == "joint" and not x.has_joint:
+        errors = "independent"
+    y, Vm, ind, comp = _long(x, errors)
+    K = x.K
+    blocks = _blocks(x, ind, errors)
+    gl = x.group[ind]
+    cells = [(g, k) for g in x.levels for k in range(K) if np.any((gl == g) & (comp == k))]
+    X1 = np.column_stack([((gl == g) & (comp == k)).astype(float) for g, k in cells])
+    X0 = np.column_stack([(comp == k).astype(float) for k in range(K)])
+    if T0 is None:
+        from .stage2 import cbme
+        T0 = cbme(x).T
+    reml = _fit(y, X1, Vm, ind, comp, blocks, K, True, T0)
+    ml1 = _fit(y, X1, Vm, ind, comp, blocks, K, False, reml["T"])
+    ml0 = _fit(y, X0, Vm, ind, comp, blocks, K, False, reml["T"])
+    means = pd.DataFrame(np.nan, index=x.levels, columns=x.components)
+    for (g, k), v in zip(cells, reml["b"]):
+        means.loc[g, x.components[k]] = v
+    lrt = 2 * (ml1["logLik"] - ml0["logLik"])
+    df = X1.shape[1] - X0.shape[1]
+    return {"T": pd.DataFrame(reml["T"], index=x.components, columns=x.components), "means": means,
+            "lrt": {"LRT": lrt, "df": df, "p": float(sps.chi2.sf(lrt, df))},
+            "logLik_reml": reml["logLik"], "logLik_ml_groups": ml1["logLik"], "logLik_ml_common": ml0["logLik"],
+            "converged": reml["converged"] and ml1["converged"] and ml0["converged"]}

canbiplotme/rng.py

@@ -0,0 +1,100 @@
+"""Random numbers identical to those of R.
+
+``RRNG(seed)`` reproduces ``set.seed(seed); rnorm(n)`` of R (Mersenne-Twister
+with R's seed scrambling, normal draws by inversion with
+``qnorm`` of Wichura, algorithm AS 241). With it the parametric bootstrap of
+the package gives, seed by seed, the replicates of the R package ``gcbiplot``
+and of the code of the thesis. About 97% of the normal draws are bitwise equal
+to R's and the rest differ by at most a few units in the last place (rounding
+inside ``qnorm``), which does not change any p-value in practice.
+"""
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+_I2_32M1 = 2.328306437080797e-10      # 1 / (2^32 - 1)
+_BIG = 134217728.0                    # 2^27
+
+
+def _r_seed_state(seed: int) -> np.ndarray:
+    s = int(seed) & 0xFFFFFFFF
+    for _ in range(50):
+        s = (69069 * s + 1) & 0xFFFFFFFF
+    st = []
+    for _ in range(625):
+        s = (69069 * s + 1) & 0xFFFFFFFF
+        st.append(s)
+    return np.array(st[1:], dtype=np.uint32)
+
+
+def qnorm_r(p):
+    """R's qnorm(p, 0, 1) (Wichura AS241, nmath/qnorm.c), vectorized."""
+    p = np.asarray(p, float); q = p - 0.5; out = np.empty_like(p)
+    c = np.abs(q) <= 0.425
+    r = 0.180625 - q[c] * q[c]
+    out[c] = q[c] * (((((((r * 2509.0809287301226727 + 33430.575583588128105) * r + 67265.770927008700853) * r
+                + 45921.953931549871457) * r + 13731.693765509461125) * r + 1971.5909503065514427) * r + 133.14166789178437745) * r
+                + 3.387132872796366608) / (((((((r * 5226.495278852545925 + 28729.085735721942674) * r + 39307.89580009271061) * r
+                + 21213.794301586595867) * r + 5394.1960214247511077) * r + 687.1870074920579083) * r + 42.313330701600911252) * r + 1.)
+    nc = ~c
+    pp = np.where(q[nc] < 0, p[nc], 1 - p[nc])
+    r = np.sqrt(-np.array([math.log(v) for v in pp]))
+    v = np.empty_like(r)
+    a = r <= 5.
+    ra = r[a] - 1.6
+    v[a] = (((((((ra * 7.7454501427834140764e-4 + .0227238449892691845833) * ra + .24178072517745061177) * ra
+              + 1.27045825245236838258) * ra + 3.64784832476320460504) * ra + 5.7694972214606914055) * ra + 4.6303378461565452959) * ra
+              + 1.42343711074968357734) / (((((((ra * 1.05075007164441684324e-9 + 5.475938084995344946e-4) * ra + .0151986665636164571966) * ra
+              + .14810397642748007459) * ra + .68976733498510000455) * ra + 1.6763848301838038494) * ra + 2.05319162663775882187) * ra + 1.)
+    b = ~a
+    rb = r[b] - 5.
+    v[b] = (((((((rb * 2.01033439929228813265e-7 + 2.71155556874348757815e-5) * rb + .0012426609473880784386) * rb
+              + .026532189526576123093) * rb + .29656057182850489123) * rb + 1.7848265399172913358) * rb + 5.4637849111641143699) * rb
+              + 6.6579046435011037772) / (((((((rb * 2.04426310338993978564e-15 + 1.4215117583164458887e-7) * rb + 1.8463183175100546818e-5) * rb
+              + 7.868691311456132591e-4) * rb + .0148753612908506148525) * rb + .13692988092273580531) * rb + .59983220655588793769) * rb + 1.)
+    v[q[nc] < 0] = -v[q[nc] < 0]
+    out[nc] = v
+    return out
+
+
+class RRNG:
+    """Generator compatible with R's default ``RNGkind`` (Mersenne-Twister, Inversion).
+
+    It exposes ``standard_normal`` and ``integers`` so that it can be passed
+    wherever the package accepts a ``numpy.random.Generator``. Draws are taken
+    in the order R takes them (``matrix(rnorm(n), ...)`` fills by columns).
+    """
+
+    def __init__(self, seed: int):
+        self.bg = np.random.MT19937()
+        self.bg.state = {"bit_generator": "MT19937", "state": {"key": _r_seed_state(seed), "pos": 624}}
+
+    def unif(self, n: int) -> np.ndarray:
+        """``runif(n)`` of R."""
+        v = self.bg.random_raw(n).astype(np.float64) * 2.3283064365386963e-10
+        v[v <= 0.0] = 0.5 * _I2_32M1
+        v[1.0 - v <= 0.0] = 1.0 - 0.5 * _I2_32M1
+        return v
+
+    def standard_normal(self, size=None) -> np.ndarray:
+        """``rnorm(n)`` of R; ``size`` may be an int or a shape (filled in C order)."""
+        n = 1 if size is None else int(np.prod(size))
+        u = self.unif(2 * n).reshape(n, 2)
+        z = qnorm_r((np.floor(_BIG * u[:, 0]) + u[:, 1]) / _BIG)
+        return z[0] if size is None else z.reshape(size)
+
+    def integers(self, high: int) -> int:
+        """``sample.int(high, 1) - 1`` of R (sample.kind = "Rejection")."""
+        bits = int(math.ceil(math.log2(high))) if high > 1 else 0
+        while True:
+            v = 0.0
+            for _ in range(0, bits + 1, 16):
+                v1 = int(math.floor(self.unif(1)[0] * 65536))
+                v = 65536 * v + v1
+            if bits > 0:
+                one64 = 1
+                v = int(v) & ((one64 << bits) - 1)
+            if v < high:
+                return int(v)

canbiplotme/simulation.py

@@ -0,0 +1,77 @@
+"""Size, power and coverage studies by simulation (section 4.6)."""
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+from scipy import stats as sps
+
+from .inference import _prepare, _rng, _with_core, cbme_test, procrustes_align, refit, simulate_coefs
+from .regions import cbme_regions
+from .stage2 import CBMEFit
+
+
+def cbme_simulate_tests(fit: CBMEFit, nsim=1000, B=199, scale=0.0, generator="model", errors="joint",
+                        level=0.95, seed=None) -> dict:
+    """Rejection rates of the bootstrap tests in ``nsim`` datasets generated from
+    the fitted model, each tested with its own bootstrap of ``B`` replicates.
+    ``scale=0``: null hypothesis (size); ``scale>0``: power. ``generator`` is
+    ``'model'`` or ``'t4-replicate'`` (non-normal validation of the thesis)."""
+    rng = _rng(seed)
+    _prepare(fit)
+    alpha = 1 - level
+    res = []
+    for _ in range(nsim):
+        means = "common" if scale == 0 else "group"
+        if generator == "model":
+            b = simulate_coefs(fit, means, errors, scale=scale, rng=rng)
+        else:
+            b = simulate_coefs(fit, means, "replicate", "t4", scale=scale, rng=rng)
+        core = refit(fit, b)
+        if core is None:
+            continue
+        fs = _with_core(fit, core, beta=b)
+        tt = cbme_test(fs, B=B, errors=errors, keep=False, seed=rng)
+        res.append({"trace": tt.global_["trace"] < alpha, "lambda1": tt.global_["lambda1"] < alpha,
+                    "any_pair": bool((tt.pairs["p_holm"] < alpha).any()),
+                    "mean_pair": float((tt.pairs["p"] < alpha).mean()),
+                    "n_pairs": int((tt.pairs["p_holm"] < alpha).sum()),
+                    "composition": tt.composition["holds"] if tt.composition else np.nan})
+    res = pd.DataFrame(res)
+    ns = len(res)
+    mc = alpha + np.array([-1, 1]) * sps.norm.ppf(0.995) * np.sqrt(alpha * (1 - alpha) / ns)
+    return {"rates": res.mean(numeric_only=True), "mc_interval_99": mc, "nsim_valid": ns, "scale": scale,
+            "generator": generator, "B": B, "results": res}
+
+
+def cbme_coverage(fit: CBMEFit, nsim=500, B_R4=199, types=("R1", "R2", "R3", "R4"), level=0.95, seed=None) -> dict:
+    """Coverage of the regions with data generated with the estimated means,
+    heterogeneity and errors; targets: projection of the true means on the
+    estimated axes (``conditional``) and population marker aligned by
+    Procrustes (``aligned``)."""
+    rng = _rng(seed)
+    _prepare(fit)
+    levs = fit.data.levels
+    G = len(levs)
+    mbar_true = (fit.mu * fit.n[:, None]).sum(0) / fit.n.sum()
+    cond = np.zeros((G, len(types))); alig = np.zeros((G, len(types))); cnt = 0
+    for _ in range(nsim):
+        b = simulate_coefs(fit, "group", rng=rng)
+        core = refit(fit, b)
+        if core is None:
+            continue
+        fs = _with_core(fit, core, beta=b)
+        rg = cbme_regions(fs, types=types, level=level, B=B_R4, seed=rng)
+        tc = (fit.mu - mbar_true) @ fs.V[:, :2]
+        ta = procrustes_align(fit.A, fs.A)
+        for ti, t in enumerate(types):
+            for g, lev in enumerate(levs):
+                r = rg["regions"][t][lev]
+                if not np.isfinite(r["q"]):
+                    continue
+                Si = np.linalg.inv(r["S"])
+                for target, acc in ((tc, cond), (ta, alig)):
+                    d = target[g] - fs.A[g]
+                    acc[g, ti] += float(d @ Si @ d <= r["q"])
+        cnt += 1
+    return {"conditional": pd.DataFrame(cond / cnt, index=levs, columns=list(types)),
+            "aligned": pd.DataFrame(alig / cnt, index=levs, columns=list(types)), "nsim_valid": cnt}

canbiplotme/stage1.py

@@ -0,0 +1,189 @@
+"""First stage for gravity data: individual PPML coefficients, pairs bootstrap and rule D6.
+
+Requires ``pyfixest`` (``pip install canbiplotme[gravity]``).
+"""
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+
+from .data import MEData
+
+
+def _pf():
+    try:
+        import pyfixest as pf
+    except ImportError as e:  # pragma: no cover
+        raise ImportError("The first stage needs pyfixest: pip install pyfixest") from e
+    return pf
+
+
+def _prepare(data: pd.DataFrame, origin, destination, frictions):
+    """Explicit origin x friction interactions (more portable than i() syntax)."""
+    d = data.copy()
+    d[origin] = d[origin].astype(str)
+    d[destination] = d[destination].astype(str)
+    d["_pair"] = d[origin] + "_" + d[destination]
+    return d
+
+
+def _design(d, origin, frictions, ex):
+    cols = {}
+    names = []
+    for f in frictions:
+        for e in ex:
+            nm = f"b__{e}__{f}"
+            cols[nm] = np.where(d[origin].values == e, d[f].values, 0.0)
+            names.append(nm)
+    return pd.concat([d.reset_index(drop=True), pd.DataFrame(cols)], axis=1), names
+
+
+def _fml(flow, names, controls, origin, destination, year, extra=()):
+    rhs = " + ".join(list(names) + list(controls or []) + list(extra))
+    return f"{flow} ~ {rhs} | {origin}^{year} + {destination}^{year}"
+
+
+def _fit(pf, fml, data, vcov):
+    return pf.fepois(fml, data=data, vcov=vcov, iwls_maxiter=100)
+
+
+def _extract(model, ex, frictions, with_vcov=True):
+    coef = model.coef()
+    K = len(frictions)
+    beta = np.full((len(ex), K), np.nan)
+    V = np.full((len(ex), K, K), np.nan)
+    vc = None
+    if with_vcov:
+        vc = pd.DataFrame(model._vcov, index=coef.index, columns=coef.index)
+    for j, e in enumerate(ex):
+        nm = [f"b__{e}__{f}" for f in frictions]
+        pres = [n in coef.index and np.isfinite(coef[n]) for n in nm]
+        idx = [k for k in range(K) if pres[k]]
+        for k in idx:
+            beta[j, k] = coef[nm[k]]
+        if with_vcov and idx:
+            sel = [nm[k] for k in idx]
+            V[j][np.ix_(idx, idx)] = vc.loc[sel, sel].values
+    return beta, V
+
+
+class GravityStage1:
+    def __init__(self, beta, V, group, id, spec):
+        self.beta, self.V, self.group, self.id, self.spec = beta, V, group, id, spec
+
+    def __repr__(self):
+        return (f"<GravityStage1> {self.beta.shape[0]} individuals in {len(set(self.group))} groups; "
+                f"frictions: {', '.join(self.spec['frictions'])}")
+
+
+def gravity_by_individual(data: pd.DataFrame, flow, origin, destination, year, frictions, group,
+                          controls=None, cluster="pair") -> GravityStage1:
+    """PPML with a coefficient vector of its own for each origin in the compared
+    frictions, common controls and origin-year and destination-year fixed
+    effects, estimated separately in each group (section 3.2.1)."""
+    pf = _pf()
+    d0 = _prepare(data, origin, destination, frictions)
+    vc = {"pair": {"CRV1": "_pair"}, "destination": {"CRV1": destination},
+          "twoway": {"CRV1": f"{origin}+{destination}"}}[cluster]
+    spec = dict(flow=flow, origin=origin, destination=destination, year=year, frictions=list(frictions),
+                group=group, controls=list(controls or []))
+    betas, Vs, grps, ids = [], [], [], []
+    for g in sorted(d0[group].dropna().astype(str).unique()):
+        d = d0[d0[group].astype(str) == g]
+        ex = sorted(d[origin].unique())
+        dd, names = _design(d, origin, frictions, ex)
+        m = _fit(pf, _fml(flow, names, controls, origin, destination, year), dd, vc)
+        b, V = _extract(m, ex, frictions)
+        betas.append(b); Vs.append(V); grps += [g] * len(ex); ids += ex
+    return GravityStage1(np.vstack(betas), np.concatenate(Vs), np.array(grps, dtype=object), ids, spec)
+
+
+def refit_on_pairs(s1: GravityStage1, data: pd.DataFrame, g: str, pairs) -> np.ndarray:
+    """Re-estimate group ``g`` on a list of drawn pairs (repeats allowed)."""
+    pf = _pf()
+    sp = s1.spec
+    d = _prepare(data, sp["origin"], sp["destination"], sp["frictions"])
+    d = d[d[sp["group"]].astype(str) == g]
+    by_pair = {k: v for k, v in d.groupby("_pair").indices.items()}
+    rows = np.concatenate([by_pair[p] for p in pairs])
+    db = d.iloc[rows]
+    ex = [i for i, gg in zip(s1.id, s1.group) if gg == g]
+    dd, names = _design(db, sp["origin"], sp["frictions"], ex)
+    try:
+        m = _fit(pf, _fml(sp["flow"], names, sp["controls"], sp["origin"], sp["destination"], sp["year"]), dd, "iid")
+    except Exception:
+        return np.full((len(ex), len(sp["frictions"])), np.nan)
+    return _extract(m, ex, sp["frictions"], with_vcov=False)[0]
+
+
+def pairs_bootstrap(s1: GravityStage1, data: pd.DataFrame, R=199, seed=20251025, min_valid=0.8) -> dict:
+    """Pairs bootstrap (section 4.4.3): resamples the origin-destination pairs of
+    each group with replacement (each with all its years) and re-estimates the
+    model with its fixed effects. Keeps the replicates (R x n_g x K per group)."""
+    sp = s1.spec
+    d = _prepare(data, sp["origin"], sp["destination"], sp["frictions"])
+    K = len(sp["frictions"])
+    reps, V, valid = {}, np.full_like(s1.V, np.nan), np.full_like(s1.beta, np.nan)
+    groups = list(dict.fromkeys(s1.group))
+    for k, g in enumerate(groups):
+        pairs = np.array(sorted(d.loc[d[sp["group"]].astype(str) == g, "_pair"].unique()))
+        ex_idx = [i for i, gg in enumerate(s1.group) if gg == g]
+        arr = np.full((R, len(ex_idx), K), np.nan)
+        for r in range(R):
+            rng = np.random.default_rng(seed + 10000 * (k + 1) + r + 1)
+            arr[r] = refit_on_pairs(s1, data, g, pairs[rng.integers(len(pairs), size=len(pairs))])
+        reps[g] = arr
+        for j, i in enumerate(ex_idx):
+            M = arr[:, j, :]
+            valid[i] = np.isfinite(M).mean(0)
+            for a in range(K):
+                for b in range(K):
+                    ok = np.isfinite(M[:, a]) & np.isfinite(M[:, b])
+                    if ok.sum() >= min_valid * R and ok.sum() > 1:
+                        V[i, a, b] = np.cov(M[ok, a], M[ok, b])[0, 1]
+    return {"replicates": reps, "V": V, "valid": valid, "R": R, "seed": seed}
+
+
+def support_counts(data: pd.DataFrame, origin, destination, binary) -> pd.DataFrame:
+    """Distinct destinations with value one and with value zero of each binary
+    friction (a destination whose value changes counts in both)."""
+    rows = []
+    for o, d in data.groupby(origin):
+        r = {"origin": str(o)}
+        for b in binary:
+            r[f"{b}_1"] = d.loc[d[b] > 0, destination].nunique()
+            r[f"{b}_0"] = d.loc[d[b] == 0, destination].nunique()
+        rows.append(r)
+    return pd.DataFrame(rows)
+
+
+def identify_d6(beta, V, ids, counts=None, binary=(), se_max=2.0, min_dest=5, valid=None, min_valid=0.8,
+                components=None) -> np.ndarray:
+    """Identification rule D6 (section 3.2.2). Returns a boolean (N, K) mask.
+
+    ``valid`` (share of valid replicates) is an optional extra filter that is
+    not part of the rule of the thesis; leave it as ``None`` to apply that rule.
+    """
+    beta = np.asarray(beta, float)
+    se = np.sqrt(np.maximum(np.einsum("ikk->ik", np.asarray(V, float)), 0))
+    ok = np.isfinite(beta) & np.isfinite(se) & (se < se_max)
+    if valid is not None:
+        ok &= np.isfinite(valid) & (np.asarray(valid) >= min_valid)
+    if len(binary) and counts is not None:
+        cs = counts.set_index("origin").reindex([str(i) for i in ids])
+        for b in binary:
+            j = components.index(b)
+            ok[:, j] &= (cs[f"{b}_1"].values >= min_dest) & (cs[f"{b}_0"].values >= min_dest)
+    return ok
+
+
+def as_me_data(s1: GravityStage1, boot: dict | None = None, identified=None) -> MEData:
+    """Second-stage data from the first stage (bootstrap covariances and
+    replicates when ``boot`` is given)."""
+    beta = s1.beta.copy()
+    if identified is not None:
+        beta[~identified] = np.nan
+    if boot is None:
+        return MEData(beta, s1.V, s1.group, id=s1.id, components=s1.spec["frictions"])
+    return MEData(beta, boot["V"], s1.group, replicates=boot["replicates"], id=s1.id,
+                  components=s1.spec["frictions"])