Skip to content

Posterior

Package

posterior

posterior

Posterior representations and diagnostics.

This subpackage provides: - ParameterPosterior: protocol for posteriors over model parameters p(θ | data) - MAPPosterior: delta distribution at θ_MAP (point estimate) - Posterior: backwards compatibility alias (deprecated, use MAPPosterior) - diagnostics: tools for checking posterior quality (ESS, R-hat, etc.)

Two-tier design
  • ParameterPosterior: represents p(θ | data) for research/diagnostics
  • PredictivePosterior: represents p(f(X*) | data) for acquisition functions (to be added in next phase)
Future extensions
  • LaplacePosterior: Gaussian approximation N(θ_MAP, Σ)
  • LangevinPosterior: MCMC samples
  • PredictivePosterior: predictions at test points

Classes:

Name Description
MAPPosterior

MAP (Maximum A Posteriori) posterior - delta distribution at θ_MAP.
ParameterPosterior

Protocol for posterior distributions over model parameters p(θ | data).
PredictivePosterior

Protocol for predictive distributions p(f(X*) | data) at test stimuli.
WPPMPredictivePosterior

Predictive posterior for WPPM models.

Functions:

Name Description
effective_sample_size

Estimate effective sample size (ESS) to calculate the number of independent
rhat

Compute R-hat convergence diagnostic.

Attributes:

Name Type Description
Posterior

Posterior 

Posterior = MAPPosterior

MAPPosterior 

MAPPosterior(params, model)

Bases: BasePosterior

MAP (Maximum A Posteriori) posterior - delta distribution at θ_MAP.

Represents a point estimate with no uncertainty.

Parameters:

Name Type Description Default
params dict

MAP parameter dictionary (θ_MAP)

 required
model WPPM

Model instance used for predictions

 required
Notes

This implements the ParameterPosterior protocol. For uncertainty quantification, use LaplacePosterior or LangevinPosterior.

Methods:

Name Description
MAP_params

Return the MAP parameters.
diagnostics

Return diagnostic information.
get_covariance_field

Create a CovarianceField object from this posterior.
log_prob

Evaluate log p(θ | data) under delta distribution.
predict_prob

Predict probability of correct response for a stimulus.
predict_thresholds

Predict discrimination threshold contour around a reference stimulus.
sample

Sample from delta distribution (returns repeated θ_MAP).

Attributes:

Name Type Description
model

Return the associated model.
params

Return the MAP parameters (θ_MAP).
Source code in src/psyphy/posterior/posterior.py 
43
44
45
def __init__(self, params, model):
    self._params = params
    self._model = model

model 

model

Return the associated model.

params 

params

Return the MAP parameters (θ_MAP).

MAP_params 

MAP_params()

Return the MAP parameters.

Returns:

Type Description
dict

Parameter dictionary.
Notes

Kept for backwards compatibility with BasePosterior. Use .params property instead.

Source code in src/psyphy/posterior/posterior.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
def MAP_params(self):
    """
    Return the MAP parameters.

    Returns
    -------
    dict
        Parameter dictionary.

    Notes
    -----
    Kept for backwards compatibility with BasePosterior.
    Use .params property instead.
    """
    return self._params

diagnostics 

diagnostics() -> dict

Return diagnostic information.

Returns:

Type Description
dict

Empty dict (no diagnostics for delta distribution)
Notes

Override this in subclasses to add optimizer convergence info.

Source code in src/psyphy/posterior/posterior.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
def diagnostics(self) -> dict:
    """
    Return diagnostic information.

    Returns
    -------
    dict
        Empty dict (no diagnostics for delta distribution)

    Notes
    -----
    Override this in subclasses to add optimizer convergence info.
    """
    return {}

get_covariance_field 

get_covariance_field()

Create a CovarianceField object from this posterior.

Returns:

Type Description
CovarianceField

Field object for evaluating Σ(x) and U(x) at arbitrary locations.

Examples:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
>>> # After fitting
>>> posterior = model.fit(data, optimizer=MAPOptimizer())
>>> field = posterior.get_covariance_field()
>>>
>>> # Evaluate at point
>>> Sigma_x = field.cov(jnp.array([0.5, 0.3]))
>>>
>>> # Evaluate over grid
>>> X_grid = jnp.stack(jnp.meshgrid(...), axis=-1)
>>> Sigmas = field.cov_batch(X_grid)
Notes

For MAP posteriors, uses θ_MAP. For variational posteriors, could use posterior mean or sample.

