EMModel & FitResult

EMModel is the high-level, scikit-learn-style wrapper around the hierarchical EM fit. You give it a list of per-subject data, an objective (fit_func), the parameter names, and (optionally) the Gaussian→natural transforms and a simulator. EMModel.fit() returns a FitResult dataclass. Parameters are fit in Gaussian (unbounded) space; subject_params() returns natural-space estimates when param_xform is provided.

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

rng = np.random.default_rng(0)
true = np.column_stack([norm2beta(rng.normal(size=20)), norm2alpha(rng.normal(size=20))])
sim = rw1a1b_sim(true, nblocks=3, ntrials=24, seed=0)
all_data = [[sim["choices"][s], sim["rewards"][s]] for s in range(20)]

model = EMModel(all_data, rw1a1b_fit, ["beta", "alpha"],
                param_xform=[norm2beta, norm2alpha], simulate_func=rw1a1b_sim)
result = model.fit(verbose=0)
print(result.m.shape, result.convergence)     # (2, 20) True/False
print(model.subject_params().shape)           # (20, 2) natural-space

EMModel

EMModel(all_data: Sequence[Sequence[Any]] | None, fit_func: Callable[..., float], param_names: Sequence[str], param_xform: Sequence[Callable] | None = None, simulate_func: Callable[..., Any] | None = None)

A high-level, sklearn-like interface: model = EMModel(all_data, fit_func, param_names, param_xform=[norm2beta, norm2alpha], simulate_func=...) result = model.fit(...) sim = model.simulate(...)

# Access parameter transformations: beta_transform = model.get_param_transform("beta") # or model.param_xform[0] transformed_value = beta_transform(0.5)

Source code in pyem/api.py

def __init__(
    self,
    all_data: Sequence[Sequence[Any]] | None,
    fit_func: Callable[..., float],
    param_names: Sequence[str],
    param_xform: Sequence[Callable] | None = None,
    simulate_func: Callable[..., Any] | None = None,
) -> None:
    self.all_data = all_data
    self.fit_func = fit_func
    self.param_names = list(param_names)
    self.simulate_func = simulate_func
    self._out: dict[str, Any] | None = None
    self._outfit: dict[str, np.ndarray] | None = None

    # Store parameter transformation functions
    if param_xform is not None:
        if len(param_xform) != len(param_names):
            raise ValueError(f"param_xform length ({len(param_xform)}) must match param_names length ({len(param_names)})")
        self.param_xform = list(param_xform)
    else:
        self.param_xform = None

outfit property

outfit: dict[str, ndarray]

Lazily-computed, cached dict of fit_func outputs for each subject.

Computed on first access after a fit/recover and cached; the cache is invalidated whenever a new fit() or recover() runs.

get_param_transform

get_param_transform(param_name_or_index)

Get parameter transformation function by parameter name or index.

Parameters:

Name	Type	Description	Default
param_name_or_index		Parameter name (str) or index (int)	required

Returns:

Type	Description
	Transformation function for the specified parameter

Raises:

Type	Description
ValueError	If param_xform was not provided or parameter not found

Source code in pyem/api.py

def get_param_transform(self, param_name_or_index):
    """
    Get parameter transformation function by parameter name or index.

    Args:
        param_name_or_index: Parameter name (str) or index (int)

    Returns:
        Transformation function for the specified parameter

    Raises:
        ValueError: If param_xform was not provided or parameter not found
    """
    if self.param_xform is None:
        raise ValueError("param_xform was not provided to the model")

    if isinstance(param_name_or_index, str):
        try:
            index = self.param_names.index(param_name_or_index)
        except ValueError:
            raise ValueError(f"Parameter '{param_name_or_index}' not found in param_names")
    else:
        index = param_name_or_index
        if not (0 <= index < len(self.param_xform)):
            raise ValueError(f"Index {index} out of range for param_xform")

    return self.param_xform[index]

EMfit

EMfit(**kwargs) -> dict[str, Any]

Deprecated legacy alias for :meth:fit; returns fit(**kwargs).__dict__ (a plain dict view of the :class:FitResult). Prefer :meth:fit.

