Skip to content

Utilities

pyem.utils collects the shared helpers used throughout the fitting pipeline: the Gaussian-space parameter transforms and objective-function helper (pyem.utils.math), group-level fit metrics for model comparison (pyem.utils.stats), and convenience plotting functions (pyem.utils.plotting).

Math transforms

These functions map between the unconstrained Gaussian space that the EM loop optimizes over and the natural (bounded) parameter ranges that models actually use. norm2alpha/alpha2norm are inverses mapping R to (0, 1) (e.g. a learning rate), and norm2beta/beta2norm are inverses mapping R to (0, max_val) (e.g. an inverse-temperature capped at max_val). softmax is a numerically stabilized choice rule (subtracting the max expected value before exponentiating) used to turn expected values into choice probabilities. calc_fval builds the scalar objective minimized during subject-level fitting — the negative log-posterior ("npl", negative log-likelihood plus negative log-prior) or just the negative log-likelihood ("nll") — and caps non-finite values at 1e7 so gradient-based optimizers keep making progress instead of hitting inf/nan.

1
2
3
4
5
6
import numpy as np
from pyem.utils.math import softmax, norm2alpha, norm2beta, alpha2norm

print(softmax(np.array([0.2, 0.8]), beta=3.0))
print(norm2alpha(0.0), norm2beta(0.0))          # 0.5, 10.0
print(alpha2norm(norm2alpha(0.7)))              # ~0.7 round-trip

softmax 

softmax(evs: ArrayLike, beta: float) -> np.ndarray
Source code in pyem/utils/math.py 
 7
 8
 9
10
11
def softmax(evs: ArrayLike, beta: float) -> np.ndarray:
    x = np.asarray(evs, dtype=float)
    x = x - np.max(x)  # stabilize
    p = np.exp(beta * x)
    return p / p.sum()

norm2alpha 

norm2alpha(x: ArrayLike) -> np.ndarray

Map R -> (0,1).

Source code in pyem/utils/math.py 
13
14
15
def norm2alpha(x: ArrayLike) -> np.ndarray:
    """Map R -> (0,1)."""
    return expit(np.asarray(x))

alpha2norm 

alpha2norm(a: ArrayLike) -> np.ndarray
Source code in pyem/utils/math.py 
17
18
19
20
21
def alpha2norm(a: ArrayLike) -> np.ndarray:
    a = np.asarray(a)
    eps = 1e-12
    a = np.clip(a, eps, 1 - eps)
    return -np.log(1.0 / a - 1.0)

norm2beta 

norm2beta(x: ArrayLike, max_val: float = 20.0) -> np.ndarray

Map R -> (0,max_val).

Source code in pyem/utils/math.py 
23
24
25
def norm2beta(x: ArrayLike, max_val: float = 20.0) -> np.ndarray:
    """Map R -> (0,max_val)."""
    return max_val / (1.0 + np.exp(-np.asarray(x)))

beta2norm 

beta2norm(b: ArrayLike, max_val: float = 20.0) -> np.ndarray
Source code in pyem/utils/math.py 
27
28
29
30
31
def beta2norm(b: ArrayLike, max_val: float = 20.0) -> np.ndarray:
    b = np.asarray(b)
    eps = 1e-12
    b = np.clip(b, eps, max_val - eps)
    return np.log(b / (max_val - b))

calc_fval 

calc_fval(negll: float, params: ArrayLike, prior=None, output: str = 'npl') -> float

Return objective value given a negative log-likelihood.

Parameters

negll : float Negative log-likelihood of the data under the model. params : array-like Parameter vector passed to the prior. Only used when prior is provided and output is "npl". prior : object or None, optional Object with a logpdf method returning the log prior density. When None the function reduces to returning negll. output : {"npl", "nll"} Indicates whether to return the negative posterior likelihood ("npl") or just the negative log-likelihood ("nll").

Returns

float Objective value suitable for minimisation. If the prior results in inf the value is capped at a large constant to keep optimisers stable.

