Core

The pyem.core subpackage implements the hierarchical Expectation-Maximization (EM) machinery that pyem.api.EMModel wraps: the outer EM loop, the subject-level (E-step) optimizer, the group-level (M-step) distribution families, prior objects, model-comparison utilities, parameter-recovery diagnostics, and the ModelSpec bundle used to register a model's identity and entry points. Most users will interact with these through EMModel rather than calling them directly, but they are documented here for anyone building custom fitting pipelines or new group-distribution families.

Hierarchical EM loop

EMfit is the low-level hierarchical-EM routine: it alternates a parallel per-subject E-step (single_subject_minimize, MAP estimation against the current group prior) with a group-level M-step (see Group M-step families) until a convergence criterion is met, returning subject-level estimates, inverse-Hessians, the final group posterior, and per-subject (negative) log-likelihood terms. EMConfig bundles the loop's tunables — iteration cap, convergence method and criterion, parallelism, RNG seed, and which M-step family to use.

EMfit

EMfit(all_data: Sequence[Sequence[Any]] | Sequence[DataFrame], objfunc: Callable[..., float], param_names: Sequence[str], *, verbose: int = 1, config: EMConfig | None = None, prior: Prior | None = None, **kwargs) -> dict[str, Any]

Hierarchical EM with MAP. Compatible with legacy signature.

list where item i contains args for objfunc for subject i

(e.g., [choices, rewards]) or a pandas DataFrame for subject i.

objfunc: callable like f(params, *subject_args, prior=None, output='npl') -> float param_names: list of parameter names (nparams)

Source code in pyem/core/em.py