Source code in pyem/api.py

def EMfit(self, **kwargs) -> dict[str, Any]:
    """Deprecated legacy alias for :meth:`fit`; returns ``fit(**kwargs).__dict__``
    (a plain dict view of the :class:`FitResult`). Prefer :meth:`fit`."""
    return self.fit(**kwargs).__dict__

compute_integrated_bic

compute_integrated_bic(nsamples: int = 500, func_output: str = 'all', nll_key: str = 'nll', ntrials_total: int | None = None) -> float

Compute integrated Bayesian Information Criterion (BICint).

Parameters:

Name	Type	Description	Default
nsamples	int	Number of samples for Monte Carlo integration	500
func_output	str	Output type to request from fit function	'all'
nll_key	str	Key to extract negative log-likelihood from fit function output	'nll'
ntrials_total	int | None	Number of trials per subject used for the BIC complexity penalty. If omitted, auto-detected from all_data (see :func:pyem.utils.stats.calc_BICint) — pass explicitly if your model's data fields aren't all trial-aligned with identical shape.	None

Returns:

Type	Description
float	Integrated BIC value

Source code in pyem/api.py

def compute_integrated_bic(
    self, nsamples: int = 500, func_output: str = "all", nll_key: str = "nll",
    ntrials_total: int | None = None,
) -> float:
    """
    Compute integrated Bayesian Information Criterion (BICint).

    Args:
        nsamples: Number of samples for Monte Carlo integration
        func_output: Output type to request from fit function
        nll_key: Key to extract negative log-likelihood from fit function output
        ntrials_total: Number of trials per subject used for the BIC complexity
            penalty. If omitted, auto-detected from ``all_data`` (see
            :func:`pyem.utils.stats.calc_BICint`) — pass explicitly if your
            model's data fields aren't all trial-aligned with identical shape.

    Returns:
        Integrated BIC value
    """
    if self._out is None:
        raise RuntimeError("Call fit() first.")

    posterior = self._out["posterior"]
    return calc_BICint(
        self.all_data,
        self.param_names,
        posterior["mu"],
        posterior["sigma"],
        self.fit_func,
        nsamples=nsamples,
        func_output=func_output,
        nll_key=nll_key,
        ntrials_total=ntrials_total,
    )

compute_lme

compute_lme() -> tuple[np.ndarray, float, np.ndarray]

Compute Laplace approximation for log model evidence (LME).

Returns:

Type	Description
tuple[ndarray, float, ndarray]	Tuple of (Laplace approximation per subject, total LME, good Hessian flags)

Source code in pyem/api.py

def compute_lme(self) -> tuple[np.ndarray, float, np.ndarray]:
    """
    Compute Laplace approximation for log model evidence (LME).

    Returns:
        Tuple of (Laplace approximation per subject, total LME, good Hessian flags)
    """
    if self._out is None:
        raise RuntimeError("Call fit() first.")

    return calc_LME(self._out["inv_h"], self._out["NPL"])

get_outfit

get_outfit() -> dict[str, np.ndarray]

Return a dictionary of outputs from the fit_func for each subject.

Source code in pyem/api.py

def get_outfit(self) -> dict[str, np.ndarray]:
    """Return a dictionary of outputs from the ``fit_func`` for each subject."""
    if self._out is None:
        raise RuntimeError("Call fit() first.")

    nsubjects = self._out["m"].shape[1]

    # determine available keys and shapes from first subject
    first_params = self._out["m"][:, 0]
    first_args = self.all_data[0]
    if not isinstance(first_args, (list, tuple)):
        first_args = (first_args,)
    first_fit = self.fit_func(first_params, *first_args, prior=None, output="all")

    arrays_dict: dict[str, np.ndarray] = {}
    for key, value in first_fit.items():
        if isinstance(value, (int, float, np.number)):
            arrays_dict[key] = np.zeros(nsubjects)
        elif isinstance(value, (list, np.ndarray)):
            arr = np.asarray(value)
            if arr.ndim == 0:
                arrays_dict[key] = np.zeros(nsubjects)
            else:
                arrays_dict[key] = np.zeros((nsubjects, *arr.shape), dtype=arr.dtype)
        else:
            arrays_dict[key] = np.empty(nsubjects, dtype=object)

    for subj_idx in range(nsubjects):
        params = self._out["m"][:, subj_idx]
        args = self.all_data[subj_idx]
        if not isinstance(args, (list, tuple)):
            args = (args,)
        subj_fit = self.fit_func(params, *args, prior=None, output="all")
        for key in arrays_dict.keys():
            if key in subj_fit:
                arrays_dict[key][subj_idx] = subj_fit[key]

    return arrays_dict