Source code in pyem/utils/math.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def calc_fval(negll: float, params: ArrayLike, prior=None, output: str = 'npl') -> float:
    """Return objective value given a negative log-likelihood.

    Parameters
    ----------
    negll : float
        Negative log-likelihood of the data under the model.
    params : array-like
        Parameter vector passed to the prior.  Only used when `prior` is
        provided and ``output`` is ``"npl"``.
    prior : object or None, optional
        Object with a ``logpdf`` method returning the log prior density.
        When ``None`` the function reduces to returning ``negll``.
    output : {"npl", "nll"}
        Indicates whether to return the negative posterior likelihood
        (``"npl"``) or just the negative log-likelihood (``"nll"``).

    Returns
    -------
    float
        Objective value suitable for minimisation.  If the prior results in
        ``inf`` the value is capped at a large constant to keep optimisers
        stable.
    """
    if output == 'npl' and prior is not None and hasattr(prior, 'logpdf'):
        # Want to minimise -log[ P(data|h) * P(h) ]
        fval = -(-negll + prior.logpdf(np.asarray(params)))
        if not np.isfinite(fval):
            # Return a very large value so gradient-based optimisers keep going
            fval = 1e7
        return fval
    else:
        return float(negll) if np.isfinite(negll) else 1e7

Statistics & model-fit metrics

These functions turn the raw output of an EM fit into the model-comparison metrics used by pyem.core.compare.ModelComparison (and are usually called through it rather than directly). calc_LME computes the Laplace-approximate log model evidence from each subject's inverse-Hessian and negative-log-posterior, summed across subjects and penalized for the number of group-level parameters. calc_BICint computes the integrated BIC by Monte Carlo integration over the group posterior — for each subject it draws parameter samples from N(mu, sigma), re-evaluates fit_func at each sample, and combines the resulting negative log-likelihoods via a numerically stable logsumexp before applying the usual k*log(n) complexity penalty. pseudo_r2_from_nll computes a McFadden-style pseudo-R² comparing a model's (median or mean) negative log-likelihood against the negative log-likelihood of chance performance among noptions alternatives.

calc_LME 

calc_LME(inv_h: ndarray, NPL: ndarray) -> tuple[np.ndarray, float, np.ndarray]
Source code in pyem/utils/stats.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def calc_LME(inv_h: np.ndarray, NPL: np.ndarray) -> tuple[np.ndarray, float, np.ndarray]:
    nparams = inv_h.shape[0]
    nsubjects = inv_h.shape[2]
    good = np.zeros(nsubjects)
    Lap = np.zeros(nsubjects)
    for i in range(nsubjects):
        try:
            sign, logdet = np.linalg.slogdet(inv_h[:, :, i])
            if sign <= 0:
                raise ValueError("non-posdef Hessian")
            Lap[i] = -NPL[i] - 0.5 * (-logdet) + (nparams/2)*np.log(2*np.pi)
            good[i] = 1
        except Exception:
            Lap[i] = np.nan
            good[i] = 0
    if np.all(np.isnan(Lap)):
        Lap[:] = 0.0
    else:
        Lap[np.isnan(Lap)] = np.nanmean(Lap)
    lme = float(np.sum(Lap) - nparams*np.log(nsubjects))
    return Lap, lme, good

calc_BICint 

calc_BICint(all_data, param_names, mu, sigma, fit_func, nsamples: int = 2000, func_output: str = 'all', nll_key: str = 'nll', ntrials_total: int | None = None, njobs: int = -1) -> float

Integrated BIC via Monte Carlo integration over the group posterior.

ntrials_total is the number of independent trials contributing to a single subject's likelihood, used for the k*log(n) complexity penalty. If not supplied, it is auto-detected from the first subject's data, assuming every array-like field in that subject's data is trial-aligned with identical shape (true for the RW/Bayes families, e.g. [choices, rewards] both shaped (nblocks, ntrials)). For families with heterogeneous per-trial shapes (e.g. a GLM's [X, Y] where X has an extra feature-count dimension), pass ntrials_total explicitly rather than relying on auto-detection.