Source code in src/psyphy/posterior/posterior.py 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def get_covariance_field(self):
    """
    Create a CovarianceField object from this posterior.

    Returns
    -------
    CovarianceField
        Field object for evaluating Σ(x) and U(x) at arbitrary locations.

    Examples
    --------
    >>> # After fitting
    >>> posterior = model.fit(data, optimizer=MAPOptimizer())
    >>> field = posterior.get_covariance_field()
    >>>
    >>> # Evaluate at point
    >>> Sigma_x = field.cov(jnp.array([0.5, 0.3]))
    >>>
    >>> # Evaluate over grid
    >>> X_grid = jnp.stack(jnp.meshgrid(...), axis=-1)
    >>> Sigmas = field.cov_batch(X_grid)

    Notes
    -----
    For MAP posteriors, uses θ_MAP.
    For variational posteriors, could use posterior mean or sample.
    """
    from psyphy.model.covariance_field import WPPMCovarianceField

    return WPPMCovarianceField.from_posterior(self)

log_prob 

log_prob(params: dict) -> ndarray

Evaluate log p(θ | data) under delta distribution.

Parameters:

Name Type Description Default
params dict

Parameters to evaluate

 required

Returns:

Type Description
ndarray

0.0 if params == θ_MAP, -∞ otherwise
Notes

Delta distribution: all mass at θ_MAP. In practice, returns 0.0 at MAP and -inf elsewhere.

Source code in src/psyphy/posterior/posterior.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def log_prob(self, params: dict) -> jnp.ndarray:
    """
    Evaluate log p(θ | data) under delta distribution.

    Parameters
    ----------
    params : dict
        Parameters to evaluate

    Returns
    -------
    jnp.ndarray
        0.0 if params == θ_MAP, -∞ otherwise

    Notes
    -----
    Delta distribution: all mass at θ_MAP.
    In practice, returns 0.0 at MAP and -inf elsewhere.
    """
    # Check if params match MAP (element-wise)
    matches = jax.tree.map(lambda x, y: jnp.allclose(x, y), params, self._params)
    all_match = jax.tree.reduce(lambda a, b: a & b, matches)

    return jnp.where(all_match, 0.0, -jnp.inf)

predict_prob 

predict_prob(stimulus)

Predict probability of correct response for a stimulus.

Parameters:

Name Type Description Default
stimulus tuple

(reference, probe).

 required

Returns:

Type Description
ndarray

Probability of correct response.
Notes

Delegates to WPPM.predict_prob(). This is not recursion: Posterior calls WPPM’s method with stored params.

Source code in src/psyphy/posterior/posterior.py 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
def predict_prob(self, stimulus):
    """
    Predict probability of correct response for a stimulus.

    Parameters
    ----------
    stimulus : tuple
        (reference, probe).

    Returns
    -------
    jnp.ndarray
        Probability of correct response.

    Notes
    -----
    Delegates to WPPM.predict_prob().
    This is not recursion: Posterior calls WPPM’s method with stored params.
    """
    return self.model.predict_prob(self.params, stimulus)

predict_thresholds 

predict_thresholds(
    reference,
    criterion: float = 0.667,
    directions: int = 16,
)

Predict discrimination threshold contour around a reference stimulus.

Parameters:

Name Type Description Default
reference ndarray

Reference point in model space.

 required
criterion float

Target performance (e.g., 2/3 for oddity).

 0.667
directions int

Number of directions to probe.

 16

Returns:

Type Description
ndarray

Contour points (MVP: unit circle).
MVP

Returns a placeholder unit circle.

Future
  • Search outward in each direction until performance crosses criterion.
  • Average over posterior samples (Laplace, MCMC) to get credible intervals.
Source code in src/psyphy/posterior/posterior.py 
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
def predict_thresholds(
    self, reference, criterion: float = 0.667, directions: int = 16
):
    """
    Predict discrimination threshold contour around a reference stimulus.

    Parameters
    ----------
    reference : jnp.ndarray
        Reference point in model space.
    criterion : float, default=0.667
        Target performance (e.g., 2/3 for oddity).
    directions : int, default=16
        Number of directions to probe.

    Returns
    -------
    jnp.ndarray
        Contour points (MVP: unit circle).

    MVP
    ---
    Returns a placeholder unit circle.

    Future
    ------
    - Search outward in each direction until performance crosses criterion.
    - Average over posterior samples (Laplace, MCMC) to get credible intervals.
    """
    angles = jnp.linspace(0, 2 * jnp.pi, directions, endpoint=False)
    contour = jnp.stack(
        [reference + jnp.array([jnp.cos(a), jnp.sin(a)]) for a in angles]
    )
    return contour

sample 

sample(n: int = 1, *, key=None)

Sample from delta distribution (returns repeated θ_MAP).

Parameters:

Name Type Description Default
n int

Number of samples

 1
key KeyArray

PRNG key (unused for delta distribution)

 None

Returns:

Type Description
dict

Parameter PyTree with leading dimension n. Each array has shape (n, ...) with identical values.
Notes

Delta distribution has no randomness - returns repeated MAP estimate.

Source code in src/psyphy/posterior/posterior.py 
 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
