Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
aryan-pan
marked this pull request as ready for review
August 14, 2026 18:48
aryan-pan
force-pushed
the
feat/pec-hierarchical-em
branch
from
September 4, 2026 19:25
9ad7649 to
8793e75
Compare
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
aryan-pan
force-pushed
the
feat/pec-hierarchical-em
branch
from
September 5, 2026 00:09
1c7be34 to
3ad66f6
Compare
This comment was marked as outdated.
This comment was marked as outdated.
aryan-pan
force-pushed
the
feat/pec-hierarchical-em
branch
2 times, most recently
from
September 5, 2026 20:02
7964202 to
5de6dc6
Compare
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
aryan-pan
force-pushed
the
feat/pec-hierarchical-em
branch
from
September 9, 2026 18:08
9be01cd to
358e47b
Compare
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
|
This PR causes the following changes to the html docs (ubuntu-latest-3.11): See CI logs for the full diff. |
|
This PR causes the following changes to the html docs (ubuntu-latest-3.11): See CI logs for the full diff. |
aryan-pan
force-pushed
the
feat/pec-hierarchical-em
branch
2 times, most recently
from
September 15, 2026 18:00
9cf11a9 to
65bed70
Compare
davidt0x
requested changes
Sep 15, 2026
|
|
||
| # parameter estimation tests with optional dependencies on linux python 3.10 | ||
| - python-version: '3.10' | ||
| os: ubuntu |
Collaborator
There was a problem hiding this comment.
Why not enable these tests on all OS?
Collaborator
Author
There was a problem hiding this comment.
added CI for windows/macOS as well. if you meant running these by default across all builds, I can make that change
| off against each other the first is the smaller, so intervals err narrow. | ||
|
|
||
| Curvature comes from central finite differences, which requires the objective to be deterministic in | ||
| `theta`. Participant models must therefore be built with common random numbers; without them the |
Collaborator
There was a problem hiding this comment.
Do we enforce this? We should throw an error if a user constructs the PEC without common random numbers.
| # mode, and one that is still infinite describes a probe that never fit. Neither is a width, | ||
| # so the prior stands in and reports that this parameter was not pinned down. | ||
| usable = np.isfinite(curvature) & (curvature > 0) | ||
| inverse = np.where(usable, 1.0 / np.where(usable, curvature, 1.0), prior_variance) |
Collaborator
There was a problem hiding this comment.
Codex caught this one:
I see the diagonal approximation is documented, but these variances also feed into the group variance update, so the effect goes beyond narrower uncertainty intervals. When parameters trade off, 1 / H[k,k] can be quite different from (H⁻¹)[k,k]. A small Gaussian test with H = [[10, 9], [9, 10]] gives variances of 0.1 here versus about 0.526 for the exact posterior.
Could we consider computing the full participant Hessian while keeping the group covariance diagonal? If we want to keep this approximation for speed, it would be useful to add a correlated-likelihood test and check how much it affects parameter recovery in the DDM example.
Collaborator
There was a problem hiding this comment.
This will of course increase the number of likelihood evaluations as number of parameters increase. However, I think we will be able to mitigate that with the GPU stuff I have been working on.
Collaborator
There was a problem hiding this comment.
Here is a example test illustrating the issue from codex:
"""Review reproduction: compare diagonal curvature with an exact Gaussian posterior.
The independent case should pass. The coupled case intentionally fails on PR #3619:
the approximation reports conditional variances instead of marginal variances.
"""
import numpy as np
import pytest
from psyneulink.core.compositions.hierarchical.laplaceem import subject_map_estep
@pytest.mark.parametrize("coupled", [False, True], ids=["independent", "parameter_tradeoff"])
def test_estep_variance_matches_exact_gaussian_posterior(coupled):
# Independent standard-normal priors on z[0] and z[1].
# Observation noise variance is 1/9, with observed values equal to zero.
# Independent case: observe each parameter separately.
# Coupled case: observe only their sum, allowing them to trade off.
def neg_log_post(z):
prior = 0.5 * np.dot(z, z)
residual_squared = np.sum(z) ** 2 if coupled else np.dot(z, z)
return float(prior + 0.5 * 9.0 * residual_squared)
post = subject_map_estep(
neg_log_post, z0=np.zeros(2), prior_variance=np.ones(2)
)
# These are exactly quadratic objectives, so neither simulation noise nor
# finite-difference truncation error explains a covariance discrepancy.
assert post.success
np.testing.assert_allclose(post.z_hat, np.zeros(2), atol=1e-6)
np.testing.assert_allclose(post.curvature, [10.0, 10.0], rtol=1e-6)
# Independent: H = 10 I, so diag(inv(H)) = [0.1, 0.1].
# Coupled: H = [[10, 9], [9, 10]], whose determinant is 19;
# inv(H) = [[10, -9], [-9, 10]] / 19.
exact_variance = np.full(2, 10.0 / 19.0 if coupled else 0.1)
np.testing.assert_allclose(post.variance, exact_variance, rtol=1e-6)| decision.output_ports[pnl.RESPONSE_TIME], | ||
| ], | ||
| data=data, | ||
| optimization_function=PECOptimizationFunction( |
Collaborator
There was a problem hiding this comment.
One possible improvement to try: could we let a participant PEC omit optimization_function when it’s only being used for log_likelihood()? We’d still create a default PECOptimizationFunction internally for the simulation and likelihood machinery, but mark the PEC as likelihood-only and raise a helpful error if someone calls run().
That would let us remove the differential-evolution settings here, which suggest a search is being performed when it isn’t. We’d just need to keep this distinct from the hierarchical group PEC, whose run() still needs to work.
Longer term, it probably makes sense to separate simulation-based likelihood evaluation from PEC entirely, so both ordinary and hierarchical fitting can use it without constructing a PEC just to score parameters. That could be a follow-up rather than something we need to tackle in this PR. Need to talk with Jon about this though.
|
This PR causes the following changes to the html docs (ubuntu-latest-3.11): See CI logs for the full diff. |
|
This PR causes the following changes to the html docs (ubuntu-latest-3.11): See CI logs for the full diff. |
|
This PR causes the following changes to the html docs (ubuntu-latest-3.11): See CI logs for the full diff. |
A participant whose objective is impossible, named in the error rather than left to be found among the rest; a participant task failing on one worker while another is still building, which is the only arrangement in which a model outlives its fit; and two participants whose identifiers differ only in type.
The documentation offers these three files as the complete worked example, and they sat under Scripts/Debug. They are now in Scripts/Examples/ParameterEstimation/hierarchical, which leaves somewhere for the other parameter-estimation examples to go. The group fit is named for what it is rather than sharing a name with the participant models the factory builds, and the comment above it no longer describes a model being passed, which has not been true since the factory became the only place a model is declared.
A participant's uncertainty comes from the curvature of their fit at its peak. It was measured one parameter at a time, moving each while the others stayed where they were, which answers how well a parameter is determined given the others rather than how well it is determined at all. Where two parameters trade off, moving one alone makes the fit worse faster than moving it while the other compensates, so the answer comes out too confident. `hierarchical_options` gains a `curvature`, still one parameter at a time by default. Set to "full" it measures the whole matrix and inverts it, at 2P squared evaluations of a participant's objective per iteration rather than 2P. This reaches further than the reported intervals. The group variance is built from these, so measuring them too small makes the population look less varied than it is: on participants whose parameters are correlated at 0.9, the group variances the default settles on are 17 and 32 percent below the values that maximize the marginal likelihood, which the full measurement reaches exactly. The group covariance stays diagonal either way. What is chosen here is how each participant is measured, not what the group is allowed to express, and the two are separate: taking the diagonal of the second moments is the correct update for a diagonal group covariance, and the quantity it calls for is the variance of a parameter with the others integrated out, which is what the full measurement gives and the one-at-a-time measurement does not. Inverting the curvature is done in the prior's own scale, where standing in for a direction the data did not pin down and refusing to report a posterior wider than the prior are both clips on an eigenvalue. The one-at-a-time measurement is the special case where the directions are the parameters themselves.
A Gaussian objective has an exact posterior covariance, so the two settings can be compared against the answer rather than against each other, at a point where neither simulation noise nor truncation error is in play. The group-level check computes what it is comparing against by direct numerical search over the marginal likelihood, rather than by any of the code under test.
What the two settings measure, what each answers, what they cost, and the figure from a posterior where the exact answer is known: a pair of parameters correlated at 0.9 reports 0.10 one at a time against a true 0.53. Also says what the setting does not do, since the two are easy to conflate: it chooses how each participant is measured, not what the group model is allowed to express.
They ran on Linux alone, where the platform-specific parts of this code -- how many threads a compiled model is given, and how a cluster starts worker processes -- are the parts least likely to be the same elsewhere.
aryan-pan
force-pushed
the
feat/pec-hierarchical-em
branch
from
September 16, 2026 21:13
d320732 to
5c1a274
Compare
The curvature this fit reports comes from finite differences, which subtract one score from another and divide by a small number. A model that simulates fresh noise on every call returns a different score for the same parameters, and the difference then measures that noise instead of the shape of the fit, amplified by the small divisor. On a drift-diffusion model scored from 60 simulations, three calls at one parameter setting returned -65.3, -45.2 and -19.9, against curvature of order one. What makes a simulated model repeat itself is same_seed_for_all_parameter_combinations. It was written down as a requirement and left to the caller to meet. Every participant model is now checked for it before the fit begins, in the process and on a worker, because nothing in the result would have shown that it was missing: the fit runs, converges, and reports intervals measured from noise. A model that does not report the setting at all is left alone, so the driver still runs against the closed-form models it is tested with.
Both ways a participant model can lack the setting, the case where it has it, and the case where there is nothing to ask -- which is how the driver stays usable against a model that has no controller.
The page asked for common random numbers and a fixed initial seed together, as one requirement for a deterministic likelihood. Only the first of them does that: with common random numbers a model repeats itself whether or not the initial seed is fixed. Fixing the seed earns its place for two other reasons, now given: a distinct seed per participant keeps one stream of simulation noise from being shared by the whole group and absorbed into its variance, and a fixed one keeps a participant's model scoring the same way after a worker rebuilds it, which a distributed fit depends on and a single-process fit never notices. The first is now checked rather than requested, so it says so.
|
This PR causes the following changes to the html docs (ubuntu-latest-3.11): See CI logs for the full diff. |
|
This PR causes the following changes to the html docs (ubuntu-latest-3.11): See CI logs for the full diff. |
The opening paragraph still described the curvature as diagonal and explained what follows from that, which stopped being the whole story when the measurement became a setting. It now names the setting and says what does not change with it, and the two properties that return a diagonal point at the same place instead of restating it.
Fifteen lines describing what a comparison cannot detect, on a helper whose body is one line and whose job is to strip a mechanism name. The limitation belongs to `ParameterSchema.check_matches`, which is what fails to separate two parameters of the same name over the same range, and reads there as a caveat on what the check promises.
The two places that ask Dask for a worker carried the same sentence. The one on the cache lookup is the one worth having, since that is where the fallback it describes is chosen.
It was offered "for unbounded parameters", which reads as something a fit can be configured to use. A ParameterEstimationComposition cannot: a parameter is specified there as the values to search over, so it carries a range and is bounded by construction. Where the transform does apply is the EM driver used directly, against a likelihood of one's own, which the driver is built to allow. The three methods lose docstrings that restated their bodies.
|
This PR causes the following changes to the html docs (ubuntu-latest-3.11): See CI logs for the full diff. |
The distributed fitting tests were reached by three jobs that installed the extra and ran only the parameter estimation directory. Everywhere else they skipped for want of the package. The extra is three pure-python packages, and five tests turn on having it; the rest of the distributed tests exist to exercise the path taken when it is absent, and are unaffected. So it is installed on every build instead, which covers three operating systems and three versions of python rather than one version on three, and needs no job of its own.
|
This PR causes the following changes to the html docs (ubuntu-latest-3.11): See CI logs for the full diff. |
Naming it in the matrix put it in the name of every job that installed it, and needed a step and an artifact suffix to carry it. It is a feature of the package like the others the build installs, so it goes where those go, and the whole of it is one word.
|
This PR causes the following changes to the html docs (ubuntu-latest-3.11): See CI logs for the full diff. |
kmantel
requested changes
Sep 17, 2026
kmantel
left a comment
kmantel
left a comment
Collaborator
There was a problem hiding this comment.
There's a lot about this/PEC that I don't know the motivations behind, so I'll defer more to Dave; I wonder about how much this is meant as a somewhat independent library-extension/plugin vs something more tightly integrated with and designed like the rest of PNL.
| id: install | ||
| with: | ||
| features: 'dev' | ||
| features: 'dev,dask' |
Collaborator
There was a problem hiding this comment.
If this will be an optional component, we should mark all tests that require dask, and don't install dask in the main jobs, skipping those tests, to make sure everything runs fine without it. Then, add a job that installs this feature and includes the dask tests. We do something similar with torch in reverse; it's a main dependency, but we have a separate job to test without it.
Collaborator
There was a problem hiding this comment.
Yeah, this is my bad. Aryan had it setup like this at first and I told him to change it. The more I think about it I think you and him are right, it should optional and tests in a smaller separate test job.
|
|
||
|
|
||
| @dataclass | ||
| class LaplaceEMResult: |
Collaborator
There was a problem hiding this comment.
There are several shared attributes with HierarchicalPECResults and EStepResult. Are they different subtypes of a similar thing?
| context: Optional[Context] = None, | ||
| distributed: bool = False, | ||
| distributed_options: Optional[Mapping] = None, | ||
| fit_method: Optional[Literal["hierarchical"]] = None, |
Collaborator
There was a problem hiding this comment.
This and similar params that have a limited number of options should use an enum/similar, maybe PNLStrEnum
| self._run_called = False | ||
|
|
||
| #: Settings accepted by `hierarchical_options`, with their defaults. | ||
| _HIERARCHICAL_OPTION_DEFAULTS = { |
| (t.to_natural(z + h * e)[k] - t.to_natural(z - h * e)[k]) / (2 * h) | ||
| for k, e in enumerate(np.eye(2)) | ||
| ]) | ||
| assert np.allclose(t.dtheta_dz(z), num, rtol=1e-5) |
Collaborator
There was a problem hiding this comment.
use np.testing.assert_allclose, because it also checks shape, here and for others
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fits a group of participants jointly rather than one at a time, so each participant's estimate is informed by the population. Estimation is by empirical-Bayes Laplace EM, with participants drawn from
z_s ~ N(beta, diag(sigma))in an unconstrained space derived from the search ranges.hierarchical_options["subject_id"]in the data identifies participants. PEC requires every column of data to be an outcome variable, so the column is removed both from the group composition's own data and from each participant's slice before it reaches the factory. pec_factory builds one participant's model from their slice, and is required because a Composition cannot be copied.Setting distributed=True fits participants across a Dask cluster, one per task, resolved through the same helpers as distributed maximum-likelihood fitting. Results are collected by participant index rather than completion order, so distributed and in-process fits agree exactly.
Documented in
HierarchicalFitting.rst, with a runnable example that fits a simulated group and reports recovery.