Source code in pyem/utils/stats.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def calc_BICint(
    all_data, param_names, mu, sigma, fit_func, nsamples: int = 2000, func_output: str = "all",
    nll_key: str = "nll", ntrials_total: int | None = None, njobs: int = -1,
) -> float:
    """Integrated BIC via Monte Carlo integration over the group posterior.

    ``ntrials_total`` is the number of independent trials contributing to a
    single subject's likelihood, used for the ``k*log(n)`` complexity
    penalty. If not supplied, it is auto-detected from the first subject's
    data, assuming every array-like field in that subject's data is
    trial-aligned with identical shape (true for the RW/Bayes families,
    e.g. ``[choices, rewards]`` both shaped ``(nblocks, ntrials)``). For
    families with heterogeneous per-trial shapes (e.g. a GLM's ``[X, Y]``
    where ``X`` has an extra feature-count dimension), pass
    ``ntrials_total`` explicitly rather than relying on auto-detection.
    """
    npar = len(param_names)
    if ntrials_total is not None:
        total_trials = ntrials_total
    elif isinstance(all_data[0], pd.DataFrame):
        total_trials = len(all_data[0])
    else:
        first = all_data[0]
        if isinstance(first, (list, tuple)) and hasattr(first[0], "size"):
            total_trials = int(first[0].size)
        else:
            raise ValueError("Unrecognized data structure in all_data")
    sigmasqrt = np.sqrt(np.asarray(sigma).reshape(-1))
    mu = np.asarray(mu).reshape(-1)
    def subj_iLog(beh):
        G = norm.rvs(loc=mu[:, None], scale=sigmasqrt[:, None], size=(len(mu), nsamples))
        subnll = []
        for k in range(nsamples):
            pars = G[:, k]
            # fit_func expected to return dict when output="all"
            info = fit_func(pars, *beh, output=func_output)
            subnll.append(info[nll_key])
        # log(sum(exp(-subnll))/nsamples), computed via logsumexp for numerical
        # stability: the naive form over/underflows when subnll (an unbounded
        # NLL, e.g. from a Gaussian-likelihood GLM fit) is very negative or
        # very positive, silently turning a valid subject into inf/-inf that
        # then gets dropped below instead of contributing its real value.
        iLog = logsumexp(-np.asarray(subnll)) - np.log(nsamples)
        return iLog
    iLogs = Parallel(n_jobs=njobs)(delayed(subj_iLog)(beh) for beh in all_data)
    iLogs = np.asarray(iLogs)
    finite = np.isfinite(iLogs)
    if not np.all(finite):
        warnings.warn(
             f"calc_BICint: dropping {int((~finite).sum())}/{len(iLogs)} subject(s) with non-finite integrated log-likelihood",
             RuntimeWarning,
         )
        iLogs = iLogs[finite]
    bicint = -2*np.sum(iLogs) + npar*np.log(total_trials)
    return float(bicint)

pseudo_r2_from_nll 

pseudo_r2_from_nll(nll: ndarray, ntrials_total: int, noptions: int, metric: str = 'median') -> float
Source code in pyem/utils/stats.py 
88
89
90
91
92
93
94
95
96
def pseudo_r2_from_nll(nll: np.ndarray, ntrials_total: int, noptions: int, metric: str = 'median') -> float:
    if metric == 'median':
        median_nll = float(np.median(nll))
        random_baseline = float(np.median(-np.log(1.0 / noptions) * ntrials_total))
        return 1.0 - (median_nll / random_baseline)
    else:
        mean_nll = float(np.mean(nll))
        random_baseline = float(np.mean(-np.log(1.0 / noptions) * ntrials_total))
        return 1.0 - (mean_nll / random_baseline)

Plotting

Convenience plotting helpers built on matplotlib and seaborn. These functions require the optional seaborn dependency — install it with the viz extra: pip install -e ".[viz]". plot_choices plots subject-averaged choice frequency over trials for a two-alternative task. plot_scatter draws an x/y scatter plot (e.g. true vs. estimated parameters from a parameter-recovery check) with an optional dashed x=y reference line and a Pearson-r annotation; it draws onto (and returns) a matplotlib Axes without calling plt.show(), so callers can compose multiple panels before displaying or saving the figure.

1
2
3
4
5
6
7
8
import matplotlib; matplotlib.use("Agg")
import numpy as np
from pyem.utils.plotting import plot_scatter

x = np.linspace(0, 1, 20)
y = x + np.random.default_rng(0).normal(scale=0.1, size=20)
ax = plot_scatter(x, "true", y, "estimated")
print(type(ax).__name__)

plot_choices 

plot_choices(choices_A, filename=None)

Plots the subject average EV over trials

Inputs
  • choices_A (np.array): subject choices for option A
  • filename (str): filename to save figure to (if None, figure is not saved)
Source code in pyem/utils/plotting.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def plot_choices(choices_A, filename=None):
    """
    Plots the subject average EV over trials

    Inputs:
        - choices_A (np.array): subject choices for option A
        - filename (str): filename to save figure to (if None, figure is not saved)
    """
    sns = _require_seaborn()
    nsubjects, nblocks, ntrials = choices_A.shape
    choices_B = np.ones_like(choices_A) - choices_A
    df_a = pd.DataFrame(np.mean(choices_A, axis=1), 
                        columns=[f'trial{i}' for i in range(ntrials)]).rename_axis('subject', axis=0).reset_index()
    df_b = pd.DataFrame(np.mean(choices_B, axis=1), 
                        columns=[f'trial{i}' for i in range(ntrials)]).rename_axis('subject', axis=0).reset_index()
    c_longA = pd.wide_to_long(df_a, stubnames='trial', i='subject', j='choices')
    c_longB = pd.wide_to_long(df_b, stubnames='trial', i='subject', j='choices')

    sns.lineplot(x='choices', y='trial',data=c_longA, label='A')
    sns.lineplot(x='choices', y='trial',data=c_longB, label='B')
    sns.despine()
    plt.xlabel('Trials')
    plt.ylabel('Choice Frequency')
    plt.legend()

    if filename:
        plt.savefig(filename, dpi=450, bbox_inches='tight')

    plt.show()