def sample(self, n: int = 1, *, key=None):
    """
    Sample from delta distribution (returns repeated θ_MAP).

    Parameters
    ----------
    n : int, default=1
        Number of samples
    key : jax.random.KeyArray, optional
        PRNG key (unused for delta distribution)

    Returns
    -------
    dict
        Parameter PyTree with leading dimension n.
        Each array has shape (n, ...) with identical values.

    Notes
    -----
    Delta distribution has no randomness - returns repeated MAP estimate.
    """
    # Stack params n times along new leading dimension
    return jax.tree.map(
        lambda x: jnp.tile(x[None, ...], (n,) + (1,) * x.ndim), self._params
    )

ParameterPosterior

Bases: Protocol

Protocol for posterior distributions over model parameters p(θ | data).

Returned by InferenceEngine.fit(model, data). Used for research workflows: diagnostics, parameter sampling, uncertainty.

Methods:

Name Description
diagnostics

Return inference-specific diagnostic information.
log_prob

Evaluate log p(θ | data) at given parameters.
predict_prob

Predict probability of correct response for a stimulus.
predict_thresholds

Predict discrimination threshold contour around a reference.
sample

Sample parameter vectors from p(θ | data).

Attributes:

Name Type Description
model

Associated generative model.
params dict

Point estimate or posterior mean parameters.

model 

model

Associated generative model.

Returns:

Type Description
Model

The WPPM or other model instance used for predictions.

params 

params: dict

Point estimate or posterior mean parameters.

Returns:

Type Description
dict

Parameter PyTree (e.g., {"log_diag": jnp.ndarray, ...})
Notes
  • MAP: θ_MAP
  • Laplace: E[θ] (Gaussian mean, equals θ_MAP)
  • MCMC: posterior mean of samples

diagnostics 

diagnostics() -> dict

Return inference-specific diagnostic information.

Returns:

Type Description
dict

Diagnostic metrics. Contents vary by inference method:

MAP: - convergence_info: optimizer status - final_loss: negative log posterior at MAP

Laplace: - log_marginal_likelihood: evidence approximation - hessian_condition_number: numerical stability - marginal_variances: parameter uncertainties

MCMC: - ess: effective sample size per parameter - rhat: Gelman-Rubin statistic - acceptance_rate: step acceptance (if applicable)
Source code in src/psyphy/posterior/parameter_posterior.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
def diagnostics(self) -> dict:
    """
    Return inference-specific diagnostic information.

    Returns
    -------
    dict
        Diagnostic metrics. Contents vary by inference method:

        MAP:
            - convergence_info: optimizer status
            - final_loss: negative log posterior at MAP

        Laplace:
            - log_marginal_likelihood: evidence approximation
            - hessian_condition_number: numerical stability
            - marginal_variances: parameter uncertainties

        MCMC:
            - ess: effective sample size per parameter
            - rhat: Gelman-Rubin statistic
            - acceptance_rate: step acceptance (if applicable)
    """
    ...

log_prob 

log_prob(params: dict) -> ndarray

Evaluate log p(θ | data) at given parameters.

Parameters:

Name Type Description Default
params dict

Parameter PyTree to evaluate

 required

Returns:

Type Description
ndarray

Log probability (scalar)

Raises:

Type Description
NotImplementedError

For MCMC posteriors (no tractable density)
Notes
  • MAP: -∞ for θ ≠ θ_MAP, 0 at θ_MAP (delta distribution)
  • Laplace: Gaussian log density
  • MCMC: raises NotImplementedError
Source code in src/psyphy/posterior/parameter_posterior.py 
 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
def log_prob(self, params: dict) -> jnp.ndarray:
    """
    Evaluate log p(θ | data) at given parameters.

    Parameters
    ----------
    params : dict
        Parameter PyTree to evaluate

    Returns
    -------
    jnp.ndarray
        Log probability (scalar)

    Raises
    ------
    NotImplementedError
        For MCMC posteriors (no tractable density)

    Notes
    -----
    - MAP: -∞ for θ ≠ θ_MAP, 0 at θ_MAP (delta distribution)
    - Laplace: Gaussian log density
    - MCMC: raises NotImplementedError
    """
    ...

predict_prob 

predict_prob(stimulus) -> ndarray

Predict probability of correct response for a stimulus.

Parameters:

Name Type Description Default
stimulus tuple

(reference, probe) in model space

 required

Returns:

Type Description
ndarray

Probability of correct response
Notes

Uses point estimate (self.params) for prediction. For uncertainty quantification, use PredictivePosterior instead.

Source code in src/psyphy/posterior/parameter_posterior.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
def predict_prob(self, stimulus) -> jnp.ndarray:
    """
    Predict probability of correct response for a stimulus.

    Parameters
    ----------
    stimulus : tuple
        (reference, probe) in model space

    Returns
    -------
    jnp.ndarray
        Probability of correct response

    Notes
    -----
    Uses point estimate (self.params) for prediction.
    For uncertainty quantification, use PredictivePosterior instead.
    """
    ...

predict_thresholds 

predict_thresholds(
    reference,
    criterion: float = 0.667,
    directions: int = 16,
) -> ndarray

Predict discrimination threshold contour around a reference.