def EMfit(
    all_data: Sequence[Sequence[Any]] | Sequence[pd.DataFrame],
    objfunc: Callable[..., float],
    param_names: Sequence[str],
    *,
    verbose: int = 1,
    config: EMConfig | None = None,
    prior: Prior | None = None,
    **kwargs,
) -> dict[str, Any]:
    """
    Hierarchical EM with MAP. Compatible with legacy signature.

    all_data: list where item i contains args for objfunc for subject i
              (e.g., [choices, rewards]) or a pandas DataFrame for subject i.
    objfunc: callable like f(params, *subject_args, prior=None, output='npl') -> float
    param_names: list of parameter names (nparams)
    """
    nsubjects = len(all_data)
    nparams = len(param_names)
    if nsubjects == 0:
        raise ValueError("EMfit received empty `all_data`: no subjects to fit.")
    if nparams == 0:
        raise ValueError("EMfit received empty `param_names`: at least one parameter is required.")
    for i, _args in enumerate(all_data):
        _empty = _args.empty if isinstance(_args, pd.DataFrame) else (len(_args) == 0)
        if _empty:
            raise ValueError(
                f"all_data[{i}] is empty; each subject's entry must contain non-empty "
                "arguments/data for objfunc."
            )
    if config is None:
        config = EMConfig()
    # Backward-compat kwargs mapping (e.g., mstep_maxit=..., njobs=..., etc.)
    if kwargs:
        if 'mstep_maxit' in kwargs: config.mstep_maxit = int(kwargs['mstep_maxit'])
        if 'convergence_method' in kwargs: config.convergence_method = kwargs['convergence_method']
        if 'convergence_custom' in kwargs: config.convergence_custom = kwargs['convergence_custom']
        if 'convergence_crit' in kwargs: config.convergence_crit = float(kwargs['convergence_crit'])
        if 'convergence_precision' in kwargs: config.convergence_precision = int(kwargs['convergence_precision'])
        if 'njobs' in kwargs: config.njobs = int(kwargs['njobs'])
        if 'seed' in kwargs: config.seed = kwargs['seed']
        if 'mstep' in kwargs: config.mstep = kwargs['mstep']
        # Optimizer knobs
        if 'optim_method' in kwargs or 'optim_options' in kwargs or 'max_restarts' in kwargs:
            from .optim import OptimConfig
            config.optim = OptimConfig(
                method=kwargs.get('optim_method', config.optim.method),
                options=kwargs.get('optim_options', config.optim.options),
                max_restarts=kwargs.get('max_restarts', config.optim.max_restarts)
            )
    if config.mstep_maxit < 1:
        raise ValueError("mstep_maxit must be >= 1")
    if prior is None:
        prior = default_prior(nparams, seed=config.seed)

    group = make_group(config.mstep)

    # initialize tracking
    NPL_hist = []
    NPL_old = np.inf
    converged = False

    # initial posterior (seeded from the user prior's own moments when available)
    if hasattr(prior, "init_moments"):
        pm, ps = prior.init_moments()
        post_mu, post_sigma = np.asarray(pm, float).copy(), np.asarray(ps, float).copy()
    elif hasattr(prior, "mu") and hasattr(prior, "sigma"):
        post_mu, post_sigma = np.asarray(prior.mu, float).copy(), np.asarray(prior.sigma, float).copy()
    else:
        raise TypeError("prior must define init_moments() or expose .mu/.sigma (GaussianPrior-like)")
    last_good_hyper = None

    with Parallel(n_jobs=config.njobs) as parallel:
        for iiter in range(config.mstep_maxit):
            # Use the user-supplied prior until a valid (ok) M-step has occurred;
            # afterwards use the chosen group family's prior built from the last
            # known-good M-step. This mirrors the pre-refactor semantics, which only
            # ever fed a validated prior into the next E-step.
            iter_prior = prior if last_good_hyper is None else group.make_prior(last_good_hyper)

            # E-step in parallel
            def _fit_subject(i: int, iter_prior=iter_prior):
                args = all_data[i]
                # normalize args into tuple for *args
                if isinstance(args, pd.DataFrame):
                    obj_args = (args,)
                else:
                    obj_args = tuple(args)
                subject_seed = None if config.seed is None else config.seed + i
                return single_subject_minimize(
                    objfunc=objfunc,
                    obj_args=obj_args,
                    nparams=nparams,
                    prior=iter_prior,
                    config=config.optim,
                    rng=np.random.default_rng(subject_seed)
                )

            results = parallel(delayed(_fit_subject)(i) for i in range(nsubjects))

            m = np.zeros((nparams, nsubjects))
            inv_h = np.zeros((nparams, nparams, nsubjects))
            NPL = np.zeros(nsubjects)
            NLP = np.zeros(nsubjects)

            for i, (q_est, hess_inv, fval, nl_prior, success, res) in enumerate(results):
                m[:, i] = q_est
                inv_h[:, :, i] = hess_inv
                NPL[i] = round(fval, config.convergence_precision)
                NLP[i] = nl_prior

            # M-step: empirical Bayes update via the chosen group-distribution family
            hyper = group.update(m, inv_h)
            mu, sigma, ok = group.moments(hyper)
            if ok:
                post_mu = mu
                post_sigma = sigma
                last_good_hyper = hyper
            # if not ok: keep last_good_hyper and post_mu/post_sigma from the last good step

            conv_val = _hier_convergence(NPL, config.convergence_method)
            if verbose:
                if (iiter == 0) or conv_val <= min(NPL_hist + [conv_val]):
                    if config.convergence_custom == "relative_npl" and iiter > 0:
                        print(f"{abs((conv_val - NPL_old) / (NPL_old if NPL_old != 0 else 1.0)):.4f} ({iiter:03d})", end=", ")
                    else:
                        print(f"{conv_val:.4f} ({iiter:03d})", end=", ")

            # convergence check
            if config.convergence_custom == "running_average" and ok and iiter > 5:
                if abs(conv_val - float(np.mean(NPL_hist[-5:]))) < config.convergence_crit:
                    converged = True
            elif config.convergence_custom == "relative_npl" and ok and iiter > 1:
                if abs((conv_val - NPL_old) / (NPL_old if NPL_old != 0 else 1.0)) < config.convergence_crit:
                    converged = True
            else:
                if iiter > 0 and abs(conv_val - NPL_old) < config.convergence_crit and ok:
                    converged = True

            NPL_hist.append(conv_val)
            NPL_old = conv_val

            if converged:
                break

    out = {
        "m": m,
        "inv_h": inv_h,
        "posterior": {"mu": post_mu, "sigma": post_sigma},
        "NPL": NPL,
        "NLPrior": NLP,
        "NLL": NPL - NLP,
        "convergence": converged,
    }
    return out

EMConfig dataclass

EMConfig(mstep_maxit: int = 200, convergence_method: ConvergenceMethod = 'sum', convergence_custom: Literal['relative_npl', 'running_average', None] = None, convergence_crit: float = 0.001, convergence_precision: int = 6, njobs: int = -2, optim: OptimConfig = OptimConfig(), seed: int | None = None, mstep: str = 'gaussian')

Subject-level optimizer

single_subject_minimize performs the E-step for a single subject: it minimizes the negative log-posterior (objfunc(...) - logprior) over that subject's data with scipy.optimize.minimize, retrying from random restarts until a successful optimization is found or the restart budget (OptimConfig.max_restarts) is exhausted. OptimConfig holds the optimizer method, solver options, restart budget, and the scale of the random restart initializations.

OptimConfig dataclass

OptimConfig(method: str = 'BFGS', options: dict | None = None, max_restarts: int = 2, x_scale: float = 0.1)

single_subject_minimize

single_subject_minimize(objfunc: ObjectiveFn, obj_args: Iterable[Any], nparams: int, prior: Prior, config: OptimConfig, rng: Generator) -> tuple[np.ndarray, np.ndarray, float, float, bool, OptimizeResult]

Returns (q_est, inv_h, fval, nl_prior, success, result)

Source code in pyem/core/optim.py