plot_scatter 

plot_scatter(x, xlabel, y, ylabel, *, ax=None, show_line=True, equal_limits=True, s=75, alpha=0.6, colorname='royalblue', annotate=True)

Scatter plot of x vs y with optional Pearson r annotation and x=y reference line. The function deliberately does not call plt.show() so that callers can aggregate multiple subplots before displaying the figure.

Parameters:

Name Type Description Default
x array - like

x-axis data

 required
xlabel str

x-axis label

 required
y array - like

y-axis data

 required
ylabel str

y-axis label

 required
ax Axes | None

Axes to draw on. If None, creates a new figure/axes.

 None
show_line bool

Whether to draw dashed x=y line.

 True
equal_limits bool

If True, sets identical limits and equal aspect.

 True
s float

Marker size.

 75
alpha float

Marker opacity.

 0.6
colorname str

Marker color.

 'royalblue'
annotate bool

If True, adds Pearson r annotation.

 True

Returns:

Type Description

matplotlib.axes.Axes: The axes the plot was drawn on.
Source code in pyem/utils/plotting.py 
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def plot_scatter(
    x,
    xlabel,
    y,
    ylabel,
    *,
    ax=None,
    show_line=True,
    equal_limits=True,
    s=75,
    alpha=0.6,
    colorname='royalblue',
    annotate=True,
):
    """
    Scatter plot of ``x`` vs ``y`` with optional Pearson r annotation and
    x=y reference line.  The function deliberately does **not** call
    ``plt.show()`` so that callers can aggregate multiple subplots before
    displaying the figure.

    Args:
        x (array-like): x-axis data
        xlabel (str): x-axis label
        y (array-like): y-axis data
        ylabel (str): y-axis label
        ax (matplotlib.axes.Axes | None): Axes to draw on. If None, creates a new figure/axes.
        show_line (bool): Whether to draw dashed x=y line.
        equal_limits (bool): If True, sets identical limits and equal aspect.
        s (float): Marker size.
        alpha (float): Marker opacity.
        colorname (str): Marker color.
        annotate (bool): If True, adds Pearson r annotation.

    Returns:
        matplotlib.axes.Axes: The axes the plot was drawn on.
    """
    sns = _require_seaborn()
    x = np.asarray(x)
    y = np.asarray(y)

    created_fig = None
    if ax is None:
        # If used standalone, give it a sensible size and turn on constrained layout
        created_fig, ax = plt.subplots(1, 1, figsize=(3.6, 3.6), constrained_layout=True)

    # Scatter
    ax.scatter(x, y, s=s, alpha=alpha, color=colorname)

    # Labels
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)

    # Pearson r annotation
    if annotate:
        mask = np.isfinite(x) & np.isfinite(y)
        if mask.any() and mask.sum() > 1:
            corr = np.corrcoef(x[mask], y[mask])[0, 1]
            ax.annotate(f'Pearson r = {corr:.2f}', xy=(0.05, 0.95),
                        xycoords='axes fraction', va='top', fontsize=11)

    # Optional x=y line and equal limits
    if show_line:
        data_min = np.nanmin([np.nanmin(x, initial=np.nan), np.nanmin(y, initial=np.nan)])
        data_max = np.nanmax([np.nanmax(x, initial=np.nan), np.nanmax(y, initial=np.nan)])
        cur_xmin, cur_xmax = ax.get_xlim()
        cur_ymin, cur_ymax = ax.get_ylim()
        lo = np.nanmin([data_min, cur_xmin, cur_ymin])
        hi = np.nanmax([data_max, cur_xmax, cur_ymax])
        ax.plot([lo, hi], [lo, hi], linestyle='--', color='k', alpha=0.6, zorder=0)

        if equal_limits:
            ax.set_xlim(lo, hi)
            ax.set_ylim(lo, hi)
            # Keep the box square without triggering excessive outer padding
            try:
                ax.set_box_aspect(1)
            except Exception:
                pass

    # A touch of margin so points/labels aren’t flush against the frame
    ax.margins(0.05)

    sns.despine()
    return ax