Parameters:

Name Type Description Default
reference ndarray

Reference point in model space

 required
criterion float

Target performance level

 0.667
directions int

Number of directions to probe

 16

Returns:

Type Description
ndarray

Threshold contour points
Notes

Uses point estimate. For credible regions, use PredictivePosterior with multiple samples.

Source code in src/psyphy/posterior/parameter_posterior.py 
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
def predict_thresholds(
    self,
    reference,
    criterion: float = 0.667,
    directions: int = 16,
) -> jnp.ndarray:
    """
    Predict discrimination threshold contour around a reference.

    Parameters
    ----------
    reference : jnp.ndarray
        Reference point in model space
    criterion : float, default=0.667
        Target performance level
    directions : int, default=16
        Number of directions to probe

    Returns
    -------
    jnp.ndarray
        Threshold contour points

    Notes
    -----
    Uses point estimate. For credible regions, use PredictivePosterior
    with multiple samples.
    """
    ...

sample 

sample(n: int, *, key: KeyArray) -> dict

Sample parameter vectors from p(θ | data).

Parameters:

Name Type Description Default
n int

Number of samples

 required
key KeyArray

PRNG key for randomness

 required

Returns:

Type Description
dict

Parameter PyTree with leading dimension n. Example: {"log_diag": jnp.ndarray with shape (n, input_dim)}
Notes
  • MAP: returns repeated θ_MAP
  • Laplace: samples from N(θ_MAP, Σ)
  • MCMC: returns stored samples (may subsample if n differs)
Source code in src/psyphy/posterior/parameter_posterior.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def sample(self, n: int, *, key: jr.KeyArray) -> dict:
    """
    Sample parameter vectors from p(θ | data).

    Parameters
    ----------
    n : int
        Number of samples
    key : jax.random.KeyArray
        PRNG key for randomness

    Returns
    -------
    dict
        Parameter PyTree with leading dimension n.
        Example: {"log_diag": jnp.ndarray with shape (n, input_dim)}

    Notes
    -----
    - MAP: returns repeated θ_MAP
    - Laplace: samples from N(θ_MAP, Σ)
    - MCMC: returns stored samples (may subsample if n differs)
    """
    ...

PredictivePosterior

Bases: Protocol

Protocol for predictive distributions p(f(X*) | data) at test stimuli.

Returned by Model.posterior(X) for use in acquisition functions.

Methods:

Name Description
cov_field

Posterior over perceptual covariance field Σ(X).
rsample

Reparameterized samples from p(f(X*) | data).

Attributes:

Name Type Description
mean ndarray

Posterior predictive mean E[f(X*) | data].
variance ndarray

Posterior predictive marginal variances Var[f(X*) | data].

mean 

mean: ndarray

Posterior predictive mean E[f(X*) | data].

Returns:

Type Description
ndarray

Shape (n_test,) for scalar outputs Shape (n_test, output_dim) for vector outputs (future)
Notes

Computed via Monte Carlo integration over parameter posterior.

variance 

variance: ndarray

Posterior predictive marginal variances Var[f(X*) | data].

Returns:

Type Description
ndarray

Shape (n_test,) for scalar outputs Shape (n_test, output_dim) for vector outputs (future)
Notes

Captures both aleatoric (model) and epistemic (parameter) uncertainty.

cov_field 

cov_field(X: ndarray) -> ndarray

Posterior over perceptual covariance field Σ(X).

Parameters:

Name Type Description Default
X ndarray

Test stimuli, shape (n_test, input_dim)

 required

Returns:

Type Description
ndarray

Posterior mean covariance E[Σ(X) | data], shape (n_test, input_dim, input_dim)
Notes

WPPM-specific method for visualizing perceptual noise structure. This is NOT the predictive covariance - it's the model's internal representation of perceptual uncertainty.

Source code in src/psyphy/posterior/predictive_posterior.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
def cov_field(self, X: jnp.ndarray) -> jnp.ndarray:
    """
    Posterior over perceptual covariance field Σ(X).

    Parameters
    ----------
    X : jnp.ndarray
        Test stimuli, shape (n_test, input_dim)

    Returns
    -------
    jnp.ndarray
        Posterior mean covariance E[Σ(X) | data],
        shape (n_test, input_dim, input_dim)

    Notes
    -----
    WPPM-specific method for visualizing perceptual noise structure.
    This is NOT the predictive covariance - it's the model's
    internal representation of perceptual uncertainty.
    """
    ...

rsample 

rsample(
    sample_shape: tuple = (), *, key: KeyArray
) -> ndarray

Reparameterized samples from p(f(X*) | data).

Parameters:

Name Type Description Default
sample_shape tuple

Shape of sample batch

 ()
key KeyArray

PRNG key

 required

Returns:

Type Description
ndarray

Shape (sample_shape, n_test) for scalar outputs Shape (sample_shape, n_test, output_dim) for vector outputs
Notes

