[AlexanderFengler]

Spine plan: https://github.com/lnccbrown/HSSMSpine/blob/main/plans/derive-aux-networks.md#pr-stack (step H2 of 8)
Depends on: main (H1, #1438, is merged)
Unblocks: CPNs on the Hub, which are promoted only once a release contains this; the pipeline's CPN hssm_missing_load gate

Closes #1324

Why

• A missing-RT row (rt == -999.0, no deadline) still has an observed choice, but the CPN was called with the parameters only. Every such row got the probability of one fixed choice, whichever choice the participant made.
• Which choice that was depended on how the network had been trained, a convention HSSM could not see or check.

What

• Breaking: the CPN input is now [θ…, choice] -> log P(choice | θ), the same layout the OPN already uses with the deadline ([θ…, deadline]). Both have input width n_params + 1.
• HSSM wires the CPN with params_only=False, like the OPN. assemble_callables splices in the response column for missing-RT rows and the deadline column for omission rows.
• An ONNX network of any other width is rejected when the model is built, with a message that links CPN input contract: pass the observed choice to the choice-probability network #1324. The check runs under hssm.track's missing-data role, so a tracked fit still records the CPN as its missing-data network.
• Missing-RT rows whose response is not in choices now raise a clear ValueError. The data validator used to skip these rows.
• tests/fixtures/ddm_cpn.onnx is regenerated at width 5 from the DDM LAN fixture with LANfactory's derive-aux. Its provenance is in tests/fixtures/ddm_cpn.provenance.json.
• The lapse term on these rows is unchanged from H1 (1 / n_choices).
• Docstrings (HSSM, HSSMBase, make_missing_data_callable, assemble_callables, make_jax_logp_funcs_from_callable) and docs/reference/models-and-likelihoods.md describe the contract. Changelog: 0.5.1 item 19, plus a "Breaking changes" note.

Migration (hand-trained CPNs only)

• A CPN trained on the old params-only input (width n_params) is rejected at build time. Retrain it on rows whose last column is the choice, [θ…, choice], labelled with P(choice | θ).
• No CPN was ever published on franklab/HSSM, so no Hub artifact breaks.
• Callables built by hand with assemble_callables(..., params_only=True) still work. They are called as f(None, *params).

Not in this PR

• Lapse choice component is not normalised on observed and choice-only rows #1323 (H3): the lapse choice factor on observed rows and for choice-only models.

Tests

• tests/test_cpn_choice_input.py (13 tests): the fixture against its provenance, the width error for a params-only CPN and a wrong-width OPN, invalid responses on missing rows, a user callable receiving the missing rows' response, and the missing-row log-likelihood against the lapse mixture of the CPN computed with onnxruntime.
• tests/test_missing_and_deadline.py: the CPN case runs again (JAX and PyTensor agree on value and gradient).
• tests/test_make_jax_single_trial_logp_from_networks.py: the input vector layout is pinned.
• tests/test_lapse_missing_rows.py and tests/integration/test_missing_data_mcmc.py: H1's closed-form CPN test and the missing-data MCMC grid run under the new contract.

uv run ruff check . && uv run ruff format --check .
uv run pytest tests/test_cpn_choice_input.py tests/test_missing_and_deadline.py \
  tests/test_lapse_missing_rows.py tests/test_make_jax_single_trial_logp_from_networks.py \
  tests/integration/test_missing_data_mcmc.py
uv run pytest tests/ -m "not slow"

🤖 Generated with Claude Code

Summary by CodeRabbit

• Behavior Changes
  • Missing-data choice-probability networks now receive the observed choice as their final input; omission-probability networks receive the deadline.
  • Parameter-only callables remain supported when configured for parameter-only inputs.
• Validation
  • Invalid responses on missing-response-time rows are rejected for choice-probability models.
  • ONNX networks with an incompatible input width are rejected during model construction. Older parameter-only choice-probability artifacts may no longer load.
• Documentation
  • Updated guidance describes the input layout for missing-data networks and callables.

[coderabbitai]

No actionable comments were generated in the recent review.

Recent review info

Run configuration

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: 21e7e75f-07c3-4e8b-a544-1fa9aa015187

Commits

Reviewing files that changed from the base of the PR and between cacdd1d and 90d922f.

Files selected for processing: 17

• docs/changelog.md
• docs/reference/models-and-likelihoods.md
• src/hssm/base.py
• src/hssm/data_validator.py
• src/hssm/distribution_utils/dist.py
• src/hssm/distribution_utils/jax.py
• src/hssm/distribution_utils/onnx.py
• src/hssm/hssm.py
• tests/conftest.py
• tests/distribution_utils/test_jax.py
• tests/fixtures/ddm_cpn.onnx
• tests/fixtures/ddm_cpn.provenance.json
• tests/integration/test_missing_data_mcmc.py
• tests/test_cpn_choice_input.py
• tests/test_lapse_missing_rows.py
• tests/test_make_jax_single_trial_logp_from_networks.py
• tests/test_missing_and_deadline.py

Included review availability: This review used your included allowance. Your plan provides up to 2 included reviews per hour; 1 remain after this review.

Walkthrough

Missing-data CPNs now receive the observed choice as their final input; OPNs continue to receive the deadline. HSSM checks static ONNX input widths during construction and validates missing-row responses when no deadline is used. Parameter-only callables remain supported with params_only=True.

Changes: Missing-Data Network Inputs

Layer / File(s)	Summary
Route data columns to missing-data networks src/hssm/distribution_utils/dist.py, src/hssm/distribution_utils/jax.py, src/hssm/distribution_utils/onnx.py, src/hssm/hssm.py, tests/distribution_utils/test_jax.py, tests/conftest.py, tests/fixtures/ddm_cpn.provenance.json, tests/integration/test_missing_data_mcmc.py, tests/test_lapse_missing_rows.py, tests/test_make_jax_single_trial_logp_from_networks.py, tests/test_missing_and_deadline.py, tests/test_cpn_choice_input.py	assemble_callables passes the response column to CPNs and the deadline column to OPNs when params_only=False. HSSM configures CPNs to use this input path. Tests cover CPN evaluation, gradients, and single-trial and vectorized calls.
Validate network widths and missing-row responses src/hssm/hssm.py, src/hssm/data_validator.py, docs/changelog.md, docs/reference/models-and-likelihoods.md, src/hssm/base.py, tests/test_cpn_choice_input.py	HSSM rejects CPN or OPN ONNX networks with a static input width other than n_params + 1. When missing data is enabled without a deadline, invalid responses on missing-RT rows raise a ValueError. Documentation describes the input contract and migration behavior.

Priority: Normal

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Bug fix

Sequence Diagram(s)

sequenceDiagram
  participant HSSM
  participant assemble_callables
  participant CPN
  HSSM->>assemble_callables: Send model parameters and missing-row data
  assemble_callables->>CPN: Send parameters followed by observed choice
  CPN-->>assemble_callables: Return choice log-probability
  assemble_callables-->>HSSM: Return missing-row log-probabilities

Loading

Merge Risk: Minimal · up to 90d92

Missing-row CPNs receive the observed choice. No material merge risk remains from the reviewed changes.

Pre-merge checks: 5

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title "Pass the observed choice to the CPN as its last input (#1324)" matches the main change in the changeset. The raw summaries show the CPN input contract change: networks receive `[θ…, choice]…
Linked Issues check		Issue #1324 is directly linked and has coding requirements. The supplied PR evidence shows that HSSM uses params_only=False for CPNs, assemble_callables appends each missing-RT row's response, a…
Out of Scope Changes check		The supplied whole-PR summary ties the implementation, validation, ONNX fixture and provenance, documentation, changelog, and tests to issue #1324. The OPN width check shares the documented `n_params …
Docstring Coverage		Docstring coverage is 86.36% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentati…

Full details: Title check

Explanation

The title "Pass the observed choice to the CPN as its last input (#1324)" matches the main change in the changeset. The raw summaries show the CPN input contract change: networks receive [θ…, choice] as the final input, the width check, and the response validation on missing-RT rows. The linked objectives describe the same change, and the issue number is referenced. The title is concise, specific, and a single sentence.

Full details: Linked Issues check

Explanation

Issue #1324 is directly linked and has coding requirements. The supplied PR evidence shows that HSSM uses params_only=False for CPNs, assemble_callables appends each missing-RT row's response, and the validator rejects invalid missing-row responses with a reference to #1324. HSSM rejects fixed-width ONNX CPN and OPN inputs that are not n_params + 1, and the check runs under missing-data tracking. The DDM CPN fixture was regenerated at width 5. The CPN test case was re-enabled, and the reported tests cover response propagation, width validation, fixture behavior, lapse mixing, and gradients. The evidence also states that the 1 / n_choices lapse term is unchanged.

Full details: Out of Scope Changes check

Explanation

The supplied whole-PR summary ties the implementation, validation, ONNX fixture and provenance, documentation, changelog, and tests to issue #1324. The OPN width check shares the documented n_params + 1 contract and supports the common missing-data network path. The params_only=True callable support documents and preserves the migration compatibility stated in the PR. No unrelated change is demonstrated. Issue #1323 is explicitly excluded.

Full details: Docstring Coverage

Explanation

Docstring coverage is 86.36% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 22 functions across 13 files. (3 skipped: 3 unsupported.)

Finishing Touches

📝 Generate docstrings

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}