def single_subject_minimize(
    objfunc: ObjectiveFn,
    obj_args: Iterable[Any],
    nparams: int,
    prior: Prior,
    config: OptimConfig,
    rng: np.random.Generator
) -> tuple[np.ndarray, np.ndarray, float, float, bool, OptimizeResult]:
    """
    Returns (q_est, inv_h, fval, nl_prior, success, result)
    """
    options = {"gtol": 1e-4, "eps": 1e-4}
    if config.options:
        options.update(config.options)

    best = None
    best_result: Optional[OptimizeResult] = None
    for attempt in range(1 + config.max_restarts):
        x0 = config.x_scale * rng.standard_normal(nparams)
        result = minimize(
            lambda x, *args: objfunc(x, *args, prior=prior),
            x0=x0,
            args=tuple(obj_args),
            method=config.method,
            options=options
        )
        if best is None or result.fun < best:
            best = float(result.fun)
            best_result = result
        if result.success:
            break  # accept successful result

    assert best_result is not None  # for type checker

    q_est = best_result.x
    fval = float(best_result.fun)
    # Hessian inverse (BFGS provides it). Fallback: diagonal approx.
    if hasattr(best_result, "hess_inv"):
        h = best_result.hess_inv
        try:
            inv_h = np.asarray(h)
        except Exception:
            n = nparams
            inv_h = np.eye(n) * max(1.0, np.linalg.norm(q_est) + 1e-6)
    else:
        n = nparams
        inv_h = np.eye(n) * max(1.0, np.linalg.norm(q_est) + 1e-6)

    nl_prior = -prior.logpdf(q_est)
    success = bool(best_result.success)
    return q_est, inv_h, fval, nl_prior, success, best_result

Model comparison

ModelComparison wraps one or more fitted EMModel instances and computes per-model fit metrics — log model evidence (LME), integrated BIC (BICint), and pseudo-R² — via compare(), and can additionally run a full simulate-and-refit identifiability analysis across models via identify() (with plot_identifiability() for visualizing the resulting confusion matrix). compare_models is the underlying function that produces the list of ComparisonRow results that ModelComparison.compare() turns into a DataFrame.

import numpy as np
from pyem import EMModel
from pyem.core.compare import ModelComparison
from pyem.models.rl_mf import rw1a1b_sim, rw1a1b_fit
from pyem.utils.math import norm2beta, norm2alpha

rng = np.random.default_rng(0)
p1 = np.column_stack([norm2beta(rng.normal(size=15)), norm2alpha(rng.normal(size=15))])
sim = rw1a1b_sim(p1, nblocks=2, ntrials=20, seed=0)
data = [[sim["choices"][s], sim["rewards"][s]] for s in range(15)]
m1 = EMModel(data, rw1a1b_fit, ["beta", "alpha"], param_xform=[norm2beta, norm2alpha], simulate_func=rw1a1b_sim)
m1.fit(verbose=0)
cmp = ModelComparison([m1], ["rw1a1b"])
print(cmp.compare(bicint_kwargs={"nsamples": 50, "func_output": "all", "nll_key": "nll"}))

compare_models

compare_models(models, model_names: List[str] | None = None, metric_order: Sequence[str] = ('LME', 'BICint', 'R2'), bicint_kwargs: dict | None = None, r2_kwargs: dict | None = None) -> List[ComparisonRow]

Source code in pyem/core/compare.py

def compare_models(
    models,  # list of EMModel (already fit) or tuples (name, FitResult, extras)
    model_names: List[str] | None = None,
    metric_order: Sequence[str] = ("LME", "BICint", "R2"),
    bicint_kwargs: dict | None = None,
    r2_kwargs: dict | None = None,
) -> List[ComparisonRow]:
    if bicint_kwargs is None:
        bicint_kwargs = {"nsamples": 2000, "func_output": "all", "nll_key": "nll"}
    rows: List[ComparisonRow] = []
    for mod_idx, item in enumerate(models):
        if hasattr(item, "fit_func") or hasattr(item, "fit"):  # EMModel instance
            out = item._out or {}
            all_data = item.all_data
            fit_func = item.fit_func
            param_names = list(item.param_names)
        else:
            _, out, all_data, fit_func = item  # explicit tuple
            param_names = list(out.get("param_names", []))
        if model_names is not None:
            name = model_names[mod_idx]
        else:
            name = f"Model_{mod_idx + 1}"
        LME = None
        if "inv_h" in out and "NPL" in out:
            _, lme, _ = calc_LME(out["inv_h"], out["NPL"])
            LME = float(lme)
        BICint = None
        if out.get("posterior") is not None and bicint_kwargs is not None:
            mu = out["posterior"]["mu"]; sigma = out["posterior"]["sigma"]
            BICint = calc_BICint(all_data, param_names, mu, sigma, fit_func, **bicint_kwargs)
        R2 = np.nan
        if r2_kwargs is not None and "NLL" in out:
            R2 = pseudo_r2_from_nll(out["NLL"], **r2_kwargs)
        rows.append(ComparisonRow(name=name, LME=LME, BICint=BICint, R2=R2))
    # sort by the first available metric
    for metric in metric_order:
        vals = [getattr(r, metric) for r in rows]
        if any(v is not None for v in vals):
            rev = True if metric in ("LME", "R2") else False  # higher better for LME,R2; lower better for BICint
            rows.sort(key=lambda r: (float("inf") if getattr(r, metric) is None else getattr(r, metric)), reverse=rev)
            break
    return rows