Enables gradient-based acquisition optimization via reparameterization trick.

Source code in src/psyphy/posterior/predictive_posterior.py 
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
def rsample(self, sample_shape: tuple = (), *, key: jr.KeyArray) -> jnp.ndarray:
    """
    Reparameterized samples from p(f(X*) | data).

    Parameters
    ----------
    sample_shape : tuple, default=()
        Shape of sample batch
    key : jax.random.KeyArray
        PRNG key

    Returns
    -------
    jnp.ndarray
        Shape (*sample_shape, n_test) for scalar outputs
        Shape (*sample_shape, n_test, output_dim) for vector outputs

    Notes
    -----
    Enables gradient-based acquisition optimization via reparameterization trick.
    """
    ...

WPPMPredictivePosterior 

WPPMPredictivePosterior(
    param_posterior: ParameterPosterior,
    X: ndarray,
    probes: ndarray | None = None,
    n_samples: int = 100,
)

Predictive posterior for WPPM models.

Computes p(f(X*) | data) via Monte Carlo integration over parameter posterior p(θ | data).

Parameters:

Name Type Description Default
param_posterior ParameterPosterior

Posterior over model parameters

 required
X (ndarray, shape(n_test, input_dim))

Test reference stimuli

 required
probes (ndarray, shape(n_test, input_dim))

Test probe stimuli. If None, predictions are over thresholds.

 None
n_samples int

Number of posterior samples for MC integration

 100

Attributes:

Name Type Description
param_posterior ParameterPosterior

Wrapped parameter posterior
X ndarray

Test stimuli
probes ndarray | None

Test probes
n_samples int

MC sample count
Notes

Uses lazy evaluation: moments computed on first access.

Methods:

Name Description
cov_field

Posterior mean covariance field E[Σ(X) | data].
rsample

Sample predictions from p(f(X*) | data).
Source code in src/psyphy/posterior/predictive_posterior.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
def __init__(
    self,
    param_posterior: ParameterPosterior,
    X: jnp.ndarray,
    probes: jnp.ndarray | None = None,
    n_samples: int = 100,
):
    self.param_posterior = param_posterior
    self.X = X
    self.probes = probes
    self.n_samples = n_samples

    # Lazy evaluation cache
    self._mean = None
    self._variance = None
    self._computed = False

X 

X = X

mean 

mean: ndarray

E[f(X*) | data], shape (n_test,).

n_samples 

n_samples = n_samples

param_posterior 

param_posterior = param_posterior

probes 

probes = probes

variance 

variance: ndarray

Var[f(X*) | data], shape (n_test,).

cov_field 

cov_field(X: ndarray) -> ndarray

Posterior mean covariance field E[Σ(X) | data].

Parameters:

Name Type Description Default
X ndarray

Test stimuli, shape (n_test, input_dim)

 required

Returns:

Type Description
ndarray

Covariance matrices, shape (n_test, input_dim, input_dim)
Notes

Averages local_covariance(x) over parameter posterior samples.

Source code in src/psyphy/posterior/predictive_posterior.py 
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
def cov_field(self, X: jnp.ndarray) -> jnp.ndarray:
    """
    Posterior mean covariance field E[Σ(X) | data].

    Parameters
    ----------
    X : jnp.ndarray
        Test stimuli, shape (n_test, input_dim)

    Returns
    -------
    jnp.ndarray
        Covariance matrices, shape (n_test, input_dim, input_dim)

    Notes
    -----
    Averages local_covariance(x) over parameter posterior samples.
    """
    key = jr.PRNGKey(0)
    param_samples = self.param_posterior.sample(self.n_samples, key=key)

    model = self.param_posterior.model

    def cov_at_x(params, x):
        """Evaluate Σ(x) with given parameters."""
        return model.local_covariance(params, x)

    # Vectorized evaluation: (n_samples, n_test, input_dim, input_dim)
    cov_samples = jax.vmap(
        lambda params: jax.vmap(lambda x: cov_at_x(params, x))(X)
    )(param_samples)

    # Return posterior mean
    return jnp.mean(cov_samples, axis=0)

rsample 

rsample(
    sample_shape: tuple = (), *, key: KeyArray
) -> ndarray

Sample predictions from p(f(X*) | data).

Parameters:

Name Type Description Default
sample_shape tuple

Batch shape

 ()
key KeyArray

PRNG key

 required

Returns:

Type Description
ndarray

Shape (*sample_shape, n_test)
Source code in src/psyphy/posterior/predictive_posterior.py 
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
def rsample(self, sample_shape: tuple = (), *, key: jr.KeyArray) -> jnp.ndarray:
    """
    Sample predictions from p(f(X*) | data).

    Parameters
    ----------
    sample_shape : tuple
        Batch shape
    key : jax.random.KeyArray
        PRNG key

    Returns
    -------
    jnp.ndarray
        Shape (*sample_shape, n_test)
    """
    n = int(jnp.prod(jnp.array(sample_shape))) if sample_shape else 1
    param_samples = self.param_posterior.sample(n, key=key)

    model = self.param_posterior.model

    if self.probes is None:
        raise NotImplementedError("Threshold sampling not yet implemented")

    def predict_one(params):
        """Predict for all test points with given params."""
        return jax.vmap(lambda r, p: model.predict_prob(params, (r, p)))(
            self.X, self.probes
        )

    samples = jax.vmap(predict_one)(param_samples)

    if sample_shape:
        return samples.reshape(*sample_shape, -1)
    return samples