scipy_minimize

scipy_minimize(seed: int | None = None) -> dict[str, Any]

Fit each subject independently using :func:scipy.optimize.minimize.

This is a non-hierarchical baseline (no population-level prior/M-step; each subject's initial guess and objective are independent of every other subject) — useful for comparing against the full EM fit.

Parameters:

Name	Type	Description	Default
seed	int | None	Seed for the :func:numpy.random.default_rng instance used to draw each subject's random initial guess (x0). Passing the same seed reproduces the same sequence of initial guesses across runs. Note that this method uses its own local Generator and no longer reads/honors the global np.random.seed() state — call with an explicit seed for reproducibility instead of relying on a prior global seed call.	None

Source code in pyem/api.py

def scipy_minimize(self, seed: int | None = None) -> dict[str, Any]:
    """Fit each subject independently using :func:`scipy.optimize.minimize`.

    This is a non-hierarchical baseline (no population-level prior/M-step;
    each subject's initial guess and objective are independent of every
    other subject) — useful for comparing against the full EM fit.

    Args:
        seed: Seed for the :func:`numpy.random.default_rng` instance used
            to draw each subject's random initial guess (``x0``). Passing
            the same seed reproduces the same sequence of initial guesses
            across runs. Note that this method uses its own local
            ``Generator`` and no longer reads/honors the global
            ``np.random.seed()`` state — call with an explicit ``seed``
            for reproducibility instead of relying on a prior global seed
            call.
    """

    if self.all_data is None:
        raise ValueError("all_data must be provided to fit the model.")

    nsubjects = len(self.all_data)
    nparams = len(self.param_names)

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

    rng = np.random.default_rng(seed)

    for subj_idx, args in enumerate(self.all_data):
        if not isinstance(args, (list, tuple)):
            args = (args,)

        def obj_func(params: np.ndarray) -> float:
            return self.fit_func(params, *args, output="npl")

        x0 = rng.standard_normal(nparams)
        res = minimize(obj_func, x0=x0, method="BFGS")
        m[:, subj_idx] = res.x
        NPL[subj_idx] = res.fun

        if hasattr(res, "hess_inv"):
            try:
                inv_h[:, :, subj_idx] = np.asarray(res.hess_inv)
            except Exception:
                inv_h[:, :, subj_idx] = np.eye(nparams) * max(1.0, np.linalg.norm(res.x) + 1e-6)
        else:
            inv_h[:, :, subj_idx] = np.eye(nparams) * max(1.0, np.linalg.norm(res.x) + 1e-6)

    posterior_sigma = np.array([np.diag(inv_h[:, :, i]) for i in range(nsubjects)]).T

    return {
        "m": m,
        "inv_h": inv_h,
        "posterior": {"mu": m, "sigma": posterior_sigma},
        "NPL": NPL,
        "convergence": True,
        "individual_fit": True,
    }

recover

recover(true_params: ndarray, pr_inputs: List[str], simulate_func: Callable = None, fit_kwargs: dict | None = None, **sim_kwargs) -> dict[str, Any]

Parameter recovery analysis given true parameters and simulation function.

Parameters:

Name	Type	Description	Default
true_params	ndarray	True parameter values (nsubjects x nparams)	required
pr_inputs	List[str]	Required. Names of the simulation output keys (from simulate_func's returned dict) to pass positionally to fit_func as subject data, in order, e.g. ["choices", "rewards"]. Every named key must be present in the simulation output and all requested series must have equal length; a per-trial row is built by zipping across them (all_data[i] = [sim[k][i] for k in pr_inputs]).	required
simulate_func	Callable	Simulation function (uses self.simulate_func if None)	None
fit_kwargs	dict | None	Additional keyword arguments forwarded to the recovery model's fit(...) call (e.g. seed, mstep, prior); verbose defaults to 0 but may be overridden via fit_kwargs. Note that end-to-end recovery reproducibility is additionally bounded by the model's simulate function RNG (currently unseeded).	None
**sim_kwargs		Additional arguments for simulation	{}

Returns:

Type	Description
dict[str, Any]	Dictionary containing true params, estimated params, and recovery metrics.
dict[str, Any]	The correlation entry provides a Pearson correlation coefficient for
dict[str, Any]	each parameter (array of length nparams).

Source code in pyem/api.py

def recover(self, true_params: np.ndarray, pr_inputs: List[str], simulate_func: Callable = None,
            fit_kwargs: dict | None = None, **sim_kwargs) -> dict[str, Any]:
    """
    Parameter recovery analysis given true parameters and simulation function.

    Args:
        true_params: True parameter values (nsubjects x nparams)
        pr_inputs: Required. Names of the simulation output keys (from
            ``simulate_func``'s returned dict) to pass positionally to
            ``fit_func`` as subject data, in order, e.g. ``["choices",
            "rewards"]``. Every named key must be present in the
            simulation output and all requested series must have equal
            length; a per-trial row is built by zipping across them
            (``all_data[i] = [sim[k][i] for k in pr_inputs]``).
        simulate_func: Simulation function (uses self.simulate_func if None)
        fit_kwargs: Additional keyword arguments forwarded to the recovery
            model's ``fit(...)`` call (e.g. seed, mstep, prior); ``verbose``
            defaults to 0 but may be overridden via fit_kwargs. Note that
            end-to-end recovery reproducibility is additionally bounded by
            the model's simulate function RNG (currently unseeded).
        **sim_kwargs: Additional arguments for simulation

    Returns:
        Dictionary containing true params, estimated params, and recovery metrics.
        The ``correlation`` entry provides a Pearson correlation coefficient for
        each parameter (array of length ``nparams``).
    """
    if simulate_func is None:
        simulate_func = self.simulate_func
    if simulate_func is None:
        raise AttributeError("No simulate_func provided.")

    # Simulate data with true parameters
    sim = simulate_func(true_params, **sim_kwargs)

    if not isinstance(sim, dict):
        raise TypeError(
            "recover() requires simulate_func to return a dict of named arrays "
            "The included GLM family's *_sim returns an (X, Y) tuple and is not compatible with "
            "the pr_inputs mechanism; wrap it to return {'X': X, 'Y': Y} or fit "
            "GLMs directly with EMModel.fit()."
        )

    # Prepare data for fitting
    missing = [k for k in pr_inputs if k not in sim]
    if missing:
        raise ValueError(
            f"Simulation output 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 pr_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 pr_inputs))]

    # Create new model instance with simulated data
    recovery_model = EMModel(
        all_data=all_data,
        fit_func=self.fit_func,
        param_names=self.param_names,
        simulate_func=simulate_func,
        param_xform=self.param_xform,
    )

    # Fit the model
    merged = {"verbose": 0, **(fit_kwargs or {})}
    fit_result = recovery_model.fit(**merged)

    # grab estimated params with transformations applied if applicable
    estimated_params = recovery_model.subject_params()

    corr = parameter_recovery(true_params, estimated_params).corr
    recovery_dict = {
        'true_params': true_params,
        'estimated_params': estimated_params,
        'sim': sim,
        'fit_result': fit_result,
        'correlation': corr,
        'recovery_model': recovery_model
    }

    # Adopt the recovery fit onto THIS model so get_outfit()/.outfit are usable
    # afterwards (the outer model is typically constructed with all_data=None for
    # recovery workflows, so we point all_data + outfit at the recovery fit).
    print(
        "recover(): this model's data and outfit now reflect the recovery fit; "
        "the recovered estimates are also in the returned dict['recovery_model']."
    )
    self.all_data = all_data
    self._out = recovery_model._out
    # assign outfit from the fit() run (reuses the recovery model's computed outfit,
    # which is built from the same all_data + fit)
    self._outfit = recovery_model.outfit
    return recovery_dict

plot_recovery

plot_recovery(recovery_dict: dict, show_line: bool = True, figsize: tuple | None = None, show: bool = True) -> plt.Figure

Plot parameter recovery as scatter plots of simulated vs estimated parameters. Creates 3 columns with as many rows as needed, with compact spacing and subplot sizes that scale with the grid.

Parameters:

Name	Type	Description	Default
recovery_dict	dict	Output from recover() method, containing: - 'true_params' (array-like, shape [n_sims, n_params]) - 'estimated_params' (array-like, shape [n_sims, n_params])	required
show_line	bool	Whether to draw x=y line	True
figsize	tuple | None	Figure size	None
show	bool	Call :func:matplotlib.pyplot.show after drawing	True

Returns:

Type	Description
Figure	matplotlib Figure object

Source code in pyem/api.py

def plot_recovery(
    self,
    recovery_dict: dict,
    show_line: bool = True,
    figsize: tuple | None = None,
    show: bool = True
) -> plt.Figure:
    """
    Plot parameter recovery as scatter plots of simulated vs estimated parameters.
    Creates 3 columns with as many rows as needed, with compact spacing and
    subplot sizes that scale with the grid.

    Args:
        recovery_dict: Output from recover() method, containing:
            - 'true_params' (array-like, shape [n_sims, n_params])
            - 'estimated_params' (array-like, shape [n_sims, n_params])
        show_line: Whether to draw x=y line
        figsize: Figure size
        show: Call :func:`matplotlib.pyplot.show` after drawing

    Returns:
        matplotlib Figure object
    """
    true_params = recovery_dict['true_params']
    estimated_params = recovery_dict['estimated_params']
    nparams = true_params.shape[1]

    # Grid: 3 columns, compute rows
    ncols = 3
    nrows = int(np.ceil(nparams / ncols))

    # Figure size: scale per-subplot to avoid tiny axes.
    # Aim for 5x5 inches per subplot (square-ish data area works well here).
    per_ax_w, per_ax_h = 3.5, 3.5
    fig_w = per_ax_w * ncols
    fig_h = per_ax_h * nrows
    if figsize is None:
        figsize = (fig_w, fig_h)

    fig, axes = plt.subplots(
        nrows, ncols,
        figsize=figsize,
        constrained_layout=True, # let Matplotlib handle spacing
        squeeze=False
    )

    # Fine-tune constrained_layout paddings (reduces big gutters)
    # w_pad/h_pad: padding around the figure edges; wspace/hspace: padding between subplots
    fig.get_layout_engine().set() #h_pad=X, w_pad=Y, hspace=Z, wspace=W

    axes = axes.ravel()
    names = list(self.param_names)[:nparams]

    for i, param_name in enumerate(names):
        ax = axes[i]
        plotting.plot_scatter(
            true_params[:, i], f'True {param_name}',
            estimated_params[:, i], f'Estimated {param_name}',
            ax=ax,
            show_line=show_line,
            equal_limits=True,     # still equalize limits (handled w/ box aspect below)
            s=100,                  # slightly smaller markers to reduce overlap
            alpha=0.6,
            colorname='royalblue',
            annotate=True,
        )
        # Title & tick/label sizing tuned so they don't collide with data
        ax.tick_params(labelsize=12)
        ax.xaxis.label.set_size(12)
        ax.yaxis.label.set_size(12)

        # Keep plots square without blowing up gutters
        # (avoid ax.set_aspect('equal', adjustable='box') here)
        try:
            ax.set_box_aspect(1)   # Matplotlib >=3.4
        except Exception:
            pass

    # Remove unused axes completely so they don't consume layout space
    for j in range(nparams, len(axes)):
        axes[j].remove()

    if show:
        plt.show()

    return fig

FitResult dataclass

FitResult(m: ndarray, inv_h: ndarray, posterior_mu: ndarray, posterior_sigma: ndarray, NPL: ndarray, NLPrior: ndarray, NLL: ndarray, convergence: bool)