ComparisonRow dataclass

ComparisonRow(name: str, LME: float | None, BICint: float | None, R2: float | None)

ModelComparison

ModelComparison(models: List[Any], model_names: Optional[List[str]] = None)

Class for performing model comparison and identifiability analysis.

Initialize ModelComparison.

Parameters:

Name	Type	Description	Default
models	List[Any]	List of EMModel instances or model tuples	required
model_names	Optional[List[str]]	Optional list of model names	None

Source code in pyem/core/compare.py

def __init__(self, models: List[Any], model_names: Optional[List[str]] = None):
    """
    Initialize ModelComparison.

    Args:
        models: List of EMModel instances or model tuples
        model_names: Optional list of model names
    """
    self.models = models
    self.model_names = model_names or [f"Model_{i+1}" for i in range(len(models))]
    self.comparison_results = None
    self.identifiability_matrix = None

compare

compare(**kwargs) -> pd.DataFrame

Run model comparison and return a pandas DataFrame indexed by model name.

Columns

• "LME (largest is best)"
• "BICint (smallest is best)"
• "pseudoR^2 (largest is best)"

Source code in pyem/core/compare.py

def compare(self, **kwargs) -> pd.DataFrame:
    """
    Run model comparison and return a pandas DataFrame indexed by model name.

    Columns:
        - "LME (largest is best)"
        - "BICint (smallest is best)"
        - "pseudoR^2 (largest is best)"
    """
    rows = compare_models(self.models, self.model_names, **kwargs)

    df = pd.DataFrame(
        [(r.name, r.LME, r.BICint, r.R2) for r in rows],
        columns=[
            "name",
            "LME (largest is best)",
            "BICint (smallest is best)",
            "pseudoR^2 (largest is best)",
        ],
    ).set_index("name").dropna(axis=1)

    self.comparison_results = df
    return df

identify

identify(mi_inputs: List[str], nrounds: int = 10, nsubjects: int = 100, *, sim_kwargs: dict | None = None, fit_kwargs: dict | None = None, bicint_kwargs: dict | None = None, r2_kwargs: dict | None = None, seed: int | None = None, verbose: int = 0) -> pd.DataFrame

For each round (outer loop), for each EMModel in self.models: - simulate using that model's simulate_func - fit every model to the simulated data using its fit_func

Returns a DataFrame with columns

['Simulated','Estimated','LME','BICint','pseudoR2','bestlme','bestbic','bestR2']

Note

"True" parameters for each round are drawn as raw standard-normal values and mapped into natural space via each model's param_xform before simulating. Construct every model with param_xform set (as you would for .recover()) — without it, models whose simulate_func validates its natural-space bounds (e.g. every RW-family model) will raise on the untransformed Gaussian values.

pseudoR2 requires ntrials_total and noptions (see :func:pyem.utils.stats.pseudo_r2_from_nll), which have no defaults. Pass them via r2_kwargs (e.g. r2_kwargs={"ntrials_total": ntrials*nblocks, "noptions": 2}) or the pseudoR2 column will be all-NaN.

Every model's simulate_func must return a dict of named arrays (as the RW and Bayes model families do). This is incompatible with the GLM family (glm_sim, logit_sim, glm_ar_sim, etc.), whose simulate_func returns an (X, Y) tuple — using a GLM model here raises ValueError.

Source code in pyem/core/compare.py