effective_sample_size 

effective_sample_size(samples: ndarray) -> float

Estimate effective sample size (ESS) to calculate the number of independent samples that a correlated MCMC chain is equivalent to.

Parameters:

Name Type Description Default
samples ndarray

Posterior samples, shape (n_samples, ...).

 required

Returns:

Type Description
float

Effective sample size (stub).
Notes

MVP: Returns the number of samples. Full WPPM mode: Compute ESS using autocorrelation structure.

Source code in src/psyphy/posterior/diagnostics.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def effective_sample_size(samples: jnp.ndarray) -> float:
    """
    Estimate effective sample size (ESS) to calculate the number of independent
    samples that a correlated MCMC chain is equivalent to.

    Parameters
    ----------
    samples : jnp.ndarray
        Posterior samples, shape (n_samples, ...).

    Returns
    -------
    float
        Effective sample size (stub).

    Notes
    -----
    MVP:
        Returns the number of samples.
    Full WPPM mode:
        Compute ESS using autocorrelation structure.
    """
    return samples.shape[0]

rhat 

rhat(chains: ndarray) -> float

Compute R-hat convergence diagnostic.

Parameters:

Name Type Description Default
chains ndarray

Posterior samples across chains, shape (n_chains, n_samples, ...).

 required

Returns:

Type Description
float

R-hat statistic (stub).
Notes

MVP: Always returns 1.0. Full WPPM mode: Implement Gelman-Rubin diagnostic [1]

References: 
1
[1] https://bookdown.org/rdpeng/advstatcomp/monitoring-convergence.html
Source code in src/psyphy/posterior/diagnostics.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
def rhat(chains: jnp.ndarray) -> float:
    """
    Compute R-hat convergence diagnostic.

    Parameters
    ----------
    chains : jnp.ndarray
        Posterior samples across chains, shape (n_chains, n_samples, ...).

    Returns
    -------
    float
        R-hat statistic (stub).

    Notes
    -----
    MVP:
        Always returns 1.0.
    Full WPPM mode:
        Implement Gelman-Rubin diagnostic [1]

    References:
    ----------
        [1] https://bookdown.org/rdpeng/advstatcomp/monitoring-convergence.html

    """
    return 1.0

Posterior Representation

posterior

posterior.py

Concrete ParameterPosterior implementations.

This module provides: - MAPPosterior: delta distribution at θ_MAP (point estimate) - Posterior: backwards compatibility alias (deprecated)

Future additions: - LaplacePosterior: Gaussian approximation - LangevinPosterior: MCMC samples

Classes:

Name Description
MAPPosterior

MAP (Maximum A Posteriori) posterior - delta distribution at θ_MAP.

Attributes:

Name Type Description
Posterior

Posterior 

Posterior = MAPPosterior

MAPPosterior 

MAPPosterior(params, model)

Bases: BasePosterior

MAP (Maximum A Posteriori) posterior - delta distribution at θ_MAP.

Represents a point estimate with no uncertainty.

Parameters:

Name Type Description Default
params dict

MAP parameter dictionary (θ_MAP)

 required
model WPPM

Model instance used for predictions

 required
Notes

This implements the ParameterPosterior protocol. For uncertainty quantification, use LaplacePosterior or LangevinPosterior.

Methods:

Name Description
MAP_params

Return the MAP parameters.
diagnostics

Return diagnostic information.
get_covariance_field

Create a CovarianceField object from this posterior.
log_prob

Evaluate log p(θ | data) under delta distribution.
predict_prob

Predict probability of correct response for a stimulus.
predict_thresholds

Predict discrimination threshold contour around a reference stimulus.
sample

Sample from delta distribution (returns repeated θ_MAP).

Attributes:

Name Type Description
model

Return the associated model.
params

Return the MAP parameters (θ_MAP).
Source code in src/psyphy/posterior/posterior.py 
43
44
45
def __init__(self, params, model):
    self._params = params
    self._model = model

model 

model

Return the associated model.

params 

params

Return the MAP parameters (θ_MAP).

MAP_params 

MAP_params()

Return the MAP parameters.

Returns:

Type Description
dict

Parameter dictionary.
Notes

Kept for backwards compatibility with BasePosterior. Use .params property instead.

Source code in src/psyphy/posterior/posterior.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
def MAP_params(self):
    """
    Return the MAP parameters.

    Returns
    -------
    dict
        Parameter dictionary.

    Notes
    -----
    Kept for backwards compatibility with BasePosterior.
    Use .params property instead.
    """
    return self._params

