[codex] Harden federated longitudinal SuStaIn

[armaneshaghi]

Summary

This review branch hardens the federated/longitudinal SuStaIn work added today. It includes Claude's exact hot-path optimisations and golden regression guard, plus review fixes for packaging and aggregate-only federation semantics.

Changes:

• Use package discovery in setup.py so pySuStaIn.federated and experiment modules are included in real wheel/sdist installs.
• Expose LongitudinalZscoreSustain from the top-level pySuStaIn package.
• Make LongitudinalZScoreSustainData explicitly satisfy AbstractSustainData and validate visit/subject shapes.
• Vectorise the longitudinal monotone-path dynamic program by visit count.
• Cache deterministic client-side stage likelihoods with a bounded LRU cache.
• Return aggregate subtype/stage summaries by default from the federated server, with explicit return_individual=True only for local simulations/debugging.
• Add golden regression tests to pin exact cross-sectional and longitudinal deterministic outputs.

Review notes

The main issue I found in today's committed branch was packaging: setup.py only listed pySuStaIn and sim, so a normal package install would omit the new federated subpackage. I also thought the previous server-side subtype_and_stage() API conflicted with the aggregate-only privacy language because it returned row-level assignments to the server by default.

Validation

• .venv/bin/python -m pytest -q -> 8 passed, 1 skipped
• rg --files -g "*.py" | xargs -r .venv/bin/python -m py_compile
• git diff --check
• Built a wheel locally and confirmed the federated package, experiments module, and longitudinal module are present.

[armaneshaghi]

@copilot review please

[armaneshaghi]

Review (Claude) — reviewed & addressed, leaving for your final review

Verdict: looks solid. CI (federated-sustain tests) is green on the latest commit.

Checked:

• Hot-path optimisations are numerically exact. The bounded per-sequence stage-likelihood cache and the vectorised longitudinal monotone-path DP reproduce the pre-optimisation outputs to machine precision — pinned by tests/test_regression_golden.py (cross-sectional + longitudinal golden loglike/sequences) plus the fed==pooled equivalence tests. 8 passed, 1 skipped locally.
• Packaging fix correct — find_packages() picks up pySuStaIn.federated{,.experiments}; .venv/tests excluded (no __init__.py at those roots). Top-level LongitudinalZscoreSustain export imports cleanly (no kde_ebm needed).
• Aggregate-only server API (subtype_and_stage(..., return_individual=False) default) — good privacy fix, covered by a test.
• Longitudinal correctness — one visit/subject is identical to cross-sectional (diff 0); full-scale 3-subtype fed-long == pooled-long (loglike diff ~7e-12, sequences identical).

Addressed:

• Added the longitudinal validation results to RESULTS.md (PR added the longitudinal model but documented only cross-sectional numbers) and refreshed the limitations.

No outstanding issues from my side. Roadmap (non-blocking): federated MCMC uncertainty, federated K-selection, further sequence-search batching. Leaving open for your review — not merging.

[armaneshaghi]

I have reviewed this PR. The mathematical approach for the longitudinal path dynamic programming and the federated cache and aggregation is entirely sound. However, I discovered an incompatibility with Numpy 2.x (TypeErrors due to array-to-scalar casting) which breaks the tests. I've created PR #2 to fix these Numpy 2.x compatibility issues on this branch. Once PR #2 is merged into this branch, this PR is good to merge!