def identify(
    self,
    mi_inputs: List[str],
    nrounds: int = 10,
    nsubjects: int = 100,
    *,
    sim_kwargs: dict | None = None,
    fit_kwargs: dict | None = None,
    bicint_kwargs: dict | None = None,
    r2_kwargs: dict | None = None,
    seed: int | None = None,
    verbose: int = 0,      # NEW: verbosity (0=silent, 1=round/model, 2=per-fit)
) -> pd.DataFrame:
    """
    For each round (outer loop), for each EMModel in self.models:
      - simulate using that model's simulate_func
      - fit every model to the simulated data using its fit_func

    Returns a DataFrame with columns:
        ['Simulated','Estimated','LME','BICint','pseudoR2','bestlme','bestbic','bestR2']

    Note:
        "True" parameters for each round are drawn as raw standard-normal
        values and mapped into natural space via each model's
        ``param_xform`` before simulating. Construct every model with
        ``param_xform`` set (as you would for ``.recover()``) — without
        it, models whose ``simulate_func`` validates its natural-space
        bounds (e.g. every RW-family model) will raise on the
        untransformed Gaussian values.

        pseudoR2 requires ``ntrials_total`` and ``noptions`` (see
        :func:`pyem.utils.stats.pseudo_r2_from_nll`), which have no
        defaults. Pass them via ``r2_kwargs`` (e.g.
        ``r2_kwargs={"ntrials_total": ntrials*nblocks, "noptions": 2}``)
        or the pseudoR2 column will be all-NaN.

        Every model's ``simulate_func`` must return a ``dict`` of named
        arrays (as the RW and Bayes model families do). This is
        incompatible with the GLM family (``glm_sim``, ``logit_sim``,
        ``glm_ar_sim``, etc.), whose ``simulate_func`` returns an
        ``(X, Y)`` tuple — using a GLM model here raises ``ValueError``.

    Raises:
        AttributeError if any simulated model has no simulate_func.
        ValueError if any simulated model's simulate_func does not return a dict.
    """
    rng = np.random.default_rng(seed)
    sim_kwargs = sim_kwargs or {}
    fit_kwargs = fit_kwargs or {}

    if bicint_kwargs is None:
        bicint_kwargs = {"nsamples": 2000, "func_output": "all", "nll_key": "nll"}
    if r2_kwargs is None:
        r2_kwargs = {}
    names = self.model_names

    # accumulators
    metrics = {
        "LME": {(s, e): [] for s in names for e in names},
        "BICint": {(s, e): [] for s in names for e in names},
        "R2": {(s, e): [] for s in names for e in names},
    }
    best_lme = {(s, e): 0 for s in names for e in names}
    best_bic = {(s, e): 0 for s in names for e in names}
    best_r2  = {(s, e): 0 for s in names for e in names}

    # ---------- OUTER LOOP OVER ROUNDS ----------
    for r in range(nrounds):
        if verbose >= 1:
            print(f"[identify] round {r+1}/{nrounds}")

        # loop over SIMULATED models
        for sm_i, sim_model in enumerate(self.models):
            sim_name = names[sm_i]

            if getattr(sim_model, "simulate_func", None) is None:
                raise AttributeError(
                    f"Model '{sim_name}' has no simulate_func; cannot run identifiability."
                )

            # choose parameters to simulate with
            nparams = len(sim_model.param_names)
            true_params = rng.normal(size=(nsubjects, nparams))
            # apply transforms if available (EMModel always defines param_xform, but
            # it's None unless the caller passed one — hasattr() alone can't detect that)
            if getattr(sim_model, "param_xform", None) is not None:
                for pi in range(nparams):
                    xform = sim_model.param_xform[pi]
                    if xform is not None and callable(xform):
                        true_params[:, pi] = xform(true_params[:, pi])

            # simulate data
            try:
                sim = sim_model.simulate_func(true_params, **sim_kwargs)
            except Exception as e:
                raise RuntimeError(f"Simulation from '{sim_name}' failed: {e}") from e

            # Validate structure
            if not isinstance(sim, dict):
                raise ValueError(f"Simulation from '{sim_name}' did not return a dict.")

            missing = [k for k in mi_inputs if k not in sim]
            if missing:
                raise ValueError(
                    f"Simulation from '{sim_name}' is missing expected keys: {missing}. "
                    f"Available keys: {list(sim.keys())}"
                )

            # Validate equal lengths for all requested inputs
            lengths = {k: len(sim[k]) for k in mi_inputs}
            if len(set(lengths.values())) != 1:
                raise ValueError(
                    f"Inconsistent lengths among requested inputs: {lengths}. "
                    "All requested series must be the same length."
                )

            # Build all_data as rows of the selected inputs (e.g., [choice, reward] per trial)
            all_data = [list(row) for row in zip(*(sim[k] for k in mi_inputs))]

            # fit EVERY model to this simulated dataset
            per_round_LME = []
            per_round_BIC = []
            per_round_R2  = []

            for em_j, est_model in enumerate(self.models):
                est_name = names[em_j]
                if verbose >= 2:
                    print(f"  - fitting Estimated='{est_name}' to Simulated='{sim_name}'")

                # build a fresh EMModel for the estimator
                tmp = EMModel(
                    all_data=all_data,
                    fit_func=est_model.fit_func,
                    param_names=est_model.param_names,
                    simulate_func=getattr(est_model, "simulate_func", None),
                )
                fit_res = tmp.fit(verbose=0, **fit_kwargs)

                # LME (higher better)
                lme = None
                if fit_res.inv_h is not None and fit_res.NPL is not None:
                    _, lme_total, _ = calc_LME(fit_res.inv_h, fit_res.NPL)  # from utils.stats
                    lme = float(lme_total)

                # BICint (lower better)
                try:
                    bic = calc_BICint(
                        all_data,
                        est_model.param_names,
                        fit_res.posterior_mu,
                        fit_res.posterior_sigma,
                        est_model.fit_func,
                        **bicint_kwargs,
                    )
                except Exception:
                    bic = np.nan

                # pseudo-R² (higher better)
                try:
                    r2 = pseudo_r2_from_nll(fit_res.NLL, **r2_kwargs)
                except Exception:
                    r2 = np.nan

                metrics["LME"][(sim_name, est_name)].append(lme)
                metrics["BICint"][(sim_name, est_name)].append(bic)
                metrics["R2"][(sim_name, est_name)].append(r2)

                per_round_LME.append((est_name, lme))
                per_round_BIC.append((est_name, bic))
                per_round_R2.append((est_name, r2))

            # determine winners for THIS (sim_name, round)
            def _winner(pairs, prefer="max"):
                vals = [(n, v) for (n, v) in pairs if v is not None and not np.isnan(v)]
                if not vals:
                    return None
                return (max if prefer == "max" else min)(vals, key=lambda t: t[1])[0]

            w_lme = _winner(per_round_LME, "max")
            w_bic = _winner(per_round_BIC, "min")
            w_r2  = _winner(per_round_R2,  "max")

            if w_lme is not None:
                best_lme[(sim_name, w_lme)] += 1
            if w_bic is not None:
                best_bic[(sim_name, w_bic)] += 1
            if w_r2 is not None:
                best_r2[(sim_name, w_r2)] += 1

    # assemble tidy dataframe
    rows = []
    for sim_name in names:
        for est_name in names:
            LME_mean = (np.nanmean(metrics["LME"][(sim_name, est_name)])
                        if metrics["LME"][(sim_name, est_name)] else np.nan)
            BIC_mean = (np.nanmean(metrics["BICint"][(sim_name, est_name)])
                        if metrics["BICint"][(sim_name, est_name)] else np.nan)
            R2_mean  = (np.nanmean(metrics["R2"][(sim_name, est_name)])
                        if metrics["R2"][(sim_name, est_name)] else np.nan)
            rows.append({
                "Simulated": sim_name,
                "Estimated": est_name,
                "LME": LME_mean,
                "BICint": BIC_mean,
                "pseudoR2": R2_mean,
                "bestlme": best_lme[(sim_name, est_name)],
                "bestbic": best_bic[(sim_name, est_name)],
                "bestR2":  best_r2[(sim_name, est_name)],
            })

    df = pd.DataFrame(rows)

    self.identifiability_matrix = df.dropna(axis=1)
    self._identify_rounds = nrounds
    if verbose >= 1:
        print("[identify] done")
    return df