diagnostics 

diagnostics() -> dict

Return diagnostic information.

Returns:

Type Description
dict

Empty dict (no diagnostics for delta distribution)
Notes

Override this in subclasses to add optimizer convergence info.

Source code in src/psyphy/posterior/posterior.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
def diagnostics(self) -> dict:
    """
    Return diagnostic information.

    Returns
    -------
    dict
        Empty dict (no diagnostics for delta distribution)

    Notes
    -----
    Override this in subclasses to add optimizer convergence info.
    """
    return {}

get_covariance_field 

get_covariance_field()

Create a CovarianceField object from this posterior.

Returns:

Type Description
CovarianceField

Field object for evaluating Σ(x) and U(x) at arbitrary locations.

Examples:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
>>> # After fitting
>>> posterior = model.fit(data, optimizer=MAPOptimizer())
>>> field = posterior.get_covariance_field()
>>>
>>> # Evaluate at point
>>> Sigma_x = field.cov(jnp.array([0.5, 0.3]))
>>>
>>> # Evaluate over grid
>>> X_grid = jnp.stack(jnp.meshgrid(...), axis=-1)
>>> Sigmas = field.cov_batch(X_grid)
Notes

For MAP posteriors, uses θ_MAP. For variational posteriors, could use posterior mean or sample.

Source code in src/psyphy/posterior/posterior.py 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def get_covariance_field(self):
    """
    Create a CovarianceField object from this posterior.

    Returns
    -------
    CovarianceField
        Field object for evaluating Σ(x) and U(x) at arbitrary locations.

    Examples
    --------
    >>> # After fitting
    >>> posterior = model.fit(data, optimizer=MAPOptimizer())
    >>> field = posterior.get_covariance_field()
    >>>
    >>> # Evaluate at point
    >>> Sigma_x = field.cov(jnp.array([0.5, 0.3]))
    >>>
    >>> # Evaluate over grid
    >>> X_grid = jnp.stack(jnp.meshgrid(...), axis=-1)
    >>> Sigmas = field.cov_batch(X_grid)

    Notes
    -----
    For MAP posteriors, uses θ_MAP.
    For variational posteriors, could use posterior mean or sample.
    """
    from psyphy.model.covariance_field import WPPMCovarianceField

    return WPPMCovarianceField.from_posterior(self)

log_prob 

log_prob(params: dict) -> ndarray

Evaluate log p(θ | data) under delta distribution.

Parameters:

Name Type Description Default
params dict

Parameters to evaluate

 required

Returns:

Type Description
ndarray

0.0 if params == θ_MAP, -∞ otherwise
Notes

Delta distribution: all mass at θ_MAP. In practice, returns 0.0 at MAP and -inf elsewhere.

Source code in src/psyphy/posterior/posterior.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def log_prob(self, params: dict) -> jnp.ndarray:
    """
    Evaluate log p(θ | data) under delta distribution.

    Parameters
    ----------
    params : dict
        Parameters to evaluate

    Returns
    -------
    jnp.ndarray
        0.0 if params == θ_MAP, -∞ otherwise

    Notes
    -----
    Delta distribution: all mass at θ_MAP.
    In practice, returns 0.0 at MAP and -inf elsewhere.
    """
    # Check if params match MAP (element-wise)
    matches = jax.tree.map(lambda x, y: jnp.allclose(x, y), params, self._params)
    all_match = jax.tree.reduce(lambda a, b: a & b, matches)

    return jnp.where(all_match, 0.0, -jnp.inf)

predict_prob 

predict_prob(stimulus)

Predict probability of correct response for a stimulus.

Parameters:

Name Type Description Default
stimulus tuple

(reference, probe).

 required

Returns:

Type Description
ndarray

Probability of correct response.
Notes

Delegates to WPPM.predict_prob(). This is not recursion: Posterior calls WPPM’s method with stored params.

Source code in src/psyphy/posterior/posterior.py 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
def predict_prob(self, stimulus):
    """
    Predict probability of correct response for a stimulus.

    Parameters
    ----------
    stimulus : tuple
        (reference, probe).

    Returns
    -------
    jnp.ndarray
        Probability of correct response.

    Notes
    -----
    Delegates to WPPM.predict_prob().
    This is not recursion: Posterior calls WPPM’s method with stored params.
    """
    return self.model.predict_prob(self.params, stimulus)

predict_thresholds 

predict_thresholds(
    reference,
    criterion: float = 0.667,
    directions: int = 16,
)

Predict discrimination threshold contour around a reference stimulus.

Parameters:

Name Type Description Default
reference ndarray

Reference point in model space.

 required
criterion float

Target performance (e.g., 2/3 for oddity).

 0.667
directions int

Number of directions to probe.

 16

Returns:

Type Description
ndarray

Contour points (MVP: unit circle).
MVP

Returns a placeholder unit circle.

