Skip to content

Creating Custom Models

The main purpose of pyEM is to make it very easy to generate new models for simulation and fitting. Every model follows the same template: a pair of _sim/_fit functions plus a ModelSpec that bundles them with a hand-authored id/description/spec. To create a custom model, follow the same template:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
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
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
from pyem.core.modelspec import ModelSpec
from pyem.utils.math import norm2alpha, norm2beta, calc_fval

def my_model_fit(params, choices, rewards, *, prior=None, output="npl"):
    """Fit function for your custom model.

    Parameters
    ----------
    params : sequence
        Parameter values in normalized space.
    choices, rewards : sequence
        Subject-specific data passed to the model.
    prior : object, optional
        Prior distribution with ``logpdf`` method.
    output : {"npl", "nll", "all"}
        Determines the function output.

    Returns
    -------
    float or dict
        Objective value or full output when ``output='all'``.
    """
    # ---- EDIT AS NEEDED  ---- 
    alpha = norm2alpha(params[0])
    beta = norm2beta(params[1])

    if not (0 <= alpha <= 1):
        return 1e7
    if not (0.001 <= beta <= 20):
        return 1e7
    # ------------------------- 

    # ---- YOUR CODE HERE  ---- 
    # Model-specific negative log-likelihood
    nll = ...
    # ------------------------- 

    if output == "all":
        return {"params": [alpha, beta], "choices": choices, "rewards": rewards, "nll": nll} # should have "params" and "nll"

    return calc_fval(nll, params, prior=prior, output=output)

def my_model_sim(params, **kwargs):
    """
    Simulation function for your custom model.

    Args:
        params: Parameter values in parameter space (n_subjects x n_params)
        **kwargs: Additional simulation parameters

    Returns:
        Dictionary with keys (CAN BE ANYTHING): "params", "choices", "rewards", etc.
    """

    # ---- YOUR CODE HERE  ---- 
    # Simulate trials/choices/rewards for your task and populate output arrays
    # ------------------------- 

    # (Optional) If you want to compute a likelihood value for debugging, define
    # it explicitly here; otherwise remove this line.
    # fval = calc_fval(nll, params, output="nll")

    return {"params": params, "choices": choices, "rewards": rewards}

# Wrapper metadata, following the same shape as every model in pyem.models.*
my_model_desc = "One-sentence description of what this model does and its free parameters."
my_model_id = "my_model"
my_model_spec = {"rl": {"softmax": ["beta"], "rw": ["alpha"]}}  # free-form, hand-authored — no fixed vocabulary to satisfy
my_model = ModelSpec(
    id=my_model_id, spec=my_model_spec, desc=my_model_desc,
    params=None, sim=my_model_sim, fit=my_model_fit,
)

ModelSpec is deliberately a plain, unopinionated container (id, spec, desc, params, sim, fit) — it imposes no naming scheme on your parameters or spec taxonomy, and nothing in EMModel/ModelComparison requires you to build one at all; it exists purely so a model can describe itself.