identifiability_analysis

identifiability_analysis(*args, **kwargs)

Backward-compatible alias for :meth:identify.

Source code in pyem/core/compare.py

def identifiability_analysis(self, *args, **kwargs):
    """Backward-compatible alias for :meth:`identify`."""
    return self.identify(*args, **kwargs)

plot_identifiability

plot_identifiability(df: DataFrame | None = None, metric: str = 'LME', *, nrounds: int | None = None, cmap: str = 'viridis', annotate: bool = True, figsize: tuple = (8, 6)) -> plt.Figure

Plot a heatmap of identifiability as the PROPORTION of rounds that the Estimated model won for each (Simulated, Estimated) pair, using the winner-count fields produced by identify().

Parameters:

Name	Type	Description	Default
df	DataFrame | None	DataFrame returned by identify(); if None, uses self.identifiability_matrix	None
metric	str	which metric's winners to visualize: {"LME","BICint","pseudoR2"}	'LME'
nrounds	int | None	total rounds per Simulated model (if None, inferred per row by sum of winners)	None
cmap	str	colormap	'viridis'
annotate	bool	whether to add cell annotations	True
figsize	tuple	figure size	(8, 6)

Returns:

Type	Description
Figure	matplotlib Figure

Source code in pyem/core/compare.py

def plot_identifiability(
    self,
    df: pd.DataFrame | None = None,
    metric: str = "LME",
    *,
    nrounds: int | None = None,
    cmap: str = "viridis",
    annotate: bool = True,
    figsize: tuple = (8, 6),
) -> plt.Figure:
    """
    Plot a heatmap of identifiability as the PROPORTION of rounds that the Estimated
    model won for each (Simulated, Estimated) pair, using the winner-count fields
    produced by identify().

    Args:
        df: DataFrame returned by identify(); if None, uses self.identifiability_matrix
        metric: which metric's winners to visualize: {"LME","BICint","pseudoR2"}
        nrounds: total rounds per Simulated model (if None, inferred per row by sum of winners)
        cmap: colormap
        annotate: whether to add cell annotations
        figsize: figure size

    Returns:
        matplotlib Figure
    """
    if df is None:
        if isinstance(self.identifiability_matrix, pd.DataFrame):
            df = self.identifiability_matrix
        else:
            raise RuntimeError("Run identify() first or pass a dataframe.")

    # Map metric -> winner-count column produced by identify()
    metric = metric.strip()
    metric_key = metric.upper()
    best_map = {
        "LME": "bestlme",
        "BICINT": "bestbic",
        "PSEUDOR2": "bestR2",
    }
    if metric_key not in best_map:
        raise ValueError("metric must be one of {'LME','BICint','pseudoR2'}")
    best_col = best_map[metric_key]

    # Pivot winner counts to matrix (rows=Simulated, cols=Estimated)
    counts = df.pivot(index="Simulated", columns="Estimated", values=best_col).astype(float)

    # Determine denominators (rounds). Prefer explicit arg, else internal, else infer per-row.
    if nrounds is None:
        nrounds = getattr(self, "_identify_rounds", None)

    if nrounds is not None:
        # Single scalar denominator for all rows
        denom = pd.Series(nrounds, index=counts.index, dtype=float)
    else:
        # Infer per Simulated model: sum of winner counts across columns
        # (handles cases with occasional no-winner rounds by yielding < 1.0 max)
        denom = counts.sum(axis=1)
        # avoid divide-by-zero
        denom[denom == 0] = np.nan

    # Compute proportions row-wise
    prop = counts.div(denom, axis=0)

    # Plot
    fig, ax = plt.subplots(figsize=figsize)
    im = ax.imshow(prop.values, cmap=cmap, aspect="auto", vmin=0.0, vmax=1.0)

    # Colorbar labeled as a proportion
    cbar = plt.colorbar(im, ax=ax)
    cbar.set_label("Proportion of rounds won")

    # Axis ticks/labels
    ax.set_xticks(range(prop.shape[1]))
    ax.set_yticks(range(prop.shape[0]))
    ax.set_xticklabels(prop.columns, rotation=45, ha="right")
    ax.set_yticklabels(prop.index)

    # Title and axis labels
    title = {
        "LME": "Log Model Evidence (proportion wins)",
        "BICINT": "Integrated BIC (proportion wins, lower is better)",
        "PSEUDOR2": "Pseudo R² (proportion wins)",
    }[metric_key]
    ax.set_title(title)
    ax.set_xlabel("Estimated")
    ax.set_ylabel("Simulated")

    # Optional per-cell annotations
    if annotate:
        for i in range(prop.shape[0]):
            for j in range(prop.shape[1]):
                val = prop.iat[i, j]
                if np.isnan(val):
                    txt = "–"
                else:
                    txt = f"{val:.2f}"
                # legible text color
                ax.text(j, i, txt,
                        ha="center", va="center",
                        color="white" if (not np.isnan(val) and val >= 0.5) else "black")

    plt.tight_layout()
    return fig