Future
  • Search outward in each direction until performance crosses criterion.
  • Average over posterior samples (Laplace, MCMC) to get credible intervals.
Source code in src/psyphy/posterior/posterior.py 
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
def predict_thresholds(
    self, reference, criterion: float = 0.667, directions: int = 16
):
    """
    Predict discrimination threshold contour around a reference stimulus.

    Parameters
    ----------
    reference : jnp.ndarray
        Reference point in model space.
    criterion : float, default=0.667
        Target performance (e.g., 2/3 for oddity).
    directions : int, default=16
        Number of directions to probe.

    Returns
    -------
    jnp.ndarray
        Contour points (MVP: unit circle).

    MVP
    ---
    Returns a placeholder unit circle.

    Future
    ------
    - Search outward in each direction until performance crosses criterion.
    - Average over posterior samples (Laplace, MCMC) to get credible intervals.
    """
    angles = jnp.linspace(0, 2 * jnp.pi, directions, endpoint=False)
    contour = jnp.stack(
        [reference + jnp.array([jnp.cos(a), jnp.sin(a)]) for a in angles]
    )
    return contour

sample 

sample(n: int = 1, *, key=None)

Sample from delta distribution (returns repeated θ_MAP).

Parameters:

Name Type Description Default
n int

Number of samples

 1
key KeyArray

PRNG key (unused for delta distribution)

 None

Returns:

Type Description
dict

Parameter PyTree with leading dimension n. Each array has shape (n, ...) with identical values.
Notes

Delta distribution has no randomness - returns repeated MAP estimate.

Source code in src/psyphy/posterior/posterior.py 
 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
def sample(self, n: int = 1, *, key=None):
    """
    Sample from delta distribution (returns repeated θ_MAP).

    Parameters
    ----------
    n : int, default=1
        Number of samples
    key : jax.random.KeyArray, optional
        PRNG key (unused for delta distribution)

    Returns
    -------
    dict
        Parameter PyTree with leading dimension n.
        Each array has shape (n, ...) with identical values.

    Notes
    -----
    Delta distribution has no randomness - returns repeated MAP estimate.
    """
    # Stack params n times along new leading dimension
    return jax.tree.map(
        lambda x: jnp.tile(x[None, ...], (n,) + (1,) * x.ndim), self._params
    )

Diagnostics

diagnostics

diagnostics.py

Posterior diagnostics.

Provides functions to check quality of posterior inference.

MVP implementation: - Stubs for effective sample size and R-hat.

Full WPPM mode: - Implement real diagnostics from posterior chains. - Include posterior predictive checks.

Functions:

Name Description
effective_sample_size

Estimate effective sample size (ESS) to calculate the number of independent
rhat

Compute R-hat convergence diagnostic.

effective_sample_size 

effective_sample_size(samples: ndarray) -> float

Estimate effective sample size (ESS) to calculate the number of independent samples that a correlated MCMC chain is equivalent to.

Parameters:

Name Type Description Default
samples ndarray

Posterior samples, shape (n_samples, ...).

 required

Returns:

Type Description
float

Effective sample size (stub).
Notes

MVP: Returns the number of samples. Full WPPM mode: Compute ESS using autocorrelation structure.

Source code in src/psyphy/posterior/diagnostics.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def effective_sample_size(samples: jnp.ndarray) -> float:
    """
    Estimate effective sample size (ESS) to calculate the number of independent
    samples that a correlated MCMC chain is equivalent to.

    Parameters
    ----------
    samples : jnp.ndarray
        Posterior samples, shape (n_samples, ...).

    Returns
    -------
    float
        Effective sample size (stub).

    Notes
    -----
    MVP:
        Returns the number of samples.
    Full WPPM mode:
        Compute ESS using autocorrelation structure.
    """
    return samples.shape[0]

rhat 

rhat(chains: ndarray) -> float

Compute R-hat convergence diagnostic.

Parameters:

Name Type Description Default
chains ndarray

Posterior samples across chains, shape (n_chains, n_samples, ...).

 required

Returns:

Type Description
float

R-hat statistic (stub).
Notes

MVP: Always returns 1.0. Full WPPM mode: Implement Gelman-Rubin diagnostic [1]

References: 
1
[1] https://bookdown.org/rdpeng/advstatcomp/monitoring-convergence.html
Source code in src/psyphy/posterior/diagnostics.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
def rhat(chains: jnp.ndarray) -> float:
    """
    Compute R-hat convergence diagnostic.

    Parameters
    ----------
    chains : jnp.ndarray
        Posterior samples across chains, shape (n_chains, n_samples, ...).

    Returns
    -------
    float
        R-hat statistic (stub).

    Notes
    -----
    MVP:
        Always returns 1.0.
    Full WPPM mode:
        Implement Gelman-Rubin diagnostic [1]

    References:
    ----------
        [1] https://bookdown.org/rdpeng/advstatcomp/monitoring-convergence.html

    """
    return 1.0