Group M-step families

Each group-distribution family (Gaussian, Laplace, Student-t, Cauchy) implements the same three-method interface — update(m, inv_h) to run the empirical-Bayes M-step over subject-level MAP means and inverse-Hessians, moments(hyper) to return the group mean/variance and a validity flag, and make_prior(hyper) to build the Prior object fed back into the next E-step — so EMfit can swap between them via the EMConfig.mstep name. make_group is the factory that dispatches a family name (and, for Student-t, a degrees-of-freedom df) to the corresponding class.

import numpy as np
from pyem.core.groupdist import make_group
g = make_group("student_t", df=8.0)
m = np.random.default_rng(0).normal(size=(2, 20))
inv_h = np.stack([np.eye(2) for _ in range(20)], axis=-1)
print(g.moments(g.update(m, inv_h)))

make_group

make_group(name: str, df: float = 8.0)

Factory for group-distribution families.

df is ignored when name == "cauchy" (Cauchy is df=1 by definition).

Source code in pyem/core/groupdist.py

def make_group(name: str, df: float = 8.0):
    """Factory for group-distribution families.

    ``df`` is ignored when ``name == "cauchy"`` (Cauchy is df=1 by definition).
    """
    if name == "gaussian":
        return GaussianGroup()
    if name == "laplace":
        return LaplaceGroup()
    if name == "student_t":
        return StudentTGroup(df)
    if name == "cauchy":
        return CauchyGroup()
    raise ValueError(f"unknown mstep {name!r}")

GaussianGroup

Gaussian group distribution (the classic MFX M-step).

LaplaceGroup

Laplace group distribution.

StudentTGroup

StudentTGroup(df: float = 8.0)

Student-t group distribution via a scale-mixture IRLS EM.

Source code in pyem/core/groupdist.py

def __init__(self, df: float = 8.0):
    self.df = float(df)

CauchyGroup

CauchyGroup()

Bases: StudentTGroup

Cauchy group distribution (Student-t with df=1).

Source code in pyem/core/groupdist.py

def __init__(self):
    super().__init__(df=1.0)

Priors

Prior is the minimal protocol (a logpdf(x) method) that every prior object satisfies; GaussianPrior is the default and is used throughout the EM loop as the per-iteration regularizer built from the group posterior. Note: GaussianPrior.sigma holds the per-parameter variance (sigma², not the standard deviation) — this matches what the empirical-Bayes group M-step reports and what EMModel.fit(prior_sigma=...) expects, so use the .variance property alias when you want the intent to read explicitly. default_prior builds a broad, weakly-informative GaussianPrior (variance 100) for a given number of parameters. UniformPrior, LaplacePrior, StudentTPrior, and CauchyPrior provide alternative independent priors, and IndependentPrior composes a list of 1-D priors (one per parameter) into a single multi-parameter prior.

import numpy as np
from pyem.core.priors import GaussianPrior, UniformPrior, default_prior
g = GaussianPrior(mu=[0.0, 0.0], sigma=[1.0, 4.0])   # sigma == variance
print(g.logpdf([0.0, 0.0]), g.variance)
print(default_prior(2).sigma)                          # broad default (variance 100)
print(UniformPrior(lo=[-1, 0], hi=[1, 1]).logpdf([0.0, 0.5]))

Prior

Bases: Protocol

Protocol for prior objects.

logpdf

logpdf(x: ndarray) -> float

Return the log probability density of x.

Source code in pyem/core/priors.py

def logpdf(self, x: np.ndarray) -> float:
    """Return the log probability density of ``x``."""

GaussianPrior dataclass

GaussianPrior(mu: ndarray, sigma: ndarray)

Independent Gaussian prior.

IMPORTANT: sigma holds the per-parameter VARIANCE (sigma^2), not the standard deviation — logpdf divides (x-mu)**2 by sigma directly. This matches the empirical-Bayes group M-step, which reports variances, and the values in EMModel.fit(prior_sigma=...) are likewise variances. Use the variance alias when you want the intent to be explicit.

variance property

variance: ndarray

Alias for sigma (which is the per-parameter variance).

default_prior

default_prior(nparams: int, seed: int | None = None) -> GaussianPrior

Source code in pyem/core/priors.py

def default_prior(nparams: int, seed: int | None = None) -> GaussianPrior:
    rng = np.random.default_rng(seed)
    mu = 0.1 * rng.standard_normal(nparams)
    sigma = np.full(nparams, 100.0)
    return GaussianPrior(mu=mu, sigma=sigma)

UniformPrior dataclass

UniformPrior(lo: ndarray, hi: ndarray)

Independent bounded-uniform prior. logpdf = -sum(log(hi-lo)) inside the box, -inf outside.

LaplacePrior dataclass

LaplacePrior(loc: ndarray, scale: ndarray)

Independent Laplace prior. scale is the diversity b (std = b*sqrt(2)).

StudentTPrior dataclass

StudentTPrior(loc: ndarray, scale: ndarray, df: ndarray)

Independent Student-t prior. scale is std-like; df = degrees of freedom.

CauchyPrior dataclass

CauchyPrior(loc: ndarray, scale: ndarray)

Independent Cauchy prior (Student-t df=1). scale is half-width gamma. Cauchy has no finite variance; init_moments returns a large finite fallback.

IndependentPrior dataclass

IndependentPrior(priors: list)

Per-parameter composer: one 1-D prior per coordinate.

Parameter recovery

parameter_recovery compares ground-truth subject-level parameters against their EM-estimated counterparts, returning per-parameter Pearson correlation and RMSE (plus the estimates themselves) bundled in a RecoveryResult dataclass — the standard diagnostic for validating that a model's parameters are identifiable from simulated data.

parameter_recovery

parameter_recovery(true_params: ndarray, est_params: ndarray) -> RecoveryResult

Compare true vs. estimated subject-level parameters: return per-parameter Pearson r and RMSE.

Source code in pyem/core/posterior.py

def parameter_recovery(true_params: np.ndarray, est_params: np.ndarray) -> RecoveryResult:
    """
    Compare true vs. estimated subject-level parameters: return per-parameter Pearson r and RMSE.
    """
    assert true_params.shape == est_params.shape
    dif = est_params - true_params
    rmse = np.sqrt(np.mean(dif**2, axis=0))
    corr = np.array([np.corrcoef(true_params[:, j], est_params[:, j])[0,1] for j in range(true_params.shape[1])])
    return RecoveryResult(corr=corr, rmse=rmse, est_params=est_params)

RecoveryResult dataclass

RecoveryResult(corr: ndarray, rmse: ndarray, est_params: ndarray)

Model specification

ModelSpec is a small, self-describing bundle (id, spec dict, description, and the model's params/sim/fit callables) used to register a model's identity and entry points in one place. It is purely additive — neither EMModel nor ModelComparison requires a ModelSpec to function — but model modules use it to expose a consistent, discoverable interface.

ModelSpec dataclass

ModelSpec(id: str, spec: dict, desc: str, params: Callable | None, sim: Callable, fit: Callable)

Self-describing bundle for a model's identity and entry points.

id and spec are hand-authored per model (there is no shared taxonomy to derive them from) — pick whatever short id and descriptive dict shape makes sense for that model family. Purely additive: nothing in :class:pyem.api.EMModel or :class:pyem.core.compare.ModelComparison requires a ModelSpec to function.